File size: 6,530 Bytes
babffc8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | from __future__ import annotations
import argparse
import sys
from pathlib import Path
import torch
from accelerate import Accelerator
from peft import PeftModel
from torch.utils.data import DataLoader
from tqdm import tqdm
sys.path.append(str(Path(__file__).resolve().parents[1]))
from src.datasets import VideoJsonlDataset
from src.io_utils import append_jsonl, apply_overrides, load_config, output_dir, save_config
from src.qwen_video_utils import load_model_for_training, load_processor, processor_inputs
from src.train_utils import (
apply_lora,
build_optimizer,
build_scheduler,
oom_help,
quantization_config_from_config,
)
class QwenVideoSFTCollator:
def __init__(self, processor):
self.processor = processor
def __call__(self, batch):
if len(batch) != 1:
raise ValueError(
"This first-version video collator supports per_device_train_batch_size=1. "
"Use gradient_accumulation_steps for larger effective batches."
)
sample = batch[0]
messages = sample["messages"]
if len(messages) < 2 or messages[-1].get("role") != "assistant":
raise ValueError(f"No assistant label in training sample id={sample.get('id')}")
full_inputs = processor_inputs(self.processor, [messages], add_generation_prompt=False)
prompt_messages = [m for m in messages if m.get("role") != "assistant"]
prompt_inputs = processor_inputs(self.processor, [prompt_messages], add_generation_prompt=True)
labels = full_inputs["input_ids"].clone()
prompt_len = min(prompt_inputs["input_ids"].shape[1], labels.shape[1])
labels[:, :prompt_len] = -100
pad_id = self.processor.tokenizer.pad_token_id
labels[full_inputs["input_ids"] == pad_id] = -100
full_inputs["labels"] = labels
return full_inputs
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", required=True)
parser.add_argument("--adapter_path", default=None)
parser.add_argument("--max_samples", type=int, default=None)
parser.add_argument("--model_name", default=None)
parser.add_argument("--output_dir", default=None)
args = parser.parse_args()
cfg = apply_overrides(load_config(args.config), model_name=args.model_name, output_dir=args.output_dir)
out_dir = output_dir(cfg)
adapter_dir = out_dir / "adapter"
out_dir.mkdir(parents=True, exist_ok=True)
accelerator = Accelerator(
gradient_accumulation_steps=int(cfg.get("gradient_accumulation_steps", 8)),
mixed_precision="bf16" if cfg.get("bf16", True) else ("fp16" if cfg.get("fp16", False) else "no"),
)
if not torch.cuda.is_available():
accelerator.print("Warning: CUDA is unavailable. QLoRA training is expected to require CUDA.")
processor = load_processor(cfg["model_name"])
dataset = VideoJsonlDataset(
cfg["train_jsonl"],
video_column=cfg["video_column"],
target_column=cfg["target_column"],
prompt=cfg["prompt"],
video_fps=cfg.get("video_fps"),
max_frames=cfg.get("max_frames"),
max_samples=args.max_samples or int(cfg.get("train_max_samples", 0)),
require_target=True,
)
if len(dataset) == 0:
raise ValueError("Training dataset is empty.")
batch_size = int(cfg.get("per_device_train_batch_size", 1))
loader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=QwenVideoSFTCollator(processor),
)
dtype = "bfloat16" if cfg.get("bf16", True) else "float16"
device_map = {"": accelerator.process_index} if torch.cuda.is_available() else None
model = load_model_for_training(
cfg["model_name"],
quantization_config=quantization_config_from_config(cfg),
device_map=device_map,
dtype=dtype,
)
if cfg.get("gradient_checkpointing", True):
model.gradient_checkpointing_enable()
if hasattr(model, "config"):
model.config.use_cache = False
if args.adapter_path:
accelerator.print(f"Continuing training from adapter: {args.adapter_path}")
model = PeftModel.from_pretrained(model, args.adapter_path, is_trainable=True)
else:
model = apply_lora(model, cfg)
if accelerator.is_main_process:
model.print_trainable_parameters()
save_config(cfg, out_dir / "training_config.yaml")
optimizer = build_optimizer(model, cfg)
steps_per_epoch = max(1, len(loader) // int(cfg.get("gradient_accumulation_steps", 8)))
scheduler = build_scheduler(optimizer, cfg, steps_per_epoch)
model, optimizer, loader, scheduler = accelerator.prepare(model, optimizer, loader, scheduler)
global_step = 0
model.train()
try:
for epoch in range(int(cfg.get("num_train_epochs", 1))):
progress = tqdm(loader, disable=not accelerator.is_main_process, desc=f"epoch {epoch + 1}")
for batch in progress:
with accelerator.accumulate(model):
outputs = model(**batch)
loss = outputs.loss
accelerator.backward(loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(model.parameters(), float(cfg.get("max_grad_norm", 1.0)))
optimizer.step()
scheduler.step()
optimizer.zero_grad()
if accelerator.sync_gradients:
global_step += 1
loss_value = float(loss.detach().float().cpu())
progress.set_postfix(loss=loss_value, step=global_step)
if accelerator.is_main_process:
append_jsonl(
out_dir / "train_log.jsonl",
{"epoch": epoch + 1, "step": global_step, "loss": loss_value},
)
except torch.cuda.OutOfMemoryError as exc:
raise RuntimeError(oom_help()) from exc
accelerator.wait_for_everyone()
if accelerator.is_main_process:
unwrapped = accelerator.unwrap_model(model)
unwrapped.save_pretrained(adapter_dir)
processor.save_pretrained(adapter_dir)
print(f"Saved adapter to {adapter_dir}")
print(f"Saved training config to {out_dir / 'training_config.yaml'}")
print(f"Saved train log to {out_dir / 'train_log.jsonl'}")
if __name__ == "__main__":
main()
|