From 8cabc0e9862ec324165747a6aa233657914772a3 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 11:12:31 -0400
Subject: [PATCH] fix: correct transformer_auto_wrap_policy API signature

---
 train.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index c9203f1..a1f8ee1 100644
--- a/train.py
+++ b/train.py
@@ -155,10 +155,15 @@ def train(config_path):
     # 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
+    from functools import partial
     
     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)
+        return partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls={Qwen3_5MoeDecoderLayer},
+            min_num_params=1,
+        )
     
     # Wrap model with FSDP on CPU first
     print("Wrapping model with FSDP on CPU...")