fix test errors

Author: ruisizhang123

examples/example_llama3.py

@@ -558,14 +558,18 @@ def forward(self, tokens: torch.Tensor, input_batch: torch.Tensor | None = None)
 
 world_size = 256
 
-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)
+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
 
@@ -618,8 +622,9 @@ def input_fn():
 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.calibrate_number = 5
 simplefsdp_autobucketing_config.save_estimation_path = (
-    "/storage/home/fmassa/work/projects/autoparallel/estimation_mast.pkl"
+    "./estimation_mast.pkl"
 )
 simple_fsdp_autobucketing_reordering_pass = partial(
     simple_fsdp_autobucketing_reordering_pass,