Import error in manager.py + switch to sync mode

Author: mreso

torchft/manager.py

@@ -36,6 +36,7 @@
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, TypeVar, cast
 
 import torch
+import torch.distributed as dist
 from torch.distributed import ReduceOp, TCPStore
 
 from torchft.checkpointing import CheckpointServer

train_fsdp.py

@@ -126,6 +126,7 @@ def state_dict():
         load_state_dict=load_state_dict,
         state_dict=state_dict,
         replica_id=f"train_fsdp_{REPLICA_GROUP_ID}",
+        use_async_quorum=False,
     )
 
     mesh = hsdp_device_mesh(NUM_REPLICA_GROUPS, NUM_REPLICAS, "cuda" if torch.cuda.is_available() else "cpu", manager=manager)
@@ -136,8 +137,6 @@ def state_dict():
 
     optimizer = Optimizer(manager, torch.optim.Adam(model.parameters(), lr=1e-5))
 
-    optimizer.zero_grad()
-
     while manager.current_step() < 500:
         model.train()
         for batch in tqdm(train_dataloader):