feat: manually wrap model with FSDP on CPU before trainer

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
+        gradient_checkpointing=False,  # FSDP handles it
-        fsdp="full_shard",  # Force FSDP1 (not FSDP2)
+        # No fsdp config - we wrap manually above
-        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",
-        },
     )
     
-    # SFT Trainer (DeepSpeed handles distributed training via config)
+    # SFT Trainer
     from trl import SFTTrainer
     
     trainer = SFTTrainer(