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

134
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) # Initialize distributed process group (required for FSDP)
if not dist.is_initialized(): if not dist.is_initialized():
dist.init_process_group(backend="nccl") dist.init_process_group(backend="nccl")
print("✓ Distributed process group initialized") print("✓ Distributed process group initialized")
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 partial( return partial(
transformer_auto_wrap_policy, transformer_auto_wrap_policy,
transformer_layer_cls={Qwen3_5MoeDecoderLayer}, transformer_layer_cls={Qwen3_5MoeDecoderLayer},
min_num_params=1, )
# 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)")
# Training arguments (no FSDP config - we handle it manually)
training_args = TrainingArguments(
output_dir=config["train_params"]["output_dir"],
num_train_epochs=config["train_params"]["num_train_epochs"],
per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
learning_rate=float(config["train_params"]["learning_rate"]),
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"]),
logging_steps=config["train_params"]["logging_steps"],
save_steps=config["train_params"]["save_steps"],
save_total_limit=config["train_params"]["save_total_limit"],
eval_strategy=config.get("eval_strategy", "steps"),
eval_steps=config.get("eval_steps", 100),
bf16=True,
gradient_checkpointing=False, # FSDP handles it
# No fsdp config - we wrap manually above
) )
# Wrap model with FSDP on CPU first # SFT Trainer
print("Wrapping model with FSDP on CPU...") from trl import SFTTrainer
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)")
# Training arguments (no FSDP config - we handle it manually) trainer = SFTTrainer(
training_args = TrainingArguments( model=model,
output_dir=config["train_params"]["output_dir"], processing_class=tokenizer,
num_train_epochs=config["train_params"]["num_train_epochs"], train_dataset=dataset["train"],
per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"], eval_dataset=dataset["test"],
gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"], args=training_args,
learning_rate=float(config["train_params"]["learning_rate"]), )
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"]),
logging_steps=config["train_params"]["logging_steps"],
save_steps=config["train_params"]["save_steps"],
save_total_limit=config["train_params"]["save_total_limit"],
eval_strategy=config.get("eval_strategy", "steps"),
eval_steps=config.get("eval_steps", 100),
bf16=True,
gradient_checkpointing=False, # FSDP handles it
# No fsdp config - we wrap manually above
)
# SFT Trainer # Train
from trl import SFTTrainer print("Starting training...")
trainer.train()
trainer = SFTTrainer( # Save
model=model, trainer.save_model(config["train_params"]["output_dir"])
processing_class=tokenizer, tokenizer.save_pretrained(config["train_params"]["output_dir"])
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
args=training_args,
)
# Train print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
print("Starting training...")
trainer.train()
# Save # Cleanup distributed process group
trainer.save_model(config["train_params"]["output_dir"]) if dist.is_initialized():
tokenizer.save_pretrained(config["train_params"]["output_dir"]) dist.destroy_process_group()
print(f"Training complete! Model saved to {config['train_params']['output_dir']}") # Return early - success
return
except Exception as e:
errors.append(("QLoRA 4-bit FSDP training", e))
print(f"✗ Failed during FSDP training: {e}")
# Cleanup distributed process group if initialized
if dist.is_initialized():
dist.destroy_process_group()
# Fall through to next strategy
def main(): def main():