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

fix: use FSDP1 with sync_module_states to avoid GPU gather spike

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 08:47:56 -04:00
parent c607888d8f
commit 70a2bf0dd6
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
train.py
View File

@@ -167,12 +167,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 removed - FSDP activation_checkpointing handles it # gradient_checkpointing removed - FSDP activation_checkpointing handles it
fsdp=True, fsdp="full_shard",
fsdp_config={ fsdp_config={
"sharding_strategy": "SHARD_GRAD_OP", "sharding_strategy": "SHARD_GRAD_OP",
"cpu_offload": False, "cpu_offload": False,
"activation_checkpointing": True, "activation_checkpointing": True,
"limit_all_gathers": True, "limit_all_gathers": True,
"sync_module_states": True,
"backward_prefetch": "backward_pre", "backward_prefetch": "backward_pre",
"forward_prefetch": "true", "forward_prefetch": "true",
"auto_wrap_policy": "TRANSFORMER_BASED_WRAP", "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",