Test Ruisi's commit

Author: fmassa

examples/example_llama3.py

@@ -558,10 +558,14 @@ def forward(self, tokens: torch.Tensor, input_batch: torch.Tensor | None = None)
 
 world_size = 256
 
-fake_store = FakeStore()
-torch.distributed.init_process_group(
-    "fake", store=fake_store, rank=0, world_size=world_size
-)
+backend = "fake"
+kwargs = {"rank": 0, "world_size": world_size}
+if True:
+    backend = "nccl"
+    fake_store = None
+    kwargs = {}
+    world_size = 8
+torch.distributed.init_process_group(backend, store=fake_store, **kwargs)
 
 use_1d_mesh = False
 
@@ -588,7 +592,7 @@ def forward(self, tokens: torch.Tensor, input_batch: torch.Tensor | None = None)
 
 def model_fn():
     model_args = TransformerModelArgs(
-        n_layers=32,
+        n_layers=2,
         vocab_size=vocab_size,
         max_seq_len=seqlen,
         multiple_of=1024,
@@ -604,6 +608,28 @@ def input_fn():
     return x
 
 
+from functools import partial
+
+from autoparallel.auto_bucketing import (
+    simple_fsdp_autobucketing_reordering_pass,
+    simplefsdp_autobucketing_config,
+)
+
+torch._inductor.config.allow_buffer_reuse = False
+torch._inductor.config.reorder_for_peak_memory = False
+torch._inductor.config.reorder_for_compute_comm_overlap = True
+simplefsdp_autobucketing_config.save_estimation_path = (
+    "/storage/home/fmassa/work/projects/autoparallel/estimation_mast.pkl"
+)
+simple_fsdp_autobucketing_reordering_pass = partial(
+    simple_fsdp_autobucketing_reordering_pass,
+    configs=simplefsdp_autobucketing_config,
+)
+torch._inductor.config.reorder_for_compute_comm_overlap_passes = [
+    simple_fsdp_autobucketing_reordering_pass
+]
+
+
 # parallelize the model
 with torch.device("meta"):
     model = model_fn()
@@ -659,3 +685,35 @@ def input_fn():
 out = parallel_mod(*x)
 out.backward(torch.randn_like(out))
 print("All good!")
+
+
+def _dump_trace(prof):
+    prof.export_chrome_trace(
+        f"/home/fmassa/work/projects/autoparallel/traces/rank{torch.distributed.get_rank()}.json"
+    )
+
+
+prof = torch.profiler.profile(
+    activities=[
+        torch.profiler.ProfilerActivity.CUDA,
+        torch.profiler.ProfilerActivity.CPU,
+    ],
+    schedule=torch.profiler.schedule(
+        warmup=2,
+        active=2,
+        wait=1,
+        repeat=1,
+        skip_first=0,
+    ),
+    record_shapes=False,
+    profile_memory=False,
+    with_stack=True,
+    with_flops=False,
+    on_trace_ready=_dump_trace,
+)
+
+with prof:
+    for i in range(10):
+        out = parallel_mod(*x)
+        out.backward(torch.randn_like(out))
+        prof.step()