diff --git a/train.py b/train.py index e3fdbed..4f459ce 100644 --- a/train.py +++ b/train.py @@ -201,7 +201,7 @@ def train(config_path): auto_wrap_policy=get_auto_wrap_policy(model), device_id=None, # Keep on CPU initially mixed_precision=None, - sync_module_states=True, + sync_module_states=False, # Model is on CPU, no sync needed use_orig_params=True, ) print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")