From 651965e844bdc1aba86edf5648a9beb75d7b28e2 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 11:03:41 -0400
Subject: [PATCH] feat: manually wrap model with FSDP on CPU before trainer

---
 train.py | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/train.py b/train.py
index 894162f..c9203f1 100644
--- a/train.py
+++ b/train.py
@@ -152,7 +152,27 @@ def train(config_path):
         },
     )
     
-    # Training arguments with FSDP
+    # Manually wrap model with FSDP before passing to trainer
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+    
+    def get_auto_wrap_policy(model):
+        from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
+        return transformer_auto_wrap_policy({Qwen3_5MoeDecoderLayer}, model)
+    
+    # Wrap model with FSDP on CPU first
+    print("Wrapping model with FSDP on CPU...")
+    model = FSDP(
+        model,
+        auto_wrap_policy=get_auto_wrap_policy(model),
+        device_id=None,  # Keep on CPU initially
+        mixed_precision=None,
+        sync_module_states=True,
+        use_orig_params=True,
+    )
+    print("✓ Model wrapped with FSDP")
+    
+    # Training arguments (no FSDP config - we handle it manually)
     training_args = TrainingArguments(
         output_dir=config["train_params"]["output_dir"],
         num_train_epochs=config["train_params"]["num_train_epochs"],
@@ -168,22 +188,11 @@ def train(config_path):
         eval_strategy=config.get("eval_strategy", "steps"),
         eval_steps=config.get("eval_steps", 100),
         bf16=True,
-        gradient_checkpointing=False,  # FSDP activation_checkpointing handles it
-        fsdp="full_shard",  # Force FSDP1 (not FSDP2)
-        fsdp_config={
-            "sharding_strategy": "SHARD_GRAD_OP",
-            "cpu_offload": False,
-            "activation_checkpointing": True,
-            "limit_all_gathers": True,
-            "sync_module_states": True,
-            "use_orig_params": True,  # Critical for LoRA/PEFT
-            # "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",
-        },
+        gradient_checkpointing=False,  # FSDP handles it
+        # No fsdp config - we wrap manually above
     )
     
-    # SFT Trainer (DeepSpeed handles distributed training via config)
+    # SFT Trainer
     from trl import SFTTrainer
     
     trainer = SFTTrainer(