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

fix: correct transformer_auto_wrap_policy API signature

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 11:12:31 -04:00
parent 651965e844
commit 8cabc0e986
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
train.py
View File

@@ -155,10 +155,15 @@ def train(config_path):
# Manually wrap model with FSDP before passing to trainer # Manually wrap model with FSDP before passing to trainer
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial
def get_auto_wrap_policy(model): def get_auto_wrap_policy(model):
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
return transformer_auto_wrap_policy({Qwen3_5MoeDecoderLayer}, model) return partial(
transformer_auto_wrap_policy,
transformer_layer_cls={Qwen3_5MoeDecoderLayer},
min_num_params=1,
)
# Wrap model with FSDP on CPU first # Wrap model with FSDP on CPU first
print("Wrapping model with FSDP on CPU...") print("Wrapping model with FSDP on CPU...")