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

fix: disable FSDP, use standard accelerate data parallelism

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 09:22:28 -04:00
parent 42d61c15b0
commit fe7df7c92f
1 changed files with 2 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
train.py
View File

@@ -168,20 +168,8 @@ def train(config_path):
eval_strategy=config.get("eval_strategy", "steps"), eval_strategy=config.get("eval_strategy", "steps"),
eval_steps=config.get("eval_steps", 100), eval_steps=config.get("eval_steps", 100),
bf16=True, bf16=True,
# gradient_checkpointing removed - FSDP activation_checkpointing handles it gradient_checkpointing=True,
fsdp=True, # No FSDP - use standard accelerate data parallelism
fsdp_config={
"sharding_strategy": "SHARD_GRAD_OP",
"cpu_offload": False,
"activation_checkpointing": True,
"limit_all_gathers": True,
"sync_module_states": False, # Don't sync from CPU, FSDP will handle placement
"backward_prefetch": "backward_pre",
"forward_prefetch": "true",
"auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
"transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
"mixed_precision": {"param_dtype": torch.bfloat16, "reduce_dtype": torch.float32, "buffer_dtype": torch.float32},
},
) )
# SFT Trainer (DeepSpeed handles distributed training via config) # SFT Trainer (DeepSpeed handles distributed training via config)