Update on "Made some stylistic changes to apply_dp"

[ghstack-poisoned]
pytorch · Jul 10, 2024 · 19cba24 · 19cba24
2 parents 5c04a9b + 67c4e9b
commit 19cba24
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchtitan/parallelisms/parallelize_llama.py b/torchtitan/parallelisms/parallelize_llama.py
@@ -459,6 +459,7 @@ def apply_dp(model, world_mesh, parallel_dims, job_config: JobConfig):
         reduce_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_reduce],
     )
     fsdp_config = {"mesh": dp_mesh, "mp_policy": mp_policy}
+
     for layer_id, transformer_block in model.layers.items():
         if parallel_dims.pp_enabled:
             # For PP, do not reshard after forward to avoid per-microbatch