Repository: NVIDIA/TransformerEngine
Pull request #1514: Verified TE2.0 with offloading
Author: login: sanandaraj5597
Created: 2025-02-27T08:50:05Z | Merged: 2025-03-09T09:38:55Z

# Description

This PR adds CPU offloading support for TE 2.0. This PR:
1) Fixes all bugs needed to support TE 2.0 with offloading
2) Writes a clean unit test to support CPU activation offloading
3) Added the CPU offloading test to TE unit test suite



Diff:
diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh
index 56d668bd12..e2fe2c0200 100644
--- a/qa/L0_pytorch_unittest/test.sh
+++ b/qa/L0_pytorch_unittest/test.sh
@@ -24,6 +24,7 @@ pytest -v -s $TE_PATH/tests/pytorch/test_multi_tensor.py || FAIL=1
 pytest -v -s $TE_PATH/tests/pytorch/test_fusible_ops.py || FAIL=1
 pytest -v -s $TE_PATH/tests/pytorch/test_permutation.py || FAIL=1
 pytest -v -s $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || FAIL=1
+pytest -v -s $TE_PATH/tests/pytorch/test_cpu_offloading.py || FAIL=1
 NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 pytest -o log_cli=true --log-cli-level=INFO -v -s $TE_PATH/tests/pytorch/fused_attn/test_fused_attn.py || FAIL=1
 
 exit $FAIL
diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py
index 61b4a2553c..ed7cdda85b 100644
--- a/tests/pytorch/test_cpu_offloading.py
+++ b/tests/pytorch/test_cpu_offloading.py
@@ -7,8 +7,12 @@
 from contextlib import nullcontext
 
 import transformer_engine.pytorch as te
+from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
 
-SIZE = 4096
+# Check if FP8 supported
+fp8_available, reason_for_no_fp8 = FP8GlobalStateManager.is_fp8_available()
+
+SIZE = 512
 
 models = {
     "linear": te.Linear,
@@ -18,40 +22,64 @@
 
 
 def _get_input():
-    return torch.empty((1, SIZE, SIZE)).cuda()  # input size - 1 * 2048 * 2048 * 4b = 16MB
+    return torch.empty((128, SIZE, SIZE)).cuda()
 
 
 def _measure_memory_between_forward_and_backward(model_cls, fp8, cpu_offload):
-    torch.cuda.empty_cache()
-    model = model_cls(SIZE, SIZE, 1)
+
+    input_layer = model_cls(SIZE, SIZE)
+    hidden_layer = model_cls(SIZE, SIZE)
+    output_layer = model_cls(SIZE, SIZE)
 
     input = _get_input()
     if cpu_offload:
-        offload_context, sync_function = te.get_cpu_offload_context(enabled=True)
+        offload_context, sync_function = te.get_cpu_offload_context(
+            enabled=True,
+            num_layers=2,
+            model_layers=3,
+            offload_activations=True,
+            offload_weights=False,
+        )
     else:
         offload_context = nullcontext()
         sync_function = lambda x: x
 
     with te.fp8_autocast(enabled=fp8), offload_context:
-        out = model(input)
+        out = input_layer(input)
+    out = sync_function(out)
+    with te.fp8_autocast(enabled=fp8), offload_context:
+        out = hidden_layer(out)
     out = sync_function(out)
-    input.data = torch.Tensor()  # delete data from input
-    out.data = torch.Tensor()  # delete data from out
+    with te.fp8_autocast(enabled=fp8), offload_context:
+        out = output_layer(out)
+    out = sync_function(out)
+
+    max_mem_used = torch.cuda.memory_allocated() / 1024**2
+
+    out.sum().backward()
+
+    del input_layer
+    del hidden_layer
+    del output_layer
     del input
     del out
-    torch.cuda.empty_cache()
-    allocated_memory_mb = torch.cuda.memory_allocated() / 1024**2
-    del model
-    return allocated_memory_mb
 
+    torch.cuda.synchronize()
+
+    return max_mem_used
 
-@pytest.mark.parametrize("fp8", [False, True])
+
+@pytest.mark.parametrize("fp8", [True, False])
 @pytest.mark.parametrize("model_key", models.keys())
 def test_cpu_offload(fp8, model_key) -> None:
+
+    if fp8 and not fp8_available:
+        pytest.skip(reason_for_no_fp8)
+
     model_cls = models[model_key]
+
     without_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, False)
