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

fix: single-process training, remove FSDP (model pre-distributed via device_map)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 15:53:51 -04:00
parent b8214f64cb
commit 9e1c5120fb
2 changed files with 53 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
train-on-this-server.sh
View File

@@ -18,9 +18,9 @@ pip install --upgrade pip
echo "Installing training dependencies..." echo "Installing training dependencies..."
pip install transformers datasets trl peft accelerate bitsandbytes deepspeed pip install transformers datasets trl peft accelerate bitsandbytes deepspeed
# Run training with accelerate data parallelism # Run training with single process (model already distributed via device_map)
echo "Starting training with accelerate..." echo "Starting training (single process, model pre-distributed)..."
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
accelerate launch --multi_gpu --num_processes 2 --mixed_precision bf16 train.py --config training/configs/ornith-35b-lora.yaml python train.py --config training/configs/ornith-35b-lora.yaml
echo "Training completed!" echo "Training completed!"

142
train.py
View File

@@ -96,105 +96,63 @@ def train(config_path):
from datasets import load_dataset from datasets import load_dataset
import os import os
# Get dataset paths - training/data/ is relative to project root # Get dataset path from config
repo_root = Path(__file__).parent dataset_path = config["dataset"][0]["path"]
train_path = str(repo_root / "training" / "data" / "train.jsonl") print(f"Loading dataset from: {dataset_path}")
test_path = str(repo_root / "training" / "data" / "test.jsonl")
print(f"Looking for dataset at: {train_path}")
dataset = load_dataset( dataset = load_dataset(
"json", "json",
data_files={ data_files={
"train": train_path, "train": dataset_path,
"test": test_path,
}, },
) )
# Manually wrap model with FSDP before passing to trainer # Model is already distributed across GPUs via device_map
try: # No FSDP needed - device_map handles distribution
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP print("✓ Model already distributed across GPUs (device_map)")
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
from functools import partial if torch.cuda.device_count() > 1:
import torch.distributed as dist print(f" GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
# Initialize distributed process group (required for FSDP) # Training arguments
if not dist.is_initialized(): training_args = TrainingArguments(
dist.init_process_group(backend="nccl") output_dir=config["train_params"]["output_dir"],
print("✓ Distributed process group initialized") num_train_epochs=config["train_params"]["num_train_epochs"],
per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
def get_auto_wrap_policy(model): gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer learning_rate=float(config["train_params"]["learning_rate"]),
return partial( lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
transformer_auto_wrap_policy, weight_decay=config["train_params"]["weight_decay"],
transformer_layer_cls={Qwen3_5MoeDecoderLayer}, 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"],
# Wrap model with FSDP on GPU save_total_limit=config["train_params"]["save_total_limit"],
print("Wrapping model with FSDP on GPU...") eval_strategy=config.get("eval_strategy", "steps"),
model = FSDP( eval_steps=config.get("eval_steps", 100),
model, bf16=True,
auto_wrap_policy=get_auto_wrap_policy(model), gradient_checkpointing=config.get("gradient_checkpointing", True),
device_id=torch.cuda.current_device(), # Keep on current GPU )
mixed_precision=None,
sync_module_states=False, # Model is already on GPU # SFT Trainer
use_orig_params=True, from trl import SFTTrainer
)
print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)") trainer = SFTTrainer(
model=model,
# Training arguments (no FSDP config - we handle it manually) processing_class=tokenizer,
training_args = TrainingArguments( train_dataset=dataset["train"],
output_dir=config["train_params"]["output_dir"], eval_dataset=dataset.get("test"),
num_train_epochs=config["train_params"]["num_train_epochs"], args=training_args,
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"]), # Train
lr_scheduler_type=config["train_params"]["lr_scheduler_type"], print("Starting training...")
weight_decay=config["train_params"]["weight_decay"], trainer.train()
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
save_steps=config["train_params"]["save_steps"], trainer.save_model(config["train_params"]["output_dir"])
save_total_limit=config["train_params"]["save_total_limit"], tokenizer.save_pretrained(config["train_params"]["output_dir"])
eval_strategy=config.get("eval_strategy", "steps"),
eval_steps=config.get("eval_steps", 100), print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
bf16=True,
gradient_checkpointing=False, # FSDP handles it
# No fsdp config - we wrap manually above
)
# SFT Trainer
from trl import SFTTrainer
trainer = SFTTrainer(
model=model,
processing_class=tokenizer,
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
args=training_args,
)
# Train
print("Starting training...")
trainer.train()
# Save
trainer.save_model(config["train_params"]["output_dir"])
tokenizer.save_pretrained(config["train_params"]["output_dir"])
print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
# Cleanup distributed process group
if dist.is_initialized():
dist.destroy_process_group()
# 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():