Add cached AV pipeline for Omni turn LoRA
This commit is contained in:
Xin Wang
634
train_omni_turn_lora.py
Normal file
634
train_omni_turn_lora.py
Normal file
@@ -0,0 +1,634 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Train Qwen2.5-Omni Thinker LoRA + EgoSpeak 10x3 head from AV caches.
|
||||
|
||||
This is a custom task runner intended to execute inside the ms-swift image.
|
||||
It uses the PEFT library shipped with ms-swift because Swift's normal SFT CLI
|
||||
optimizes generated text tokens, whereas this experiment reads one continuous
|
||||
Thinker hidden state and applies a multi-horizon classification loss.
|
||||
|
||||
The cached audio and vision towers remain frozen. The reconstructed sequence
|
||||
is forwarded with ``use_cache=False`` so all retained 2-second blocks
|
||||
participate in autograd through the Thinker LoRA adapters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import importlib.metadata
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
import streaming_turn_model as core
|
||||
from easycom_head import (
|
||||
NUM_CLASSES,
|
||||
NUM_HORIZONS,
|
||||
TurnPredictionHead,
|
||||
checkpoint_head_config,
|
||||
egospeak_multihorizon_loss,
|
||||
)
|
||||
from omni_cached_av import (
|
||||
SCHEMA_VERSION,
|
||||
CachedBlockStore,
|
||||
collate_cached_sequences,
|
||||
forward_last_state,
|
||||
load_json,
|
||||
reconstruct_open_stream_sequence,
|
||||
save_json,
|
||||
torch_load,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_LORA_TARGETS = (
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
)
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-id", required=True)
|
||||
parser.add_argument(
|
||||
"--cache-dir",
|
||||
type=Path,
|
||||
action="append",
|
||||
required=True,
|
||||
help="Cache directory from cache_omni_av_pieces.py; repeat for sessions",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target",
|
||||
type=Path,
|
||||
help="Override target path for a single --cache-dir",
|
||||
)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument("--head-checkpoint", type=Path)
|
||||
parser.add_argument("--max-examples", type=int, default=0)
|
||||
parser.add_argument("--max-steps", type=int, default=200)
|
||||
parser.add_argument("--batch-size", type=int, default=1)
|
||||
parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
|
||||
parser.add_argument("--num-workers", type=int, default=0)
|
||||
parser.add_argument("--block-cache-size", type=int, default=16)
|
||||
parser.add_argument("--learning-rate", type=float, default=2e-5)
|
||||
parser.add_argument("--head-learning-rate", type=float, default=1e-4)
|
||||
parser.add_argument("--weight-decay", type=float, default=0.01)
|
||||
parser.add_argument("--warmup-steps", type=int, default=10)
|
||||
parser.add_argument("--max-grad-norm", type=float, default=1.0)
|
||||
parser.add_argument("--lora-rank", type=int, default=8)
|
||||
parser.add_argument("--lora-alpha", type=int, default=16)
|
||||
parser.add_argument("--lora-dropout", type=float, default=0.05)
|
||||
parser.add_argument(
|
||||
"--lora-target-modules",
|
||||
default=",".join(DEFAULT_LORA_TARGETS),
|
||||
)
|
||||
parser.add_argument("--head-dropout", type=float, default=0.1)
|
||||
parser.add_argument("--gradient-checkpointing", action="store_true")
|
||||
parser.add_argument("--seed", type=int, default=1337)
|
||||
parser.add_argument("--log-every", type=int, default=5)
|
||||
parser.add_argument("--eval-max-batches", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
choices=("bfloat16", "float16"),
|
||||
default="bfloat16",
|
||||
)
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument("--flash-attention-2", action="store_true")
|
||||
parser.add_argument("--overwrite-output", action="store_true")
|
||||
parser.add_argument("--no-progress", action="store_true")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def _seed_everything(seed: int) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def _package_version(name: str) -> str | None:
|
||||
try:
|
||||
return importlib.metadata.version(name)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
return None
|
||||
|
||||
|
||||
def _load_targets(path: Path) -> np.ndarray:
|
||||
targets = np.load(path, mmap_mode="r")
|
||||
if targets.ndim != 2 or targets.shape[1] != NUM_CLASSES:
|
||||
raise ValueError(f"Expected [T,{NUM_CLASSES}] target, got {targets.shape}: {path}")
|
||||
return targets
|
||||
|
||||
|
||||
class OmniTurnCacheDataset(Dataset[dict[str, Any]]):
|
||||
"""One sample per 200-ms anchor, reconstructed lazily on CPU."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cache_dirs: Sequence[Path],
|
||||
*,
|
||||
target_override: Path | None,
|
||||
max_examples: int,
|
||||
seed: int,
|
||||
block_cache_size: int,
|
||||
) -> None:
|
||||
if target_override is not None and len(cache_dirs) != 1:
|
||||
raise ValueError("--target can override only one --cache-dir")
|
||||
self.sources: list[dict[str, Any]] = []
|
||||
self.examples: list[tuple[int, dict[str, Any]]] = []
|
||||
hidden_size: int | None = None
|
||||
reference_contract: tuple[Any, ...] | None = None
|
||||
|
||||
for source_index, raw_cache_dir in enumerate(cache_dirs):
|
||||
cache_dir = raw_cache_dir.expanduser().resolve()
|
||||
manifest_path = cache_dir / "manifest.json"
|
||||
manifest = load_json(manifest_path)
|
||||
if int(manifest.get("schema_version", -1)) != SCHEMA_VERSION:
|
||||
raise ValueError(
|
||||
f"Unsupported cache schema in {manifest_path}: "
|
||||
f"{manifest.get('schema_version')}"
|
||||
)
|
||||
static = torch_load(cache_dir / manifest["static_path"])
|
||||
target_path = (
|
||||
target_override.expanduser().resolve()
|
||||
if target_override is not None
|
||||
else Path(manifest["target_path"])
|
||||
)
|
||||
targets = _load_targets(target_path)
|
||||
if len(targets) != int(manifest["target_steps"]):
|
||||
raise ValueError(
|
||||
f"Target length changed for {manifest_path}: "
|
||||
f"{len(targets)} vs {manifest['target_steps']}"
|
||||
)
|
||||
source_hidden = int(manifest["hidden_size"])
|
||||
hidden_size = source_hidden if hidden_size is None else hidden_size
|
||||
if source_hidden != hidden_size:
|
||||
raise ValueError("Cache directories have different hidden sizes")
|
||||
contract = (
|
||||
manifest["model_id"],
|
||||
manifest["step_seconds"],
|
||||
manifest["native_block_seconds"],
|
||||
manifest["memory_seconds"],
|
||||
manifest["video_fps"],
|
||||
manifest["positions_per_chunk"],
|
||||
)
|
||||
reference_contract = contract if reference_contract is None else reference_contract
|
||||
if contract != reference_contract:
|
||||
raise ValueError(
|
||||
"Cache directories do not share the same model/timing contract: "
|
||||
f"{manifest_path}"
|
||||
)
|
||||
self.sources.append(
|
||||
{
|
||||
"cache_dir": cache_dir,
|
||||
"manifest": manifest,
|
||||
"static": static,
|
||||
"targets": targets,
|
||||
"store": CachedBlockStore(cache_dir, max_blocks=block_cache_size),
|
||||
}
|
||||
)
|
||||
for anchor in manifest["anchors"]:
|
||||
target_index = int(anchor["target_index"])
|
||||
if not bool(anchor.get("has_future_target", True)):
|
||||
continue
|
||||
if target_index + NUM_HORIZONS >= len(targets):
|
||||
continue
|
||||
self.examples.append((source_index, anchor))
|
||||
|
||||
if not self.examples:
|
||||
raise ValueError("No cache anchors have ten future target steps")
|
||||
if max_examples < 0:
|
||||
raise ValueError("--max-examples must be non-negative")
|
||||
if max_examples and max_examples < len(self.examples):
|
||||
generator = random.Random(seed)
|
||||
selected = generator.sample(range(len(self.examples)), max_examples)
|
||||
self.examples = [self.examples[index] for index in sorted(selected)]
|
||||
self.hidden_size = int(hidden_size or 0)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.examples)
|
||||
|
||||
def __getitem__(self, index: int) -> dict[str, Any]:
|
||||
source_index, anchor = self.examples[index]
|
||||
source = self.sources[source_index]
|
||||
embeddings, positions = reconstruct_open_stream_sequence(
|
||||
store=source["store"],
|
||||
static=source["static"],
|
||||
anchor=anchor,
|
||||
)
|
||||
target_index = int(anchor["target_index"])
|
||||
targets = torch.from_numpy(
|
||||
np.array(
|
||||
source["targets"][
|
||||
target_index + 1 : target_index + NUM_HORIZONS + 1
|
||||
],
|
||||
dtype=np.float32,
|
||||
copy=True,
|
||||
)
|
||||
)
|
||||
if targets.shape != (NUM_HORIZONS, NUM_CLASSES):
|
||||
raise RuntimeError(
|
||||
f"Unexpected future target shape {tuple(targets.shape)} at {target_index}"
|
||||
)
|
||||
return {
|
||||
"inputs_embeds": embeddings,
|
||||
"position_ids": positions,
|
||||
"targets": targets,
|
||||
"step_index": target_index,
|
||||
}
|
||||
|
||||
|
||||
def _load_initial_head(
|
||||
checkpoint: Path | None,
|
||||
*,
|
||||
hidden_size: int,
|
||||
dropout: float,
|
||||
) -> TurnPredictionHead:
|
||||
if checkpoint is None:
|
||||
return TurnPredictionHead(hidden_size=hidden_size, dropout=dropout)
|
||||
try:
|
||||
payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
|
||||
except (TypeError, RuntimeError):
|
||||
payload = torch.load(checkpoint, map_location="cpu")
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError(f"Expected dict head checkpoint: {checkpoint}")
|
||||
config = checkpoint_head_config(payload)
|
||||
if int(config["hidden_size"]) != hidden_size:
|
||||
raise ValueError(
|
||||
f"Head hidden size {config['hidden_size']} != cache hidden size {hidden_size}"
|
||||
)
|
||||
head = TurnPredictionHead(
|
||||
hidden_size=hidden_size,
|
||||
num_horizons=int(config["num_horizons"]),
|
||||
num_classes=int(config["num_classes"]),
|
||||
dropout=float(config["dropout"]),
|
||||
)
|
||||
head.load_state_dict(payload.get("head", payload))
|
||||
print(f"Initialized classification head from {checkpoint}")
|
||||
return head
|
||||
|
||||
|
||||
def _prepare_language_lora(
|
||||
model: torch.nn.Module,
|
||||
*,
|
||||
rank: int,
|
||||
alpha: int,
|
||||
dropout: float,
|
||||
target_modules: Sequence[str],
|
||||
gradient_checkpointing: bool,
|
||||
) -> torch.nn.Module:
|
||||
try:
|
||||
from peft import LoraConfig, get_peft_model
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"PEFT is required. The ms-swift 3.12.5 image should include it."
|
||||
) from exc
|
||||
|
||||
model.requires_grad_(False)
|
||||
language_model = model.model
|
||||
matched = [
|
||||
name
|
||||
for name, module in language_model.named_modules()
|
||||
if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in target_modules
|
||||
]
|
||||
if not matched:
|
||||
raise RuntimeError(
|
||||
"No language-model linear layers matched LoRA targets "
|
||||
f"{list(target_modules)}"
|
||||
)
|
||||
config = LoraConfig(
|
||||
r=rank,
|
||||
lora_alpha=alpha,
|
||||
lora_dropout=dropout,
|
||||
target_modules=list(target_modules),
|
||||
bias="none",
|
||||
)
|
||||
language_model = get_peft_model(language_model, config)
|
||||
language_model.config.use_cache = False
|
||||
if gradient_checkpointing:
|
||||
try:
|
||||
language_model.gradient_checkpointing_enable(
|
||||
gradient_checkpointing_kwargs={"use_reentrant": False}
|
||||
)
|
||||
except TypeError:
|
||||
language_model.gradient_checkpointing_enable()
|
||||
language_model.print_trainable_parameters()
|
||||
print(
|
||||
f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
|
||||
f"last={matched[-1]}"
|
||||
)
|
||||
return language_model
|
||||
|
||||
|
||||
def _lr_lambda(step: int, *, total_steps: int, warmup_steps: int) -> float:
|
||||
if warmup_steps > 0 and step < warmup_steps:
|
||||
return float(step + 1) / float(warmup_steps)
|
||||
remaining = max(1, total_steps - warmup_steps)
|
||||
progress = min(1.0, max(0.0, (step - warmup_steps) / remaining))
|
||||
return 0.5 * (1.0 + math.cos(math.pi * progress))
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _evaluate_loss(
|
||||
language_model: torch.nn.Module,
|
||||
head: TurnPredictionHead,
|
||||
loader: DataLoader,
|
||||
*,
|
||||
device: torch.device,
|
||||
model_dtype: torch.dtype,
|
||||
max_batches: int,
|
||||
) -> float:
|
||||
language_was_training = language_model.training
|
||||
head_was_training = head.training
|
||||
language_model.eval()
|
||||
head.eval()
|
||||
losses: list[float] = []
|
||||
for batch_index, batch in enumerate(loader):
|
||||
if max_batches > 0 and batch_index >= max_batches:
|
||||
break
|
||||
inputs_embeds = batch["inputs_embeds"].to(
|
||||
device=device, dtype=model_dtype, non_blocking=True
|
||||
)
|
||||
position_ids = batch["position_ids"].to(device=device, non_blocking=True)
|
||||
attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
|
||||
readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
|
||||
targets = batch["targets"].to(device=device, dtype=torch.float32)
|
||||
states = forward_last_state(
|
||||
language_model,
|
||||
inputs_embeds=inputs_embeds,
|
||||
position_ids=position_ids,
|
||||
attention_mask=attention_mask,
|
||||
readout_indices=readout_indices,
|
||||
)
|
||||
logits = head(states.float())
|
||||
losses.append(float(egospeak_multihorizon_loss(logits, targets).item()))
|
||||
language_model.train(language_was_training)
|
||||
head.train(head_was_training)
|
||||
return float(np.mean(losses))
|
||||
|
||||
|
||||
def _jsonable_args(args: argparse.Namespace) -> dict[str, Any]:
|
||||
output: dict[str, Any] = {}
|
||||
for key, value in vars(args).items():
|
||||
if isinstance(value, Path):
|
||||
output[key] = str(value)
|
||||
elif isinstance(value, list) and value and isinstance(value[0], Path):
|
||||
output[key] = [str(item) for item in value]
|
||||
else:
|
||||
output[key] = value
|
||||
return output
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> None:
|
||||
args = parse_args(argv)
|
||||
if args.max_steps <= 0:
|
||||
raise ValueError("--max-steps must be positive")
|
||||
if args.batch_size <= 0 or args.gradient_accumulation_steps <= 0:
|
||||
raise ValueError("Batch size and gradient accumulation must be positive")
|
||||
if args.num_workers != 0:
|
||||
print(
|
||||
"Warning: each worker owns an independent block-file LRU; use 0 for "
|
||||
"the 16-example pilot and increase only after profiling."
|
||||
)
|
||||
output_dir = args.output_dir.expanduser().resolve()
|
||||
known_outputs = (
|
||||
output_dir / "adapter_config.json",
|
||||
output_dir / "adapter",
|
||||
output_dir / "turn_head.pt",
|
||||
output_dir / "train_summary.json",
|
||||
)
|
||||
if not args.overwrite_output and any(path.exists() for path in known_outputs):
|
||||
raise FileExistsError(
|
||||
f"Training output already exists in {output_dir}; pass --overwrite-output"
|
||||
)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
_seed_everything(args.seed)
|
||||
core._configure_qwen_logging()
|
||||
device = torch.device(args.device)
|
||||
if device.type != "cuda":
|
||||
raise ValueError("This LoRA pilot is intended for one CUDA GPU")
|
||||
|
||||
dataset = OmniTurnCacheDataset(
|
||||
args.cache_dir,
|
||||
target_override=args.target,
|
||||
max_examples=args.max_examples,
|
||||
seed=args.seed,
|
||||
block_cache_size=args.block_cache_size,
|
||||
)
|
||||
generator = torch.Generator().manual_seed(args.seed)
|
||||
train_loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=True,
|
||||
persistent_workers=args.num_workers > 0,
|
||||
collate_fn=collate_cached_sequences,
|
||||
generator=generator,
|
||||
)
|
||||
eval_loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=0,
|
||||
pin_memory=True,
|
||||
collate_fn=collate_cached_sequences,
|
||||
)
|
||||
print(
|
||||
f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, "
|
||||
f"batch={args.batch_size} x accum={args.gradient_accumulation_steps}"
|
||||
)
|
||||
|
||||
base_model, _ = core.load_qwen_omni(
|
||||
args.model_id,
|
||||
dtype=args.dtype,
|
||||
device_map=args.device,
|
||||
use_flash_attention_2=args.flash_attention_2,
|
||||
)
|
||||
target_modules = tuple(
|
||||
value.strip() for value in args.lora_target_modules.split(",") if value.strip()
|
||||
)
|
||||
if not target_modules:
|
||||
raise ValueError("--lora-target-modules is empty")
|
||||
language_model = _prepare_language_lora(
|
||||
base_model,
|
||||
rank=args.lora_rank,
|
||||
alpha=args.lora_alpha,
|
||||
dropout=args.lora_dropout,
|
||||
target_modules=target_modules,
|
||||
gradient_checkpointing=args.gradient_checkpointing,
|
||||
)
|
||||
# Cached embeddings make both modality towers unnecessary during training.
|
||||
# Keeping only the language model materially reduces L20 memory use.
|
||||
del base_model
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
model_dtype = next(language_model.parameters()).dtype
|
||||
head = _load_initial_head(
|
||||
args.head_checkpoint,
|
||||
hidden_size=dataset.hidden_size,
|
||||
dropout=args.head_dropout,
|
||||
).to(device=device, dtype=torch.float32)
|
||||
|
||||
lora_parameters = [parameter for parameter in language_model.parameters() if parameter.requires_grad]
|
||||
if not lora_parameters:
|
||||
raise RuntimeError("LoRA preparation produced no trainable parameters")
|
||||
optimizer = torch.optim.AdamW(
|
||||
[
|
||||
{"params": lora_parameters, "lr": args.learning_rate},
|
||||
{"params": head.parameters(), "lr": args.head_learning_rate},
|
||||
],
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(
|
||||
optimizer,
|
||||
lr_lambda=lambda step: _lr_lambda(
|
||||
step, total_steps=args.max_steps, warmup_steps=args.warmup_steps
|
||||
),
|
||||
)
|
||||
|
||||
initial_loss = _evaluate_loss(
|
||||
language_model,
|
||||
head,
|
||||
eval_loader,
|
||||
device=device,
|
||||
model_dtype=model_dtype,
|
||||
max_batches=args.eval_max_batches,
|
||||
)
|
||||
print(f"Initial cached loss: {initial_loss:.6f}")
|
||||
language_model.train()
|
||||
head.train()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
iterator = iter(train_loader)
|
||||
history: list[dict[str, float | int]] = []
|
||||
progress = tqdm(
|
||||
total=args.max_steps,
|
||||
desc="Thinker LoRA",
|
||||
unit="update",
|
||||
disable=args.no_progress,
|
||||
)
|
||||
for update_step in range(1, args.max_steps + 1):
|
||||
accumulated_loss = 0.0
|
||||
for _ in range(args.gradient_accumulation_steps):
|
||||
try:
|
||||
batch = next(iterator)
|
||||
except StopIteration:
|
||||
iterator = iter(train_loader)
|
||||
batch = next(iterator)
|
||||
inputs_embeds = batch["inputs_embeds"].to(
|
||||
device=device, dtype=model_dtype, non_blocking=True
|
||||
)
|
||||
position_ids = batch["position_ids"].to(device=device, non_blocking=True)
|
||||
attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
|
||||
readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
|
||||
targets = batch["targets"].to(device=device, dtype=torch.float32)
|
||||
states = forward_last_state(
|
||||
language_model,
|
||||
inputs_embeds=inputs_embeds,
|
||||
position_ids=position_ids,
|
||||
attention_mask=attention_mask,
|
||||
readout_indices=readout_indices,
|
||||
)
|
||||
logits = head(states.float())
|
||||
loss = egospeak_multihorizon_loss(logits, targets)
|
||||
(loss / args.gradient_accumulation_steps).backward()
|
||||
accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps
|
||||
|
||||
trainable_parameters = [*lora_parameters, *head.parameters()]
|
||||
grad_norm = torch.nn.utils.clip_grad_norm_(
|
||||
trainable_parameters, args.max_grad_norm
|
||||
)
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
record = {
|
||||
"step": update_step,
|
||||
"loss": accumulated_loss,
|
||||
"grad_norm": float(grad_norm),
|
||||
"lora_lr": float(optimizer.param_groups[0]["lr"]),
|
||||
"head_lr": float(optimizer.param_groups[1]["lr"]),
|
||||
}
|
||||
history.append(record)
|
||||
progress.update(1)
|
||||
progress.set_postfix(
|
||||
loss=f"{accumulated_loss:.4f}",
|
||||
lr=f"{optimizer.param_groups[0]['lr']:.2g}",
|
||||
)
|
||||
if update_step == 1 or update_step % args.log_every == 0:
|
||||
print(
|
||||
f"step={update_step} loss={accumulated_loss:.6f} "
|
||||
f"grad_norm={float(grad_norm):.4f} "
|
||||
f"lora_lr={optimizer.param_groups[0]['lr']:.3g}",
|
||||
flush=True,
|
||||
)
|
||||
progress.close()
|
||||
|
||||
final_loss = _evaluate_loss(
|
||||
language_model,
|
||||
head,
|
||||
eval_loader,
|
||||
device=device,
|
||||
model_dtype=model_dtype,
|
||||
max_batches=args.eval_max_batches,
|
||||
)
|
||||
adapter_dir = output_dir / "adapter"
|
||||
language_model.save_pretrained(adapter_dir)
|
||||
torch.save(
|
||||
{
|
||||
"head": head.state_dict(),
|
||||
"head_config": {
|
||||
"hidden_size": dataset.hidden_size,
|
||||
"num_horizons": NUM_HORIZONS,
|
||||
"num_classes": NUM_CLASSES,
|
||||
"dropout": head.dropout,
|
||||
},
|
||||
"initial_head_checkpoint": (
|
||||
str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
|
||||
),
|
||||
"training_step": args.max_steps,
|
||||
},
|
||||
output_dir / "turn_head.pt",
|
||||
)
|
||||
summary = {
|
||||
"schema_version": 1,
|
||||
"status": "pilot" if args.max_examples else "training",
|
||||
"model_id": args.model_id,
|
||||
"cache_dirs": [str(path.expanduser().resolve()) for path in args.cache_dir],
|
||||
"num_examples": len(dataset),
|
||||
"initial_loss": initial_loss,
|
||||
"final_loss": final_loss,
|
||||
"loss_reduction_ratio": (
|
||||
(initial_loss - final_loss) / initial_loss if initial_loss else None
|
||||
),
|
||||
"peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / (1024 ** 3),
|
||||
"packages": {
|
||||
"torch": torch.__version__,
|
||||
"transformers": _package_version("transformers"),
|
||||
"peft": _package_version("peft"),
|
||||
"ms-swift": _package_version("ms-swift"),
|
||||
},
|
||||
"args": _jsonable_args(args),
|
||||
"history": history,
|
||||
}
|
||||
save_json(output_dir / "train_summary.json", summary)
|
||||
print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})")
|
||||
print(f"Saved LoRA adapter to {adapter_dir}")
|
||||
print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
|
||||
print(f"Peak CUDA allocated: {summary['peak_cuda_allocated_gib']:.2f} GiB")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user