From 70a2bf0dd6feb66245d5da3fc2a8b5a879f0786d Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 08:47:56 -0400
Subject: [PATCH] fix: use FSDP1 with sync_module_states to avoid GPU gather
 spike

---
 train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index a6ef3ce..560973e 100644
--- 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",