feat: restore FSDP with SHARD_GRAD_OP + sync_module_states

This commit is contained in:
Christian Medina
2026-07-02 09:36:08 -04:00
parent 7b39fc3a1b
commit b7d680966e

View File

@@ -169,7 +169,17 @@ def train(config_path):
eval_steps=config.get("eval_steps", 100),
bf16=True,
gradient_checkpointing=True,
# No FSDP - use standard accelerate data parallelism
fsdp=True,
fsdp_config={
"sharding_strategy": "SHARD_GRAD_OP",
"cpu_offload": False,
"activation_checkpointing": True,
"limit_all_gathers": True,
"sync_module_states": True,
"mixed_precision": {"param_dtype": torch.bfloat16, "reduce_dtype": torch.float32, "buffer_dtype": torch.float32},
"auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
"transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
},
)
# SFT Trainer (DeepSpeed handles distributed training via config)