fix: fsdp=True instead of list + add activation_checkpointing

This commit is contained in:
Christian Medina
2026-07-02 00:16:33 -04:00
parent 8747111671
commit dfe85c3c82

View File

@@ -167,13 +167,14 @@ 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=True,
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",
"activation_checkpointing": True,
},
)