cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 23:45:05 -04:00
parent 05cbd7b6b2
commit 8a5edb25e0
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@@ -219,11 +219,13 @@ def train(config_path):
eval_steps=config.get("eval_steps", 100), eval_steps=config.get("eval_steps", 100),
bf16=True, bf16=True,
gradient_checkpointing=config.get("gradient_checkpointing", True), gradient_checkpointing=config.get("gradient_checkpointing", True),
fsdp="full_shard auto_wrap", fsdp=["full_shard", "auto_wrap"],
fsdp_config={ fsdp_config={
"backward_prefetch": "backward_pre", "backward_prefetch": "backward_pre",
"forward_prefetch": "true", "forward_prefetch": "true",
"cpu_offload": "false", "cpu_offload": "false",
"auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
"transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
}, },
) )