#!/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()
#!/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,