fix float8 after the HSDP PR (pytorch#575)

C-TC · Sep 12, 2024 · 2a25f4d · 2a25f4d
1 parent f2a1551
commit 2a25f4d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/torchtitan/float8.py b/torchtitan/float8.py
@@ -49,8 +49,7 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
 
         # Mutates the model inplace replacing instances of torch.nn.Linear with Float8Linear
         enable_fsdp_float8_all_gather = (
-            parallel_dims.dp_enabled
-            and parallel_dims.dp_type == "fsdp"
+            parallel_dims.dp_shard_enabled
             and float8_config.enable_fsdp_float8_all_gather
         )
         scaling_type_input = ScalingType(float8_config.scaling_type_input)