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!"

130
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"],
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=config.get("gradient_checkpointing", True),
)
def get_auto_wrap_policy(model): # SFT Trainer
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer from trl import SFTTrainer
return partial(
transformer_auto_wrap_policy,
transformer_layer_cls={Qwen3_5MoeDecoderLayer},
)
# Wrap model with FSDP on GPU trainer = SFTTrainer(
print("Wrapping model with FSDP on GPU...") model=model,
model = FSDP( processing_class=tokenizer,
model, train_dataset=dataset["train"],
auto_wrap_policy=get_auto_wrap_policy(model), eval_dataset=dataset.get("test"),
device_id=torch.cuda.current_device(), # Keep on current GPU args=training_args,
mixed_precision=None, )
sync_module_states=False, # Model is already on GPU
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) # Train
training_args = TrainingArguments( print("Starting training...")
output_dir=config["train_params"]["output_dir"], trainer.train()
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
)
# SFT Trainer # Save
from trl import SFTTrainer trainer.save_model(config["train_params"]["output_dir"])
tokenizer.save_pretrained(config["train_params"]["output_dir"])
trainer = SFTTrainer( print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
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():