-    torch.cuda.empty_cache()
+
     with_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, True)
 
-    assert without_offloading > 30
-    assert with_offloading < 10
+    assert with_offloading < without_offloading
diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py
index c47130fe78..8294a76742 100644
--- a/transformer_engine/pytorch/cpu_offload.py
+++ b/transformer_engine/pytorch/cpu_offload.py
@@ -137,9 +137,7 @@ def __init__(
         super().__init__()
 
     def on_save_for_backward(self, tensor: torch.Tensor) -> Any:
-        retrieve_identifier = self.offload_handler.tensor_push(
-            tensor.data, **self.handler_extra_kwargs
-        )
+        retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs)
         return retrieve_identifier
 
     def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor:
@@ -235,19 +233,15 @@ def on_group_commit_backward(self):
     @staticmethod
     def offload(src_tensor, pin_memory=True):
         """Offload."""
-        fp8_offload = isinstance(src_tensor, Float8Tensor)
 
         cpu_backup = torch.empty(
             src_tensor.size(),
-            dtype=torch.uint8 if fp8_offload else src_tensor.dtype,
+            dtype=src_tensor.dtype,
             layout=src_tensor.layout,
             device="cpu",
             pin_memory=pin_memory,
         )
 
-        if fp8_offload:
-            cpu_backup = Float8Tensor.make_like(src_tensor, data=cpu_backup)
-
         cpu_backup.copy_(src_tensor, non_blocking=pin_memory)
         state = (src_tensor.device, cpu_backup)
         return state
@@ -311,6 +305,9 @@ def __init__(
         self.num_layers = num_model_group
         # Data Structure to maintain reference to activation tensors
         self.tensor_tag_to_buf = {}
+        # Data structure to hold the FP8/MXFP8 tensor objects
+        self.fp8_tensor_object_map = {}
+        self.float8_transpose_cache_valid = {}
         # Tracking the number of layers offloaded
         self.offloaded_group_count = 0
         # Core data structure that decides the window for offloading
@@ -341,18 +338,46 @@ def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
             ),
         )
 
+        is_quantized_tensor = callable(getattr(tensor, "prepare_for_saving", None))
+
         if not torch_stray_tensor:
+
             # obtain a unique tensor tag
             tensor_tag = (self.current_group, self.tensor_count_current_group)
             self.tensor_count_current_group += 1
+
             assert tensor_tag not in self.tensor_tag_to_state
 
-            self.tensor_tag_to_state[tensor_tag] = tensor
+            if is_quantized_tensor:
+                tensor_list, _ = tensor.prepare_for_saving()
+
+                self.tensor_tag_to_state[tensor_tag] = []
+                self.tensor_tag_to_buf[tensor_tag] = []
 
-            if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(
-                tensor
-            ):
-                self.tensor_tag_to_buf[tensor_tag] = tensor
+                self.fp8_tensor_object_map[tensor_tag] = tensor
+                if isinstance(tensor, Float8Tensor):
+                    self.float8_transpose_cache_valid[tensor_tag] = getattr(
+                        tensor, "_transpose_invalid"
+                    )
+            else:
+                tensor_list = [tensor]
+
+            for t in tensor_list:
+                if is_quantized_tensor:
+                    self.tensor_tag_to_state[tensor_tag].append(t)
+                else:
+                    self.tensor_tag_to_state[tensor_tag] = t
+
+                if (
+                    self.current_group < self.num_offload_group
+                    and self.tensor_need_offloading_checker(t)
+                ):
+                    if is_quantized_tensor:
+                        self.tensor_tag_to_buf[tensor_tag].append(t)
+                        # Need to clear the internal data reference for the quantized tensors
+                        tensor.clear()
+                    else:
+                        self.tensor_tag_to_buf[tensor_tag] = t
         else:
             tensor_tag = (-1, self.torch_tensor_count)
             self.torch_tensor_count += 1
