Activation offloading to CPU's for the Linear, Layernorm Linear and the Layernorm MLP modules

docs/api/pytorch.rst

@@ -40,3 +40,5 @@ pyTorch
 .. autoapifunction:: transformer_engine.pytorch.checkpoint
 
 .. autoapifunction:: transformer_engine.pytorch.onnx_export
+
+.. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context

tests/pytorch/test_sanity.py

@@ -4,6 +4,7 @@
 
 from dataclasses import dataclass
 from typing import Optional
+from contextlib import nullcontext
 
 import torch
 import pytest
@@ -20,6 +21,7 @@
     TransformerLayer,
     RMSNorm,
     LayerNorm,
+    get_cpu_offload_context,
 )
 from transformer_engine.common import recipe
 
@@ -215,17 +217,24 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, dtype, config, fp8_reci
             assert torch.count_nonzero(p.main_grad) > 0, "Gradient not accumulated."
 
 
-def _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad):
+def _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad, cpu_offload):
     te_inp_hidden_states = torch.randn(
         config.seq_len, config.batch_size, config.hidden_size, dtype=dtype, requires_grad=True
     ).cuda()
 
     if skip_wgrad:
         _disable_wgrads(block)
 
+    if cpu_offload:
+        offload_context, sync_function = get_cpu_offload_context(enabled=True)
+    else:
+        offload_context = nullcontext()
+        sync_function = lambda x: x
+
     use_fp8 = fp8_recipe is not None
-    with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
+    with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe), offload_context:
         te_out = block(te_inp_hidden_states)
+    te_out = sync_function(te_out)
     loss = te_out.sum()
     loss.backward()
     torch.cuda.synchronize()
@@ -449,9 +458,11 @@ def test_sanity_layernorm_mlp(dtype, fp8_recipe, model, skip_wgrad,
 @pytest.mark.parametrize("activation", all_activations)
 @pytest.mark.parametrize("normalization", all_normalizations)
 @pytest.mark.parametrize("parallel_attention_mlp", all_boolean)
+@pytest.mark.parametrize("cpu_offload", all_boolean)
 def test_sanity_gpt(dtype, fp8_recipe, model, skip_wgrad,
                     zero_centered_gamma, bias, activation,
-                    normalization, parallel_attention_mlp):
+                    normalization, parallel_attention_mlp,
+                    cpu_offload):
     config = model_configs[model]
 
     if fp8_recipe is not None:
@@ -489,7 +500,7 @@ def test_sanity_gpt(dtype, fp8_recipe, model, skip_wgrad,
         .cuda()
     )
 
-    _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad)
+    _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad, cpu_offload)
 
 
 def test_sanity_gpt_126m():
@@ -512,6 +523,7 @@ def test_sanity_gpt_126m():
         activation="gelu",
         normalization="LayerNorm",
         parallel_attention_mlp=False,
+        cpu_offload=False,
     )
 
 
@@ -713,7 +725,7 @@ def test_sanity_drop_path(dtype, fp8_recipe, model, skip_wgrad):
         .cuda()
     )
 
-    _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad)
+    _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad, False)
 
 
 @pytest.mark.parametrize("dtype", param_types)
@@ -751,7 +763,7 @@ def test_sanity_fused_qkv_params(dtype, fp8_recipe, model, skip_wgrad):
         .cuda()
     )
 
-    _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad)
+    _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad, False)
 
 
 @pytest.mark.parametrize("dtype", param_types)

transformer_engine/pytorch/__init__.py

@@ -17,6 +17,7 @@
 from .export import onnx_export
 from .distributed import checkpoint
 from .distributed import CudaRNGStatesTracker
+from .cpu_offload import get_cpu_offload_context
 # Register custom op symbolic ONNX functions
 from .te_onnx_extensions import (
     onnx_cast_to_fp8,