scripts/omni/train_qwen3_omni_lora.py

#!/usr/bin/env python3
"""Conservative LoRA SFT for Qwen3-Omni action/subtask label generation."""

from __future__ import annotations

import argparse
import json
import math
import random
import time
from pathlib import Path
from types import MethodType

import torch
import torch.nn.functional as F

from qwen3_omni_dataset_utils import (
    build_messages,
    DEFAULT_MODEL_ID,
    has_empty_audio_items,
    is_empty_audio_exception,
    load_jsonl,
    sample_has_audio,
    sample_without_audio,
)


def parse_args() -> argparse.Namespace:
    workspace_default = Path(__file__).resolve().parents[2]
    parser = argparse.ArgumentParser(description="Train Qwen3-Omni LoRA on exported Ropedia windows.")
    parser.add_argument("--dataset-jsonl", type=Path, required=True)
    parser.add_argument("--run-id", default="qwen_lora_text_video_audio")
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--results-dir", type=Path)
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID)
    parser.add_argument(
        "--backbone-config",
        type=Path,
        default=workspace_default / "configs" / "omni_backbones" / "qwen3_omni_lora.json",
        help="Backbone contract JSON recorded with the run for model-extension tracking.",
    )
    parser.add_argument("--train-split", default="train")
    parser.add_argument("--val-split", default="val")
    parser.add_argument("--include-unspecified-in-train", action="store_true")
    parser.add_argument("--max-train-samples", type=int, default=0)
    parser.add_argument("--max-val-samples", type=int, default=64)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=8)
    parser.add_argument("--learning-rate", type=float, default=1e-4)
    parser.add_argument("--weight-decay", type=float, default=0.0)
    parser.add_argument("--max-grad-norm", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--dtype", default="bfloat16", choices=["auto", "bfloat16", "float16", "float32"])
    parser.add_argument("--local-files-only", action="store_true")
    parser.add_argument("--trust-remote-code", action="store_true")
    parser.add_argument("--use-audio-in-video", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument("--gradient-checkpointing", action="store_true")
    parser.add_argument("--progress-every", type=int, default=1)
    parser.add_argument(
        "--loss-logit-tail-only",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="For SFT, project only the supervised assistant-answer tail through lm_head before CE loss.",
    )
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--lora-dropout", type=float, default=0.05)
    parser.add_argument(
        "--lora-target-modules",
        default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
        help="Comma-separated module names passed to PEFT LoRAConfig.",
    )
    return parser.parse_args()


def dtype_arg(value: str):
    if value == "auto":
        return "auto"
    return {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[value]


def select_samples(samples: list[dict], split: str, include_unspecified: bool) -> list[dict]:
    rows = [sample for sample in samples if sample.get("split") == split]
    if include_unspecified:
        rows.extend(sample for sample in samples if sample.get("split") == "unspecified")
    return rows


def patch_rotary_position_device(model) -> bool:
    """Keep Qwen3-Omni rotary position ids aligned under model-parallel device maps."""
    inner_model = getattr(model, "model", None)
    rotary = getattr(inner_model, "rotary_emb", None)
    if rotary is None or getattr(rotary, "_ropedia_position_device_patch", False):
        return False

    original_forward = rotary.forward

    def forward_with_aligned_position_ids(self, x, position_ids, *args, **kwargs):
        if hasattr(self, "inv_freq") and hasattr(x, "device") and self.inv_freq.device != x.device:
            self._buffers["inv_freq"] = self.inv_freq.to(x.device)
        if hasattr(position_ids, "to") and hasattr(x, "device") and position_ids.device != x.device:
            position_ids = position_ids.to(x.device)
        return original_forward(x, position_ids, *args, **kwargs)

    rotary.forward = MethodType(forward_with_aligned_position_ids, rotary)
    rotary._ropedia_position_device_patch = True
    return True


def patch_qwen3_omni_rotary_classes() -> None:
    """Patch Qwen3-Omni MRoPE classes before Accelerate installs device hooks."""
    from transformers.models.qwen3_omni_moe import modeling_qwen3_omni_moe as qwen3_omni_moe

    def patch_mrope_class(class_name: str) -> None:
        rotary_cls = getattr(qwen3_omni_moe, class_name, None)
        if rotary_cls is None or getattr(rotary_cls, "_ropedia_class_device_patch", False):
            return

        @torch.no_grad()
        def forward(self, x, position_ids):
            if position_ids.ndim == 2:
                position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
            target_device = x.device
            inv_freq = self.inv_freq.to(target_device)
            position_ids = position_ids.to(target_device)
            inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
            position_ids_expanded = position_ids[:, :, None, :].float()

            device_type = target_device.type if isinstance(target_device.type, str) and target_device.type != "mps" else "cpu"
            with qwen3_omni_moe.maybe_autocast(device_type=device_type, enabled=False):
                freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
                freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
                emb = torch.cat((freqs, freqs), dim=-1)
                cos = emb.cos() * self.attention_scaling
                sin = emb.sin() * self.attention_scaling

            return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

        rotary_cls.forward = forward
        rotary_cls._ropedia_class_device_patch = True

    patch_mrope_class("Qwen3OmniMoeThinkerTextRotaryEmbedding")
    patch_mrope_class("Qwen3OmniMoeTalkerRotaryEmbedding")


def patch_qwen3_omni_norm_classes() -> None:
    """Patch Qwen3-Omni RMSNorm classes for model-parallel device maps."""
    from transformers.models.qwen3_omni_moe import modeling_qwen3_omni_moe as qwen3_omni_moe

    def patch_norm_class(class_name: str) -> None:
        norm_cls = getattr(qwen3_omni_moe, class_name, None)
        if norm_cls is None or getattr(norm_cls, "_ropedia_class_device_patch", False):
            return

        def forward(self, hidden_states):
            input_dtype = hidden_states.dtype
            norm_states = hidden_states.to(torch.float32)
            variance = norm_states.pow(2).mean(-1, keepdim=True)
            norm_states = norm_states * torch.rsqrt(variance + self.variance_epsilon)
            weight = self.weight.to(hidden_states.device)
            return weight * norm_states.to(input_dtype)

        norm_cls.forward = forward
        norm_cls._ropedia_class_device_patch = True

    patch_norm_class("Qwen3OmniMoeRMSNorm")
    patch_norm_class("Qwen3OmniMoeTextRMSNorm")
    patch_norm_class("Qwen3OmniMoeThinkerTextRMSNorm")
    patch_norm_class("Qwen3OmniMoeCode2WavRMSNorm")


def cast_floating_parameters(model, target_dtype) -> None:
    if isinstance(target_dtype, str):
        return
    for param in model.parameters():
        if param.is_floating_point() and param.dtype != target_dtype:
            param.data = param.data.to(target_dtype)


def build_trainable_cpu_state_dict(model) -> dict[str, torch.Tensor]:
    state_dict = {}
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        clean_name = name
        if clean_name.startswith("module."):
            clean_name = clean_name[len("module.") :]
        state_dict[clean_name] = param.detach().to("cpu", copy=True)
    return state_dict


def strip_module_prefix(name: str) -> str:
    while name.startswith("module."):
        name = name[len("module.") :]
    return name


def thinker_adapter_state_from_full_state(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    adapter_state: dict[str, torch.Tensor] = {}
    for name, tensor in state_dict.items():
        clean_name = strip_module_prefix(name)
        if clean_name.startswith("thinker."):
            clean_name = clean_name[len("thinker.") :]
        elif clean_name.startswith("base_model.model.thinker."):
            clean_name = clean_name[len("base_model.model.thinker.") :]
        if "lora_" not in clean_name:
            continue
        if tensor.numel() == 0:
            raise ValueError(f"Gathered LoRA state still has an empty tensor: {name}")
        adapter_state[clean_name] = tensor.detach().to("cpu", copy=True)
    if not adapter_state:
        raise ValueError("Gathered state dict did not contain any LoRA tensors.")
    return adapter_state


def adapter_shape_summary(state_dict: dict[str, torch.Tensor]) -> dict[str, object]:
    prefixes = {}
    for name, tensor in state_dict.items():
        prefix = name.split(".")[2] if name.startswith("base_model.model.") and len(name.split(".")) > 2 else name.split(".")[0]
        row = prefixes.setdefault(prefix, {"tensors": 0, "numel": 0})
        row["tensors"] += 1
        row["numel"] += int(tensor.numel())
    return {
        "adapter_tensors": len(state_dict),
        "adapter_bytes": sum(t.numel() * t.element_size() for t in state_dict.values()),
        "prefixes": prefixes,
    }


class TailSlicingLmHead(torch.nn.Module):
    """Wrap lm_head so SFT can avoid full-prompt vocab logits."""

    def __init__(self, base: torch.nn.Module) -> None:
        super().__init__()
        self.base = base
        self.tail_start: int | None = None
        self._ropedia_tail_slicing_lm_head = True

    @property
    def weight(self):
        return self.base.weight

    @property
    def bias(self):
        return getattr(self.base, "bias", None)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        if self.tail_start is not None:
            hidden_states = hidden_states[:, self.tail_start :, :]
        return self.base(hidden_states)


def install_tail_slicing_lm_head(model) -> bool:
    base = model.get_base_model() if hasattr(model, "get_base_model") else model
    lm_head = getattr(base, "lm_head", None)
    if lm_head is None:
        return False
    if getattr(lm_head, "_ropedia_tail_slicing_lm_head", False):
        return True
    setattr(base, "lm_head", TailSlicingLmHead(lm_head))
    return True


def set_tail_slicing_lm_head(model, tail_start: int | None) -> bool:
    updated = False
    for module in model.modules():
        if getattr(module, "_ropedia_tail_slicing_lm_head", False):
            module.tail_start = tail_start
            updated = True
    return updated


def first_supervised_label(labels: torch.Tensor) -> int | None:
    active = labels[0].ne(-100).nonzero(as_tuple=False)
    if active.numel() == 0:
        return None
    return int(active[0].item())


def load_backbone_profile(path: Path | None) -> dict:
    if path is None:
        return {
            "id": "qwen3_omni_lora",
            "display_name": "Qwen3-Omni LoRA",
            "dataset_contract": "xperience10m_episode_json_qa_v1",
            "training_objective": "structured_episode_understanding_json_qa",
            "primary_metrics": [],
        }
    path = path.expanduser()
    if not path.is_absolute():
        path = Path(__file__).resolve().parents[2] / path
    if not path.exists():
        raise FileNotFoundError(f"Backbone config not found: {path}")
    payload = json.loads(path.read_text(encoding="utf-8"))
    return {
        "id": payload.get("id"),
        "display_name": payload.get("display_name"),
        "status": payload.get("status"),
        "model_family": payload.get("model_family"),
        "dataset_contract": payload.get("dataset_contract"),
        "training_objective": payload.get("training_objective"),
        "split_policy": payload.get("split_policy", {}),
        "modalities": payload.get("modalities", {}),
        "primary_metrics": payload.get("primary_metrics", []),
        "config_path": str(path),
    }


def load_model_processor(args: argparse.Namespace):
    from qwen3_omni_compat import patch_qwen3_omni_config

    patch_qwen3_omni_config()
    from peft import LoraConfig, get_peft_model
    from transformers import Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoeProcessor

    patch_qwen3_omni_rotary_classes()
    patch_qwen3_omni_norm_classes()

    model_kwargs = {
        "dtype": dtype_arg(args.dtype),
        "local_files_only": args.local_files_only,
    }
    if args.device_map and args.device_map.lower() != "none":
        model_kwargs["device_map"] = args.device_map
    if args.trust_remote_code:
        model_kwargs["trust_remote_code"] = True
    omni_model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(args.model_id, **model_kwargs)
    if hasattr(omni_model, "disable_talker"):
        omni_model.disable_talker()
    model = omni_model.thinker
    if args.device_map and args.device_map.lower() != "none":
        patch_rotary_position_device(model)
    if args.gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()

    processor_kwargs = {"local_files_only": args.local_files_only}
    if args.trust_remote_code:
        processor_kwargs["trust_remote_code"] = True
    processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_id, **processor_kwargs)

    config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        bias="none",
        target_modules=[item.strip() for item in args.lora_target_modules.split(",") if item.strip()],
    )
    model = get_peft_model(model, config)
    if args.loss_logit_tail_only:
        install_tail_slicing_lm_head(model)
    cast_floating_parameters(model, dtype_arg(args.dtype))
    model.print_trainable_parameters()
    return model, processor


def move_inputs(inputs, device, dtype=None):
    for key, value in list(inputs.items()):
        if hasattr(value, "to"):
            if dtype is not None and getattr(value, "is_floating_point", lambda: False)():
                inputs[key] = value.to(device=device, dtype=dtype)
            else:
                inputs[key] = value.to(device)
    return inputs


def compute_answer_token_loss(model, inputs: dict, tail_only: bool = True) -> torch.Tensor:
    """Compute CE only on supervised answer tokens to avoid full-logit fp32 casts."""
    labels = inputs.pop("labels")
    tail_start = 0
    if tail_only:
        first_label = first_supervised_label(labels)
        if first_label is None:
            return model(**inputs).logits.sum() * 0.0
        tail_start = max(first_label - 1, 0)
        tail_only = set_tail_slicing_lm_head(model, tail_start)
    try:
        output = model(**inputs)
    finally:
        if tail_only:
            set_tail_slicing_lm_head(model, None)
    logits = output.logits
    labels = labels.to(logits.device)
    if tail_only:
        if logits.shape[1] == labels.shape[1]:
            logits = logits[:, tail_start:, :]
        shift_logits = logits[..., :-1, :]
        shift_labels = labels[..., tail_start + 1 : tail_start + 1 + shift_logits.shape[1]]
    else:
        shift_logits = logits[..., :-1, :]
        shift_labels = labels[..., 1:]
    active = shift_labels != -100
    if not active.any().item():
        return logits.sum() * 0.0
    active_logits = shift_logits[active]
    active_labels = shift_labels[active]
    return F.cross_entropy(active_logits.float(), active_labels, reduction="mean")


def prepare_sample(processor, sample: dict, use_audio_in_video: bool, device, dtype=None) -> dict:
    from qwen_omni_utils import process_mm_info

    active_sample = sample
    for attempt in range(2):
        full_messages = build_messages(active_sample, active_sample["label_options"], include_answer=True)
        prompt_messages = build_messages(active_sample, active_sample["label_options"], include_answer=False)
        full_text = processor.apply_chat_template(full_messages, tokenize=False)
        prompt_text = processor.apply_chat_template(prompt_messages, add_generation_prompt=True, tokenize=False)
        audios, images, videos = process_mm_info(full_messages, use_audio_in_video=use_audio_in_video)
        if attempt == 0 and sample_has_audio(active_sample) and has_empty_audio_items(audios):
            active_sample = sample_without_audio(active_sample)
            continue
        try:
            inputs = processor(
                text=full_text,
                audio=audios,
                images=images,
                videos=videos,
                return_tensors="pt",
                padding=True,
                use_audio_in_video=use_audio_in_video,
            )
            break
        except RuntimeError as exc:
            if attempt == 0 and sample_has_audio(active_sample) and is_empty_audio_exception(exc):
                active_sample = sample_without_audio(active_sample)
                continue
            raise
    else:
        raise RuntimeError("Unable to prepare multimodal sample after dropping empty audio.")
    labels = inputs["input_ids"].clone()
    prompt_ids = processor.tokenizer(prompt_text, add_special_tokens=False, return_tensors="pt")["input_ids"]
    prompt_len = min(prompt_ids.shape[1], labels.shape[1])
    labels[:, :prompt_len] = -100
    pad_id = processor.tokenizer.pad_token_id
    if pad_id is not None:
        labels[inputs["input_ids"] == pad_id] = -100
    inputs["labels"] = labels
    return move_inputs(inputs, device, dtype=dtype)


def write_progress(path: Path, row: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as fp:
        fp.write(json.dumps(row, ensure_ascii=False) + "\n")


def distributed_slice(samples: list[dict], process_index: int, num_processes: int) -> list[dict]:
    if num_processes <= 1:
        return list(samples)
    shard = list(samples[process_index::num_processes])
    max_len = math.ceil(len(samples) / num_processes)
    if not samples:
        return []
    if not shard:
        shard = [samples[process_index % len(samples)]]
    while len(shard) < max_len:
        shard.append(random.choice(shard))
    return shard


def evaluate_loss(model, processor, samples: list[dict], args: argparse.Namespace, device, dtype=None, accelerator=None) -> float | None:
    if not samples:
        return None
    losses = []
    model.eval()
    with torch.no_grad():
        for sample in samples:
            inputs = prepare_sample(processor, sample, args.use_audio_in_video, device, dtype=dtype)
            loss = compute_answer_token_loss(model, inputs, tail_only=args.loss_logit_tail_only)
            losses.append(float(loss.detach().cpu()))
    model.train()
    local = torch.tensor([sum(losses), len(losses)], dtype=torch.float32, device=device)
    if accelerator is not None:
        gathered = accelerator.gather(local)
        total_loss = float(gathered[0::2].sum().detach().cpu())
        total_count = float(gathered[1::2].sum().detach().cpu())
        return total_loss / total_count if total_count else None
    return sum(losses) / len(losses) if losses else None


def main() -> int:
    args = parse_args()
    from accelerate import Accelerator

    accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps)
    workspace_default = Path(__file__).resolve().parents[2]
    if args.output_dir is None:
        args.output_dir = workspace_default / "checkpoints" / args.run_id / "adapter_lora"
    if args.results_dir is None:
        args.results_dir = workspace_default / "results" / "omni_finetune" / args.run_id
    args.output_dir.mkdir(parents=True, exist_ok=True)
    args.results_dir.mkdir(parents=True, exist_ok=True)
    progress_path = args.results_dir / "progress.jsonl"
    if accelerator.is_main_process and progress_path.exists():
        progress_path.unlink()
    backbone_profile = load_backbone_profile(args.backbone_config)
    torch.manual_seed(args.seed + accelerator.process_index)
    random.seed(args.seed + accelerator.process_index)

    samples = load_jsonl(args.dataset_jsonl)
    train_samples = select_samples(samples, args.train_split, args.include_unspecified_in_train)
    val_samples = [sample for sample in samples if sample.get("split") == args.val_split]
    if args.max_train_samples > 0:
        train_samples = train_samples[: args.max_train_samples]
    if args.max_val_samples > 0:
        val_samples = val_samples[: args.max_val_samples]
    if not train_samples:
        raise ValueError("No training samples selected. Check --train-split or use --include-unspecified-in-train.")
    rank_train_samples = distributed_slice(train_samples, accelerator.process_index, accelerator.num_processes)
    rank_val_samples = distributed_slice(val_samples, accelerator.process_index, accelerator.num_processes) if val_samples else []

    if accelerator.is_main_process:
        write_progress(progress_path, {
            "event": "setup_done",
            "run_id": args.run_id,
            "dataset_jsonl": str(args.dataset_jsonl),
            "num_processes": accelerator.num_processes,
            "num_train_samples": len(train_samples),
            "num_val_samples": len(val_samples),
            "rank0_samples_per_epoch": len(rank_train_samples),
            "backbone_id": backbone_profile.get("id"),
            "dataset_contract": backbone_profile.get("dataset_contract"),
            "training_objective": backbone_profile.get("training_objective"),
            "loss_mode": "answer_token_ce",
            "loss_logit_tail_only": args.loss_logit_tail_only,
            "timestamp": time.time(),
        })
    if accelerator.num_processes > 1 and args.device_map == "auto":
        args.device_map = "none"
    if accelerator.is_main_process:
        write_progress(progress_path, {
            "event": "model_load_start",
            "run_id": args.run_id,
            "model_id": args.model_id,
            "backbone_id": backbone_profile.get("id"),
            "device_map": args.device_map,
            "dtype": args.dtype,
            "timestamp": time.time(),
        })
    model, processor = load_model_processor(args)
    if accelerator.is_main_process:
        write_progress(progress_path, {
            "event": "model_load_done",
            "run_id": args.run_id,
            "timestamp": time.time(),
        })
    optimizer = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.learning_rate, weight_decay=args.weight_decay)
    if accelerator.is_main_process:
        write_progress(progress_path, {
            "event": "accelerator_prepare_start",
            "run_id": args.run_id,
            "timestamp": time.time(),
        })
    model, optimizer = accelerator.prepare(model, optimizer)
    if accelerator.is_main_process:
        write_progress(progress_path, {
            "event": "accelerator_prepare_done",
            "run_id": args.run_id,
            "timestamp": time.time(),
        })
    device = accelerator.device
    model_dtype = next(model.parameters()).dtype

    history = []
    global_step = 0
    optimizer.zero_grad(set_to_none=True)
    model.train()
    if accelerator.is_main_process:
        write_progress(progress_path, {
            "event": "train_loop_start",
            "run_id": args.run_id,
            "model_id": args.model_id,
            "dataset_jsonl": str(args.dataset_jsonl),
            "num_processes": accelerator.num_processes,
            "num_train_samples": len(train_samples),
            "num_val_samples": len(val_samples),
            "rank_samples_per_epoch": len(rank_train_samples),
            "epochs": args.epochs,
            "timestamp": time.time(),
        })
    for epoch in range(1, args.epochs + 1):
        random.shuffle(rank_train_samples)
        epoch_loss = 0.0
        seen = 0
        steps_in_epoch = math.ceil(len(rank_train_samples) / max(args.batch_size, 1))
        for batch_start in range(0, len(rank_train_samples), args.batch_size):
            batch = rank_train_samples[batch_start : batch_start + args.batch_size]
            batch_loss = 0.0
            for sample in batch:
                with accelerator.accumulate(model):
                    inputs = prepare_sample(processor, sample, args.use_audio_in_video, device, dtype=model_dtype)
                    loss = compute_answer_token_loss(model, inputs, tail_only=args.loss_logit_tail_only)
                    accelerator.backward(loss)
                    batch_loss += float(loss.detach().cpu())
                    if accelerator.sync_gradients:
                        accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                    optimizer.step()
                    optimizer.zero_grad(set_to_none=True)
            seen += len(batch)
            epoch_loss += batch_loss
            global_step += 1
            if accelerator.is_main_process and (global_step % args.progress_every == 0 or batch_start // max(args.batch_size, 1) == steps_in_epoch - 1):
                write_progress(progress_path, {
                    "event": "train_step",
                    "epoch": epoch,
                    "global_step": global_step,
                    "rank0_seen": seen,
                    "rank0_samples_per_epoch": len(rank_train_samples),
                    "rank0_batch_loss": batch_loss / max(len(batch), 1),
                    "timestamp": time.time(),
                })
        val_loss = evaluate_loss(model, processor, rank_val_samples, args, device, dtype=model_dtype, accelerator=accelerator)
        epoch_row = {
            "epoch": epoch,
            "train_loss": epoch_loss / max(len(rank_train_samples), 1),
            "val_loss": val_loss,
            "global_step": global_step,
        }
        history.append(epoch_row)
        if accelerator.is_main_process:
            print(json.dumps(epoch_row, indent=2))
            write_progress(progress_path, {"event": "epoch_end", **epoch_row, "timestamp": time.time()})

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        write_progress(progress_path, {
            "event": "save_start",
            "checkpoint_dir": str(args.output_dir),
            "save_mode": "trainable_lora_state_dict",
            "timestamp": time.time(),
        })
    accelerator.wait_for_everyone()
    unwrapped = accelerator.unwrap_model(model)
    gathered_state = accelerator.get_state_dict(model) if accelerator.num_processes > 1 else None
    if accelerator.is_main_process:
        if gathered_state is not None:
            adapter_state = thinker_adapter_state_from_full_state(gathered_state)
            state_source = "accelerator_full_state_dict"
        else:
            adapter_state = thinker_adapter_state_from_full_state(build_trainable_cpu_state_dict(unwrapped))
            state_source = "local_trainable_state_dict"
        shape_summary = adapter_shape_summary(adapter_state)
        write_progress(progress_path, {
            "event": "save_state_dict_built",
            "checkpoint_dir": str(args.output_dir),
            "state_source": state_source,
            "trainable_tensors": shape_summary["adapter_tensors"],
            "trainable_bytes": shape_summary["adapter_bytes"],
            "shape_summary": shape_summary,
            "timestamp": time.time(),
        })
        peft_model = getattr(unwrapped, "thinker", unwrapped)
        peft_model.save_pretrained(args.output_dir, state_dict=adapter_state, is_main_process=True)
        processor.save_pretrained(args.output_dir)
        write_progress(progress_path, {
            "event": "save_done",
            "checkpoint_dir": str(args.output_dir),
            "timestamp": time.time(),
        })
    metadata = {
        "run_id": args.run_id,
        "model_id": args.model_id,
        "backbone": backbone_profile,
        "dataset_jsonl": str(args.dataset_jsonl),
        "checkpoint_dir": str(args.output_dir),
        "num_processes": accelerator.num_processes,
        "num_train_samples": len(train_samples),
        "num_val_samples": len(val_samples),
        "history": history,
        "lora": {
            "r": args.lora_r,
            "alpha": args.lora_alpha,
            "dropout": args.lora_dropout,
            "target_modules": [item.strip() for item in args.lora_target_modules.split(",") if item.strip()],
        },
        "use_audio_in_video": args.use_audio_in_video,
        "loss_mode": "answer_token_ce",
        "loss_logit_tail_only": args.loss_logit_tail_only,
    }
    if accelerator.is_main_process:
        (args.output_dir / "training_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
        (args.results_dir / "config.yaml").write_text(
            "\n".join([
                f"run_id: {args.run_id}",
                "stage: qwen_lora_text_video_audio",
                f"backbone_id: {backbone_profile.get('id')}",
                f"dataset_contract: {backbone_profile.get('dataset_contract')}",
                f"model_id: {args.model_id}",
                f"dataset_jsonl: {args.dataset_jsonl}",
                f"checkpoint_dir: {args.output_dir}",
                f"num_processes: {accelerator.num_processes}",
                f"epochs: {args.epochs}",
                f"learning_rate: {args.learning_rate}",
                f"lora_r: {args.lora_r}",
                f"lora_alpha: {args.lora_alpha}",
                "loss_mode: answer_token_ce",
                f"loss_logit_tail_only: {args.loss_logit_tail_only}",
            ]) + "\n",
            encoding="utf-8",
        )
        (args.results_dir / "training_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
        report = [
            "# Qwen3-Omni LoRA Training",
            "",
            f"- Backbone profile: `{backbone_profile.get('display_name')}`",
            f"- Dataset contract: `{backbone_profile.get('dataset_contract')}`",
            f"- Training objective: `{backbone_profile.get('training_objective')}`",
            f"- Base model: `{args.model_id}`",
            f"- Dataset: `{args.dataset_jsonl}`",
            f"- Train samples: `{len(train_samples)}`",
            f"- Validation samples: `{len(val_samples)}`",
            f"- Processes: `{accelerator.num_processes}`",
            f"- Epochs: `{args.epochs}`",
            "- Loss: answer-token cross entropy over supervised JSON tokens",
            f"- Logit projection: `{'assistant-answer tail only' if args.loss_logit_tail_only else 'full sequence'}`",
            f"- Final train loss: `{history[-1]['train_loss']:.6f}`",
            "",
            "Only LoRA parameters are trained; the base Qwen3-Omni weights remain frozen.",
        ]
        if history[-1]["val_loss"] is not None:
            report.append(f"- Final val loss: `{history[-1]['val_loss']:.6f}`")
        (args.results_dir / "RUN_REPORT.md").write_text("\n".join(report) + "\n", encoding="utf-8")
        write_progress(progress_path, {"event": "complete", "checkpoint_dir": str(args.output_dir), "timestamp": time.time()})
        print(f"Wrote LoRA adapter to {args.output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())