@@ -364,7 +389,14 @@ def tensor_pop(self, tensor_tag, **kwargs):
         """Tensor pop."""
         assert tensor_tag in self.tensor_tag_to_state
         tensor = self.tensor_tag_to_state.pop(tensor_tag)
+
+        # Handling the quantized tensor case specially here
+        if isinstance(tensor, list):
+            self.fp8_tensor_object_map[tensor_tag].restore_from_saved(tensor)
+            tensor = self.fp8_tensor_object_map.pop(tensor_tag)
+
         self.tensor_tag_to_buf.pop(tensor_tag, None)
+
         # the tensor should have been copied back in on_group_commit_backward()
         # which invokes bulk_reload_group.
         assert not isinstance(tensor, tuple)
@@ -377,13 +409,23 @@ def bulk_offload_group(self, group_to_offload):
                 group_id, _ = tensor_tag
                 if group_id == group_to_offload:
                     assert not isinstance(state, tuple)
-                    tensor_on_device = state
 
-                    # if offload, return the reference to cpu copy
-                    if self.tensor_need_offloading_checker(tensor_on_device):
-                        state = SynchronizedGroupOffloadHandler.offload(tensor_on_device)
-                        self.tensor_tag_to_state[tensor_tag] = state
-                        tensor_on_device.data = torch.Tensor()  # Force to release memory
+                    is_quantized_tensor = isinstance(state, list)
+
+                    if is_quantized_tensor:
+                        tensor_list = state
+                        self.tensor_tag_to_state[tensor_tag] = []
+                    else:
+                        tensor_list = [state]
+
+                    for tensor_on_device in tensor_list:
+                        # if offload, return the reference to cpu copy
+                        if self.tensor_need_offloading_checker(tensor_on_device):
+                            state = SynchronizedGroupOffloadHandler.offload(tensor_on_device)
+                        if is_quantized_tensor:
+                            self.tensor_tag_to_state[tensor_tag].append(state)
+                        else:
+                            self.tensor_tag_to_state[tensor_tag] = state
 
     def synchronize_on_group_commit_forward(self, current_group):
         """Synchronize on group commit forward."""
@@ -433,6 +475,23 @@ def bulk_reload_group(self, group_to_reload):
                     if isinstance(state, tuple):
                         recovered_tensor = SynchronizedGroupOffloadHandler.reload(state)
                         self.tensor_tag_to_state[tensor_label] = recovered_tensor
+                    elif isinstance(state, list):
+                        tensor_list = []
+                        for state_tuple in state:
+                            if isinstance(state_tuple, tuple):
+                                tensor_list.append(
+                                    SynchronizedGroupOffloadHandler.reload(state_tuple)
+                                )
+                            else:
+                                tensor_list.append(state_tuple)
+                        _ = self.fp8_tensor_object_map[tensor_label].restore_from_saved(tensor_list)
+                        if isinstance(self.fp8_tensor_object_map[tensor_label], Float8Tensor):
+                            self.fp8_tensor_object_map[tensor_label]._transpose_invalid = (
+                                self.float8_transpose_cache_valid.pop(tensor_label)
+                            )
+                        self.tensor_tag_to_state[tensor_label] = self.fp8_tensor_object_map.pop(
+                            tensor_label
+                        )
 
     def on_group_commit_backward(self):
         # first decrement the current group.
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index 5944039cf0..c9e65bd93a 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -348,10 +348,6 @@ def clear(self):
         self._transpose = torch.Tensor() if self._transpose is not None else None
         self._transpose_invalid = True
 
-    def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], Float8TensorBase]:
-        """Prepare the tensor base for saving for backward"""
-        return [self], None
-
     @classmethod
     def __torch_dispatch__(cls, func, types, args, kwargs=None):
 
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index db369de803..6e3835fbef 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -285,10 +285,6 @@ def clear(self):
         self._rowwise_data = torch.Tensor() if self._rowwise_data is not None else None
         self._columnwise_data = torch.Tensor() if self._columnwise_data is not None else None
 
