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

feat: FSDP training failures now fallback to next strategy

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 11:48:03 -04:00
parent ac1417567c
commit 5a13ca1d1c
1 changed files with 83 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
train.py
View File

@@ -176,76 +176,90 @@ 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 try:
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from functools import partial from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import torch.distributed as dist from functools import partial
import torch.distributed as dist
# Initialize distributed process group (required for FSDP)
if not dist.is_initialized(): # Initialize distributed process group (required for FSDP)
dist.init_process_group(backend="nccl") if not dist.is_initialized():
print("✓ Distributed process group initialized") dist.init_process_group(backend="nccl")
print("✓ Distributed process group initialized")
def get_auto_wrap_policy(model):
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer def get_auto_wrap_policy(model):
return partial( from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
transformer_auto_wrap_policy, return partial(
transformer_layer_cls={Qwen3_5MoeDecoderLayer}, transformer_auto_wrap_policy,
min_num_params=1, transformer_layer_cls={Qwen3_5MoeDecoderLayer},
)
# Wrap model with FSDP on CPU first
print("Wrapping model with FSDP on CPU...")
model = FSDP(
model,
auto_wrap_policy=get_auto_wrap_policy(model),
device_id=None, # Keep on CPU initially
mixed_precision=None,
sync_module_states=True,
use_orig_params=True,
) )
print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")
# Wrap model with FSDP on CPU first
print("Wrapping model with FSDP on CPU...") # Training arguments (no FSDP config - we handle it manually)
model = FSDP( training_args = TrainingArguments(
model, output_dir=config["train_params"]["output_dir"],
auto_wrap_policy=get_auto_wrap_policy(model), num_train_epochs=config["train_params"]["num_train_epochs"],
device_id=None, # Keep on CPU initially per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
mixed_precision=None, gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
sync_module_states=True, learning_rate=float(config["train_params"]["learning_rate"]),
use_orig_params=True, lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
) weight_decay=config["train_params"]["weight_decay"],
print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)") warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]),
logging_steps=config["train_params"]["logging_steps"],
# Training arguments (no FSDP config - we handle it manually) save_steps=config["train_params"]["save_steps"],
training_args = TrainingArguments( save_total_limit=config["train_params"]["save_total_limit"],
output_dir=config["train_params"]["output_dir"], eval_strategy=config.get("eval_strategy", "steps"),
num_train_epochs=config["train_params"]["num_train_epochs"], eval_steps=config.get("eval_steps", 100),
per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"], bf16=True,
gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"], gradient_checkpointing=False, # FSDP handles it
learning_rate=float(config["train_params"]["learning_rate"]), # No fsdp config - we wrap manually above
lr_scheduler_type=config["train_params"]["lr_scheduler_type"], )
weight_decay=config["train_params"]["weight_decay"],
warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]), # SFT Trainer
logging_steps=config["train_params"]["logging_steps"], from trl import SFTTrainer
save_steps=config["train_params"]["save_steps"],
save_total_limit=config["train_params"]["save_total_limit"], trainer = SFTTrainer(
eval_strategy=config.get("eval_strategy", "steps"), model=model,
eval_steps=config.get("eval_steps", 100), processing_class=tokenizer,
bf16=True, train_dataset=dataset["train"],
gradient_checkpointing=False, # FSDP handles it eval_dataset=dataset["test"],
# No fsdp config - we wrap manually above args=training_args,
) )
# SFT Trainer # Train
from trl import SFTTrainer print("Starting training...")
trainer.train()
trainer = SFTTrainer(
model=model, # Save
processing_class=tokenizer, trainer.save_model(config["train_params"]["output_dir"])
train_dataset=dataset["train"], tokenizer.save_pretrained(config["train_params"]["output_dir"])
eval_dataset=dataset["test"],
args=training_args, print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
)
# Cleanup distributed process group
# Train if dist.is_initialized():
print("Starting training...") dist.destroy_process_group()
trainer.train()
# Return early - success
# Save return
trainer.save_model(config["train_params"]["output_dir"]) except Exception as e:
tokenizer.save_pretrained(config["train_params"]["output_dir"]) errors.append(("QLoRA 4-bit FSDP training", e))
print(f"✗ Failed during FSDP training: {e}")
print(f"Training complete! Model saved to {config['train_params']['output_dir']}") # Cleanup distributed process group if initialized
if dist.is_initialized():
dist.destroy_process_group()
# Fall through to next strategy
def main(): def main():