fix float8 after the HSDP PR (#575)

tianyu-l · web-flow · commit 2a25f4d7868c · 2024-09-12T16:30:43.000-07:00
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)

Original file line number	Diff line number	Diff line change
`@@ -49,8 +49,7 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):`
`49`	`49`
`50`	`50`	`# Mutates the model inplace replacing instances of torch.nn.Linear with Float8Linear`
`51`	`51`	`enable_fsdp_float8_all_gather = (`
`52`		`- parallel_dims.dp_enabled`
`53`		`- and parallel_dims.dp_type == "fsdp"`
	`52`	`+ parallel_dims.dp_shard_enabled`
`54`	`53`	`and float8_config.enable_fsdp_float8_all_gather`
`55`	`54`	`)`
`56`	`55`	`scaling_type_input = ScalingType(float8_config.scaling_type_input)`