-    def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorBase]:
-        """Prepare the tensor base for saving for backward"""
-        return [self], None
-
     @classmethod
     def __torch_dispatch__(cls, func, types, args, kwargs=None):
 
diff --git a/transformer_engine/pytorch/tensor/quantized_tensor.py b/transformer_engine/pytorch/tensor/quantized_tensor.py
index b540cd91a1..00815452da 100644
--- a/transformer_engine/pytorch/tensor/quantized_tensor.py
+++ b/transformer_engine/pytorch/tensor/quantized_tensor.py
@@ -27,7 +27,7 @@ def prepare_for_saving(
         if tensor is None:
             tensor_list.append(None)
             tensor_objects_list.append(None)
-        elif type(tensor) in (torch.Tensor, torch.nn.Parameter):
+        elif isinstance(tensor, torch.Tensor):
             tensor_list.append(tensor)
             tensor_objects_list.append(None)
         else:


Reviews:
user: login: ptrendx
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-27T23:26:41Z
user: login: sanandaraj5597
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-02-27T23:33:31Z
user: login: ptrendx
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-05T00:17:01Z
user: login: sudhakarsingh27
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-06T00:15:50Z
user: login: ptrendx
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-06T01:24:17Z
user: login: ksivaman
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-06T04:38:37Z
user: login: sanandaraj5597
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T06:39:33Z
user: login: sanandaraj5597
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T06:39:42Z
user: login: sanandaraj5597
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T06:39:57Z
user: login: ksivaman
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-06T08:52:20Z
user: login: ksivaman
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-06T13:24:33Z
user: login: ksivaman
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-03-07T06:32:18Z
user: login: sanandaraj5597
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T06:38:54Z
user: login: ksivaman
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-07T08:27:48Z

Review comments:
diff_hunk: @@ -105,8 +105,8 @@ def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], Float8Tensor
 
         """
         tensors = [self._data, self._transpose]
-        # self._data = None
-        # self._transpose = None
+        self._data = None
+        self._transpose = None
path: transformer_engine/pytorch/tensor/_internal/float8_tensor_base.py
user: login: ptrendx
body: This is being handled by #1500.
created_at: 2025-02-27T23:26:41Z
author_association: MEMBER
diff_hunk: @@ -105,8 +105,8 @@ def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], Float8Tensor
 
         """
         tensors = [self._data, self._transpose]
-        # self._data = None
-        # self._transpose = None
+        self._data = None
+        self._transpose = None
path: transformer_engine/pytorch/tensor/_internal/float8_tensor_base.py
user: login: sanandaraj5597
body: Sounds good, this was creating a nested reference making offloading unable to free up the GPU memory. I can do a merge conflict resolution once the other PR is merged.

Unit tests will fail if we revert this for MXFP8.
created_at: 2025-02-27T23:33:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,7 +27,7 @@ def prepare_for_saving(
         if tensor is None:
             tensor_list.append(None)
             tensor_objects_list.append(None)
-        elif type(tensor) in (torch.Tensor, torch.nn.Parameter):
+        elif isinstance(tensor, (torch.Tensor, torch.nn.Parameter)):
path: transformer_engine/pytorch/tensor/quantized_tensor.py
user: login: sudhakarsingh27
body: `torch.nn.Parameter` are `torch.Tensor` subclasses: https://pytorch.org/docs/stable/generated/torch.nn.parameter.Parameter.html
created_at: 2025-03-06T00:15:46Z
author_association: COLLABORATOR
diff_hunk: @@ -341,18 +342,44 @@ def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
             ),
         )
 
+        fp8_tensor = isinstance(
+            tensor,
+            (
+                MXFP8Tensor,
+                Float8Tensor,
+            ),
+        )
+
         if not torch_stray_tensor:
+
             # obtain a unique tensor tag
             tensor_tag = (self.current_group, self.tensor_count_current_group)
             self.tensor_count_current_group += 1
-            assert tensor_tag not in self.tensor_tag_to_state
 
