Blame for code/experiments/train_qwen3vl_sft.py

#!/usr/bin/env python3
"""Audited full-parameter Qwen3-VL trainer for the pure-VLM backbone."""

from __future__ import annotations

import argparse
import inspect
import json
import math
import os
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List

import torch
from PIL import Image
from torch.utils.data import Dataset, SequentialSampler

try:
    from transformers import (
        AutoModelForImageTextToText,
        AutoProcessor,
        Trainer,
        TrainingArguments,
        set_seed,
    )
except ModuleNotFoundError:
    AutoModelForImageTextToText = None
    AutoProcessor = None
    TrainingArguments = None
    set_seed = None

    class Trainer:  # type: ignore[no-redef]
        """Import-only fallback for dependency-light contract tests."""

        pass

from viss_core.data_v0.trainer_adapter import (
    infer_release_root,
    load_training_rows,
    normalize_messages,
)
class MediaSemanticDataset(Dataset):
    """Load only a sealed pure-semantic JSONL release."""

    def __init__(
        self,
        path: Path,
        processor: Any,
        *,
        max_samples: int = 0,
        media_root: Path | None = None,
        development_unsealed_overlay: bool = False,
    ) -> None:
        self.development_unsealed_overlay = development_unsealed_overlay
        if development_unsealed_overlay:
            self.rows = [
                json.loads(line)
                for line in path.read_text(encoding="utf-8").splitlines()
                if line.strip()
            ]
        else:
            self.rows = _load_training_rows(path)
        if max_samples > 0:
            self.rows = self.rows[:max_samples]
        self.processor = processor
        inferred_root = _infer_release_root(path)
        if development_unsealed_overlay:
            if media_root is None:
                raise RuntimeError("development overlay requires --media-root")
            self.media_root = media_root.resolve(strict=True)
        else:
            if inferred_root is None:
                raise RuntimeError("semantic release root could not be inferred")
            if media_root is not None and media_root.resolve() != inferred_root:
                raise RuntimeError("--media-root may not override sealed release media")
            self.media_root = inferred_root

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, index: int) -> Dict[str, Any]:
        row = self.rows[index]
        raw_weight = row.get("sample_loss_weight")
        sample_loss_weight = float(1.0 if raw_weight is None else raw_weight)
        if not math.isfinite(sample_loss_weight) or sample_loss_weight <= 0:
            raise RuntimeError(
                "train-ready contract violated: "
                f"uid={row.get('uid') or f'row_{index}'} "
                f"sample_loss_weight={sample_loss_weight}"
            )
        return {
            "uid": str(row.get("uid") or f"row_{index}"),
            "messages": (
                _normalize_development_messages(row["messages"])
                if self.development_unsealed_overlay
                else _normalize_messages(row["messages"], media_root=self.media_root)
            ),
            "sample_loss_weight": sample_loss_weight,
        }


@dataclass
class MediaSemanticCollator:
    processor: Any
    max_length: int

    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        encoded = []
        labels_list = []
        weights_list = []
        for feature in features:
            messages = feature["messages"]
            item = _encode_processor_item(self.processor, messages)
            if item["input_ids"].numel() > self.max_length:
                raise RuntimeError(
                    f"train-ready contract violated: uid={feature['uid']} "
                    f"tokens={item['input_ids'].numel()} max_length={self.max_length}"
                )
            labels, loss_weights = _assistant_only_labels_and_weights(
                getattr(self.processor, "tokenizer", self.processor),
                item["input_ids"],
                messages,
            )
            if not torch.any(labels != -100):
                raise RuntimeError(
                    f"train-ready contract violated: uid={feature['uid']} "
                    "has no assistant labels"
                )
            encoded.append(item)
            labels_list.append(labels)
            weights_list.append(loss_weights)

        tokenizer = getattr(self.processor, "tokenizer", self.processor)
        if tokenizer.pad_token_id is None:
            raise RuntimeError("tokenizer has no pad token")
        maximum = max(item["input_ids"].numel() for item in encoded)
        batch: Dict[str, list[torch.Tensor]] = {
            "input_ids": [],
            "attention_mask": [],
            "labels": [],
            "loss_weights": [],
            "sample_loss_weight": [],
        }
        pixel_values = []
        image_grid_thw = []
        for item, labels, loss_weights, feature in zip(
            encoded, labels_list, weights_list, features, strict=True
        ):
            pad = maximum - item["input_ids"].numel()
            batch["input_ids"].append(
                torch.nn.functional.pad(
                    item["input_ids"], (0, pad), value=tokenizer.pad_token_id
                )
            )
            batch["attention_mask"].append(
                torch.nn.functional.pad(item["attention_mask"], (0, pad), value=0)
            )
            batch["labels"].append(
                torch.nn.functional.pad(labels, (0, pad), value=-100)
            )
            batch["loss_weights"].append(
                torch.nn.functional.pad(loss_weights, (0, pad), value=0.0)
            )
            batch["sample_loss_weight"].append(
                torch.tensor(
                    float(feature["sample_loss_weight"]), dtype=torch.float32
                )
            )
            if "pixel_values" in item:
                pixel_values.append(item["pixel_values"])
            if "image_grid_thw" in item:
                image_grid_thw.append(item["image_grid_thw"])
        output = {key: torch.stack(value) for key, value in batch.items()}
        if pixel_values:
            output["pixel_values"] = torch.cat(pixel_values, dim=0)
        if image_grid_thw:
            output["image_grid_thw"] = torch.cat(image_grid_thw, dim=0)
        return output


