fix bug for non first microbatch in te

Azure · Feb 21, 2024 · eadca11 · eadca11
1 parent 0a28e0f
commit eadca11
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/msamp/te/modules.py b/msamp/te/modules.py
@@ -62,7 +62,7 @@ def set_fp8_weights(self):
                 weight_cast_attr = f'weight{i}_fp8'
                 weight_transpose_attr = f'weight{i}_t_fp8'
 
-                if (hasattr(self, weight_cast_attr) and getattr(self, weight_cast_attr).shape == shape):
+                if (hasattr(self, weight_cast_attr) and getattr(self, weight_cast_attr)._data.shape == shape):
                     return
 
                 setattr(

diff --git a/tests/te/test_te_replacer.py b/tests/te/test_te_replacer.py
@@ -164,5 +164,8 @@ def test_fp8_ddp_with_te(self):
         fp8_format = Format.HYBRID
         fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo='max')
         with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
-            output = model(x, attention_mask=None)
+            output = model(x, attention_mask=None, is_first_microbatch=True)
+            output.sum().backward()
+        with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+            output = model(x, attention_mask=None, is_first_microbatch=False)
             output.sum().backward()