fix: specify correct transformer layer class for FSDP auto-wrap

2026-07-01 23:45:05 -04:00
parent 05cbd7b6b2
commit 8a5edb25e0
1 changed files with 3 additions and 1 deletions
--- a/train.py
+++ b/train.py
@@ -219,11 +219,13 @@ def train(config_path):
        eval_steps=config.get("eval_steps", 100),
        bf16=True,
        gradient_checkpointing=config.get("gradient_checkpointing", True),
-        fsdp="full_shard auto_wrap",
+        fsdp=["full_shard", "auto_wrap"],
        fsdp_config={
            "backward_prefetch": "backward_pre",
            "forward_prefetch": "true",
            "cpu_offload": "false",
+            "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+            "transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
        },
    )