def main() -> None:
    if any(
        value is None
        for value in (
            AutoModelForImageTextToText,
            AutoProcessor,
            TrainingArguments,
            set_seed,
        )
    ):
        raise RuntimeError(
            "training requires transformers; the dependency-light import path "
            "supports only data and loss-mask contract tests"
        )
    parser = argparse.ArgumentParser(
        description="Full-parameter Qwen3-VL pure-semantic SFT."
    )
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--train", type=Path, required=True)
    parser.add_argument("--val", type=Path)
    parser.add_argument("--media-root", type=Path)
    parser.add_argument("--val-media-root", type=Path)
    parser.add_argument(
        "--development-unsealed-overlay",
        action="store_true",
        help="Development only: read an unsealed JSONL overlay with absolute media paths.",
    )
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--max-length", type=int, default=8192)
    parser.add_argument("--max-pixels", type=int, default=0)
    parser.add_argument("--per-device-batch-size", type=int, default=1)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=8)
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument(
        "--vision-projector-lr-scale",
        type=float,
        default=1.0,
        help=(
            "Multiply the optimizer learning rate for trainable visual "
            "merger/deepstack projector parameters."
        ),
    )
    parser.add_argument(
        "--lr-scheduler-type",
        choices=("cosine", "constant_with_warmup"),
        default="cosine",
    )
    parser.add_argument(
        "--train-order", choices=("shuffled", "sequential"), default="shuffled"
    )
    parser.add_argument(
        "--stage-lr-plan",
        type=Path,
        help="Audited row-boundary learning-rate plan for one ordered run.",
    )
    parser.add_argument("--warmup-ratio", type=float, default=0.03)
    parser.add_argument("--deepspeed", type=Path, required=True)
    parser.add_argument("--freeze-vision-tower", action="store_true")
    parser.add_argument("--train-vision-projector", action="store_true")
    parser.add_argument("--logging-steps", type=int, default=10)
    parser.add_argument("--save-steps", type=int, default=250)
    parser.add_argument("--save-total-limit", type=int, default=2)
    parser.add_argument("--skip-final-save", action="store_true")
    parser.add_argument("--max-train-samples", type=int, default=0)
    parser.add_argument("--max-val-samples", type=int, default=0)
    parser.add_argument("--max-steps", type=int, default=-1)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--attn-implementation", default="flash_attention_2")
    parser.add_argument("--dataloader-num-workers", type=int, default=0)
    args = parser.parse_args()

    if args.max_length <= 0 or args.max_pixels < 0:
        raise ValueError("length must be positive and max-pixels non-negative")
    if (
        not math.isfinite(args.vision_projector_lr_scale)
        or args.vision_projector_lr_scale <= 0
    ):
        raise ValueError("vision-projector-lr-scale must be finite and positive")
    if args.train_vision_projector and not args.freeze_vision_tower:
        raise ValueError("--train-vision-projector requires --freeze-vision-tower")
    if args.vision_projector_lr_scale != 1.0 and not args.train_vision_projector:
        raise ValueError(
            "a non-default vision-projector LR scale requires a trainable projector"
        )
    if not args.deepspeed.is_file():
        raise FileNotFoundError(args.deepspeed)

    set_seed(args.seed)
    has_eval = args.val is not None and args.val.exists()
    training_args = _training_args(args, has_eval)
    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
    _configure_processor_pixel_budget(processor, args.max_pixels)
    tokenizer = getattr(processor, "tokenizer", processor)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForImageTextToText.from_pretrained(
        args.model,
        dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation=args.attn_implementation,
    )
    model.config.use_cache = False
    model.gradient_checkpointing_enable()
    if args.freeze_vision_tower:
        model.visual.requires_grad_(False)
        if args.train_vision_projector:
            for name in ("merger", "deepstack_merger_list"):
                module = getattr(model.visual, name, None)
                if module is None:
                    raise RuntimeError(
                        f"Qwen3-VL visual module lacks projector component {name!r}"
                    )
                module.requires_grad_(True)
    _audit_trainable_parameters(model, args)

    if _is_main_process():
        trainable = sum(
            _logical_parameter_numel(parameter)
            for parameter in model.parameters()
            if parameter.requires_grad
        )
        total = sum(
            _logical_parameter_numel(parameter) for parameter in model.parameters()
        )
        if total <= 0:
            raise RuntimeError("could not determine logical parameter count")
        print(
            {
                "training_mode": "full_parameter_no_adapter",
                "vision_tower_frozen": args.freeze_vision_tower,
                "trainable_parameters": trainable,
                "total_parameters": total,
                "trainable_fraction": trainable / total,
            },
            flush=True,
        )

    train_dataset = MediaSemanticDataset(
        args.train,
        processor,
        max_samples=args.max_train_samples,
        media_root=args.media_root,
        development_unsealed_overlay=args.development_unsealed_overlay,
    )
    global_batch = (
        args.per_device_batch_size
        * args.gradient_accumulation_steps
        * int(os.environ.get("WORLD_SIZE", "1"))
    )
    stage_lr_plan = (
        _load_stage_lr_plan(
            args.stage_lr_plan,
            release_manifest=inferred_root_manifest(train_dataset.media_root),
            dataset_rows=len(train_dataset),
            global_batch=global_batch,
            optimizer_learning_rate=args.lr,
            warmup_ratio=args.warmup_ratio,
        )
        if args.stage_lr_plan is not None
        else None
    )
    eval_dataset = (
        MediaSemanticDataset(
            args.val,
            processor,
            max_samples=args.max_val_samples,
            media_root=args.val_media_root,
        )
        if has_eval
        else None
    )
    trainer = SemanticTokenTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=MediaSemanticCollator(processor, args.max_length),
        train_order=args.train_order,
        stage_lr_plan=stage_lr_plan,
        vision_projector_lr_scale=args.vision_projector_lr_scale,
        vision_projector_parameter_ids={
            id(parameter)
            for name in ("merger", "deepstack_merger_list")
            for parameter in getattr(model.visual, name).parameters()
        }
        if args.train_vision_projector
        else set(),
    )
    trainer.train()
    if args.skip_final_save:
        if _is_main_process():
            print(
                {
                    "output": str(args.output),
                    "training_mode": "full_parameter_no_adapter",
                    "final_model": None,
                    "systems_smoke_skipped_final_save": True,
                },
                flush=True,
            )
        return
    _release_save_reserve()
    final_dir = args.output / "final_model"
    trainer.save_model(final_dir)
    processor.save_pretrained(final_dir)
    if _is_main_process():
        print(
            {
                "output": str(args.output),
                "training_mode": "full_parameter_no_adapter",
                "final_model": str(final_dir),
            },
            flush=True,
        )