-            self.tensor_tag_to_state[tensor_tag] = tensor
+            assert tensor_tag not in self.tensor_tag_to_state
 
-            if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(
-                tensor
-            ):
-                self.tensor_tag_to_buf[tensor_tag] = tensor
+            if fp8_tensor:
path: transformer_engine/pytorch/cpu_offload.py
user: login: ptrendx
body: I don't like the fact that we encode a list of possible quantized tensors here as that is going to be a surprise place to fix when we add a new tensor type. Could we maybe just check if the tensor has `prepare_for_saving` attribute instead of looking its class in a hardcoded list?
created_at: 2025-03-06T01:24:17Z
author_association: MEMBER
diff_hunk: @@ -27,7 +27,7 @@ def prepare_for_saving(
         if tensor is None:
             tensor_list.append(None)
             tensor_objects_list.append(None)
-        elif type(tensor) in (torch.Tensor, torch.nn.Parameter):
+        elif isinstance(tensor, (torch.Tensor, torch.nn.Parameter)):
path: transformer_engine/pytorch/tensor/quantized_tensor.py
user: login: ksivaman
body: ```suggestion
        elif isinstance(tensor, torch.Tensor):
```
created_at: 2025-03-06T04:38:36Z
author_association: MEMBER
diff_hunk: @@ -27,7 +27,7 @@ def prepare_for_saving(
         if tensor is None:
             tensor_list.append(None)
             tensor_objects_list.append(None)
-        elif type(tensor) in (torch.Tensor, torch.nn.Parameter):
+        elif isinstance(tensor, (torch.Tensor, torch.nn.Parameter)):
path: transformer_engine/pytorch/tensor/quantized_tensor.py
user: login: sanandaraj5597
body: Fixed it.
created_at: 2025-03-06T06:39:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -341,18 +342,44 @@ def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
             ),
         )
 
+        fp8_tensor = isinstance(
+            tensor,
+            (
+                MXFP8Tensor,
+                Float8Tensor,
+            ),
+        )
+
         if not torch_stray_tensor:
+
             # obtain a unique tensor tag
             tensor_tag = (self.current_group, self.tensor_count_current_group)
             self.tensor_count_current_group += 1
-            assert tensor_tag not in self.tensor_tag_to_state
 
-            self.tensor_tag_to_state[tensor_tag] = tensor
+            assert tensor_tag not in self.tensor_tag_to_state
 
-            if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(
-                tensor
-            ):
-                self.tensor_tag_to_buf[tensor_tag] = tensor
+            if fp8_tensor:
path: transformer_engine/pytorch/cpu_offload.py
user: login: sanandaraj5597
body: Agreed, updated the code.
created_at: 2025-03-06T06:39:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,7 +27,7 @@ def prepare_for_saving(
         if tensor is None:
             tensor_list.append(None)
             tensor_objects_list.append(None)
-        elif type(tensor) in (torch.Tensor, torch.nn.Parameter):
+        elif isinstance(tensor, (torch.Tensor, torch.nn.Parameter)):
path: transformer_engine/pytorch/tensor/quantized_tensor.py
user: login: sanandaraj5597
body: Nice catch, fixed.
created_at: 2025-03-06T06:39:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -341,18 +342,38 @@ def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
             ),
         )
 
+        quantized_tensor = callable(getattr(tensor, "prepare_for_saving", None))
path: transformer_engine/pytorch/cpu_offload.py
user: login: ksivaman
body: Nitpick: `quantized_tensor` &rarr; `is_quantized_tensor`
created_at: 2025-03-06T08:52:20Z
author_association: MEMBER
diff_hunk: @@ -100,8 +100,8 @@ def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorB
 
         """
         tensors = [self._rowwise_data, self._columnwise_data]
-        self._rowwise_data = None
-        self._columnwise_data = None
+        # self._rowwise_data = None
+        # self._columnwise_data = None
path: transformer_engine/pytorch/tensor/_internal/mxfp8_tensor_base.py
user: login: ksivaman
body: Wait, why are we removing these again?
created_at: 2025-03-07T06:32:11Z
author_association: MEMBER
diff_hunk: @@ -100,8 +100,8 @@ def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorB
 
         """
         tensors = [self._rowwise_data, self._columnwise_data]
