Add the option to turn on async-TP

[ghstack-poisoned]

Author: Yifu Wang

torchtitan/config_manager.py

@@ -241,6 +241,12 @@ def __init__(self):
             action="store_true",
             help="Whether to apply loss parallel when sequence parallel is enabled",
         )
+        self.parser.add_argument(
+            "--experimental.enable_async_tensor_parallel",
+            default=False,
+            action="store_true",
+            help="Whether to apply async tensor parallel (currently only effective when compile is enabled)",
+        )
         self.parser.add_argument(
             "--experimental.pipeline_parallel_degree",
             type=int,

torchtitan/parallelisms/parallelize_llama.py

@@ -394,6 +394,12 @@ def apply_tp(model, world_mesh, parallel_dims, job_config: JobConfig):
             parallelize_plan=layer_plan,
         )
 
+    if job_config.experimental.enable_async_tensor_parallel:
+        from torch.distributed._symmetric_memory import enable_symm_mem_for_group
+
+        torch._inductor.config._micro_pipeline_tp = True
+        enable_symm_mem_for_group(tp_mesh.get_group().group_name)
+
     logger.info("Applied Tensor Parallelism to the model")
     return model
 

train.py

@@ -252,6 +252,11 @@ def loss_fn(pred, labels):
         for m in model_parts
     ]
 
+    # for ease of testing TP in lieu of FSDP
+    if job_config.training.tensor_parallel_degree == world_size:
+        for model in model_parts:
+            model.to(torch.bfloat16)
+
     init_device = "cpu" if job_config.checkpoint.create_seed_checkpoint else "cuda"
     for model in model_parts:
         model.to_empty(device=init_device)