def _release_save_reserve() -> None:
    raw = os.environ.get("STREAMEVA_SAVE_RESERVE_PATH")
    if raw and _is_main_process():
        path = Path(raw).resolve()
        if (
            path.name != "save_reserve.bin"
            or path.parent.parent != Path("/tmp")
            or not path.parent.name.startswith("streameva_")
        ):
            raise RuntimeError("save reserve path is outside the short runtime root")
        path.unlink(missing_ok=True)
        os.sync()
        print(
            {
                "save_reserve_released": str(path),
                "free_bytes": shutil.disk_usage(path.parent).free,
            },
            flush=True,
        )
    if torch.distributed.is_available() and torch.distributed.is_initialized():
        torch.distributed.barrier()


def _audit_trainable_parameters(model: Any, args: argparse.Namespace) -> None:
    if not all(
        parameter.requires_grad for parameter in model.language_model.parameters()
    ):
        raise RuntimeError("full-parameter contract froze a language parameter")
    if not all(parameter.requires_grad for parameter in model.lm_head.parameters()):
        raise RuntimeError("full-parameter contract froze an LM-head parameter")
    if not args.freeze_vision_tower:
        return
    projector_ids: set[int] = set()
    if args.train_vision_projector:
        for name in ("merger", "deepstack_merger_list"):
            projector_ids.update(
                id(parameter)
                for parameter in getattr(model.visual, name).parameters()
            )
    for parameter in model.visual.parameters():
        if parameter.requires_grad != (id(parameter) in projector_ids):
            raise RuntimeError("frozen-vision/trainable-projector contract violated")


def _load_training_rows(
    path: Path, *, training_schema: str = "semantic"
) -> List[Dict[str, Any]]:
    if training_schema != "semantic":
        raise ValueError(f"unsupported training schema: {training_schema}")
    return load_training_rows(path)


def _normalize_development_messages(messages: Any) -> List[Dict[str, Any]]:
    """Normalize an explicit development overlay without weakening release loading."""
    if not isinstance(messages, list):
        raise TypeError("development messages must be a list")
    roles = [message.get("role") for message in messages]
    if roles != ["user", "assistant"]:
        raise ValueError("development overlay accepts pure-VLM rows only")
    output: List[Dict[str, Any]] = []
    for message in messages:
        content: List[Dict[str, Any]] = []
        for item in message.get("content") or []:
            if item.get("type") == "text":
                content.append({"type": "text", "text": str(item.get("text") or "")})
                continue
            if item.get("type") != "image" or message["role"] != "user":
                raise ValueError("development overlay contains unsupported content")
            image_path = Path(str(item.get("image") or ""))
            if not image_path.is_absolute() or not image_path.is_file():
                raise RuntimeError("development overlay image must be an existing absolute path")
            content.append(
                {"type": "image", "image": Image.open(image_path).convert("RGB")}
            )
        if not content:
            raise ValueError("development overlay message content is empty")
        output.append(
            {
                "role": str(message["role"]),
                "content": content,
                "loss": bool(message.get("loss")),
            }
        )
    return output


def _infer_release_root(
    path: Path, *, training_schema: str = "semantic"
) -> Path | None:
    if training_schema != "semantic":
        raise ValueError(f"unsupported training schema: {training_schema}")
    return infer_release_root(path)


def _normalize_messages(
    messages: Any,
    *,
    media_root: Path | None = None,
    training_schema: str = "semantic",
) -> List[Dict[str, Any]]:
    if training_schema != "semantic":
        raise ValueError(f"unsupported training schema: {training_schema}")
    return normalize_messages(messages, media_root=media_root)


def _load_images(messages: List[Dict[str, Any]]) -> List[Image.Image]:
    images = []
    for message in messages:
        for item in message.get("content", []):
            if not isinstance(item, dict) or item.get("type") != "image":
                continue
            image = item.get("image")
            if isinstance(image, Image.Image):
                images.append(image)
            elif image:
                images.append(Image.open(image).convert("RGB"))
    return images


def _configure_processor_pixel_budget(
    processor: Any, max_pixels: int
) -> Dict[str, int]:
    image_processor = getattr(processor, "image_processor", None)
    size = getattr(image_processor, "size", None)
    if image_processor is None or not isinstance(size, dict):
        raise RuntimeError("processor has no configurable image size contract")
    shortest = int(size.get("shortest_edge", 0))
    longest = int(size.get("longest_edge", 0))
    if shortest <= 0 or longest <= 0 or shortest > longest:
        raise RuntimeError("processor image size contract is invalid")
    requested = longest if max_pixels == 0 else int(max_pixels)
    if requested < shortest:
        raise ValueError(f"max_pixels={requested} is below shortest_edge={shortest}")
    image_processor.size = {
        "shortest_edge": shortest,
        "longest_edge": requested,
    }
    return dict(image_processor.size)


def _encode_processor_item(
    processor: Any, messages: List[Dict[str, Any]]
) -> Dict[str, torch.Tensor]:
    rendered = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False,
    )
    images = _load_images(messages)
    if rendered.count("<|image_pad|>") != len(images):
        raise RuntimeError("chat-template image slots and message images differ")
    kwargs: Dict[str, Any] = {
        "text": rendered,
        "return_dict": True,
        "return_tensors": "pt",
    }
    if images:
        kwargs["images"] = images
    item = processor(**kwargs)
    return {
        key: (
            value.squeeze(0)
            if hasattr(value, "dim")
            and value.dim() > 1
            and key in {"input_ids", "attention_mask"}
            else value
        )
        for key, value in item.items()
    }


def _locate_rendered_text_token_positions(
    tokenizer: Any,
    input_ids: torch.Tensor,
    messages: List[Dict[str, Any]],
) -> Dict[tuple[int, int], List[int]] | None:
    """Align message text spans through chat rendering and visual expansion."""

    if not hasattr(tokenizer, "apply_chat_template"):
        return None
    try:
        rendered = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False,
        )
        encoded = tokenizer(
            rendered,
            add_special_tokens=False,
            return_offsets_mapping=True,
        )
    except (KeyError, NotImplementedError, TypeError, ValueError):
        return None
    template_ids = encoded.get("input_ids")
    offsets = encoded.get("offset_mapping")
    if template_ids is None or offsets is None:
        return None
    if hasattr(template_ids, "tolist"):
        template_ids = template_ids.tolist()
    if hasattr(offsets, "tolist"):
        offsets = offsets.tolist()
    if template_ids and isinstance(template_ids[0], list):
        template_ids = template_ids[0]
    if offsets and isinstance(offsets[0][0], list):
        offsets = offsets[0]

    actual_ids = input_ids.tolist()
    visual_pad_ids = {
        token_id
        for token in ("<|image_pad|>", "<|video_pad|>")
        for token_id in [tokenizer.convert_tokens_to_ids(token)]
        if isinstance(token_id, int) and token_id >= 0
    }
    template_to_actual: Dict[int, int] = {}
    actual_index = 0
    for template_index, token_id in enumerate(template_ids):
        if actual_index >= len(actual_ids) or actual_ids[actual_index] != token_id:
            raise RuntimeError("rendered template and processor input IDs diverged")
        template_to_actual[template_index] = actual_index
        actual_index += 1
        if token_id in visual_pad_ids:
            while actual_index < len(actual_ids) and actual_ids[actual_index] == token_id:
                actual_index += 1
    if actual_index != len(actual_ids):
        raise RuntimeError("processor input IDs contain non-template tokens")

    output: Dict[tuple[int, int], List[int]] = {}
    character_cursor = 0
    for message_index, message in enumerate(messages):
        text_items = [
            item
            for item in message.get("content", [])
            if isinstance(item, dict) and item.get("type") == "text"
        ]
        for text_index, item in enumerate(text_items):
            text = str(item.get("text", ""))
            start = rendered.find(text, character_cursor)
            if start < 0:
                raise RuntimeError("message text is absent from the rendered template")
            end = start + len(text)
            positions = [
                template_to_actual[index]
                for index, pair in enumerate(offsets)
                if int(pair[0]) < end and int(pair[1]) > start
            ]
            if text and not positions:
                raise RuntimeError("rendered message text has no token positions")
            output[(message_index, text_index)] = positions
            character_cursor = end
    return output


def _assistant_only_labels(
    tokenizer: Any,
    input_ids: torch.Tensor,
    messages: List[Dict[str, Any]],
) -> torch.Tensor:
    labels, _ = _assistant_only_labels_and_weights(tokenizer, input_ids, messages)
    return labels


def _assistant_only_labels_and_weights(
    tokenizer: Any,
    input_ids: torch.Tensor,
    messages: List[Dict[str, Any]],
) -> tuple[torch.Tensor, torch.Tensor]:
    """Supervise assistant semantic text and its turn terminator only."""

    labels = torch.full_like(input_ids, -100)
    weights = torch.zeros_like(input_ids, dtype=torch.float32)
    ids = input_ids.tolist()
    cursor = 0
    rendered_positions = _locate_rendered_text_token_positions(
        tokenizer, input_ids, messages
    )
    for message_index, message in enumerate(messages):
        is_assistant = message.get("role") == "assistant"
        supervise = is_assistant and message.get("loss", True) is not False
        text_items = [
            item
            for item in message.get("content", [])
            if isinstance(item, dict) and item.get("type") == "text"
        ]
        for text_index, item in enumerate(text_items):
            text = str(item.get("text", ""))
            if rendered_positions is not None:
                positions = rendered_positions[(message_index, text_index)]
                start = positions[0] if positions else cursor
                semantic_end = positions[-1] + 1 if positions else cursor
            else:
                target = tokenizer(text, add_special_tokens=False)["input_ids"]
                start = _find_subsequence(ids, target, cursor)
                if start < 0:
                    raise RuntimeError("could not locate message text in role order")
                semantic_end = min(start + len(target), labels.numel())
            end = semantic_end
            im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
            if (
                is_assistant
                and text_index == len(text_items) - 1
                and im_end_id is not None
                and end < labels.numel()
                and ids[end] == im_end_id
            ):
                end += 1
            cursor = end
            if supervise:
                labels[start:end] = input_ids[start:end]
                weights[start:end] = 1.0
    return labels, weights


def _find_subsequence(haystack: List[int], needle: List[int], start: int) -> int:
    if not needle:
        return -1
    for index in range(max(start, 0), len(haystack) - len(needle) + 1):
        if haystack[index : index + len(needle)] == needle:
            return index
    return -1


def _split_optimizer_groups_for_projector(
    param_groups: list[dict[str, Any]],
    *,
    projector_parameter_ids: set[int],
    learning_rate_scale: float,
) -> tuple[list[dict[str, Any]], int]:
    """Preserve group options while isolating/scaling projector parameters."""

    replacement_groups: list[dict[str, Any]] = []
    projector_parameters = 0
    for original in param_groups:
        projector = [
            parameter
            for parameter in original["params"]
            if id(parameter) in projector_parameter_ids
        ]
        backbone = [
            parameter
            for parameter in original["params"]
            if id(parameter) not in projector_parameter_ids
        ]
        if backbone:
            group = dict(original)
            group["params"] = backbone
            replacement_groups.append(group)
        if projector:
            group = dict(original)
            group["params"] = projector
            group["lr"] = float(group["lr"]) * learning_rate_scale
            if "initial_lr" in group:
                group["initial_lr"] = (
                    float(group["initial_lr"]) * learning_rate_scale
                )
            replacement_groups.append(group)
            projector_parameters += len(projector)
    return replacement_groups, projector_parameters


class SemanticTokenTrainer(Trainer):
    def __init__(
        self,
        *args: Any,
        train_order: str = "shuffled",
        stage_lr_plan: dict[str, Any] | None = None,
        vision_projector_lr_scale: float = 1.0,
        vision_projector_parameter_ids: set[int] | None = None,
        **kwargs: Any,
    ) -> None:
        if train_order not in {"shuffled", "sequential"}:
            raise ValueError("unsupported train order")
        if stage_lr_plan is not None and train_order != "sequential":
            raise ValueError("stage learning-rate plan requires sequential stage order")
        self.train_order = train_order
        self.stage_lr_plan = stage_lr_plan
        self.vision_projector_lr_scale = float(vision_projector_lr_scale)
        self.vision_projector_parameter_ids = set(
            vision_projector_parameter_ids or set()
        )
        self._projector_lr_groups_applied = False
        if (
            self.vision_projector_lr_scale != 1.0
            and not self.vision_projector_parameter_ids
        ):
            raise ValueError("projector LR scaling requires projector parameters")
        super().__init__(*args, **kwargs)

    def create_optimizer(self) -> Any:
        """Split Trainer's established decay groups and scale projector LR only."""

        optimizer = super().create_optimizer()
        if self.vision_projector_lr_scale == 1.0:
            return optimizer
        if self._projector_lr_groups_applied:
            return optimizer
        if optimizer is None:
            raise RuntimeError("optimizer is absent while applying projector LR")

        replacement_groups, projector_parameters = (
            _split_optimizer_groups_for_projector(
                list(optimizer.param_groups),
                projector_parameter_ids=self.vision_projector_parameter_ids,
                learning_rate_scale=self.vision_projector_lr_scale,
            )
        )
        if projector_parameters != len(self.vision_projector_parameter_ids):
            raise RuntimeError(
                "optimizer projector parameter coverage diverged: "
                f"{projector_parameters} != "
                f"{len(self.vision_projector_parameter_ids)}"
            )
        optimizer.param_groups[:] = replacement_groups
        self._projector_lr_groups_applied = True
        if _is_main_process():
            print(
                {
                    "optimizer_lr_policy": "split_visual_projector",
                    "base_learning_rate": self.args.learning_rate,
                    "vision_projector_lr_scale": self.vision_projector_lr_scale,
                    "vision_projector_learning_rate": (
                        self.args.learning_rate
                        * self.vision_projector_lr_scale
                    ),
                    "projector_parameter_tensors": projector_parameters,
                    "optimizer_parameter_groups": len(replacement_groups),
                },
                flush=True,
            )
        return optimizer

    def create_scheduler(
        self, num_training_steps: int, optimizer: Any = None
    ) -> Any:
        if self.stage_lr_plan is None:
            try:
                return super().create_scheduler(num_training_steps, optimizer)
            except TypeError:
                return super().create_scheduler(num_training_steps)
        if num_training_steps != int(self.stage_lr_plan["total_steps"]):
            raise RuntimeError(
                "stage LR plan and Trainer optimization steps diverged: "
                f"{self.stage_lr_plan['total_steps']} != {num_training_steps}"
            )
        if self.lr_scheduler is None:
            active_optimizer = self.optimizer if optimizer is None else optimizer
            if active_optimizer is None:
                raise RuntimeError("optimizer is absent while creating stage scheduler")
            self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
                active_optimizer,
                lr_lambda=lambda step: _stage_lr_multiplier(
                    self.stage_lr_plan, int(step)
                ),
            )
        return self.lr_scheduler

    def _get_train_sampler(self, train_dataset: Any = None) -> Any:
        dataset = self.train_dataset if train_dataset is None else train_dataset
        if self.train_order == "sequential":
            return SequentialSampler(dataset)
        try:
            return super()._get_train_sampler(train_dataset)
        except TypeError:
            return super()._get_train_sampler()

    def compute_loss(
        self,
        model: Any,
        inputs: Dict[str, Any],
        return_outputs: bool = False,
        num_items_in_batch: Any = None,
    ) -> Any:
        del num_items_in_batch
        labels = inputs.pop("labels")
        loss_weights = inputs.pop("loss_weights")
        sample_loss_weight = inputs.pop("sample_loss_weight").float()
        outputs = model(**inputs)
        shift_logits = outputs.logits[..., :-1, :]
        shift_labels = labels[..., 1:].contiguous()
        shift_weights = loss_weights[..., 1:].contiguous().float()
        loss = _sparse_weighted_token_loss(
            shift_logits, shift_labels, shift_weights, sample_loss_weight
        )
        return (loss, outputs) if return_outputs else loss


def _sparse_weighted_token_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
    loss_weights: torch.Tensor,
    sample_loss_weight: torch.Tensor,
) -> torch.Tensor:
    """Compute the same weighted CE without densifying every prompt token.

    Qwen emits logits for the visual tokens, prompt and assistant answer, while
    the semantic SFT contract supervises only assistant answer positions.  The
    previous implementation cast the complete ``batch x sequence x vocab``
    tensor to fp32 and then discarded all ignored positions.  Selecting active
    positions first is mathematically equivalent and avoids the transient
    13--20 GiB allocation seen on long examples.
    """

    if logits.ndim != 3 or labels.shape != logits.shape[:2]:
        raise RuntimeError("logits and labels do not align")
    if loss_weights.shape != labels.shape:
        raise RuntimeError("token loss weights do not align with labels")
    if sample_loss_weight.shape != (logits.shape[0],):
        raise RuntimeError("sample loss weights do not align with the batch")

    active = labels.ne(-100) & loss_weights.gt(0)
    active_batch, _ = torch.nonzero(active, as_tuple=True)
    if active_batch.numel() == 0:
        raise RuntimeError("batch has no positively weighted assistant labels")
    active_logits = logits[active].float()
    active_labels = labels[active]
    active_weights = loss_weights[active].float()
    token_loss = torch.nn.functional.cross_entropy(
        active_logits, active_labels, reduction="none"
    )
    numerator = torch.zeros(
        logits.shape[0], device=logits.device, dtype=torch.float32
    ).scatter_add_(0, active_batch, token_loss * active_weights)
    denominator = torch.zeros_like(numerator).scatter_add_(
        0, active_batch, active_weights
    )
    per_sample_loss = numerator / denominator.clamp_min(1e-8)
    return (per_sample_loss * sample_loss_weight.float()).mean()


def inferred_root_manifest(release_root: Path) -> dict[str, Any]:
    path = release_root / "manifest.json"
    if not path.is_file():
        raise RuntimeError("sealed release manifest is absent")
    return json.loads(path.read_text(encoding="utf-8"))


def _load_stage_lr_plan(
    path: Path,
    *,
    release_manifest: dict[str, Any],
    dataset_rows: int,
    global_batch: int,
    optimizer_learning_rate: float,
    warmup_ratio: float,
) -> dict[str, Any]:
    """Validate and compile row boundaries into optimizer-step boundaries."""

    raw = json.loads(path.read_text(encoding="utf-8"))
    if raw.get("format") != "iclr_video.stage_lr_plan.v1":
        raise RuntimeError("unsupported stage LR plan format")
    if int(raw.get("dataset_rows", -1)) != dataset_rows or global_batch <= 0:
        raise RuntimeError("stage LR plan dataset/global-batch contract diverged")
    if int(raw.get("global_batch", -1)) != global_batch:
        raise RuntimeError("stage LR plan global batch diverged")
    manifest_stages = (release_manifest.get("inputs") or {}).get(
        "curriculum_stages"
    )
    stages = raw.get("stages")
    if not isinstance(stages, list) or not stages or not isinstance(manifest_stages, list):
        raise RuntimeError("stage LR plan or release stages are absent")
    if len(stages) != len(manifest_stages):
        raise RuntimeError("stage LR plan count diverged from release")
    compiled = []
    previous_row = 0
    previous_step = 0
    maximum_lr = 0.0
    for index, (stage, manifest_stage) in enumerate(
        zip(stages, manifest_stages, strict=True)
    ):
        end_row = int(stage.get("end_row_exclusive", -1))
        learning_rate = float(stage.get("learning_rate", 0.0))
        if (
            int(stage.get("index", -1)) != index
            or int(manifest_stage.get("index", -1)) != index
            or str(stage.get("release_id")) != str(manifest_stage.get("release_id"))
            or end_row != int(manifest_stage.get("end_row_exclusive", -1))
            or end_row <= previous_row
            or not math.isfinite(learning_rate)
            or learning_rate <= 0
        ):
            raise RuntimeError("stage LR entry diverged from release manifest")
        end_step = math.ceil(end_row / global_batch)
        if end_step <= previous_step:
            raise RuntimeError("stage has no optimizer steps")
        stage_steps = end_step - previous_step
        warmup_steps = math.ceil(stage_steps * warmup_ratio)
        compiled.append(
            {
                "end_step_exclusive": end_step,
                "index": index,
                "learning_rate": learning_rate,
                "release_id": str(stage["release_id"]),
                "start_step": previous_step,
                "warmup_steps": warmup_steps,
            }
        )
        maximum_lr = max(maximum_lr, learning_rate)
        previous_row = end_row
        previous_step = end_step
    if previous_row != dataset_rows:
        raise RuntimeError("stage LR plan does not cover the complete dataset")
    if not math.isclose(
        maximum_lr, optimizer_learning_rate, rel_tol=1e-9, abs_tol=0.0
    ):
        raise RuntimeError("optimizer LR must equal the maximum stage LR")
    return {
        "base_learning_rate": optimizer_learning_rate,
        "format": raw["format"],
        "stages": compiled,
        "total_steps": previous_step,
    }


