feat: restore FSDP with SHARD_GRAD_OP + sync_module_states
This commit is contained in:
12
train.py
12
train.py
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user