fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Author: pggPL

tests/pytorch/test_cpu_offloading.py

@@ -119,7 +119,6 @@ def test_cpu_offload(fp8_recipe, model_key) -> None:
     the difference being the size of the FP8 cache that is not offloaded to the CPU.
     We also expect this memory consumption to be smaller than in scenario (1).
     """
-
     model_cls = model_types[model_key]
     models_list = [model_cls() for _ in range(NUM_LAYERS)]
 

transformer_engine/pytorch/cpu_offload.py

@@ -22,7 +22,7 @@ def mark_activation_offload(*tensors):
         if isinstance(tensor, torch.Tensor):
             tensor.activation_offloading = True
         else:
-            data_tensors = tensor.get_data_tensors()
+            data_tensors = tensor.get_data_tensors(scaling_factors=True)
             for tensor in data_tensors:
                 if tensor is not None:
                     tensor.activation_offloading = True

transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py

@@ -112,8 +112,10 @@ def restore_from_saved(
         self._columnwise_scale_inv = tensors[3]
         return tensors[4:]
 
-    def get_data_tensors(self):
+    def get_data_tensors(self, scaling_factors=False):
         """Get this Tensor's data."""
+        if scaling_factors:
+            return self._rowwise_data, self._columnwise_data, self._rowwise_scale_inv, self._columnwise_scale_inv
         return self._rowwise_data, self._columnwise_data
 
     def _transpose_dq_columnwise_output(self, columnwise_dq: torch.Tensor) -> torch.Tensor:

transformer_engine/pytorch/tensor/_internal/float8_tensor_base.py

@@ -128,8 +128,10 @@ def restore_from_saved(
         self._scale_inv = tensors[2]
         return tensors[3:]
 
-    def get_data_tensors(self):
+    def get_data_tensors(self, scaling_factors=False):
         """Get this Tensor's data."""
+        if scaling_factors:
+            return self._data, self._transpose, self._scale_inv
         return self._data, self._transpose
 
     def dequantize(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor:

transformer_engine/pytorch/tensor/_internal/mxfp8_tensor_base.py

@@ -131,8 +131,10 @@ def restore_from_saved(
         self._columnwise_scale_inv = tensors[3]
         return tensors[4:]
 
-    def get_data_tensors(self):
+    def get_data_tensors(self, scaling_factors=False):
         """Get this Tensor's data."""
+        if scaling_factors:
+            return self._rowwise_data, self._columnwise_data, self._rowwise_scale_inv, self._columnwise_scale_inv
         return self._rowwise_data, self._columnwise_data
 
     def dequantize(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: