1007 lines
37 KiB
Python
1007 lines
37 KiB
Python
#!/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 queue
|
|
import random
|
|
import threading
|
|
import time
|
|
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,
|
|
pieces_to_cache,
|
|
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(
|
|
"--online-current-tail",
|
|
action="store_true",
|
|
help=(
|
|
"Benchmark complete-2s-block caching: read completed history from "
|
|
"cache but recompute the current 0.2--1.8s AV tail from source media"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--easycom-session-dir",
|
|
type=Path,
|
|
help="Original Session_N media; required by --online-current-tail",
|
|
)
|
|
parser.add_argument(
|
|
"--producer-device",
|
|
help=(
|
|
"Optional second CUDA device for asynchronous frozen-AV production, "
|
|
"for example cuda:1 while --device is cuda:0"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--prefetch-depth",
|
|
type=int,
|
|
default=4,
|
|
help="Maximum producer-ready CPU batches queued ahead of training",
|
|
)
|
|
parser.add_argument("--frame-cache-size", type=int, default=64)
|
|
parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord")
|
|
parser.add_argument("--decoder-gpu-id", type=int, default=0)
|
|
parser.add_argument("--decoder-threads", type=int, default=1)
|
|
parser.add_argument("--max-open-readers", type=int, default=1)
|
|
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,
|
|
defer_current_tail: bool = False,
|
|
) -> 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]]] = []
|
|
self.defer_current_tail = bool(defer_current_tail)
|
|
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]
|
|
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}"
|
|
)
|
|
item: dict[str, Any] = {
|
|
"targets": targets,
|
|
"step_index": target_index,
|
|
}
|
|
if self.defer_current_tail:
|
|
item["source_index"] = source_index
|
|
item["anchor"] = anchor
|
|
else:
|
|
embeddings, positions = reconstruct_open_stream_sequence(
|
|
store=source["store"],
|
|
static=source["static"],
|
|
anchor=anchor,
|
|
)
|
|
item["inputs_embeds"] = embeddings
|
|
item["position_ids"] = positions
|
|
return item
|
|
|
|
|
|
def _collate_deferred_tail(batch: Sequence[dict[str, Any]]) -> dict[str, Any]:
|
|
if len(batch) != 1:
|
|
raise ValueError("--online-current-tail benchmark currently requires batch size 1")
|
|
return batch[0]
|
|
|
|
|
|
class OnlineCurrentTailMaterializer:
|
|
"""Recompute only the current partial block while history stays cached."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
dataset: OmniTurnCacheDataset,
|
|
model: torch.nn.Module,
|
|
processor: Any,
|
|
session_dir: Path,
|
|
frame_cache_size: int,
|
|
decoder: str,
|
|
decoder_gpu_id: int,
|
|
decoder_threads: int,
|
|
max_open_readers: int,
|
|
) -> None:
|
|
if len(dataset.sources) != 1:
|
|
raise ValueError("Online-tail benchmark supports exactly one cache directory")
|
|
from streaming_turn_model_directly import (
|
|
CausalSourceAV,
|
|
_easycom_directory_specs,
|
|
)
|
|
|
|
self.dataset = dataset
|
|
self.model = model
|
|
source = dataset.sources[0]
|
|
manifest = source["manifest"]
|
|
config = core.StreamConfig(
|
|
step_seconds=float(manifest["step_seconds"]),
|
|
native_block_seconds=float(manifest["native_block_seconds"]),
|
|
memory_seconds=float(manifest["memory_seconds"]),
|
|
video_fps=float(manifest["video_fps"]),
|
|
video_min_pixels=int(manifest["video_min_pixels"]),
|
|
video_max_pixels=int(manifest["video_max_pixels"]),
|
|
)
|
|
media = CausalSourceAV(
|
|
_easycom_directory_specs(session_dir),
|
|
output_fps=config.video_fps,
|
|
required_duration=float(manifest["cached_end_seconds"]),
|
|
frame_cache_size=frame_cache_size,
|
|
decoder_backend=decoder,
|
|
decoder_gpu_id=decoder_gpu_id,
|
|
decoder_threads=decoder_threads,
|
|
max_open_readers=max_open_readers,
|
|
)
|
|
self.extractor = core.OpenStreamReplayExtractor(
|
|
model, processor, config, fixed_grid_av=media
|
|
)
|
|
self.step_seconds = config.step_seconds
|
|
self.block_seconds = config.native_block_seconds
|
|
self.steps_per_block = int(source["static"]["steps_per_block"])
|
|
self.calls = 0
|
|
self.online_av_calls = 0
|
|
self.total_seconds = 0.0
|
|
self.online_av_seconds = 0.0
|
|
|
|
def reset_timing(self) -> None:
|
|
self.calls = 0
|
|
self.online_av_calls = 0
|
|
self.total_seconds = 0.0
|
|
self.online_av_seconds = 0.0
|
|
|
|
@torch.no_grad()
|
|
def __call__(self, item: dict[str, Any]) -> dict[str, Any]:
|
|
started = time.perf_counter()
|
|
source_index = int(item["source_index"])
|
|
source = self.dataset.sources[source_index]
|
|
anchor = item["anchor"]
|
|
tail_steps = int(anchor["tail_steps"])
|
|
current_piece = None
|
|
if tail_steps < self.steps_per_block:
|
|
av_started = time.perf_counter()
|
|
block_start = int(anchor["block_index"]) * self.block_seconds
|
|
interval_end = block_start + tail_steps * self.step_seconds
|
|
pieces = self.extractor._prepare_pieces(
|
|
video_path=None,
|
|
audio_path=None,
|
|
interval=(block_start, interval_end),
|
|
use_audio_in_video=True,
|
|
)
|
|
closure_embeddings = self.model.get_input_embeddings()(
|
|
pieces.closure_ids
|
|
)
|
|
current_piece = pieces_to_cache(
|
|
pieces,
|
|
closure_embeddings=closure_embeddings,
|
|
floating_dtype=torch.bfloat16,
|
|
)
|
|
self.online_av_seconds += time.perf_counter() - av_started
|
|
self.online_av_calls += 1
|
|
embeddings, positions = reconstruct_open_stream_sequence(
|
|
store=source["store"],
|
|
static=source["static"],
|
|
anchor=anchor,
|
|
current_piece=current_piece,
|
|
)
|
|
batch = collate_cached_sequences(
|
|
[
|
|
{
|
|
"inputs_embeds": embeddings,
|
|
"position_ids": positions,
|
|
"targets": item["targets"],
|
|
"step_index": item["step_index"],
|
|
}
|
|
]
|
|
)
|
|
self.total_seconds += time.perf_counter() - started
|
|
self.calls += 1
|
|
return batch
|
|
|
|
|
|
class AsyncTailBatchProducer:
|
|
"""Prepare online-tail batches on a second GPU behind a bounded queue."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
loader: DataLoader,
|
|
materializer: OnlineCurrentTailMaterializer,
|
|
total_batches: int,
|
|
prefetch_depth: int,
|
|
) -> None:
|
|
if total_batches <= 0:
|
|
raise ValueError("total_batches must be positive")
|
|
if prefetch_depth <= 0:
|
|
raise ValueError("--prefetch-depth must be positive")
|
|
self.loader = loader
|
|
self.materializer = materializer
|
|
self.total_batches = int(total_batches)
|
|
self.queue: queue.Queue[Any] = queue.Queue(maxsize=int(prefetch_depth))
|
|
self.stop_event = threading.Event()
|
|
self.thread = threading.Thread(
|
|
target=self._run,
|
|
name="omni-av-tail-producer",
|
|
daemon=True,
|
|
)
|
|
self.produced = 0
|
|
self.wait_seconds = 0.0
|
|
|
|
def _put(self, value: Any) -> bool:
|
|
while not self.stop_event.is_set():
|
|
try:
|
|
self.queue.put(value, timeout=0.1)
|
|
return True
|
|
except queue.Full:
|
|
continue
|
|
return False
|
|
|
|
def _run(self) -> None:
|
|
try:
|
|
iterator = iter(self.loader)
|
|
for _ in range(self.total_batches):
|
|
if self.stop_event.is_set():
|
|
return
|
|
try:
|
|
deferred = next(iterator)
|
|
except StopIteration:
|
|
iterator = iter(self.loader)
|
|
deferred = next(iterator)
|
|
prepared = self.materializer(deferred)
|
|
if not self._put(prepared):
|
|
return
|
|
self.produced += 1
|
|
except BaseException as error:
|
|
self._put(error)
|
|
|
|
def start(self) -> None:
|
|
self.thread.start()
|
|
|
|
def next_batch(self) -> dict[str, Any]:
|
|
started = time.perf_counter()
|
|
value = self.queue.get()
|
|
self.wait_seconds += time.perf_counter() - started
|
|
if isinstance(value, BaseException):
|
|
raise RuntimeError("Asynchronous AV producer failed") from value
|
|
if not isinstance(value, dict):
|
|
raise TypeError(f"Producer returned {type(value)!r}, expected dict")
|
|
return value
|
|
|
|
def close(self) -> None:
|
|
self.stop_event.set()
|
|
self.thread.join(timeout=30.0)
|
|
if self.thread.is_alive():
|
|
raise RuntimeError("Asynchronous AV producer did not stop")
|
|
|
|
|
|
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,
|
|
materialize_batch: Any | None = None,
|
|
) -> 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
|
|
if materialize_batch is not None:
|
|
batch = materialize_batch(batch)
|
|
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.online_current_tail:
|
|
if args.easycom_session_dir is None:
|
|
raise ValueError(
|
|
"--easycom-session-dir is required by --online-current-tail"
|
|
)
|
|
if len(args.cache_dir) != 1:
|
|
raise ValueError("Online-tail benchmark requires exactly one --cache-dir")
|
|
if args.batch_size != 1:
|
|
raise ValueError("Online-tail benchmark currently requires --batch-size 1")
|
|
if args.num_workers != 0:
|
|
raise ValueError("Online-tail benchmark currently requires --num-workers 0")
|
|
if args.producer_device is not None:
|
|
if not args.online_current_tail:
|
|
raise ValueError("--producer-device requires --online-current-tail")
|
|
if args.producer_device == args.device:
|
|
raise ValueError("Producer and Thinker trainer must use different devices")
|
|
if not args.producer_device.startswith("cuda:"):
|
|
raise ValueError("--producer-device must be a CUDA device such as cuda:1")
|
|
if args.prefetch_depth <= 0:
|
|
raise ValueError("--prefetch-depth 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,
|
|
defer_current_tail=args.online_current_tail,
|
|
)
|
|
collate_fn = (
|
|
_collate_deferred_tail
|
|
if args.online_current_tail
|
|
else collate_cached_sequences
|
|
)
|
|
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_fn,
|
|
generator=generator,
|
|
)
|
|
eval_loader = DataLoader(
|
|
dataset,
|
|
batch_size=args.batch_size,
|
|
shuffle=False,
|
|
num_workers=0,
|
|
pin_memory=True,
|
|
collate_fn=collate_fn,
|
|
)
|
|
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, processor = 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,
|
|
)
|
|
materialize_batch: OnlineCurrentTailMaterializer | None = None
|
|
producer_model: torch.nn.Module | None = None
|
|
if args.online_current_tail:
|
|
assert args.easycom_session_dir is not None
|
|
materializer_model = base_model
|
|
materializer_processor = processor
|
|
if args.producer_device is not None:
|
|
producer_model, materializer_processor = core.load_qwen_omni(
|
|
args.model_id,
|
|
dtype=args.dtype,
|
|
device_map=args.producer_device,
|
|
use_flash_attention_2=args.flash_attention_2,
|
|
)
|
|
producer_model.requires_grad_(False)
|
|
materializer_model = producer_model
|
|
materialize_batch = OnlineCurrentTailMaterializer(
|
|
dataset=dataset,
|
|
model=materializer_model,
|
|
processor=materializer_processor,
|
|
session_dir=args.easycom_session_dir,
|
|
frame_cache_size=args.frame_cache_size,
|
|
decoder=args.decoder,
|
|
decoder_gpu_id=args.decoder_gpu_id,
|
|
decoder_threads=args.decoder_threads,
|
|
max_open_readers=args.max_open_readers,
|
|
)
|
|
print(
|
|
"Cache policy benchmark: complete 2s history cached; "
|
|
"current 0.2--1.8s tail computed online"
|
|
)
|
|
if args.producer_device is not None:
|
|
print(
|
|
f"Producer-consumer enabled: frozen AV={args.producer_device}, "
|
|
f"Thinker LoRA={args.device}, queue={args.prefetch_depth}"
|
|
)
|
|
# The consumer retains the same language object through its PEFT
|
|
# wrapper, so the outer multimodal container and towers can go.
|
|
del base_model
|
|
del processor
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
else:
|
|
# All AV phases are cached, so the modality towers can be discarded.
|
|
del base_model
|
|
del processor
|
|
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,
|
|
materialize_batch=materialize_batch,
|
|
)
|
|
print(f"Initial cached loss: {initial_loss:.6f}")
|
|
if materialize_batch is not None:
|
|
materialize_batch.reset_timing()
|
|
language_model.train()
|
|
head.train()
|
|
optimizer.zero_grad(set_to_none=True)
|
|
training_started = time.perf_counter()
|
|
batch_producer: AsyncTailBatchProducer | None = None
|
|
if materialize_batch is not None and args.producer_device is not None:
|
|
batch_producer = AsyncTailBatchProducer(
|
|
loader=train_loader,
|
|
materializer=materialize_batch,
|
|
total_batches=args.max_steps * args.gradient_accumulation_steps,
|
|
prefetch_depth=args.prefetch_depth,
|
|
)
|
|
batch_producer.start()
|
|
iterator = iter(train_loader) if batch_producer is None else None
|
|
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):
|
|
if batch_producer is not None:
|
|
batch = batch_producer.next_batch()
|
|
else:
|
|
assert iterator is not None
|
|
try:
|
|
batch = next(iterator)
|
|
except StopIteration:
|
|
iterator = iter(train_loader)
|
|
batch = next(iterator)
|
|
if materialize_batch is not None and batch_producer is None:
|
|
batch = materialize_batch(batch)
|
|
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()
|
|
training_seconds = time.perf_counter() - training_started
|
|
if batch_producer is not None:
|
|
batch_producer.close()
|
|
online_timing = None
|
|
if materialize_batch is not None:
|
|
online_timing = {
|
|
"materialize_calls": materialize_batch.calls,
|
|
"materialize_seconds": materialize_batch.total_seconds,
|
|
"online_av_calls": materialize_batch.online_av_calls,
|
|
"online_av_seconds": materialize_batch.online_av_seconds,
|
|
"consumer_queue_wait_seconds": (
|
|
batch_producer.wait_seconds if batch_producer is not None else None
|
|
),
|
|
"producer_batches": (
|
|
batch_producer.produced if batch_producer is not None else None
|
|
),
|
|
}
|
|
|
|
final_loss = _evaluate_loss(
|
|
language_model,
|
|
head,
|
|
eval_loader,
|
|
device=device,
|
|
model_dtype=model_dtype,
|
|
max_batches=args.eval_max_batches,
|
|
materialize_batch=materialize_batch,
|
|
)
|
|
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,
|
|
"training_seconds": training_seconds,
|
|
"updates_per_second": args.max_steps / training_seconds,
|
|
"online_tail_timing": online_timing,
|
|
"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),
|
|
"producer_peak_cuda_allocated_gib": (
|
|
torch.cuda.max_memory_allocated(torch.device(args.producer_device))
|
|
/ (1024 ** 3)
|
|
if args.producer_device is not None
|
|
else None
|
|
),
|
|
"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"Training speed: {training_seconds:.3f}s / {args.max_steps} updates = "
|
|
f"{args.max_steps / training_seconds:.3f} update/s"
|
|
)
|
|
if online_timing is not None:
|
|
print(
|
|
"Online current-tail AV: "
|
|
f"{online_timing['online_av_seconds']:.3f}s / "
|
|
f"{online_timing['online_av_calls']} calls"
|
|
)
|
|
if online_timing["consumer_queue_wait_seconds"] is not None:
|
|
print(
|
|
"Consumer queue wait: "
|
|
f"{online_timing['consumer_queue_wait_seconds']:.3f}s; "
|
|
f"producer batches={online_timing['producer_batches']}"
|
|
)
|
|
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 summary["producer_peak_cuda_allocated_gib"] is not None:
|
|
print(
|
|
"Producer peak CUDA allocated: "
|
|
f"{summary['producer_peak_cuda_allocated_gib']:.2f} GiB"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|