fix: correct transformer_auto_wrap_policy API signature

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...")