-        self._rowwise_data = None
-        self._columnwise_data = None
+        # self._rowwise_data = None
+        # self._columnwise_data = None
path: transformer_engine/pytorch/tensor/_internal/mxfp8_tensor_base.py
user: login: sanandaraj5597
body: Do you still think we need this?
created_at: 2025-03-07T06:38:54Z
author_association: CONTRIBUTOR

PR comments:
user: login: ptrendx
created_at: 2025-03-03T22:35:06Z
body: @sanandaraj5597 Could you resolve the merge conflicts?
author_association: MEMBER
user: login: sanandaraj5597
created_at: 2025-03-03T23:12:00Z
body: @ptrendx Fixed merge conflicts. Please trigger CI if you are fine with everything.
author_association: CONTRIBUTOR
user: login: ksivaman
created_at: 2025-03-04T16:32:25Z
body: /te-ci pytorch
author_association: MEMBER
user: login: ptrendx
created_at: 2025-03-05T00:19:57Z
body: @sanandaraj5597 You need to skip the FP8 test on the devices that do not support FP8 (like Ampere).
author_association: MEMBER
user: login: ksivaman
created_at: 2025-03-06T13:41:39Z
body: /te-ci pytorch
author_association: MEMBER
user: login: ksivaman
created_at: 2025-03-07T06:30:12Z
body: /te-ci pytorch
author_association: MEMBER
user: login: sanandaraj5597
created_at: 2025-03-08T07:40:13Z
body: /te-ci pytorch
author_association: CONTRIBUTOR
user: login: DXZDXZ
created_at: 2025-03-15T03:56:51Z
body: Hello,

I found some issues when testing the CPU Offload feature, and the cuda memory was not truly released. As shown below:

