fixes

Signed-off-by: Pawel Gadzinski <[email protected]>

Author: pggPL

tests/pytorch/test_cpu_offloading.py

@@ -26,10 +26,10 @@
     recipe.Float8BlockScaling(),
 ]
 
-SIZE = 512
+SIZE = 64
 NUM_HEADS = 8
 NUM_LAYERS = 5
-EPSILON = 0.1
+EPSILON = 0.05
 
 # Flash attention saves some internal tensor for the backward pass
 # that cannot be offloaded to CPU.
@@ -48,7 +48,7 @@
         SIZE, NUM_HEADS, params_dtype=torch.bfloat16
     ),
     "transformer_layer": lambda: te.TransformerLayer(
-        SIZE, SIZE, NUM_HEADS, params_dtype=torch.bfloat16, hidden_dropout=0.0
+         SIZE, SIZE, NUM_HEADS, params_dtype=torch.bfloat16, hidden_dropout=0.0
     ),
 }
 
@@ -97,7 +97,8 @@ def _measure_memory_between_forward_and_backward(models, fp8_recipe, cpu_offload
         ), offload_context:
             tensor = model(tensor)
         tensor = sync_function(tensor)
-
+    
+    import gc; gc.collect()
     max_mem_used = torch.cuda.memory_allocated() / (1024**2)
     torch.cuda.synchronize()
 
@@ -119,7 +120,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: