1627 lines
59 KiB
Python
1627 lines
59 KiB
Python
#!/usr/bin/env python3
|
|
"""Train a frozen or LoRA Qwen2.5-Omni Thinker + 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 inspect
|
|
import json
|
|
import math
|
|
import os
|
|
import queue
|
|
import random
|
|
import re
|
|
import threading
|
|
import time
|
|
from contextlib import ExitStack
|
|
from dataclasses import dataclass
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch import nn
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
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",
|
|
)
|
|
|
|
_TRANSFORMER_LAYER_PATTERN = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DistributedRun:
|
|
enabled: bool
|
|
rank: int
|
|
local_rank: int
|
|
world_size: int
|
|
device: torch.device
|
|
|
|
@property
|
|
def is_main(self) -> bool:
|
|
return self.rank == 0
|
|
|
|
|
|
def _initialize_distributed(args: argparse.Namespace) -> DistributedRun:
|
|
"""Initialize torchrun/NCCL when WORLD_SIZE is greater than one."""
|
|
|
|
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
|
if world_size <= 1:
|
|
device = torch.device(args.device)
|
|
if device.type != "cuda":
|
|
raise ValueError("This LoRA trainer requires a CUDA device")
|
|
torch.cuda.set_device(device)
|
|
return DistributedRun(
|
|
enabled=False,
|
|
rank=0,
|
|
local_rank=device.index or 0,
|
|
world_size=1,
|
|
device=device,
|
|
)
|
|
|
|
if not dist.is_available():
|
|
raise RuntimeError("torch.distributed is unavailable in this PyTorch build")
|
|
if "LOCAL_RANK" not in os.environ or "RANK" not in os.environ:
|
|
raise RuntimeError("Multi-GPU training must be launched with torchrun")
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
torch.cuda.set_device(local_rank)
|
|
dist.init_process_group(
|
|
backend="nccl",
|
|
timeout=timedelta(minutes=float(args.ddp_timeout_minutes)),
|
|
)
|
|
return DistributedRun(
|
|
enabled=True,
|
|
rank=dist.get_rank(),
|
|
local_rank=local_rank,
|
|
world_size=dist.get_world_size(),
|
|
device=torch.device("cuda", local_rank),
|
|
)
|
|
|
|
|
|
def _distributed_barrier(run: DistributedRun) -> None:
|
|
if run.enabled:
|
|
dist.barrier()
|
|
|
|
|
|
def _distributed_average(value: float, run: DistributedRun) -> float:
|
|
if not run.enabled:
|
|
return float(value)
|
|
tensor = torch.tensor(value, dtype=torch.float64, device=run.device)
|
|
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
|
|
return float((tensor / run.world_size).item())
|
|
|
|
|
|
def _synchronize_trainable_parameters(
|
|
modules: Sequence[torch.nn.Module], run: DistributedRun
|
|
) -> None:
|
|
"""Broadcast only LoRA/head tensors, avoiding a second 3B-model broadcast."""
|
|
|
|
if not run.enabled:
|
|
return
|
|
for module in modules:
|
|
for parameter in module.parameters():
|
|
if parameter.requires_grad:
|
|
dist.broadcast(parameter.data, src=0)
|
|
|
|
|
|
def _wrap_ddp(module: torch.nn.Module, run: DistributedRun) -> torch.nn.Module:
|
|
if not run.enabled:
|
|
return module
|
|
kwargs: dict[str, Any] = {
|
|
"device_ids": [run.local_rank],
|
|
"output_device": run.local_rank,
|
|
"broadcast_buffers": False,
|
|
"find_unused_parameters": False,
|
|
}
|
|
# PyTorch >= 2.7 can skip the expensive initial broadcast of all frozen
|
|
# base-model parameters. The checkpoint is identical on every rank and
|
|
# trainable LoRA/head tensors were synchronized explicitly above.
|
|
if "init_sync" in inspect.signature(DDP.__init__).parameters:
|
|
kwargs["init_sync"] = False
|
|
return DDP(module, **kwargs)
|
|
|
|
|
|
def _load_base_model(
|
|
args: argparse.Namespace, run: DistributedRun
|
|
) -> tuple[torch.nn.Module, Any]:
|
|
"""Load one rank at a time so four shard readers do not hit storage together."""
|
|
|
|
if not run.enabled:
|
|
return core.load_qwen_omni(
|
|
args.model_id,
|
|
dtype=args.dtype,
|
|
device_map=str(run.device),
|
|
use_flash_attention_2=args.flash_attention_2,
|
|
)
|
|
|
|
loaded: tuple[torch.nn.Module, Any] | None = None
|
|
for loading_rank in range(run.world_size):
|
|
if run.rank == loading_rank:
|
|
print(
|
|
f"[rank {run.rank}/{run.world_size}] loading model on {run.device}",
|
|
flush=True,
|
|
)
|
|
loaded = core.load_qwen_omni(
|
|
args.model_id,
|
|
dtype=args.dtype,
|
|
device_map=str(run.device),
|
|
use_flash_attention_2=args.flash_attention_2,
|
|
)
|
|
dist.barrier()
|
|
if loaded is None:
|
|
raise RuntimeError(f"Rank {run.rank} did not load a model")
|
|
return loaded
|
|
|
|
|
|
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(
|
|
"--adapter-checkpoint",
|
|
type=Path,
|
|
help=(
|
|
"Existing PEFT adapter directory to continue from; optimizer and "
|
|
"scheduler are reinitialized"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--freeze-thinker",
|
|
action="store_true",
|
|
help=(
|
|
"Do not create or update LoRA parameters; run the same cached "
|
|
"Thinker sequences under no_grad and train only the turn head"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--freeze-head",
|
|
action="store_true",
|
|
help=(
|
|
"Keep the initialized turn head fixed and update only Thinker "
|
|
"LoRA parameters; requires --head-checkpoint"
|
|
),
|
|
)
|
|
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(
|
|
"--audio-feature-device",
|
|
default=None,
|
|
help=(
|
|
"Device used by the native Whisper log-Mel extractor for online "
|
|
"current tails, for example model, cuda:0, or cpu; omitted "
|
|
"preserves processor defaults"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--audio-right-padding-samples",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For online current tails, replace model-maximum audio padding "
|
|
"with actual samples plus this many zeros (hop-aligned); omitted "
|
|
"preserves the original maximum-length padding"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--decoder",
|
|
choices=("decord", "torchcodec", "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(
|
|
"--lora-last-n-layers",
|
|
type=int,
|
|
default=0,
|
|
help=(
|
|
"Apply a newly created LoRA adapter only to the final N Thinker "
|
|
"language layers; 0 keeps all matched layers"
|
|
),
|
|
)
|
|
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(
|
|
"--save-every-steps",
|
|
type=int,
|
|
default=0,
|
|
help=(
|
|
"Save an evaluation-ready adapter/head checkpoint every N optimizer "
|
|
"updates under checkpoint-step-N; 0 disables intermediate saves"
|
|
),
|
|
)
|
|
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",
|
|
help="Single-GPU device; torchrun overrides this with LOCAL_RANK",
|
|
)
|
|
parser.add_argument(
|
|
"--ddp-timeout-minutes",
|
|
type=float,
|
|
default=60.0,
|
|
help="NCCL process-group timeout used by torchrun multi-GPU training",
|
|
)
|
|
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,
|
|
audio_feature_device: str | torch.device | None,
|
|
audio_right_padding_samples: int | None,
|
|
) -> 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"])
|
|
if audio_feature_device == "model":
|
|
try:
|
|
audio_feature_device = next(model.parameters()).device
|
|
except StopIteration as exc:
|
|
raise RuntimeError("Online-tail AV model has no parameters") from exc
|
|
self.audio_feature_device = audio_feature_device
|
|
self.audio_right_padding_samples = audio_right_padding_samples
|
|
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,
|
|
audio_feature_device=self.audio_feature_device,
|
|
audio_right_padding_samples=self.audio_right_padding_samples,
|
|
)
|
|
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,
|
|
verbose: bool = True,
|
|
) -> 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))
|
|
if verbose:
|
|
print(f"Initialized classification head from {checkpoint}")
|
|
return head
|
|
|
|
|
|
def _prepare_language_lora(
|
|
model: torch.nn.Module,
|
|
*,
|
|
adapter_checkpoint: Path | None,
|
|
rank: int,
|
|
alpha: int,
|
|
dropout: float,
|
|
target_modules: Sequence[str],
|
|
last_n_layers: int,
|
|
gradient_checkpointing: bool,
|
|
verbose: bool = True,
|
|
) -> torch.nn.Module:
|
|
try:
|
|
from peft import LoraConfig, PeftModel, 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
|
|
if adapter_checkpoint is not None:
|
|
adapter_checkpoint = adapter_checkpoint.expanduser().resolve()
|
|
config_path = adapter_checkpoint / "adapter_config.json"
|
|
if not config_path.is_file():
|
|
raise FileNotFoundError(
|
|
f"Missing PEFT adapter config: {config_path}"
|
|
)
|
|
language_model = PeftModel.from_pretrained(
|
|
language_model,
|
|
str(adapter_checkpoint),
|
|
is_trainable=True,
|
|
)
|
|
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()
|
|
if verbose:
|
|
language_model.print_trainable_parameters()
|
|
trainable_lora = sum(
|
|
1
|
|
for name, parameter in language_model.named_parameters()
|
|
if parameter.requires_grad and "lora_" in name
|
|
)
|
|
print(
|
|
f"Continued trainable LoRA from {adapter_checkpoint}; "
|
|
f"trainable LoRA tensors={trainable_lora}"
|
|
)
|
|
return language_model
|
|
|
|
matched_all = [
|
|
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_all:
|
|
raise RuntimeError(
|
|
"No language-model linear layers matched LoRA targets "
|
|
f"{list(target_modules)}"
|
|
)
|
|
layer_indices = sorted(
|
|
{
|
|
int(match.group(1))
|
|
for name in matched_all
|
|
if (match := _TRANSFORMER_LAYER_PATTERN.search(name)) is not None
|
|
}
|
|
)
|
|
layers_to_transform: list[int] | None = None
|
|
matched = matched_all
|
|
if last_n_layers:
|
|
if not layer_indices:
|
|
raise RuntimeError(
|
|
"Could not identify Thinker layer indices in matched LoRA modules"
|
|
)
|
|
if last_n_layers > len(layer_indices):
|
|
raise ValueError(
|
|
f"--lora-last-n-layers={last_n_layers} exceeds the "
|
|
f"{len(layer_indices)} matched Thinker layers"
|
|
)
|
|
layers_to_transform = layer_indices[-last_n_layers:]
|
|
selected = set(layers_to_transform)
|
|
matched = [
|
|
name
|
|
for name in matched_all
|
|
if (
|
|
(match := _TRANSFORMER_LAYER_PATTERN.search(name)) is not None
|
|
and int(match.group(1)) in selected
|
|
)
|
|
]
|
|
if not matched:
|
|
raise RuntimeError(
|
|
"Final-layer selection removed every matched LoRA module"
|
|
)
|
|
config_kwargs: dict[str, Any] = {}
|
|
if layers_to_transform is not None:
|
|
config_kwargs.update(
|
|
layers_to_transform=layers_to_transform,
|
|
layers_pattern="layers",
|
|
)
|
|
config = LoraConfig(
|
|
r=rank,
|
|
lora_alpha=alpha,
|
|
lora_dropout=dropout,
|
|
target_modules=list(target_modules),
|
|
bias="none",
|
|
**config_kwargs,
|
|
)
|
|
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()
|
|
if verbose:
|
|
language_model.print_trainable_parameters()
|
|
print(
|
|
f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
|
|
f"last={matched[-1]}"
|
|
)
|
|
if layers_to_transform is not None:
|
|
print(
|
|
"LoRA layer restriction: "
|
|
f"{layers_to_transform[0]}--{layers_to_transform[-1]} "
|
|
f"({len(layers_to_transform)} of {len(layer_indices)} layers)"
|
|
)
|
|
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,
|
|
distributed_run: DistributedRun,
|
|
materialize_batch: Any | None = None,
|
|
) -> float:
|
|
language_was_training = language_model.training
|
|
head_was_training = head.training
|
|
language_model.eval()
|
|
head.eval()
|
|
loss_sum = 0.0
|
|
example_count = 0
|
|
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())
|
|
batch_size = int(targets.shape[0])
|
|
loss_sum += float(egospeak_multihorizon_loss(logits, targets).item()) * batch_size
|
|
example_count += batch_size
|
|
language_model.train(language_was_training)
|
|
head.train(head_was_training)
|
|
totals = torch.tensor(
|
|
[loss_sum, float(example_count)], dtype=torch.float64, device=device
|
|
)
|
|
if distributed_run.enabled:
|
|
dist.all_reduce(totals, op=dist.ReduceOp.SUM)
|
|
if float(totals[1].item()) <= 0:
|
|
raise RuntimeError("Evaluation loader produced no examples")
|
|
return float((totals[0] / totals[1]).item())
|
|
|
|
|
|
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 _save_weight_checkpoint(
|
|
checkpoint_dir: Path,
|
|
*,
|
|
raw_language_model: torch.nn.Module,
|
|
raw_head: TurnPredictionHead,
|
|
args: argparse.Namespace,
|
|
hidden_size: int,
|
|
training_step: int,
|
|
) -> None:
|
|
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
|
if not args.freeze_thinker:
|
|
raw_language_model.save_pretrained(checkpoint_dir / "adapter")
|
|
torch.save(
|
|
{
|
|
"head": raw_head.state_dict(),
|
|
"head_config": {
|
|
"hidden_size": hidden_size,
|
|
"num_horizons": NUM_HORIZONS,
|
|
"num_classes": NUM_CLASSES,
|
|
"dropout": raw_head.dropout,
|
|
},
|
|
"initial_head_checkpoint": (
|
|
str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
|
|
),
|
|
"head_frozen": bool(args.freeze_head),
|
|
"training_step": training_step,
|
|
},
|
|
checkpoint_dir / "turn_head.pt",
|
|
)
|
|
|
|
|
|
def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
|
|
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.save_every_steps < 0:
|
|
raise ValueError("--save-every-steps must be non-negative")
|
|
if args.online_current_tail:
|
|
if run.enabled:
|
|
raise ValueError(
|
|
"DDP currently supports fully cached AV training only; remove "
|
|
"--online-current-tail when launching with torchrun"
|
|
)
|
|
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.ddp_timeout_minutes <= 0:
|
|
raise ValueError("--ddp-timeout-minutes must be positive")
|
|
if args.adapter_checkpoint is not None and args.head_checkpoint is None:
|
|
if run.is_main:
|
|
print(
|
|
"Warning: --adapter-checkpoint was provided without "
|
|
"--head-checkpoint; the turn head will be reinitialized."
|
|
)
|
|
if args.freeze_thinker and args.adapter_checkpoint is not None:
|
|
raise ValueError(
|
|
"--freeze-thinker cannot be combined with --adapter-checkpoint"
|
|
)
|
|
if args.freeze_thinker and args.freeze_head:
|
|
raise ValueError(
|
|
"--freeze-thinker and --freeze-head would leave no trainable parameters"
|
|
)
|
|
if args.freeze_head and args.head_checkpoint is None:
|
|
raise ValueError("--freeze-head requires --head-checkpoint")
|
|
if args.lora_last_n_layers < 0:
|
|
raise ValueError("--lora-last-n-layers must be non-negative")
|
|
if args.adapter_checkpoint is not None and args.lora_last_n_layers:
|
|
raise ValueError(
|
|
"Layer selection is stored in an existing adapter config; omit "
|
|
"--lora-last-n-layers when using --adapter-checkpoint"
|
|
)
|
|
if args.freeze_thinker and args.lora_last_n_layers:
|
|
raise ValueError(
|
|
"--lora-last-n-layers has no effect with --freeze-thinker"
|
|
)
|
|
if args.freeze_thinker and args.gradient_checkpointing:
|
|
raise ValueError(
|
|
"--gradient-checkpointing is unnecessary with --freeze-thinker"
|
|
)
|
|
if (
|
|
args.audio_right_padding_samples is not None
|
|
and args.audio_right_padding_samples < 0
|
|
):
|
|
raise ValueError("--audio-right-padding-samples must be non-negative")
|
|
if not args.online_current_tail and (
|
|
args.audio_feature_device is not None
|
|
or args.audio_right_padding_samples is not None
|
|
):
|
|
raise ValueError(
|
|
"Audio preprocessing overrides apply only with --online-current-tail"
|
|
)
|
|
if args.num_workers != 0 and run.is_main:
|
|
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" / "adapter_config.json",
|
|
output_dir / "turn_head.pt",
|
|
output_dir / "train_summary.json",
|
|
)
|
|
intermediate_outputs = list(output_dir.glob("checkpoint-step-*"))
|
|
if not args.overwrite_output and (
|
|
any(path.exists() for path in known_outputs) or intermediate_outputs
|
|
):
|
|
raise FileExistsError(
|
|
f"Training output already exists in {output_dir}; pass --overwrite-output"
|
|
)
|
|
if run.is_main:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
_distributed_barrier(run)
|
|
_seed_everything(args.seed)
|
|
core._configure_qwen_logging()
|
|
device = run.device
|
|
|
|
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 + run.rank)
|
|
train_sampler: DistributedSampler | None = None
|
|
eval_sampler: DistributedSampler | None = None
|
|
if run.enabled:
|
|
train_sampler = DistributedSampler(
|
|
dataset,
|
|
num_replicas=run.world_size,
|
|
rank=run.rank,
|
|
shuffle=True,
|
|
seed=args.seed,
|
|
drop_last=False,
|
|
)
|
|
eval_sampler = DistributedSampler(
|
|
dataset,
|
|
num_replicas=run.world_size,
|
|
rank=run.rank,
|
|
shuffle=False,
|
|
drop_last=False,
|
|
)
|
|
train_loader = DataLoader(
|
|
dataset,
|
|
batch_size=args.batch_size,
|
|
shuffle=train_sampler is None,
|
|
sampler=train_sampler,
|
|
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,
|
|
sampler=eval_sampler,
|
|
num_workers=0,
|
|
pin_memory=True,
|
|
collate_fn=collate_fn,
|
|
)
|
|
if run.is_main:
|
|
print(
|
|
f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, "
|
|
f"world={run.world_size}, per_gpu_batch={args.batch_size}, "
|
|
f"accum={args.gradient_accumulation_steps}, "
|
|
"global_batch="
|
|
f"{run.world_size * args.batch_size * args.gradient_accumulation_steps}"
|
|
)
|
|
|
|
base_model, processor = _load_base_model(args, run)
|
|
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")
|
|
if args.freeze_thinker:
|
|
base_model.requires_grad_(False)
|
|
language_model = base_model.model
|
|
language_model.config.use_cache = False
|
|
language_model.eval()
|
|
if run.is_main:
|
|
print(
|
|
"Frozen-Thinker control: no LoRA adapter; only the turn head "
|
|
"is trainable"
|
|
)
|
|
else:
|
|
language_model = _prepare_language_lora(
|
|
base_model,
|
|
adapter_checkpoint=args.adapter_checkpoint,
|
|
rank=args.lora_rank,
|
|
alpha=args.lora_alpha,
|
|
dropout=args.lora_dropout,
|
|
target_modules=target_modules,
|
|
last_n_layers=args.lora_last_n_layers,
|
|
gradient_checkpointing=args.gradient_checkpointing,
|
|
verbose=run.is_main,
|
|
)
|
|
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,
|
|
audio_feature_device=args.audio_feature_device,
|
|
audio_right_padding_samples=args.audio_right_padding_samples,
|
|
)
|
|
print(
|
|
"Cache policy benchmark: complete 2s history cached; "
|
|
"current 0.2--1.8s tail computed online"
|
|
)
|
|
print(
|
|
"Online-tail audio preprocessing: "
|
|
f"feature_device={args.audio_feature_device or 'processor default'}, "
|
|
"padding="
|
|
+ (
|
|
"model maximum"
|
|
if args.audio_right_padding_samples is None
|
|
else f"actual + >= {args.audio_right_padding_samples} samples"
|
|
)
|
|
)
|
|
if args.producer_device is not None:
|
|
print(
|
|
f"Producer-consumer enabled: frozen AV={args.producer_device}, "
|
|
f"Thinker LoRA={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,
|
|
verbose=run.is_main,
|
|
).to(device=device, dtype=torch.float32)
|
|
if args.freeze_head:
|
|
head.requires_grad_(False)
|
|
head.eval()
|
|
if run.is_main:
|
|
print("Frozen-head control: only Thinker LoRA parameters are trainable")
|
|
|
|
raw_language_model = language_model
|
|
raw_head = head
|
|
synchronization_modules = tuple(
|
|
module
|
|
for module, trainable in (
|
|
(raw_language_model, not args.freeze_thinker),
|
|
(raw_head, not args.freeze_head),
|
|
)
|
|
if trainable
|
|
)
|
|
_synchronize_trainable_parameters(synchronization_modules, run)
|
|
if not args.freeze_thinker:
|
|
language_model = _wrap_ddp(raw_language_model, run)
|
|
if not args.freeze_head:
|
|
head = _wrap_ddp(raw_head, run)
|
|
# Model initialization is synchronized above; rank-specific RNG streams
|
|
# avoid identical dropout masks for different distributed samples.
|
|
_seed_everything(args.seed + run.rank)
|
|
|
|
lora_parameters = [
|
|
parameter
|
|
for parameter in raw_language_model.parameters()
|
|
if parameter.requires_grad
|
|
]
|
|
if not args.freeze_thinker and not lora_parameters:
|
|
raise RuntimeError("LoRA preparation produced no trainable parameters")
|
|
if args.freeze_thinker:
|
|
optimizer = torch.optim.AdamW(
|
|
[{"params": raw_head.parameters(), "lr": args.head_learning_rate}],
|
|
weight_decay=args.weight_decay,
|
|
)
|
|
head_group_index = 0
|
|
elif args.freeze_head:
|
|
optimizer = torch.optim.AdamW(
|
|
[{"params": lora_parameters, "lr": args.learning_rate}],
|
|
weight_decay=args.weight_decay,
|
|
)
|
|
head_group_index = None
|
|
else:
|
|
optimizer = torch.optim.AdamW(
|
|
[
|
|
{"params": lora_parameters, "lr": args.learning_rate},
|
|
{"params": raw_head.parameters(), "lr": args.head_learning_rate},
|
|
],
|
|
weight_decay=args.weight_decay,
|
|
)
|
|
head_group_index = 1
|
|
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,
|
|
distributed_run=run,
|
|
materialize_batch=materialize_batch,
|
|
)
|
|
if run.is_main:
|
|
print(f"Initial cached loss: {initial_loss:.6f}")
|
|
if materialize_batch is not None:
|
|
materialize_batch.reset_timing()
|
|
if args.freeze_thinker:
|
|
language_model.eval()
|
|
else:
|
|
language_model.train()
|
|
if args.freeze_head:
|
|
head.eval()
|
|
else:
|
|
head.train()
|
|
optimizer.zero_grad(set_to_none=True)
|
|
_distributed_barrier(run)
|
|
torch.cuda.synchronize(device)
|
|
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()
|
|
sampler_epoch = 0
|
|
if train_sampler is not None:
|
|
train_sampler.set_epoch(sampler_epoch)
|
|
iterator = iter(train_loader) if batch_producer is None else None
|
|
history: list[dict[str, float | int]] = []
|
|
saved_checkpoint_steps: list[int] = []
|
|
progress = tqdm(
|
|
total=args.max_steps,
|
|
desc=("Frozen Thinker head" if args.freeze_thinker else "Thinker LoRA"),
|
|
unit="update",
|
|
disable=args.no_progress or not run.is_main,
|
|
)
|
|
for update_step in range(1, args.max_steps + 1):
|
|
accumulated_loss = 0.0
|
|
for accumulation_index 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:
|
|
sampler_epoch += 1
|
|
if train_sampler is not None:
|
|
train_sampler.set_epoch(sampler_epoch)
|
|
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)
|
|
synchronize_gradients = (
|
|
accumulation_index == args.gradient_accumulation_steps - 1
|
|
)
|
|
with ExitStack() as stack:
|
|
if run.enabled and not synchronize_gradients:
|
|
if not args.freeze_thinker:
|
|
assert isinstance(language_model, DDP)
|
|
stack.enter_context(language_model.no_sync())
|
|
if not args.freeze_head:
|
|
assert isinstance(head, DDP)
|
|
stack.enter_context(head.no_sync())
|
|
if args.freeze_thinker:
|
|
with torch.no_grad():
|
|
states = forward_last_state(
|
|
language_model,
|
|
inputs_embeds=inputs_embeds,
|
|
position_ids=position_ids,
|
|
attention_mask=attention_mask,
|
|
readout_indices=readout_indices,
|
|
)
|
|
else:
|
|
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 = (
|
|
list(raw_head.parameters())
|
|
if args.freeze_thinker
|
|
else (
|
|
lora_parameters
|
|
if args.freeze_head
|
|
else [*lora_parameters, *raw_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)
|
|
reported_loss = _distributed_average(accumulated_loss, run)
|
|
lora_lr = (
|
|
0.0 if args.freeze_thinker else float(optimizer.param_groups[0]["lr"])
|
|
)
|
|
head_lr = (
|
|
0.0
|
|
if head_group_index is None
|
|
else float(optimizer.param_groups[head_group_index]["lr"])
|
|
)
|
|
record = {
|
|
"step": update_step,
|
|
"loss": reported_loss,
|
|
"grad_norm": float(grad_norm),
|
|
"lora_lr": lora_lr,
|
|
"head_lr": head_lr,
|
|
}
|
|
if run.is_main:
|
|
history.append(record)
|
|
progress.update(1)
|
|
progress.set_postfix(
|
|
loss=f"{reported_loss:.4f}",
|
|
lr=f"{head_lr:.2g}" if args.freeze_thinker else f"{lora_lr:.2g}",
|
|
)
|
|
if run.is_main and (update_step == 1 or update_step % args.log_every == 0):
|
|
print(
|
|
f"step={update_step} loss={reported_loss:.6f} "
|
|
f"grad_norm={float(grad_norm):.4f} "
|
|
+ (
|
|
f"head_lr={head_lr:.3g}"
|
|
if args.freeze_thinker
|
|
else f"lora_lr={lora_lr:.3g} head_lr={head_lr:.3g}"
|
|
),
|
|
flush=True,
|
|
)
|
|
if args.save_every_steps and update_step % args.save_every_steps == 0:
|
|
torch.cuda.synchronize(device)
|
|
_distributed_barrier(run)
|
|
if run.is_main:
|
|
checkpoint_dir = output_dir / f"checkpoint-step-{update_step}"
|
|
_save_weight_checkpoint(
|
|
checkpoint_dir,
|
|
raw_language_model=raw_language_model,
|
|
raw_head=raw_head,
|
|
args=args,
|
|
hidden_size=dataset.hidden_size,
|
|
training_step=update_step,
|
|
)
|
|
saved_checkpoint_steps.append(update_step)
|
|
print(f"Saved intermediate checkpoint to {checkpoint_dir}", flush=True)
|
|
_distributed_barrier(run)
|
|
progress.close()
|
|
torch.cuda.synchronize(device)
|
|
_distributed_barrier(run)
|
|
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,
|
|
distributed_run=run,
|
|
materialize_batch=materialize_batch,
|
|
)
|
|
local_peak_gib = torch.cuda.max_memory_allocated(device) / (1024 ** 3)
|
|
peak_tensor = torch.tensor(local_peak_gib, dtype=torch.float64, device=device)
|
|
if run.enabled:
|
|
gathered_peaks = [torch.zeros_like(peak_tensor) for _ in range(run.world_size)]
|
|
dist.all_gather(gathered_peaks, peak_tensor)
|
|
per_rank_peak_gib = [float(value.item()) for value in gathered_peaks]
|
|
else:
|
|
per_rank_peak_gib = [local_peak_gib]
|
|
|
|
if run.is_main:
|
|
_save_weight_checkpoint(
|
|
output_dir,
|
|
raw_language_model=raw_language_model,
|
|
raw_head=raw_head,
|
|
args=args,
|
|
hidden_size=dataset.hidden_size,
|
|
training_step=args.max_steps,
|
|
)
|
|
global_batch_size = (
|
|
run.world_size * args.batch_size * args.gradient_accumulation_steps
|
|
)
|
|
summary = {
|
|
"schema_version": 1,
|
|
"status": "pilot" if args.max_examples else "training",
|
|
"model_id": args.model_id,
|
|
"model_variant": (
|
|
"frozen_base_thinker_head_only"
|
|
if args.freeze_thinker
|
|
else (
|
|
"thinker_lora_frozen_head"
|
|
if args.freeze_head
|
|
else "thinker_lora"
|
|
)
|
|
),
|
|
"thinker_frozen": bool(args.freeze_thinker),
|
|
"head_frozen": bool(args.freeze_head),
|
|
"lora_last_n_layers": int(args.lora_last_n_layers),
|
|
"initial_adapter_checkpoint": (
|
|
str(args.adapter_checkpoint.expanduser().resolve())
|
|
if args.adapter_checkpoint is not None
|
|
else None
|
|
),
|
|
"optimizer_state_resumed": False,
|
|
"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,
|
|
"estimated_examples_per_second": (
|
|
args.max_steps * global_batch_size / training_seconds
|
|
),
|
|
"distributed": {
|
|
"enabled": run.enabled,
|
|
"world_size": run.world_size,
|
|
"per_gpu_batch_size": args.batch_size,
|
|
"gradient_accumulation_steps": args.gradient_accumulation_steps,
|
|
"global_batch_size": global_batch_size,
|
|
"per_rank_peak_cuda_allocated_gib": per_rank_peak_gib,
|
|
},
|
|
"online_tail_timing": online_timing,
|
|
"loss_reduction_ratio": (
|
|
(initial_loss - final_loss) / initial_loss if initial_loss else None
|
|
),
|
|
"peak_cuda_allocated_gib": max(per_rank_peak_gib),
|
|
"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,
|
|
"intermediate_checkpoint_steps": saved_checkpoint_steps,
|
|
}
|
|
if run.is_main:
|
|
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; "
|
|
f"estimated {summary['estimated_examples_per_second']:.2f} examples/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']}"
|
|
)
|
|
if args.freeze_thinker:
|
|
print("Frozen Thinker: no LoRA adapter was created")
|
|
else:
|
|
print(f"Saved LoRA adapter to {adapter_dir}")
|
|
print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
|
|
print(
|
|
"Peak CUDA allocated across ranks: "
|
|
f"{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"
|
|
)
|
|
_distributed_barrier(run)
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
run = _initialize_distributed(args)
|
|
try:
|
|
_run_training(args, run)
|
|
finally:
|
|
if run.enabled and dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|