def _stage_lr_multiplier(plan: dict[str, Any], step: int) -> float:
    """Return a per-stage warmup/cosine multiplier for one optimizer step."""

    base_lr = float(plan["base_learning_rate"])
    for stage in plan["stages"]:
        start = int(stage["start_step"])
        end = int(stage["end_step_exclusive"])
        if step >= end:
            continue
        offset = max(step - start, 0)
        duration = end - start
        warmup = int(stage["warmup_steps"])
        scale = float(stage["learning_rate"]) / base_lr
        if warmup > 0 and offset < warmup:
            return scale * (offset / warmup)
        decay_steps = max(duration - warmup, 1)
        progress = min(max((offset - warmup) / decay_steps, 0.0), 1.0)
        return scale * 0.5 * (1.0 + math.cos(math.pi * progress))
    return 0.0


def _training_args(args: argparse.Namespace, has_eval: bool) -> Any:
    kwargs: Dict[str, Any] = {
        "output_dir": str(args.output),
        "num_train_epochs": args.epochs,
        "per_device_train_batch_size": args.per_device_batch_size,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "learning_rate": args.lr,
        "warmup_ratio": args.warmup_ratio,
        "bf16": True,
        "logging_steps": args.logging_steps,
        "save_steps": args.save_steps,
        "save_total_limit": args.save_total_limit,
        "max_steps": args.max_steps,
        "gradient_checkpointing": True,
        "lr_scheduler_type": args.lr_scheduler_type,
        "optim": "adamw_torch",
        "weight_decay": 0.0,
        "max_grad_norm": 1.0,
        "tf32": True,
        "save_only_model": True,
        "report_to": "none",
        "remove_unused_columns": False,
        "ddp_find_unused_parameters": False,
        "dataloader_num_workers": args.dataloader_num_workers,
        "dataloader_pin_memory": True,
        "seed": args.seed,
        "deepspeed": str(args.deepspeed),
    }
    signature = inspect.signature(TrainingArguments.__init__)
    if args.save_steps >= 999999 and "save_strategy" in signature.parameters:
        kwargs["save_strategy"] = "no"
    if args.dataloader_num_workers > 0:
        if "dataloader_persistent_workers" in signature.parameters:
            kwargs["dataloader_persistent_workers"] = True
        if "dataloader_prefetch_factor" in signature.parameters:
            kwargs["dataloader_prefetch_factor"] = 2
    if "data_seed" in signature.parameters:
        kwargs["data_seed"] = args.seed
    evaluation_key = (
        "eval_strategy"
        if "eval_strategy" in signature.parameters
        else "evaluation_strategy"
    )
    kwargs[evaluation_key] = "steps" if has_eval else "no"
    if has_eval:
        kwargs["eval_steps"] = max(args.save_steps, args.logging_steps)
    return TrainingArguments(**kwargs)


def _is_main_process() -> bool:
    return int(os.environ.get("RANK", "0")) == 0


def _logical_parameter_numel(parameter: torch.nn.Parameter) -> int:
    return int(getattr(parameter, "ds_numel", parameter.numel()))


if __name__ == "__main__":
    main()