fix: use FSDP1 with sync_module_states to avoid GPU gather spike

2026-07-02 08:47:56 -04:00
parent c607888d8f
commit 70a2bf0dd6
1 changed files with 2 additions and 1 deletions
--- a/train.py
+++ b/train.py
@@ -167,12 +167,13 @@ def train(config_path):
        eval_steps=config.get("eval_steps", 100),
        bf16=True,
        # gradient_checkpointing removed - FSDP activation_checkpointing handles it
-        fsdp=True,
+        fsdp="full_shard",
        fsdp_config={
            "sharding_strategy": "SHARD_GRAD_OP",
            "cpu_offload": False,
            "activation_checkpointing": True,
            "limit_all_gathers": True,
+            "sync_module_states": True,
            "backward_prefetch": "backward_pre",
            "forward_prefetch": "true",
            "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",