From b7d680966e16efa7cfbd7863c3ee7920b95dc031 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 09:36:08 -0400
Subject: [PATCH] feat: restore FSDP with SHARD_GRAD_OP + sync_module_states

---
 train.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 8cbfe19..13421d9 100644
--- a/train.py
+++ b/train.py
@@ -169,7 +169,17 @@ def train(config_path):
         eval_steps=config.get("eval_steps", 100),
         bf16=True,
         gradient_checkpointing=True,
-        # No FSDP - use standard accelerate data parallelism
+        fsdp=True,
+        fsdp_config={
+            "sharding_strategy": "SHARD_GRAD_OP",
+            "cpu_offload": False,
+            "activation_checkpointing": True,
+            "limit_all_gathers": True,
+            "sync_module_states": True,
+            "mixed_precision": {"param_dtype": torch.bfloat16, "reduce_dtype": torch.float32, "buffer_dtype": torch.float32},
+            "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+            "transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
+        },
     )
     
     # SFT Trainer (DeepSpeed handles distributed training via config)