Add cached AV pipeline for Omni turn LoRA
This commit is contained in:
Xin Wang
426
cache_omni_av_pieces.py
Normal file
426
cache_omni_av_pieces.py
Normal file
@@ -0,0 +1,426 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Cache exact Qwen2.5-Omni AV pieces for Thinker LoRA training.
|
||||
|
||||
For every source-aligned 2-second block this script caches ten causal phases:
|
||||
0.2, 0.4, ..., 2.0 seconds. The 2.0-second phase is reused as completed
|
||||
history; phases 0.2--1.8 are used only for the current partial block. No
|
||||
future samples enter a phase.
|
||||
|
||||
The optional verification compares a full-sequence reconstruction from the
|
||||
cache with the existing open-stream KV replay. Run this 24-second pilot before
|
||||
caching a full EasyCom session.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import platform
|
||||
from pathlib import Path
|
||||
from typing import Any, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
import streaming_turn_model as core
|
||||
from easycom_head import NUM_HORIZONS
|
||||
from omni_cached_av import (
|
||||
SCHEMA_VERSION,
|
||||
CachedBlockStore,
|
||||
anchor_intervals,
|
||||
cache_block_filename,
|
||||
cache_size_bytes,
|
||||
forward_last_state,
|
||||
make_anchor,
|
||||
pieces_to_cache,
|
||||
reconstruct_open_stream_sequence,
|
||||
save_json,
|
||||
)
|
||||
from streaming_turn_model_directly import CausalSourceAV, _easycom_directory_specs
|
||||
|
||||
|
||||
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("--easycom-session-dir", type=Path, required=True)
|
||||
parser.add_argument("--target", type=Path, required=True)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--end",
|
||||
type=float,
|
||||
default=24.0,
|
||||
help="Dense session endpoint in seconds; use 24 for the first pilot",
|
||||
)
|
||||
parser.add_argument("--step-seconds", type=float, default=0.2)
|
||||
parser.add_argument("--native-block-seconds", type=float, default=2.0)
|
||||
parser.add_argument("--memory-seconds", type=float, default=8.0)
|
||||
parser.add_argument("--video-fps", type=float, default=2.0)
|
||||
parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
|
||||
parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
|
||||
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(
|
||||
"--dtype",
|
||||
choices=("auto", "bfloat16", "float16", "float32"),
|
||||
default="bfloat16",
|
||||
)
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument("--flash-attention-2", action="store_true")
|
||||
parser.add_argument(
|
||||
"--verify-every",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Compare every Nth 5-Hz state with exact KV replay; 0 disables",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--minimum-cosine",
|
||||
type=float,
|
||||
default=0.999,
|
||||
help="Fail if cached full-forward/reference KV cosine is below this",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite",
|
||||
action="store_true",
|
||||
help="Overwrite this cache's known block/static/manifest files",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def _integer_ratio(numerator: float, denominator: float, name: str) -> int:
|
||||
ratio = numerator / denominator
|
||||
rounded = round(ratio)
|
||||
if not math.isclose(ratio, rounded, rel_tol=0.0, abs_tol=1e-6):
|
||||
raise ValueError(f"{name} must be integral, got {ratio}")
|
||||
return int(rounded)
|
||||
|
||||
|
||||
def _load_targets(path: Path) -> np.ndarray:
|
||||
targets = np.load(path, mmap_mode="r")
|
||||
if targets.ndim != 2 or targets.shape[1] != 3:
|
||||
raise ValueError(f"Expected target [T,3], got {targets.shape}: {path}")
|
||||
return targets
|
||||
|
||||
|
||||
def _safe_output_path(path: Path, overwrite: bool) -> None:
|
||||
if path.exists() and not overwrite:
|
||||
raise FileExistsError(f"Refusing to overwrite {path}; pass --overwrite")
|
||||
|
||||
|
||||
def _model_device(model: torch.nn.Module) -> torch.device:
|
||||
try:
|
||||
return next(model.parameters()).device
|
||||
except StopIteration as exc:
|
||||
raise RuntimeError("Model has no parameters") from exc
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _cache_blocks(
|
||||
*,
|
||||
extractor: core.OpenStreamReplayExtractor,
|
||||
model: torch.nn.Module,
|
||||
output_dir: Path,
|
||||
num_steps: int,
|
||||
steps_per_block: int,
|
||||
step_seconds: float,
|
||||
block_seconds: float,
|
||||
overwrite: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
num_blocks = math.ceil(num_steps / steps_per_block)
|
||||
records: list[dict[str, Any]] = []
|
||||
input_embeddings = model.get_input_embeddings()
|
||||
for block_index in tqdm(range(num_blocks), desc="cache AV blocks", unit="block"):
|
||||
first_step = block_index * steps_per_block
|
||||
phase_count = min(steps_per_block, num_steps - first_step)
|
||||
block_start = block_index * block_seconds
|
||||
variants: dict[str, dict[str, torch.Tensor]] = {}
|
||||
token_lengths: dict[str, int] = {}
|
||||
for tail_steps in range(1, phase_count + 1):
|
||||
interval_end = block_start + tail_steps * step_seconds
|
||||
pieces = extractor._prepare_pieces(
|
||||
video_path=None,
|
||||
audio_path=None,
|
||||
interval=(block_start, interval_end),
|
||||
use_audio_in_video=True,
|
||||
)
|
||||
closure_embeddings = input_embeddings(pieces.closure_ids)
|
||||
cached = pieces_to_cache(
|
||||
pieces,
|
||||
closure_embeddings=closure_embeddings,
|
||||
floating_dtype=torch.bfloat16,
|
||||
)
|
||||
variants[str(tail_steps)] = cached
|
||||
token_lengths[str(tail_steps)] = int(
|
||||
cached["prefix_embeddings"].shape[1]
|
||||
+ cached["payload_embeddings"].shape[1]
|
||||
+ cached["closure_embeddings"].shape[1]
|
||||
)
|
||||
path = output_dir / cache_block_filename(block_index)
|
||||
_safe_output_path(path, overwrite)
|
||||
torch.save(
|
||||
{
|
||||
"schema_version": SCHEMA_VERSION,
|
||||
"block_index": block_index,
|
||||
"start_seconds": block_start,
|
||||
"end_seconds": min(
|
||||
block_start + block_seconds, num_steps * step_seconds
|
||||
),
|
||||
"variants": variants,
|
||||
},
|
||||
path,
|
||||
)
|
||||
records.append(
|
||||
{
|
||||
"block_index": block_index,
|
||||
"path": path.name,
|
||||
"phase_count": phase_count,
|
||||
"token_lengths_without_query": token_lengths,
|
||||
"size_bytes": path.stat().st_size,
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _verify_cache(
|
||||
*,
|
||||
model: torch.nn.Module,
|
||||
extractor: core.OpenStreamReplayExtractor,
|
||||
cache_dir: Path,
|
||||
static: dict[str, Any],
|
||||
anchors: list[dict[str, Any]],
|
||||
verify_every: int,
|
||||
minimum_cosine: float,
|
||||
block_seconds: float,
|
||||
step_seconds: float,
|
||||
) -> list[dict[str, float]]:
|
||||
if verify_every <= 0:
|
||||
return []
|
||||
selected = [
|
||||
anchor
|
||||
for anchor in anchors
|
||||
if (int(anchor["step_index"]) + 1) % verify_every == 0
|
||||
]
|
||||
if anchors and anchors[-1] not in selected:
|
||||
selected.append(anchors[-1])
|
||||
store = CachedBlockStore(cache_dir, max_blocks=max(2, int(static["memory_blocks"])))
|
||||
device = _model_device(model)
|
||||
checks: list[dict[str, float]] = []
|
||||
for anchor in tqdm(selected, desc="verify cached replay", unit="state"):
|
||||
embeddings, positions = reconstruct_open_stream_sequence(
|
||||
store=store, static=static, anchor=anchor
|
||||
)
|
||||
embeddings = embeddings.unsqueeze(0).to(device=device, dtype=model.dtype)
|
||||
positions = positions.unsqueeze(1).to(device=device)
|
||||
attention_mask = torch.ones(
|
||||
1, embeddings.shape[1], dtype=torch.long, device=device
|
||||
)
|
||||
cached_state = forward_last_state(
|
||||
model.model,
|
||||
inputs_embeds=embeddings,
|
||||
position_ids=positions,
|
||||
attention_mask=attention_mask,
|
||||
readout_indices=torch.tensor(
|
||||
[embeddings.shape[1] - 1], dtype=torch.long, device=device
|
||||
),
|
||||
)
|
||||
exact_state = extractor.state_from_scratch(
|
||||
video_path=None,
|
||||
audio_path=None,
|
||||
intervals=anchor_intervals(
|
||||
anchor,
|
||||
block_seconds=block_seconds,
|
||||
step_seconds=step_seconds,
|
||||
),
|
||||
use_audio_in_video=True,
|
||||
)
|
||||
cached_float = cached_state.float()
|
||||
exact_float = exact_state.float()
|
||||
check = {
|
||||
"time_seconds": float(anchor["time_seconds"]),
|
||||
"cosine": float(
|
||||
F.cosine_similarity(cached_float, exact_float, dim=-1).item()
|
||||
),
|
||||
"mean_abs": float((cached_float - exact_float).abs().mean().item()),
|
||||
"max_abs": float((cached_float - exact_float).abs().max().item()),
|
||||
}
|
||||
checks.append(check)
|
||||
print(
|
||||
f"verify t={check['time_seconds']:.1f}s "
|
||||
f"cosine={check['cosine']:.8f} "
|
||||
f"mean_abs={check['mean_abs']:.6f} "
|
||||
f"max_abs={check['max_abs']:.6f}",
|
||||
flush=True,
|
||||
)
|
||||
if check["cosine"] < minimum_cosine:
|
||||
raise RuntimeError(
|
||||
"Cached full-sequence reconstruction disagrees with open-stream "
|
||||
f"KV replay at t={check['time_seconds']:.1f}s: "
|
||||
f"cosine={check['cosine']:.8f} < {minimum_cosine:.8f}"
|
||||
)
|
||||
return checks
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> None:
|
||||
args = parse_args(argv)
|
||||
core._configure_qwen_logging()
|
||||
if args.end <= 0:
|
||||
raise ValueError("--end must be positive")
|
||||
steps_per_block = _integer_ratio(
|
||||
args.native_block_seconds, args.step_seconds, "native block / step"
|
||||
)
|
||||
memory_blocks = _integer_ratio(
|
||||
args.memory_seconds, args.native_block_seconds, "memory / native block"
|
||||
)
|
||||
num_steps = _integer_ratio(args.end, args.step_seconds, "end / step")
|
||||
targets = _load_targets(args.target)
|
||||
if num_steps > len(targets):
|
||||
raise ValueError(
|
||||
f"Requested {num_steps} cached steps but target has only {len(targets)}"
|
||||
)
|
||||
|
||||
output_dir = args.output_dir.expanduser().resolve()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
manifest_path = output_dir / "manifest.json"
|
||||
static_path = output_dir / "static.pt"
|
||||
_safe_output_path(manifest_path, args.overwrite)
|
||||
_safe_output_path(static_path, args.overwrite)
|
||||
|
||||
config = core.StreamConfig(
|
||||
step_seconds=args.step_seconds,
|
||||
native_block_seconds=args.native_block_seconds,
|
||||
memory_seconds=args.memory_seconds,
|
||||
video_fps=args.video_fps,
|
||||
video_min_pixels=args.video_min_pixels,
|
||||
video_max_pixels=args.video_max_pixels,
|
||||
)
|
||||
config.validate()
|
||||
specs = _easycom_directory_specs(args.easycom_session_dir)
|
||||
media = CausalSourceAV(
|
||||
specs,
|
||||
output_fps=args.video_fps,
|
||||
required_duration=args.end,
|
||||
frame_cache_size=args.frame_cache_size,
|
||||
decoder_backend=args.decoder,
|
||||
decoder_gpu_id=args.decoder_gpu_id,
|
||||
decoder_threads=args.decoder_threads,
|
||||
max_open_readers=args.max_open_readers,
|
||||
)
|
||||
print(
|
||||
f"Direct media ready: duration={args.end:.1f}s, "
|
||||
f"grid={media.expected_grid_frames} frames, audio={len(media.audio)} samples"
|
||||
)
|
||||
|
||||
model, processor = core.load_qwen_omni(
|
||||
args.model_id,
|
||||
dtype=args.dtype,
|
||||
device_map=args.device,
|
||||
use_flash_attention_2=args.flash_attention_2,
|
||||
)
|
||||
model.requires_grad_(False)
|
||||
extractor = core.OpenStreamReplayExtractor(
|
||||
model, processor, config, fixed_grid_av=media
|
||||
)
|
||||
with torch.inference_mode():
|
||||
query_ids = extractor.query_ids
|
||||
query_embeddings = model.get_input_embeddings()(query_ids)
|
||||
positions_per_chunk = core._tmrope_positions_per_chunk(model)
|
||||
static = {
|
||||
"schema_version": SCHEMA_VERSION,
|
||||
"steps_per_block": steps_per_block,
|
||||
"memory_blocks": memory_blocks,
|
||||
"positions_per_chunk": positions_per_chunk,
|
||||
"query_ids": query_ids.detach().cpu().long(),
|
||||
"query_embeddings": query_embeddings.detach().cpu().to(torch.bfloat16),
|
||||
}
|
||||
torch.save(static, static_path)
|
||||
|
||||
block_records = _cache_blocks(
|
||||
extractor=extractor,
|
||||
model=model,
|
||||
output_dir=output_dir,
|
||||
num_steps=num_steps,
|
||||
steps_per_block=steps_per_block,
|
||||
step_seconds=args.step_seconds,
|
||||
block_seconds=args.native_block_seconds,
|
||||
overwrite=args.overwrite,
|
||||
)
|
||||
anchors = [
|
||||
make_anchor(
|
||||
step_index=step_index,
|
||||
steps_per_block=steps_per_block,
|
||||
memory_blocks=memory_blocks,
|
||||
step_seconds=args.step_seconds,
|
||||
num_target_steps=len(targets),
|
||||
num_horizons=NUM_HORIZONS,
|
||||
)
|
||||
for step_index in range(num_steps)
|
||||
]
|
||||
checks = _verify_cache(
|
||||
model=model,
|
||||
extractor=extractor,
|
||||
cache_dir=output_dir,
|
||||
static=static,
|
||||
anchors=anchors,
|
||||
verify_every=args.verify_every,
|
||||
minimum_cosine=args.minimum_cosine,
|
||||
block_seconds=args.native_block_seconds,
|
||||
step_seconds=args.step_seconds,
|
||||
)
|
||||
|
||||
manifest = {
|
||||
"schema_version": SCHEMA_VERSION,
|
||||
"cache_kind": "qwen2.5-omni-frozen-av-open-stream-phases",
|
||||
"session": args.easycom_session_dir.resolve().name,
|
||||
"model_id": args.model_id,
|
||||
"target_path": str(args.target.expanduser().resolve()),
|
||||
"target_steps": int(len(targets)),
|
||||
"cached_steps": num_steps,
|
||||
"cached_end_seconds": args.end,
|
||||
"hidden_size": int(query_embeddings.shape[-1]),
|
||||
"floating_dtype": "bfloat16",
|
||||
"step_seconds": args.step_seconds,
|
||||
"native_block_seconds": args.native_block_seconds,
|
||||
"memory_seconds": args.memory_seconds,
|
||||
"steps_per_block": steps_per_block,
|
||||
"memory_blocks": memory_blocks,
|
||||
"positions_per_chunk": positions_per_chunk,
|
||||
"video_fps": args.video_fps,
|
||||
"video_min_pixels": args.video_min_pixels,
|
||||
"video_max_pixels": args.video_max_pixels,
|
||||
"use_audio_in_video": True,
|
||||
"interval_convention": "half_open_[start,end)",
|
||||
"phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history",
|
||||
"label_policy": "target[i+1:i+11]",
|
||||
"static_path": static_path.name,
|
||||
"blocks": block_records,
|
||||
"anchors": anchors,
|
||||
"verification": checks,
|
||||
"media": media.metadata(),
|
||||
"software": {
|
||||
"python": platform.python_version(),
|
||||
"torch": torch.__version__,
|
||||
},
|
||||
}
|
||||
save_json(manifest_path, manifest)
|
||||
size = cache_size_bytes(
|
||||
[static_path, *[output_dir / record["path"] for record in block_records]]
|
||||
)
|
||||
print(
|
||||
f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}"
|
||||
)
|
||||
print(f"Cache tensor size: {size / (1024 ** 3):.3f} GiB")
|
||||
if checks:
|
||||
print(
|
||||
"Verification: "
|
||||
f"lowest cosine={min(item['cosine'] for item in checks):.8f}, "
|
||||
f"worst max_abs={max(item['max_abs'] for item in checks):.6f}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
365
omni_cached_av.py
Normal file
365
omni_cached_av.py
Normal file
@@ -0,0 +1,365 @@
|
||||
"""Cached open-stream AV pieces for Qwen2.5-Omni turn-taking LoRA.
|
||||
|
||||
The modality towers are frozen for the first fine-tuning experiment. This
|
||||
module serializes their exact output embeddings and Qwen's local TMRoPE
|
||||
coordinates, then reconstructs the same still-open media turn without using a
|
||||
KV cache. The resulting sequence can be passed through the Thinker with
|
||||
``use_cache=False`` so gradients reach LoRA adapters in every retained block.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Sequence
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
PIECE_TENSOR_NAMES = (
|
||||
"prefix_embeddings",
|
||||
"prefix_position_ids",
|
||||
"payload_embeddings",
|
||||
"payload_position_ids",
|
||||
"closure_ids",
|
||||
"closure_position_ids",
|
||||
"closure_embeddings",
|
||||
)
|
||||
|
||||
|
||||
def load_json(path: str | Path) -> dict[str, Any]:
|
||||
return json.loads(Path(path).read_text())
|
||||
|
||||
|
||||
def save_json(path: str | Path, payload: dict[str, Any]) -> None:
|
||||
Path(path).write_text(json.dumps(payload, indent=2) + "\n")
|
||||
|
||||
|
||||
def torch_load(path: str | Path) -> Any:
|
||||
"""Load tensor-only cache files across old and new PyTorch defaults."""
|
||||
|
||||
try:
|
||||
return torch.load(Path(path), map_location="cpu", weights_only=True)
|
||||
except TypeError:
|
||||
return torch.load(Path(path), map_location="cpu")
|
||||
|
||||
|
||||
def cache_block_filename(block_index: int) -> str:
|
||||
return f"block_{block_index:06d}.pt"
|
||||
|
||||
|
||||
def pieces_to_cache(
|
||||
pieces: Any,
|
||||
*,
|
||||
closure_embeddings: torch.Tensor,
|
||||
floating_dtype: torch.dtype = torch.bfloat16,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Detach one ``OpenMediaPieces`` object into a portable CPU payload."""
|
||||
|
||||
output: dict[str, torch.Tensor] = {}
|
||||
for name in PIECE_TENSOR_NAMES:
|
||||
value = closure_embeddings if name == "closure_embeddings" else getattr(pieces, name)
|
||||
value = value.detach().cpu().contiguous()
|
||||
if value.is_floating_point():
|
||||
value = value.to(floating_dtype)
|
||||
output[name] = value
|
||||
validate_piece(output)
|
||||
return output
|
||||
|
||||
|
||||
def validate_piece(piece: dict[str, torch.Tensor]) -> None:
|
||||
missing = [name for name in PIECE_TENSOR_NAMES if name not in piece]
|
||||
if missing:
|
||||
raise ValueError(f"Cached AV piece is missing {missing}")
|
||||
prefix = piece["prefix_embeddings"]
|
||||
payload = piece["payload_embeddings"]
|
||||
closure = piece["closure_embeddings"]
|
||||
if prefix.ndim != 3 or prefix.shape[0] != 1:
|
||||
raise ValueError(f"Invalid prefix embeddings: {tuple(prefix.shape)}")
|
||||
if payload.ndim != 3 or payload.shape[0] != 1:
|
||||
raise ValueError(f"Invalid payload embeddings: {tuple(payload.shape)}")
|
||||
if closure.ndim != 3 or closure.shape[0] != 1:
|
||||
raise ValueError(f"Invalid closure embeddings: {tuple(closure.shape)}")
|
||||
hidden_sizes = {int(prefix.shape[-1]), int(payload.shape[-1]), int(closure.shape[-1])}
|
||||
if len(hidden_sizes) != 1:
|
||||
raise ValueError(f"Cached piece hidden sizes disagree: {sorted(hidden_sizes)}")
|
||||
expected_positions = (
|
||||
("prefix_position_ids", prefix.shape[1]),
|
||||
("payload_position_ids", payload.shape[1]),
|
||||
("closure_position_ids", closure.shape[1]),
|
||||
)
|
||||
for name, length in expected_positions:
|
||||
value = piece[name]
|
||||
if value.shape != (3, 1, length):
|
||||
raise ValueError(
|
||||
f"{name} must have shape {(3, 1, length)}, got {tuple(value.shape)}"
|
||||
)
|
||||
if piece["closure_ids"].shape != (1, closure.shape[1]):
|
||||
raise ValueError(
|
||||
"closure_ids/closure_embeddings mismatch: "
|
||||
f"{tuple(piece['closure_ids'].shape)} vs {tuple(closure.shape)}"
|
||||
)
|
||||
|
||||
|
||||
def make_anchor(
|
||||
*,
|
||||
step_index: int,
|
||||
steps_per_block: int,
|
||||
memory_blocks: int,
|
||||
step_seconds: float,
|
||||
num_target_steps: int,
|
||||
num_horizons: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Describe the exact completed-history/current-tail selection at time t."""
|
||||
|
||||
if step_index < 0:
|
||||
raise ValueError("step_index must be non-negative")
|
||||
block_index = step_index // steps_per_block
|
||||
tail_steps = step_index % steps_per_block + 1
|
||||
first_history = max(0, block_index - (memory_blocks - 1))
|
||||
history_blocks = list(range(first_history, block_index))
|
||||
return {
|
||||
"step_index": int(step_index),
|
||||
"time_seconds": float((step_index + 1) * step_seconds),
|
||||
"block_index": int(block_index),
|
||||
"tail_steps": int(tail_steps),
|
||||
"history_blocks": history_blocks,
|
||||
"target_index": int(step_index),
|
||||
"has_future_target": bool(step_index + num_horizons < num_target_steps),
|
||||
}
|
||||
|
||||
|
||||
class CachedBlockStore:
|
||||
"""Small CPU LRU for block shard files shared by adjacent samples."""
|
||||
|
||||
def __init__(self, cache_dir: str | Path, max_blocks: int = 16) -> None:
|
||||
if max_blocks <= 0:
|
||||
raise ValueError("max_blocks must be positive")
|
||||
self.cache_dir = Path(cache_dir)
|
||||
self.max_blocks = int(max_blocks)
|
||||
self._blocks: OrderedDict[int, dict[str, Any]] = OrderedDict()
|
||||
|
||||
def get(self, block_index: int) -> dict[str, Any]:
|
||||
cached = self._blocks.get(block_index)
|
||||
if cached is not None:
|
||||
self._blocks.move_to_end(block_index)
|
||||
return cached
|
||||
path = self.cache_dir / cache_block_filename(block_index)
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(path)
|
||||
payload = torch_load(path)
|
||||
if int(payload.get("block_index", -1)) != block_index:
|
||||
raise ValueError(f"Block index mismatch in {path}")
|
||||
variants = payload.get("variants")
|
||||
if not isinstance(variants, dict) or not variants:
|
||||
raise ValueError(f"No variants in {path}")
|
||||
self._blocks[block_index] = payload
|
||||
self._blocks.move_to_end(block_index)
|
||||
while len(self._blocks) > self.max_blocks:
|
||||
self._blocks.popitem(last=False)
|
||||
return payload
|
||||
|
||||
def variant(self, block_index: int, tail_steps: int) -> dict[str, torch.Tensor]:
|
||||
block = self.get(block_index)
|
||||
variants = block["variants"]
|
||||
piece = variants.get(str(tail_steps), variants.get(tail_steps))
|
||||
if piece is None:
|
||||
raise KeyError(
|
||||
f"Block {block_index} has no tail_steps={tail_steps}; "
|
||||
f"available={sorted(str(key) for key in variants)}"
|
||||
)
|
||||
validate_piece(piece)
|
||||
return piece
|
||||
|
||||
|
||||
def _strip_embedding_batch(value: torch.Tensor) -> torch.Tensor:
|
||||
if value.ndim != 3 or value.shape[0] != 1:
|
||||
raise ValueError(f"Expected [1,L,D], got {tuple(value.shape)}")
|
||||
return value[0]
|
||||
|
||||
|
||||
def _strip_position_batch(value: torch.Tensor) -> torch.LongTensor:
|
||||
if value.ndim != 3 or value.shape[:2] != (3, 1):
|
||||
raise ValueError(f"Expected [3,1,L], got {tuple(value.shape)}")
|
||||
return value[:, 0, :].long()
|
||||
|
||||
|
||||
def reconstruct_open_stream_sequence(
|
||||
*,
|
||||
store: CachedBlockStore,
|
||||
static: dict[str, Any],
|
||||
anchor: dict[str, Any],
|
||||
) -> tuple[torch.Tensor, torch.LongTensor]:
|
||||
"""Materialize one differentiable Thinker input from frozen AV embeddings.
|
||||
|
||||
Complete history blocks use their 2.0-second variant. Only the current
|
||||
block uses its 0.2--2.0-second phase. The first retained block keeps the
|
||||
chat/media prefix; later blocks append payload only. The final block's
|
||||
media closure and the artificial assistant readout query are appended once.
|
||||
"""
|
||||
|
||||
steps_per_block = int(static["steps_per_block"])
|
||||
positions_per_chunk = int(static["positions_per_chunk"])
|
||||
history_blocks = [int(value) for value in anchor["history_blocks"]]
|
||||
current_block = int(anchor["block_index"])
|
||||
selected: list[tuple[int, int]] = [
|
||||
(block_index, steps_per_block) for block_index in history_blocks
|
||||
]
|
||||
selected.append((current_block, int(anchor["tail_steps"])))
|
||||
if not selected:
|
||||
raise ValueError("An anchor must select at least its current block")
|
||||
|
||||
embedding_parts: list[torch.Tensor] = []
|
||||
position_parts: list[torch.LongTensor] = []
|
||||
first_payload_base: int | None = None
|
||||
last_piece: dict[str, torch.Tensor] | None = None
|
||||
last_shift = 0
|
||||
|
||||
for sequence_index, (block_index, tail_steps) in enumerate(selected):
|
||||
piece = store.variant(block_index, tail_steps)
|
||||
local_payload_positions = _strip_position_batch(piece["payload_position_ids"])
|
||||
local_payload_base = int(local_payload_positions.min().item())
|
||||
if first_payload_base is None:
|
||||
first_payload_base = local_payload_base
|
||||
desired_payload_base = first_payload_base + sequence_index * positions_per_chunk
|
||||
position_shift = desired_payload_base - local_payload_base
|
||||
|
||||
if sequence_index == 0:
|
||||
embedding_parts.append(_strip_embedding_batch(piece["prefix_embeddings"]))
|
||||
position_parts.append(_strip_position_batch(piece["prefix_position_ids"]))
|
||||
embedding_parts.append(_strip_embedding_batch(piece["payload_embeddings"]))
|
||||
position_parts.append(local_payload_positions + position_shift)
|
||||
last_piece = piece
|
||||
last_shift = position_shift
|
||||
|
||||
if last_piece is None:
|
||||
raise RuntimeError("No cached AV piece was selected")
|
||||
closure_embeddings = _strip_embedding_batch(last_piece["closure_embeddings"])
|
||||
closure_positions = (
|
||||
_strip_position_batch(last_piece["closure_position_ids"]) + last_shift
|
||||
)
|
||||
embedding_parts.append(closure_embeddings)
|
||||
position_parts.append(closure_positions)
|
||||
|
||||
query_embeddings = static["query_embeddings"]
|
||||
if query_embeddings.ndim != 3 or query_embeddings.shape[0] != 1:
|
||||
raise ValueError(
|
||||
f"static query_embeddings must be [1,Q,D], got {tuple(query_embeddings.shape)}"
|
||||
)
|
||||
query_embeddings = query_embeddings[0]
|
||||
query_start = int(closure_positions.max().item()) + 1
|
||||
query_scalar = torch.arange(
|
||||
query_start,
|
||||
query_start + query_embeddings.shape[0],
|
||||
dtype=torch.long,
|
||||
)
|
||||
query_positions = query_scalar.view(1, -1).expand(3, -1)
|
||||
embedding_parts.append(query_embeddings)
|
||||
position_parts.append(query_positions)
|
||||
|
||||
embeddings = torch.cat(embedding_parts, dim=0).contiguous()
|
||||
position_ids = torch.cat(position_parts, dim=1).contiguous()
|
||||
if position_ids.shape != (3, embeddings.shape[0]):
|
||||
raise RuntimeError(
|
||||
"Reconstructed position/embedding mismatch: "
|
||||
f"{tuple(position_ids.shape)} vs {tuple(embeddings.shape)}"
|
||||
)
|
||||
return embeddings, position_ids
|
||||
|
||||
|
||||
def collate_cached_sequences(
|
||||
batch: Sequence[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Right-pad variable token sequences while preserving each readout index."""
|
||||
|
||||
if not batch:
|
||||
raise ValueError("Cannot collate an empty batch")
|
||||
hidden_size = int(batch[0]["inputs_embeds"].shape[-1])
|
||||
max_length = max(int(item["inputs_embeds"].shape[0]) for item in batch)
|
||||
dtype = batch[0]["inputs_embeds"].dtype
|
||||
inputs_embeds = torch.zeros(
|
||||
len(batch), max_length, hidden_size, dtype=dtype
|
||||
)
|
||||
position_ids = torch.zeros(3, len(batch), max_length, dtype=torch.long)
|
||||
attention_mask = torch.zeros(len(batch), max_length, dtype=torch.long)
|
||||
readout_indices = torch.empty(len(batch), dtype=torch.long)
|
||||
targets = torch.stack([item["targets"] for item in batch])
|
||||
|
||||
for batch_index, item in enumerate(batch):
|
||||
embeddings = item["inputs_embeds"]
|
||||
positions = item["position_ids"]
|
||||
length = int(embeddings.shape[0])
|
||||
if embeddings.shape[-1] != hidden_size:
|
||||
raise ValueError("Mixed hidden sizes in one batch")
|
||||
if positions.shape != (3, length):
|
||||
raise ValueError(
|
||||
f"Invalid position_ids {tuple(positions.shape)} for length {length}"
|
||||
)
|
||||
inputs_embeds[batch_index, :length] = embeddings
|
||||
position_ids[:, batch_index, :length] = positions
|
||||
attention_mask[batch_index, :length] = 1
|
||||
readout_indices[batch_index] = length - 1
|
||||
|
||||
return {
|
||||
"inputs_embeds": inputs_embeds,
|
||||
"position_ids": position_ids,
|
||||
"attention_mask": attention_mask,
|
||||
"readout_indices": readout_indices,
|
||||
"targets": targets,
|
||||
"step_indices": torch.tensor(
|
||||
[int(item["step_index"]) for item in batch], dtype=torch.long
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def anchor_intervals(
|
||||
anchor: dict[str, Any],
|
||||
*,
|
||||
block_seconds: float,
|
||||
step_seconds: float,
|
||||
) -> list[tuple[float, float]]:
|
||||
"""Return source intervals used by the reference KV replay for an anchor."""
|
||||
|
||||
intervals = [
|
||||
(index * block_seconds, (index + 1) * block_seconds)
|
||||
for index in anchor["history_blocks"]
|
||||
]
|
||||
block_index = int(anchor["block_index"])
|
||||
current_start = block_index * block_seconds
|
||||
current_end = current_start + int(anchor["tail_steps"]) * step_seconds
|
||||
intervals.append((current_start, current_end))
|
||||
return intervals
|
||||
|
||||
|
||||
def forward_last_state(
|
||||
language_model: torch.nn.Module,
|
||||
*,
|
||||
inputs_embeds: torch.Tensor,
|
||||
position_ids: torch.LongTensor,
|
||||
attention_mask: torch.LongTensor,
|
||||
readout_indices: torch.LongTensor,
|
||||
) -> torch.Tensor:
|
||||
"""Run full-sequence causal Thinker attention and gather each query state."""
|
||||
|
||||
outputs = language_model(
|
||||
inputs_embeds=inputs_embeds,
|
||||
position_ids=position_ids,
|
||||
attention_mask=attention_mask,
|
||||
use_cache=False,
|
||||
output_hidden_states=False,
|
||||
return_dict=True,
|
||||
cache_position=torch.arange(
|
||||
inputs_embeds.shape[1], device=inputs_embeds.device
|
||||
),
|
||||
)
|
||||
batch_indices = torch.arange(
|
||||
inputs_embeds.shape[0], device=inputs_embeds.device
|
||||
)
|
||||
return outputs.last_hidden_state[batch_indices, readout_indices]
|
||||
|
||||
|
||||
def cache_size_bytes(paths: Iterable[Path]) -> int:
|
||||
return sum(path.stat().st_size for path in paths if path.is_file())
|
||||
@@ -8,5 +8,8 @@ Pillow
|
||||
tqdm
|
||||
scikit-learn
|
||||
|
||||
# Used by train_omni_turn_lora.py; already included in the ms-swift image.
|
||||
peft
|
||||
|
||||
# Optional on compatible NVIDIA GPUs: flash-attn
|
||||
# Optional NVDEC backend: PyNvVideoCodec
|
||||
|
||||
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