```
============================================================================================= test session starts =============================================================================================
platform linux -- Python 3.10.12, pytest-8.1.1, pluggy-1.5.0
rootdir: /path/to/my/TransformerEngine
plugins: hydra-core-1.3.2, typeguard-4.4.1, anyio-4.6.0, hypothesis-5.35.1, flakefinder-1.1.0, rerunfailures-14.0, shard-0.1.2, xdist-3.6.1, xdoctest-1.0.2
collected 6 items                                                                                                                                                                                             
Running 6 items in this shard

tests/pytorch/test_cpu_offloading.py FFFFFF                                                                                                                                                             [100%]

================================================================================================== FAILURES ===================================================================================================
________________________________________________________________________________________ test_cpu_offload[linear-True] ________________________________________________________________________________________

fp8 = True, model_key = 'linear'

    @pytest.mark.parametrize("fp8", [True, False])
    @pytest.mark.parametrize("model_key", models.keys())
    def test_cpu_offload(fp8, model_key) -> None:
    
        if fp8 and not fp8_available:
            pytest.skip(reason_for_no_fp8)
    
        model_cls = models[model_key]
    
        without_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, False)
    
        with_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, True)
    
>       assert with_offloading < without_offloading
E       assert 484.384765625 < 484.3232421875

tests/pytorch/test_cpu_offloading.py:85: AssertionError
_______________________________________________________________________________________ test_cpu_offload[linear-False] ________________________________________________________________________________________

fp8 = False, model_key = 'linear'

    @pytest.mark.parametrize("fp8", [True, False])
    @pytest.mark.parametrize("model_key", models.keys())
    def test_cpu_offload(fp8, model_key) -> None:
    
        if fp8 and not fp8_available:
            pytest.skip(reason_for_no_fp8)
    
        model_cls = models[model_key]
    
        without_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, False)
    
        with_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, True)
    
>       assert with_offloading < without_offloading
E       assert 547.12890625 < 547.12890625

tests/pytorch/test_cpu_offloading.py:85: AssertionError
____________________________________________________________________________________ test_cpu_offload[layernorm_mlp-True] _____________________________________________________________________________________

fp8 = True, model_key = 'layernorm_mlp'

    @pytest.mark.parametrize("fp8", [True, False])
    @pytest.mark.parametrize("model_key", models.keys())
    def test_cpu_offload(fp8, model_key) -> None:
    
        if fp8 and not fp8_available:
            pytest.skip(reason_for_no_fp8)
    
        model_cls = models[model_key]
    
        without_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, False)
    
        with_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, True)
    
>       assert with_offloading < without_offloading
E       assert 1226.8984375 < 1226.7783203125

tests/pytorch/test_cpu_offloading.py:85: AssertionError
____________________________________________________________________________________ test_cpu_offload[layernorm_mlp-False] ____________________________________________________________________________________

fp8 = False, model_key = 'layernorm_mlp'

    @pytest.mark.parametrize("fp8", [True, False])
    @pytest.mark.parametrize("model_key", models.keys())
    def test_cpu_offload(fp8, model_key) -> None:
    
        if fp8 and not fp8_available:
            pytest.skip(reason_for_no_fp8)
    
        model_cls = models[model_key]
    
        without_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, False)
    
        with_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, True)
    
>       assert with_offloading < without_offloading
E       assert 1703.88671875 < 1703.88671875

tests/pytorch/test_cpu_offloading.py:85: AssertionError
___________________________________________________________________________________ test_cpu_offload[layernorm_linear-True] ___________________________________________________________________________________

fp8 = True, model_key = 'layernorm_linear'

    @pytest.mark.parametrize("fp8", [True, False])
    @pytest.mark.parametrize("model_key", models.keys())
    def test_cpu_offload(fp8, model_key) -> None:
    
        if fp8 and not fp8_available:
            pytest.skip(reason_for_no_fp8)
    
        model_cls = models[model_key]
    
        without_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, False)
    
        with_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, True)
    
>       assert with_offloading < without_offloading
E       assert 742.509765625 < 742.4482421875

tests/pytorch/test_cpu_offloading.py:85: AssertionError
__________________________________________________________________________________ test_cpu_offload[layernorm_linear-False] ___________________________________________________________________________________

fp8 = False, model_key = 'layernorm_linear'

    @pytest.mark.parametrize("fp8", [True, False])
    @pytest.mark.parametrize("model_key", models.keys())
    def test_cpu_offload(fp8, model_key) -> None:
    
        if fp8 and not fp8_available:
            pytest.skip(reason_for_no_fp8)
    
        model_cls = models[model_key]
    
        without_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, False)
    
        with_offloading = _measure_memory_between_forward_and_backward(model_cls, fp8, True)
    
>       assert with_offloading < without_offloading
E       assert 933.00390625 < 933.00390625

tests/pytorch/test_cpu_offloading.py:85: AssertionError
============================================================================================== warnings summary ===============================================================================================
tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-False]
tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-False]
tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-False]
  /usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/jit.py:113: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
    with torch.cuda.amp.autocast(enabled=False):

tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-False]
tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-False]
tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-False]
  /usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/jit.py:123: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
    with torch.cuda.amp.autocast(enabled=False):

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=========================================================================================== short test summary info ===========================================================================================
FAILED tests/pytorch/test_cpu_offloading.py::test_cpu_offload[linear-True] - assert 484.384765625 < 484.3232421875
FAILED tests/pytorch/test_cpu_offloading.py::test_cpu_offload[linear-False] - assert 547.12890625 < 547.12890625
FAILED tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-True] - assert 1226.8984375 < 1226.7783203125
FAILED tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_mlp-False] - assert 1703.88671875 < 1703.88671875
FAILED tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_linear-True] - assert 742.509765625 < 742.4482421875
FAILED tests/pytorch/test_cpu_offloading.py::test_cpu_offload[layernorm_linear-False] - assert 933.00390625 < 933.00390625
======================================================================================== 6 failed, 6 warnings in 8.24s ========================================================================================
```

Also, I noticed the following code changes in this pr:

<img width="1533" alt="image" src="https://github.com/user-attachments/assets/93d81f39-6338-47aa-b106-b15ebc71db9d" />

So where is the cuda memory released?




author_association: NONE