From dbc97138e9d0dabd4e5d75c7e570f8dea61bc20e Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 00:58:41 -0400
Subject: [PATCH] fix: use SHARD_GRAD_OP instead of FULL_SHARD for 2-GPU setup

---
 train.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index 427a083..c7d3175 100644
--- a/train.py
+++ b/train.py
@@ -169,12 +169,14 @@ def train(config_path):
         # gradient_checkpointing removed - FSDP activation_checkpointing handles it
         fsdp=True,
         fsdp_config={
+            "sharding_strategy": "SHARD_GRAD_OP",
+            "cpu_offload": False,
+            "activation_checkpointing": True,
+            "limit_all_gathers": True,
             "backward_prefetch": "backward_pre",
             "forward_prefetch": "true",
-            "cpu_offload": "false",
             "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
             "transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
-            "activation_checkpointing": True,
         },
     )