233 lines
8.8 KiB
Python
233 lines
8.8 KiB
Python
"""
|
|
SFT VLM Training with LoRA Support
|
|
Extension of train_sft_vlm.py with LoRA for parameter-efficient fine-tuning
|
|
"""
|
|
import os
|
|
import sys
|
|
|
|
__package__ = "trainer"
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|
|
|
import argparse
|
|
import time
|
|
import warnings
|
|
import torch
|
|
import torch.distributed as dist
|
|
from contextlib import nullcontext
|
|
from torch import optim, nn
|
|
from torch.nn.parallel import DistributedDataParallel
|
|
from torch.utils.data import DataLoader, DistributedSampler
|
|
from transformers import AutoTokenizer
|
|
from model.model_vlm import MiniMindVLM, VLMConfig
|
|
from dataset.lm_dataset import VLMDataset
|
|
from trainer.trainer_utils import get_lr, Logger, is_main_process, init_distributed_mode, setup_seed, init_vlm_model, vlm_checkpoint, SkipBatchSampler
|
|
from trainer.lora_utils import get_lora_config, apply_lora
|
|
|
|
warnings.filterwarnings('ignore')
|
|
|
|
|
|
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
|
start_time = time.time()
|
|
for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1):
|
|
input_ids = input_ids.to(args.device)
|
|
labels = labels.to(args.device)
|
|
pixel_values = pixel_values.to(args.device)
|
|
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
|
|
for param_group in optimizer.param_groups:
|
|
param_group['lr'] = lr
|
|
|
|
with autocast_ctx:
|
|
res = model(input_ids, labels=labels, pixel_values=pixel_values)
|
|
loss = res.loss + res.aux_loss
|
|
loss = loss / args.accumulation_steps
|
|
|
|
scaler.scale(loss).backward()
|
|
|
|
if (step + 1) % args.accumulation_steps == 0:
|
|
scaler.unscale_(optimizer)
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
|
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
|
|
optimizer.zero_grad(set_to_none=True)
|
|
|
|
if step % args.log_interval == 0 or step == iters - 1:
|
|
spend_time = time.time() - start_time
|
|
logger.log(
|
|
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.8f} epoch_time:{}min:'.format(
|
|
epoch + 1,
|
|
args.epochs,
|
|
step,
|
|
iters,
|
|
loss.item() * args.accumulation_steps,
|
|
optimizer.param_groups[-1]['lr'],
|
|
spend_time / (step + 1) * iters // 60 - spend_time // 60))
|
|
|
|
if (wandb is not None) and is_main_process():
|
|
wandb.log({"loss": loss.item() * args.accumulation_steps,
|
|
"lr": optimizer.param_groups[-1]['lr'],
|
|
"epoch": epoch + 1,
|
|
"step": step})
|
|
|
|
if step % args.save_interval == 0 or step == iters - 1:
|
|
if is_main_process():
|
|
vlm_checkpoint(model.module if hasattr(model, 'module') else model,
|
|
epoch,
|
|
step,
|
|
args,
|
|
suffix=f"_lora" if args.use_lora else "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="MiniMind VLM SFT with LoRA")
|
|
parser.add_argument("--out_dir", type=str, default="out", help="Output directory")
|
|
parser.add_argument("--epochs", type=int, default=1, help="Number of epochs")
|
|
parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
|
|
parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate")
|
|
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
|
parser.add_argument("--dtype", type=str, default="bfloat16" if torch.cuda.is_bf16_supported() else "float16")
|
|
parser.add_argument("--use_wandb", action="store_true", help="Use Weights & Biases")
|
|
parser.add_argument("--num_workers", type=int, default=0, help="Data loader workers")
|
|
parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation")
|
|
parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping")
|
|
parser.add_argument("--warmup_iters", type=int, default=0, help="Warmup iterations")
|
|
parser.add_argument("--log_interval", type=int, default=100, help="Log interval")
|
|
parser.add_argument("--save_interval", type=int, default=1000, help="Save interval")
|
|
parser.add_argument("--local_rank", type=int, default=-1, help="DDP local rank")
|
|
parser.add_argument("--use_moe", action="store_true", help="Use Mixture of Experts")
|
|
parser.add_argument("--resume_step", type=int, default=0, help="Resume from step")
|
|
parser.add_argument("--hidden_size", type=int, default=512, help="Hidden size")
|
|
parser.add_argument("--num_heads", type=int, default=16, help="Number of attention heads")
|
|
parser.add_argument("--num_layers", type=int, default=8, help="Number of layers")
|
|
|
|
# LoRA-specific arguments
|
|
parser.add_argument("--use_lora", action="store_true", help="Use LoRA for parameter-efficient training")
|
|
parser.add_argument("--lora_r", type=int, default=16, help="LoRA rank")
|
|
parser.add_argument("--lora_alpha", type=int, default=32, help="LoRA alpha")
|
|
parser.add_argument("--lora_dropout", type=float, default=0.05, help="LoRA dropout")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Setup
|
|
init_distributed_mode(args)
|
|
setup_seed(42)
|
|
logger = Logger(args.out_dir, is_main_process())
|
|
|
|
# Load tokenizer and model
|
|
tokenizer = AutoTokenizer.from_pretrained('./model')
|
|
|
|
# Create VLM config
|
|
lm_config = VLMConfig(
|
|
hidden_size=args.hidden_size,
|
|
num_attention_heads=args.num_heads,
|
|
num_hidden_layers=args.num_layers,
|
|
use_moe=args.use_moe,
|
|
)
|
|
|
|
# Initialize model
|
|
model = init_vlm_model(lm_config, args.device)
|
|
|
|
# Apply LoRA if requested
|
|
if args.use_lora:
|
|
print("\n" + "=" * 50)
|
|
print("🚀 Applying LoRA for Parameter-Efficient Training")
|
|
print("=" * 50)
|
|
|
|
lora_config = get_lora_config(
|
|
r=args.lora_r,
|
|
lora_alpha=args.lora_alpha,
|
|
lora_dropout=args.lora_dropout,
|
|
)
|
|
|
|
model = apply_lora(model, lora_config, verbose=True)
|
|
print("=" * 50 + "\n")
|
|
|
|
# DDP if needed
|
|
if args.ddp:
|
|
model = DistributedDataParallel(model, device_ids=[args.local_rank])
|
|
|
|
# Optimizer
|
|
optimizer = optim.Adam(
|
|
filter(lambda p: p.requires_grad, model.parameters()), # Only train LoRA params if using LoRA
|
|
lr=args.learning_rate
|
|
)
|
|
|
|
# AMP scaler
|
|
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
|
|
autocast_ctx = (
|
|
torch.cuda.amp.autocast(dtype=torch.bfloat16)
|
|
if args.dtype == 'bfloat16'
|
|
else torch.cuda.amp.autocast(dtype=torch.float16)
|
|
if args.dtype == 'float16'
|
|
else nullcontext()
|
|
)
|
|
|
|
# Load dataset
|
|
dataset = VLMDataset(
|
|
parquet_path='./dataset/sft_vlm.parquet',
|
|
tokenizer=tokenizer,
|
|
preprocess=model.processor if not args.ddp else model.module.processor,
|
|
max_length=512
|
|
)
|
|
|
|
# Data loader
|
|
if args.ddp:
|
|
sampler = DistributedSampler(dataset)
|
|
loader = DataLoader(
|
|
dataset,
|
|
batch_size=args.batch_size,
|
|
sampler=sampler,
|
|
num_workers=args.num_workers,
|
|
pin_memory=True
|
|
)
|
|
else:
|
|
loader = DataLoader(
|
|
dataset,
|
|
batch_size=args.batch_size,
|
|
shuffle=True,
|
|
num_workers=args.num_workers,
|
|
pin_memory=True
|
|
)
|
|
|
|
# WandB
|
|
wandb = None
|
|
if args.use_wandb and is_main_process():
|
|
import wandb as wandb_lib
|
|
wandb = wandb_lib
|
|
wandb.init(
|
|
project="drivemind-v",
|
|
name=f"sft_lora_r{args.lora_r}" if args.use_lora else "sft_full",
|
|
config=vars(args)
|
|
)
|
|
|
|
# Training loop
|
|
iters = len(loader)
|
|
logger.log(f"Starting SFT training {'with LoRA' if args.use_lora else 'full fine-tuning'}...")
|
|
logger.log(f"Total iterations: {iters} per epoch")
|
|
|
|
for epoch in range(args.epochs):
|
|
train_epoch(epoch, loader, iters, start_step=0, wandb=wandb)
|
|
|
|
# Save final model
|
|
if is_main_process():
|
|
if args.use_lora:
|
|
# Save LoRA adapters (tiny file ~2-5MB)
|
|
save_path = f"{args.out_dir}/lora_adapters_sft_vlm_{args.hidden_size}"
|
|
model_to_save = model.module if hasattr(model, 'module') else model
|
|
model_to_save.save_pretrained(save_path)
|
|
logger.log(f"✅ LoRA adapters saved to {save_path}")
|
|
else:
|
|
vlm_checkpoint(
|
|
model.module if hasattr(model, 'module') else model,
|
|
args.epochs - 1,
|
|
iters - 1,
|
|
args,
|
|
suffix="_final"
|
|
)
|
|
|
|
if wandb:
|
|
wandb.finish()
|
|
|
|
logger.log("Training complete!")
|