diff --git "a/code/rocm.patch" "b/code/rocm.patch" new file mode 100644--- /dev/null +++ "b/code/rocm.patch" @@ -0,0 +1,3179 @@ +diff --git a/configuration_bailingmm2.py b/configuration_bailingmm2.py +index 5ab2542..b20eca3 100644 +--- a/configuration_bailingmm2.py ++++ b/configuration_bailingmm2.py +@@ -20,6 +20,11 @@ from configuration_bailing_moe_v2 import BailingMoeV2Config + + class BailingMM2Config(PretrainedConfig): + model_type = "bailingmm_moe_v2_lite" ++ # Declared so transformers' `_attn_implementation` setter recurses into both towers. ++ # Without it an explicit attn_implementation (e.g. "eager" on ROCm, which has no ++ # flash-attn) never reaches them, and their "flash_attention_2" defaults raise at ++ # model construction. ++ sub_configs = {"vision_config": Qwen2_5_VLVisionConfig, "llm_config": BailingMoeV2Config} + + def __init__( + self, +diff --git a/diffusion/transformer.py b/diffusion/transformer.py +index d47ca9f..89a2845 100644 +--- a/diffusion/transformer.py ++++ b/diffusion/transformer.py +@@ -37,6 +37,20 @@ ADALN_EMBED_DIM = 256 + SEQ_MULTI_OF = 32 + + ++def _native_sdpa_is_active(processor) -> bool: ++ """True when attention would go to diffusers' default native SDPA backend (no per-model backend, ++ no context parallelism, and the active global backend is NATIVE).""" ++ if processor._attention_backend is not None or processor._parallel_config is not None: ++ return False ++ try: ++ from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry ++ ++ name, _ = _AttentionBackendRegistry.get_active_backend() ++ return name == AttentionBackendName.NATIVE ++ except Exception: ++ return False ++ ++ + class TimestepEmbedder(nn.Module): + def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): + super().__init__() +@@ -130,16 +144,26 @@ class SingleStreamAttentionProcessor: + attention_mask = attention_mask[:, None, None, :] + + # Compute joint attention +- hidden_states = dispatch_attention_fn( +- query, +- key, +- value, +- attn_mask=attention_mask, +- dropout_p=0.0, +- is_causal=False, +- backend=self._attention_backend, +- parallel_config=self._parallel_config, +- ) ++ if _native_sdpa_is_active(self): ++ # What diffusers' default "native" backend computes, but SDPA receives contiguous ++ # [B, H, L, D] tensors instead of permuted views. PyTorch's math SDPA (the only SDPA ++ # kernel that runs on ROCm gfx1151) is ~2x faster on contiguous inputs, with ++ # bit-identical output. ++ q, k, v = (x.transpose(1, 2).contiguous() for x in (query, key, value)) ++ hidden_states = F.scaled_dot_product_attention( ++ q, k, v, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ++ ).transpose(1, 2) ++ else: ++ hidden_states = dispatch_attention_fn( ++ query, ++ key, ++ value, ++ attn_mask=attention_mask, ++ dropout_p=0.0, ++ is_causal=False, ++ backend=self._attention_backend, ++ parallel_config=self._parallel_config, ++ ) + + # Reshape back + hidden_states = hidden_states.flatten(2, 3) +diff --git a/generate_paired.sh b/generate_paired.sh +new file mode 100755 +index 0000000..864f462 +--- /dev/null ++++ b/generate_paired.sh +@@ -0,0 +1,162 @@ ++#!/usr/bin/env bash ++# Paired pipeline: Ling-3.0-flash-VL prompt enhancement -> Ming-Image text-to-image. ++# ++# Stage 1 pe_ling.py caption -> validated structured JSON prompt ++# (system prompt: assets/t2i_rewriter_system_prompt.txt) ++# Stage 2 infer.py --task text-to-image --prompt -> PNG ++# (infer.py reads --prompt as a file when the path exists) ++# ++# Artifacts land in --output-dir: enhanced_prompt.json (overwritten per run) ++# plus the PNG(s) infer.py writes (image_00.png for text-to-image). ++# Fails loudly at every stage (set -Eeuo pipefail + ERR trap + stage checks). ++set -Eeuo pipefail ++trap 'printf "generate_paired: FAILED at line %d (exit %d)\n" "$LINENO" "$?" >&2' ERR ++ ++SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" ++PYTHON="${PYTHON:-python3}" ++ ++# Local llama-server seat serving Ling-3.0-flash-VL on the target box. ++DEFAULT_BASE_URL="http://127.0.0.1:8090/v1" ++DEFAULT_PE_MODEL="ling-3.0-flash-vl-mtp-halo-STRIX_LEAN" ++ ++usage() { ++ cat <<'EOF' ++Usage: generate_paired.sh --model DIR_OR_REPO CAPTION [options] [-- EXTRA_INFER_ARGS...] ++ ++Enhances CAPTION with Ling-3.0-flash-VL (pe_ling.py), validates the structured ++JSON rewrite, then renders it with infer.py --task text-to-image. ++ ++Required: ++ CAPTION free-form design caption (positional) ++ --model DIR_OR_REPO Ming checkpoint directory or HF repo id ++ (may also be set via the MING_MODEL environment variable) ++ ++Passthrough to infer.py (all optional; infer.py defaults in parentheses): ++ --resolution N resolution bucket, 1024 or 2048 for text-to-image; ++ other positive values snap to the nearest bucket (2048) ++ --seed N generation seed (42) ++ --steps N diffusion steps (12) ++ -- everything after this is passed to infer.py verbatim ++ (e.g. -- --validate-only --dtype float16) ++ ++Prompt-enhancement endpoint: ++ --base-url URL OpenAI-compatible base URL (http://127.0.0.1:8090/v1) ++ --pe-model ID chat model id served there ++ (ling-3.0-flash-vl-mtp-halo-STRIX_LEAN) ++ LITELLM_API_KEY env exported key is sent as a Bearer token (for a gated ++ OpenAI-compatible gateway such as LiteLLM) ++ ++Other: ++ --output-dir DIR artifact directory (outputs/paired) ++ -h, --help this help ++ ++Examples: ++ ./generate_paired.sh --model /models/Ming-Image-0.1-Design \ ++ "espresso machine product poster, warm morning light" --resolution 2048 ++ ++ LITELLM_API_KEY=sk-... ./generate_paired.sh \ ++ --base-url http://:4000/v1 --pe-model \ ++ --model /models/Ming-Image-0.1-Design "a caption" --seed 7 ++EOF ++} ++ ++die() { ++ printf 'generate_paired: %s\n' "$*" >&2 ++ exit 1 ++} ++ ++model="${MING_MODEL:-}" ++base_url="$DEFAULT_BASE_URL" ++pe_model="$DEFAULT_PE_MODEL" ++output_dir="outputs/paired" ++resolution="" ++seed="" ++steps="" ++caption="" ++extra_infer_args=() ++ ++while [[ $# -gt 0 ]]; do ++ case "$1" in ++ --model) [[ $# -ge 2 ]] || die "--model requires a value"; model="$2"; shift 2 ;; ++ --base-url) [[ $# -ge 2 ]] || die "--base-url requires a value"; base_url="$2"; shift 2 ;; ++ --pe-model) [[ $# -ge 2 ]] || die "--pe-model requires a value"; pe_model="$2"; shift 2 ;; ++ --output-dir) [[ $# -ge 2 ]] || die "--output-dir requires a value"; output_dir="$2"; shift 2 ;; ++ --resolution) [[ $# -ge 2 ]] || die "--resolution requires a value"; resolution="$2"; shift 2 ;; ++ --seed) [[ $# -ge 2 ]] || die "--seed requires a value"; seed="$2"; shift 2 ;; ++ --steps) [[ $# -ge 2 ]] || die "--steps requires a value"; steps="$2"; shift 2 ;; ++ -h|--help) usage; exit 0 ;; ++ --) shift; extra_infer_args+=("$@"); break ;; ++ -*) usage >&2; die "unknown option: $1" ;; ++ *) ++ if [[ -n "$caption" ]]; then ++ usage >&2 ++ die "unexpected extra argument: $1 (CAPTION was already given)" ++ fi ++ caption="$1" ++ shift ++ ;; ++ esac ++done ++ ++[[ -n "$caption" ]] || { usage >&2; die "CAPTION is required"; } ++[[ -n "$model" ]] || { usage >&2; die "--model DIR_OR_REPO is required (or set MING_MODEL)"; } ++if [[ -n "$resolution" && ! "$resolution" =~ ^[0-9]+$ ]]; then ++ die "--resolution must be a positive integer, got: $resolution" ++fi ++if [[ -n "$seed" && ! "$seed" =~ ^-?[0-9]+$ ]]; then ++ die "--seed must be an integer, got: $seed" ++fi ++if [[ -n "$steps" && ! "$steps" =~ ^[0-9]+$ ]]; then ++ die "--steps must be a positive integer, got: $steps" ++fi ++[[ -f "$SCRIPT_DIR/pe_ling.py" ]] || die "missing stage-1 script: $SCRIPT_DIR/pe_ling.py" ++[[ -f "$SCRIPT_DIR/infer.py" ]] || die "missing stage-2 script: $SCRIPT_DIR/infer.py" ++command -v "$PYTHON" >/dev/null 2>&1 || die "python interpreter not found: $PYTHON (override with PYTHON=...)" ++ ++mkdir -p -- "$output_dir" || die "cannot create output directory: $output_dir" ++prompt_json="$output_dir/enhanced_prompt.json" ++ ++printf '== stage 1/2: prompt enhancement (pe_ling.py, model %s @ %s)\n' "$pe_model" "$base_url" >&2 ++"$PYTHON" "$SCRIPT_DIR/pe_ling.py" "$caption" \ ++ --out "$prompt_json" \ ++ --base-url "$base_url" \ ++ --model "$pe_model" ++[[ -s "$prompt_json" ]] || die "prompt enhancement produced no prompt file: $prompt_json" ++ ++printf '== stage 2/2: Ming-Image text-to-image (infer.py, model %s)\n' "$model" >&2 ++infer_args=( ++ --model "$model" ++ --task text-to-image ++ --prompt "$prompt_json" ++ --output-dir "$output_dir" ++) ++if [[ -n "$resolution" ]]; then infer_args+=(--resolution "$resolution"); fi ++if [[ -n "$seed" ]]; then infer_args+=(--seed "$seed"); fi ++if [[ -n "$steps" ]]; then infer_args+=(--steps "$steps"); fi ++if [[ ${#extra_infer_args[@]} -gt 0 ]]; then infer_args+=("${extra_infer_args[@]}"); fi ++validate_only=0 ++for arg in ${extra_infer_args[@]+"${extra_infer_args[@]}"}; do ++ if [[ "$arg" == "--validate-only" ]]; then validate_only=1; fi ++done ++"$PYTHON" "$SCRIPT_DIR/infer.py" "${infer_args[@]}" ++ ++if [[ "$validate_only" -eq 1 ]]; then ++ printf 'generate_paired: --validate-only dry run, no PNG expected; enhanced prompt: %s\n' \ ++ "$prompt_json" >&2 ++ exit 0 ++fi ++ ++# infer.py exits non-zero on failure (set -e above); additionally verify the ++# promised PNG artifacts actually exist so a silent no-write still fails ++# loudly. -newer pins the check to THIS run: stage 2 always writes its PNG ++# after stage 1 wrote enhanced_prompt.json, so stale PNGs do not satisfy it. ++pngs=() ++while IFS= read -r png; do ++ pngs+=("$png") ++done < <(find "$output_dir" -maxdepth 1 -name '*.png' -type f -newer "$prompt_json" | sort) ++if [[ ${#pngs[@]} -eq 0 ]]; then ++ die "infer.py exited 0 but wrote no PNG under $output_dir in this run" ++fi ++printf 'generate_paired: enhanced prompt: %s\n' "$prompt_json" >&2 ++printf 'generate_paired: %d PNG(s):\n' "${#pngs[@]}" >&2 ++printf '%s\n' "${pngs[@]}" +diff --git a/infer.py b/infer.py +index 9dda84a..814e8f3 100644 +--- a/infer.py ++++ b/infer.py +@@ -100,6 +100,22 @@ def parse_args() -> argparse.Namespace: + action="store_true", + help="Validate model profile and task arguments without loading weights", + ) ++ parser.add_argument( ++ "--attention-bf16-reduction", ++ action="store_true", ++ help=( ++ "Let PyTorch's math attention kernel (the only SDPA kernel that runs on ROCm gfx1151) " ++ "stay in bf16 instead of upcasting to fp32: faster, less precise." ++ ), ++ ) ++ parser.add_argument( ++ "--release-mllm-after-conditioning", ++ action="store_true", ++ help=( ++ "Free the MLLM, vision tower and connector as soon as the conditioning is computed, " ++ "before the diffusion steps. Lowers peak memory; one image per process." ++ ), ++ ) + return parser.parse_args() + + +@@ -320,6 +336,10 @@ def load_model_and_processor(model_directory: Path, args): + ) + processor = load_bailingmm2_processor(processor_directory) + ++ if getattr(args, "attention_bf16_reduction", False): ++ # The math SDPA kernel upcasts bf16 inputs to fp32 by default; this keeps it in bf16. ++ torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True) ++ + dtype = _dtype(args.dtype) + load_kwargs = { + "torch_dtype": dtype, +@@ -351,9 +371,37 @@ def load_model_and_processor(model_directory: Path, args): + model = model.to(device=args.device, dtype=dtype) + elif device_plan is not None: + _validate_balanced_placement(model, device_plan, torch) ++ if getattr(args, "release_mllm_after_conditioning", False): ++ _release_mllm_before_sampling(model) + return model, processor + + ++def _release_mllm_before_sampling(model) -> None: ++ """Free the MLLM-side modules once the conditioning exists (--release-mllm-after-conditioning). ++ ++ Wraps the diffusion sampler: by the time it is called the conditioning tensors are computed, ++ so the language model, vision tower and connector are moved to the meta device (releasing ++ their memory) before the diffusion steps start. The model cannot generate again afterwards. ++ """ ++ import gc ++ ++ import torch ++ ++ original_sample = model.diffusion_loss.sample ++ ++ def sample_after_release(*args, **kwargs): ++ for name in ("model", "vision", "linear_proj", "connector"): ++ module = getattr(model, name, None) ++ if module is not None: ++ module.to("meta") ++ gc.collect() ++ if torch.cuda.is_available(): ++ torch.cuda.empty_cache() ++ return original_sample(*args, **kwargs) ++ ++ model.diffusion_loss.sample = sample_after_release ++ ++ + def run_generation( + model, + processor, +diff --git a/modeling_bailing_moe_v2.py b/modeling_bailing_moe_v2.py +index a608f45..b62b66c 100644 +--- a/modeling_bailing_moe_v2.py ++++ b/modeling_bailing_moe_v2.py +@@ -28,7 +28,6 @@ import torch.nn.functional as F + import torch.utils.checkpoint + from torch import nn + from torch.nn import CrossEntropyLoss +-import transformer_engine.pytorch as te + from transformers.activations import ACT2FN + from transformers.cache_utils import Cache, DynamicCache + from transformers.modeling_attn_mask_utils import ( +diff --git a/modeling_bailingmm2.py b/modeling_bailingmm2.py +index fc0ecfa..27552f3 100644 +--- a/modeling_bailingmm2.py ++++ b/modeling_bailingmm2.py +@@ -444,6 +444,49 @@ class BailingMM2NativeForConditionalGeneration(PreTrainedModel): + self.diffusion_loss.to(device) + self.loaded_image_gen_modules = True + @classmethod ++ def _from_int8_checkpoint(cls, vlm_directory, device, **kwargs): ++ """Load an mllm/ component written by quant/quantize_stream.py (weight-only int8). ++ ++ The model is built with its parameters on the meta device, the Linear modules listed ++ in int8_manifest.json become Int8Linear shells, and every stored tensor is loaded ++ straight onto `device`, so BF16 weights for the quantized modules never exist in memory. ++ """ ++ from accelerate import init_empty_weights ++ from quant.load_int8 import load_int8_mllm_ ++ ++ device_map = kwargs.pop("device_map", None) ++ if device_map is not None: ++ # infer.py's default "balanced" plan on a single-GPU box maps every module to GPU 0; ++ # that is honoured. Splitting the int8 model across devices is not supported. ++ targets = set(device_map.values()) if isinstance(device_map, dict) else {device_map} ++ if len(targets) != 1 or not isinstance(next(iter(targets)), int): ++ raise ValueError( ++ "the int8 mllm checkpoint loads onto a single GPU; device_map targets " ++ f"{sorted(map(str, targets))} (use --device-map none)" ++ ) ++ device = torch.device("cuda", next(iter(targets))) ++ supported = ("torch_dtype", "dtype", "attn_implementation") ++ unsupported = sorted(key for key in kwargs if key not in supported) ++ if unsupported: ++ raise ValueError( ++ f"the int8 mllm checkpoint loads onto a single device; unsupported arguments: {unsupported}" ++ ) ++ device = torch.device(device) if device is not None else torch.device("cpu") ++ if device.type == "cuda" and device.index is None: ++ device = torch.device("cuda", torch.cuda.current_device()) ++ config = BailingMM2Config.from_pretrained(vlm_directory) ++ with init_empty_weights(): ++ model = cls._from_config(config, **kwargs) ++ report = load_int8_mllm_(model, vlm_directory, device) ++ # Buffers built in __init__ (rotary inv_freq) are not stored in the checkpoint; they follow ++ # the weights. Int8Linear keeps its fp32 scales through this and any later dtype cast. ++ model.to(device) ++ logger.info(f"int8 mllm loaded from {vlm_directory}: {report}") ++ model.tie_weights() ++ model.eval() ++ return model ++ ++ @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], +@@ -488,7 +531,9 @@ class BailingMM2NativeForConditionalGeneration(PreTrainedModel): + f"{vlm_directory}. Migrate the package to the component " + "layout before loading." + ) +- if load_vlm: ++ if load_vlm and os.path.exists(os.path.join(vlm_directory, "int8_manifest.json")): ++ model = cls._from_int8_checkpoint(vlm_directory, image_gen_device, **kwargs) ++ elif load_vlm: + model = super().from_pretrained( + vlm_directory, + *model_args, +diff --git a/pe_ling.py b/pe_ling.py +new file mode 100644 +index 0000000..88da86f +--- /dev/null ++++ b/pe_ling.py +@@ -0,0 +1,446 @@ ++#!/usr/bin/env python3 ++"""Prompt enhancement (PE) for Ming-Image text-to-image via a Ling-3.0-flash-VL seat. ++ ++Per the README, PE is a pre-processing step *outside* ``infer.py``: an ++instruction-following VLM rewrites a short caption into the structured ++Figma-style JSON prompt that the text-to-image pipeline consumes, and the ++result is passed to ``infer.py --prompt`` as raw text or via a file. ++ ++This module drives any OpenAI-compatible ``/chat/completions`` endpoint using ++only the standard library (``urllib``): by default the local llama-server seat ++serving Ling-3.0-flash-VL, optionally the LiteLLM lab gateway (Bearer auth via ++``--api-key`` or the ``LITELLM_API_KEY`` environment variable). The rewriter ++system prompt is read verbatim from ``assets/t2i_rewriter_system_prompt.txt``. ++ ++The reply is parsed robustly (```json fences and surrounding prose are ++tolerated), then validated against the schema the system prompt demands. On a ++parse or validation failure the request is retried exactly once with the ++errors appended to the user turn; if that still fails, PromptEnhancementError ++is raised with the errors. Invalid JSON is never passed through silently. ++ ++CLI: ++ python pe_ling.py "a caption" --out prompt.json \ ++ [--base-url http://127.0.0.1:8090/v1] \ ++ [--model ling-3.0-flash-vl-mtp-halo-STRIX_LEAN] ++""" ++ ++from __future__ import annotations ++ ++import argparse ++import json ++import os ++import re ++import sys ++import time ++import urllib.error ++import urllib.request ++from pathlib import Path ++from typing import Any, Dict, List, Optional, Tuple ++ ++CODE_DIRECTORY = Path(__file__).resolve().parent ++SYSTEM_PROMPT_PATH = CODE_DIRECTORY / "assets" / "t2i_rewriter_system_prompt.txt" ++ ++# The Ling-3.0-flash-VL seat already served on the target box (llama-server, ++# OpenAI-compatible, thinking disabled); both endpoints speak the same ++# /chat/completions protocol. ++DEFAULT_BASE_URL = "http://127.0.0.1:8090/v1" ++DEFAULT_MODEL = "ling-3.0-flash-vl-mtp-halo-STRIX_LEAN" ++API_KEY_ENV = "LITELLM_API_KEY" ++ ++# Low temperature: the rewrite is a deterministic schema transformation, not ++# creative sampling. ++DEFAULT_TEMPERATURE = 0.2 ++# The upstream example rewrite (assets/t2i_four_seasons_cabin_prompt.json) is ++# ~5 KB (~2k tokens); dense multi-layer infographic rewrites run several times ++# longer, so leave generous headroom for a complete JSON object. ++DEFAULT_MAX_TOKENS = 16384 ++# A multi-thousand-token completion on the local seat can take minutes. ++DEFAULT_TIMEOUT_SECONDS = 600.0 ++ ++REPAIR_INSTRUCTION = "Return only the corrected JSON object: no prose, no code fences." ++ ++CANVAS_SETTINGS_KEYS = ("aspect_ratio", "ambient_lighting", "image_style") ++LAYER_KEYS = ("description", "coordinates", "hierarchy_and_relation", "color_specs") ++COORDINATE_FIELDS = ("cx", "cy", "w", "h") ++ ++# `coordinates` must be ONE string of the form ++# "cx: 0.500, cy: 0.500, w: 1.000, h: 1.000". The upstream example also uses ++# bare integers ("h: 1"), so accept any decimal spelling and enforce the ++# [0, 1] range on the parsed value. Whitespace around ':' and ',' is ++# tolerated; the key order is fixed. ++_COORDINATE_NUMBER = r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)" ++COORDINATES_RE = re.compile( ++ rf"^\s*cx:\s*(?P{_COORDINATE_NUMBER})\s*," ++ rf"\s*cy:\s*(?P{_COORDINATE_NUMBER})\s*," ++ rf"\s*w:\s*(?P{_COORDINATE_NUMBER})\s*," ++ rf"\s*h:\s*(?P{_COORDINATE_NUMBER})\s*$" ++) ++ ++# Hex colors: #RGB, #RGBA, #RRGGBB, #RRGGBBAA (the upstream example uses ++# #RRGGBB; the alpha forms keep RGBA-design outputs from failing validation). ++HEX_COLOR_RE = re.compile( ++ r"^#(?:[0-9a-fA-F]{3}|[0-9a-fA-F]{4}|[0-9a-fA-F]{6}|[0-9a-fA-F]{8})$" ++) ++ ++ ++class PromptEnhancementError(RuntimeError): ++ """PE failed: transport/protocol error, or schema failure after the retry.""" ++ ++ def __init__( ++ self, ++ message: str, ++ errors: Optional[List[str]] = None, ++ reply: Optional[str] = None, ++ ): ++ super().__init__(message) ++ self.errors = list(errors or []) ++ self.reply = reply ++ ++ ++def load_system_prompt(path: Path = SYSTEM_PROMPT_PATH) -> str: ++ """Return the released rewriter system prompt, verbatim.""" ++ return path.read_text(encoding="utf-8") ++ ++ ++def extract_json_object(text: str) -> Dict[str, Any]: ++ """Return the first complete top-level JSON object found in ``text``. ++ ++ Models sometimes wrap JSON in ```json fences or add prose around it. ++ Scanning every ``{`` position with ``JSONDecoder.raw_decode`` (which ++ decodes a document at an offset and ignores trailing data) recovers the ++ object in all of those shapes. Raises ValueError when no complete JSON ++ object is present, e.g. a reply truncated mid-object. ++ """ ++ decoder = json.JSONDecoder() ++ position = text.find("{") ++ while position != -1: ++ try: ++ document, _ = decoder.raw_decode(text, position) ++ except ValueError: ++ position = text.find("{", position + 1) ++ continue ++ return document ++ snippet = text.strip() ++ if len(snippet) > 300: ++ snippet = snippet[:300] + "..." ++ raise ValueError( ++ f"reply contains no complete top-level JSON object " ++ f"({len(text)} characters); starts with: {snippet!r}" ++ ) ++ ++ ++def _check_exact_keys( ++ mapping: Dict[str, Any], expected: Tuple[str, ...], path: str, errors: List[str] ++) -> None: ++ missing = [key for key in expected if key not in mapping] ++ unexpected = [key for key in mapping if key not in expected] ++ if missing: ++ errors.append(f"{path}: missing required key(s): {', '.join(missing)}") ++ if unexpected: ++ errors.append( ++ f"{path}: unexpected key(s): {', '.join(unexpected)} " ++ f"(exactly {', '.join(expected)} are required)" ++ ) ++ ++ ++def _check_non_empty_string(value: Any, path: str, errors: List[str]) -> None: ++ if not isinstance(value, str): ++ errors.append(f"{path}: expected a string, got {type(value).__name__}") ++ elif not value.strip(): ++ errors.append(f"{path}: string is empty") ++ ++ ++def _check_coordinates(value: Any, path: str, errors: List[str]) -> None: ++ if not isinstance(value, str): ++ errors.append( ++ f"{path}: must be ONE string of the form " ++ f"'cx: 0.500, cy: 0.500, w: 1.000, h: 1.000', got {type(value).__name__}" ++ ) ++ return ++ match = COORDINATES_RE.match(value) ++ if match is None: ++ errors.append( ++ f"{path}: {value!r} is not of the form " ++ f"'cx: 0.500, cy: 0.500, w: 1.000, h: 1.000'" ++ ) ++ return ++ for field in COORDINATE_FIELDS: ++ number = float(match.group(field)) ++ if not 0.0 <= number <= 1.0: ++ errors.append(f"{path}: {field}={match.group(field)} is outside [0, 1]") ++ ++ ++def _check_color_specs(value: Any, path: str, errors: List[str]) -> None: ++ if not isinstance(value, list): ++ errors.append( ++ f"{path}: expected a list of hex colors, got {type(value).__name__}" ++ ) ++ return ++ for index, color in enumerate(value): ++ if not isinstance(color, str) or HEX_COLOR_RE.match(color) is None: ++ errors.append( ++ f"{path}[{index}]: {color!r} is not a hex color " ++ f"(expected #RGB, #RGBA, #RRGGBB, or #RRGGBBAA)" ++ ) ++ ++ ++def validate_enhanced_prompt(document: Any) -> List[str]: ++ """Return schema errors for a rewritten prompt; an empty list means valid. ++ ++ Schema demanded by assets/t2i_rewriter_system_prompt.txt: exactly two ++ top-level keys ``canvas_settings`` (exactly ``aspect_ratio``, ++ ``ambient_lighting``, ``image_style``) and ``layers`` (each layer exactly ++ ``description``, ``coordinates``, ``hierarchy_and_relation``, ++ ``color_specs``); ``coordinates`` is a string "cx: 0.500, cy: 0.500, ++ w: 1.000, h: 1.000" with values in [0, 1]; ``color_specs`` is a list of ++ hex colors. ``layers`` must hold at least one visible layer -- an empty ++ list means the rewrite failed even though it is type-correct. ++ """ ++ if not isinstance(document, dict): ++ return [f"top level: expected a JSON object, got {type(document).__name__}"] ++ errors: List[str] = [] ++ _check_exact_keys(document, ("canvas_settings", "layers"), "top level", errors) ++ ++ if "canvas_settings" in document: ++ canvas = document["canvas_settings"] ++ if not isinstance(canvas, dict): ++ errors.append( ++ f"canvas_settings: expected a JSON object, got {type(canvas).__name__}" ++ ) ++ else: ++ _check_exact_keys(canvas, CANVAS_SETTINGS_KEYS, "canvas_settings", errors) ++ for key in CANVAS_SETTINGS_KEYS: ++ if key in canvas: ++ _check_non_empty_string( ++ canvas[key], f"canvas_settings.{key}", errors ++ ) ++ ++ if "layers" in document: ++ layers = document["layers"] ++ if not isinstance(layers, list): ++ errors.append(f"layers: expected a list, got {type(layers).__name__}") ++ elif not layers: ++ errors.append("layers: expected at least one visible layer") ++ else: ++ for index, layer in enumerate(layers): ++ path = f"layers[{index}]" ++ if not isinstance(layer, dict): ++ errors.append( ++ f"{path}: expected a JSON object, got {type(layer).__name__}" ++ ) ++ continue ++ _check_exact_keys(layer, LAYER_KEYS, path, errors) ++ for key in ("description", "hierarchy_and_relation"): ++ if key in layer: ++ _check_non_empty_string(layer[key], f"{path}.{key}", errors) ++ if "coordinates" in layer: ++ _check_coordinates( ++ layer["coordinates"], f"{path}.coordinates", errors ++ ) ++ if "color_specs" in layer: ++ _check_color_specs(layer["color_specs"], f"{path}.color_specs", errors) ++ return errors ++ ++ ++def _chat_completion( ++ base_url: str, ++ model: str, ++ messages: List[Dict[str, str]], ++ *, ++ temperature: float, ++ max_tokens: int, ++ api_key: Optional[str], ++ timeout: float, ++) -> Tuple[str, Optional[str]]: ++ """POST one chat completion; return (content, finish_reason).""" ++ url = base_url.rstrip("/") + "/chat/completions" ++ payload = json.dumps( ++ { ++ "model": model, ++ "messages": messages, ++ "temperature": temperature, ++ "max_tokens": max_tokens, ++ "stream": False, ++ } ++ ).encode("utf-8") ++ headers = {"Content-Type": "application/json"} ++ if api_key: ++ headers["Authorization"] = f"Bearer {api_key}" ++ request = urllib.request.Request(url, data=payload, headers=headers, method="POST") ++ try: ++ with urllib.request.urlopen(request, timeout=timeout) as response: ++ body = response.read().decode("utf-8", errors="replace") ++ except urllib.error.HTTPError as error: ++ detail = error.read().decode("utf-8", errors="replace") ++ raise PromptEnhancementError( ++ f"HTTP {error.code} from {url}: {detail[:2000]}" ++ ) from error ++ except urllib.error.URLError as error: ++ raise PromptEnhancementError(f"cannot reach {url}: {error.reason}") from error ++ except OSError as error: # includes socket timeouts during the read ++ raise PromptEnhancementError(f"request to {url} failed: {error}") from error ++ ++ try: ++ envelope = json.loads(body) ++ choice = envelope["choices"][0] ++ content = choice["message"]["content"] ++ except (json.JSONDecodeError, KeyError, IndexError, TypeError) as error: ++ raise PromptEnhancementError( ++ f"malformed chat completion response from {url}: {body[:500]}" ++ ) from error ++ finish_reason = choice.get("finish_reason") ++ if not isinstance(content, str) or not content.strip(): ++ raise PromptEnhancementError( ++ f"empty completion content from {url} (finish_reason={finish_reason!r})" ++ ) ++ return content, finish_reason ++ ++ ++def enhance( ++ caption: str, ++ base_url: str, ++ model: str, ++ api_key: Optional[str] = None, ++ timeout: float = DEFAULT_TIMEOUT_SECONDS, ++ temperature: float = DEFAULT_TEMPERATURE, ++ max_tokens: int = DEFAULT_MAX_TOKENS, ++) -> Dict[str, Any]: ++ """Return the validated structured rewrite of ``caption``. ++ ++ Sends the verbatim rewriter system prompt plus the caption to ++ ``{base_url}/chat/completions``. On a parse or schema failure, retries ++ exactly once with the validation errors appended to the user turn; if ++ that also fails, raises PromptEnhancementError carrying the errors. ++ """ ++ system_prompt = load_system_prompt() ++ messages = [ ++ {"role": "system", "content": system_prompt}, ++ {"role": "user", "content": caption}, ++ ] ++ request_kwargs = { ++ "temperature": temperature, ++ "max_tokens": max_tokens, ++ "api_key": api_key, ++ "timeout": timeout, ++ } ++ errors: List[str] = [] ++ content = "" ++ for attempt in (1, 2): ++ content, finish_reason = _chat_completion( ++ base_url, model, messages, **request_kwargs ++ ) ++ document: Optional[Dict[str, Any]] = None ++ try: ++ document = extract_json_object(content) ++ except ValueError as error: ++ errors = [str(error)] ++ if document is not None: ++ errors = validate_enhanced_prompt(document) ++ if not errors: ++ assert document is not None # errors empty implies extraction succeeded ++ return document ++ if finish_reason == "length": ++ errors.append( ++ "the reply was cut off (finish_reason='length'): the complete " ++ f"JSON object must fit within max_tokens={max_tokens}" ++ ) ++ print(f"pe_ling: attempt {attempt}/2 failed validation:", file=sys.stderr) ++ for error in errors: ++ print(f"pe_ling: - {error}", file=sys.stderr) ++ if attempt == 1: ++ retry_content = ( ++ f"{caption}\n\n" ++ "Your previous reply failed schema validation:\n" ++ + "".join(f"- {error}\n" for error in errors) ++ + "\n" ++ + REPAIR_INSTRUCTION ++ ) ++ messages = [ ++ {"role": "system", "content": system_prompt}, ++ {"role": "user", "content": retry_content}, ++ ] ++ raise PromptEnhancementError( ++ "prompt enhancement failed schema validation after 2 attempts:\n" ++ + "".join(f" - {error}\n" for error in errors).rstrip(), ++ errors=errors, ++ reply=content, ++ ) ++ ++ ++def main() -> None: ++ parser = argparse.ArgumentParser( ++ description=( ++ "Enhance a Ming-Image text-to-image caption into the validated " ++ "structured JSON prompt via an OpenAI-compatible Ling-3.0-flash-VL " ++ "endpoint." ++ ) ++ ) ++ parser.add_argument("caption", help="free-form design caption to enhance") ++ parser.add_argument( ++ "--out", ++ type=Path, ++ help="write the validated JSON here (default: stdout, summary on stderr)", ++ ) ++ parser.add_argument( ++ "--base-url", ++ default=DEFAULT_BASE_URL, ++ help=f"OpenAI-compatible base URL (default: {DEFAULT_BASE_URL})", ++ ) ++ parser.add_argument( ++ "--model", ++ default=DEFAULT_MODEL, ++ help=f"chat model id served at the endpoint (default: {DEFAULT_MODEL})", ++ ) ++ parser.add_argument( ++ "--api-key", ++ default=os.environ.get(API_KEY_ENV), ++ help=f"Bearer token for gated endpoints; defaults to ${API_KEY_ENV} when set", ++ ) ++ parser.add_argument( ++ "--timeout", ++ type=float, ++ default=DEFAULT_TIMEOUT_SECONDS, ++ help=f"per-request timeout in seconds (default: {DEFAULT_TIMEOUT_SECONDS})", ++ ) ++ parser.add_argument( ++ "--temperature", ++ type=float, ++ default=DEFAULT_TEMPERATURE, ++ help=f"sampling temperature (default: {DEFAULT_TEMPERATURE})", ++ ) ++ parser.add_argument( ++ "--max-tokens", ++ type=int, ++ default=DEFAULT_MAX_TOKENS, ++ help=f"completion token budget (default: {DEFAULT_MAX_TOKENS})", ++ ) ++ args = parser.parse_args() ++ ++ started = time.perf_counter() ++ try: ++ document = enhance( ++ args.caption, ++ args.base_url, ++ args.model, ++ api_key=args.api_key, ++ timeout=args.timeout, ++ temperature=args.temperature, ++ max_tokens=args.max_tokens, ++ ) ++ except PromptEnhancementError as error: ++ print(f"pe_ling: {error}", file=sys.stderr) ++ raise SystemExit(1) ++ elapsed = time.perf_counter() - started ++ layer_count = len(document["layers"]) ++ payload = json.dumps(document, indent=2, ensure_ascii=False) + "\n" ++ if args.out is not None: ++ args.out.parent.mkdir(parents=True, exist_ok=True) ++ args.out.write_text(payload, encoding="utf-8") ++ print(f"pe_ling: {elapsed:.1f}s, {layer_count} layer(s) -> {args.out}") ++ else: ++ sys.stdout.write(payload) ++ print(f"pe_ling: {elapsed:.1f}s, {layer_count} layer(s)", file=sys.stderr) ++ ++ ++if __name__ == "__main__": ++ main() +diff --git a/quant/__init__.py b/quant/__init__.py +new file mode 100644 +index 0000000..2e60ff4 +--- /dev/null ++++ b/quant/__init__.py +@@ -0,0 +1 @@ ++"""Weight-only INT8 for the Ming-Image MLLM: quantize_stream.py writes it, load_int8.py loads it.""" +diff --git a/quant/int8_linear.py b/quant/int8_linear.py +new file mode 100644 +index 0000000..5939ea5 +--- /dev/null ++++ b/quant/int8_linear.py +@@ -0,0 +1,221 @@ ++"""Weight-only symmetric per-output-channel INT8 linear. ++ ++Scales stay float32 across dtype casts. ``module.to(dtype=torch.bfloat16)`` ++(and ``.bfloat16()`` / ``.half()`` / ``.to(device, dtype)``) must not touch them; ++device moves still do. The int8 weight codes are likewise dtype-stable. ++""" ++ ++from __future__ import annotations ++ ++import torch ++import torch.nn.functional as F ++from torch import nn ++ ++# Leaf names of Linear modules whose 2-D weights are quantized. ++# Exact match: the routers are `gate` / `image_gate` / `audio_gate`, NOT `gate_proj`. ++QUANT_LEAVES = frozenset( ++ {"query_key_value", "dense", "gate_proj", "up_proj", "down_proj"} ++) ++ ++QUANT_RULE = ( ++ "Quantize ONLY 2-D .weight tensors under model.model.layers. whose owning " ++ "module's leaf name is exactly one of query_key_value, dense, gate_proj, " ++ "up_proj, down_proj. Everything else stays byte-identical BF16: embeddings, " ++ "lm_head, all norms, the vision tower, linear_proj, and the three routers " ++ "(modules named gate, image_gate, audio_gate — leaf match, not a substring). " ++ "Per-output-channel symmetric: scale = absmax/127, " ++ "q = clamp(round(w/scale), -127, 127). All-zero rows: scale 1.0, q 0." ++) ++ ++# Real checkpoint keys look like `model.model.layers.N...`. A bare ++# `layers.N...` name is the same stack with the root prefix omitted (tests). ++_DECODER_LAYER_PREFIXES = ((), ("model", "model")) ++ ++ ++def _weight_leaf(tensor_name: str) -> str | None: ++ """Owning module's leaf name if `tensor_name` ends in `.weight`, else None.""" ++ if not isinstance(tensor_name, str) or not tensor_name.endswith(".weight"): ++ return None ++ module = tensor_name[: -len(".weight")] ++ if not module: ++ return None ++ return module.rsplit(".", 1)[-1] ++ ++ ++def _under_decoder_layers(tensor_name: str) -> bool: ++ """True when the tensor lives under the MLLM decoder `model.model.layers` stack. ++ ++ `layers` must be its own path component, followed by a layer index. The ++ components before it must be empty or end in `model.model` — so a vision ++ tower that happens to contain the substring "layers" is not selected, and ++ `gate` is never selected just because `gate_proj` contains those letters. ++ """ ++ parts = tensor_name.split(".") ++ for i, part in enumerate(parts): ++ if part != "layers": ++ continue ++ if i + 1 >= len(parts) or not parts[i + 1].isdigit(): ++ continue ++ prefix = tuple(parts[:i]) ++ if prefix in _DECODER_LAYER_PREFIXES: ++ return True ++ if len(prefix) >= 2 and prefix[-2:] == ("model", "model"): ++ return True ++ return False ++ ++ ++def quant_rule_leaf(tensor_name: str) -> str | None: ++ """Leaf name if the name matches the quantize rule, ignoring rank. ++ ++ Returns None when the tensor is not a candidate. A candidate whose rank is ++ not 2 is a hard error for the stream (see quantize_stream); ``is_quantizable`` ++ itself returns False for that case. ++ """ ++ leaf = _weight_leaf(tensor_name) ++ if leaf not in QUANT_LEAVES: ++ return None ++ if not _under_decoder_layers(tensor_name): ++ return None ++ return leaf ++ ++ ++def is_quantizable(tensor_name: str, shape) -> bool: ++ """True only for 2-D quantize-rule weights. See ``QUANT_RULE``.""" ++ if quant_rule_leaf(tensor_name) is None: ++ return False ++ try: ++ rank = len(shape) ++ except TypeError: ++ return False ++ return rank == 2 ++ ++ ++def quantize_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ++ """Per-output-channel symmetric int8. ++ ++ ``scale = absmax(row) / 127``, ``q = clamp(round(w / scale), -127, 127)``. ++ An all-zero row gets scale 1.0 and q 0 (no div-by-zero, no NaN/Inf). ++ """ ++ if weight.ndim != 2: ++ raise ValueError( ++ f"quantize_weight expects a 2-D weight, got shape {tuple(weight.shape)}" ++ ) ++ wf = weight.detach().to(dtype=torch.float32) ++ absmax = wf.abs().amax(dim=1) ++ scale = absmax / 127.0 ++ zero = scale == 0 ++ # All-zero rows would divide by 0. Force scale 1 and q 0 instead of NaN. ++ scale = torch.where(zero, torch.ones_like(scale), scale) ++ q = torch.round(wf / scale[:, None]).clamp(-127, 127).to(dtype=torch.int8) ++ q = torch.where(zero[:, None], torch.zeros_like(q), q) ++ return q.contiguous(), scale.to(dtype=torch.float32).contiguous() ++ ++ ++def _scale_name(weight_name: str) -> str: ++ if not weight_name.endswith(".weight"): ++ raise ValueError(f"not a weight tensor name: {weight_name}") ++ return weight_name[: -len("weight")] + "scale" ++ ++ ++class Int8Linear(nn.Module): ++ """``F.linear`` on a weight dequantized from int8 + per-row float32 scale. ++ ++ ``weight`` is int8 ``[out, in]``, ``scale`` is float32 ``[out]``, ``bias`` ++ (optional) keeps the source dtype. All three are buffers. ++ """ ++ ++ def __init__(self, weight: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None): ++ super().__init__() ++ if weight.dtype != torch.int8 or weight.ndim != 2: ++ raise ValueError( ++ f"weight must be int8 [out, in], got dtype={weight.dtype} shape={tuple(weight.shape)}" ++ ) ++ if scale.dtype != torch.float32 or tuple(scale.shape) != (weight.shape[0],): ++ raise ValueError( ++ f"scale must be float32 [{weight.shape[0]}], got dtype={scale.dtype} shape={tuple(scale.shape)}" ++ ) ++ if bias is not None: ++ if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: ++ raise ValueError( ++ f"bias must be [{weight.shape[0]}], got shape={tuple(bias.shape)}" ++ ) ++ self.in_features = int(weight.shape[1]) ++ self.out_features = int(weight.shape[0]) ++ self.register_buffer("weight", weight) ++ self.register_buffer("scale", scale) ++ self.register_buffer("bias", bias) ++ ++ def _apply(self, fn, *args, **kwargs): ++ # Pull dtype-stable buffers out before Module._apply. Putting them back ++ # with only a device move (never fn's dtype cast) keeps scale float32 ++ # and weight int8. Bias is left in the dict so it follows the cast. ++ saved: dict[str, torch.Tensor] = {} ++ for name in ("weight", "scale"): ++ buf = self._buffers.get(name, None) ++ if buf is not None: ++ saved[name] = buf ++ self._buffers[name] = None ++ try: ++ out = super()._apply(fn, *args, **kwargs) ++ finally: ++ for name, buf in saved.items(): ++ self._buffers[name] = _move_device_keep_dtype(buf, fn) ++ return out ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ # One dequant in fp32, one cast to the activation dtype, then linear. ++ w = (self.weight.float() * self.scale[:, None]).to(dtype=x.dtype) ++ return F.linear(x, w, self.bias) ++ ++ @classmethod ++ def from_linear(cls, linear: nn.Linear) -> "Int8Linear": ++ if not isinstance(linear, nn.Linear): ++ raise TypeError(f"from_linear expects nn.Linear, got {type(linear).__name__}") ++ q, scale = quantize_weight(linear.weight.data) ++ if linear.bias is None: ++ bias = None ++ else: ++ bias = linear.bias.detach().clone() ++ return cls(q, scale, bias) ++ ++ @classmethod ++ def shell( ++ cls, ++ in_features: int, ++ out_features: int, ++ bias: bool, ++ bias_dtype: torch.dtype, ++ device, ++ ) -> "Int8Linear": ++ """Empty buffers (for ``meta``). Does not read or write any weight values.""" ++ dev = torch.device(device) if not isinstance(device, torch.device) else device ++ weight = torch.empty((out_features, in_features), dtype=torch.int8, device=dev) ++ scale = torch.empty((out_features,), dtype=torch.float32, device=dev) ++ if bias: ++ bias_t: torch.Tensor | None = torch.empty( ++ (out_features,), dtype=bias_dtype, device=dev ++ ) ++ else: ++ bias_t = None ++ return cls(weight, scale, bias_t) ++ ++ def extra_repr(self) -> str: ++ return ( ++ f"in_features={self.in_features}, out_features={self.out_features}, " ++ f"bias={self.bias is not None}" ++ ) ++ ++ ++def _move_device_keep_dtype(buf: torch.Tensor, fn) -> torch.Tensor: ++ """Apply only the device change implied by ``fn``, preserving ``buf``'s dtype and values. ++ ++ Probed with a 0-element tensor so a dtype cast cannot round the real scale. ++ """ ++ try: ++ probe = torch.empty((), dtype=buf.dtype, device=buf.device) ++ moved = fn(probe) ++ except Exception: ++ return buf ++ if not torch.is_tensor(moved) or moved.device == buf.device: ++ return buf ++ return buf.to(device=moved.device) +diff --git a/quant/load_int8.py b/quant/load_int8.py +new file mode 100644 +index 0000000..1117892 +--- /dev/null ++++ b/quant/load_int8.py +@@ -0,0 +1,172 @@ ++"""Load a streamed INT8 Ming MLLM checkpoint onto a meta-initialized model. ++ ++``model`` must already exist with parameters on ``meta`` (for example under ++``accelerate.init_empty_weights()``). Quantized modules listed in ++``int8_manifest.json`` are swapped from ``nn.Linear`` to ``Int8Linear.shell`` ++before the shards are assigned in. ++""" ++ ++from __future__ import annotations ++ ++import json ++from pathlib import Path ++ ++import torch ++from safetensors.torch import load_file ++from torch import nn ++ ++try: # imported as the `quant` package (modeling_bailingmm2.py) ++ from .int8_linear import Int8Linear ++except ImportError: # run from inside quant/ (CLI, tests) ++ from int8_linear import Int8Linear ++ ++MANIFEST_NAME = "int8_manifest.json" ++INDEX_NAME = "model.safetensors.index.json" ++ ++ ++def load_int8_mllm_(model: nn.Module, int8_dir, device) -> dict: ++ """Swap quantize-rule linears for INT8 shells and assign shard tensors. ++ ++ Returns ``{"modules_swapped", "tensors_loaded", "bytes_loaded"}``. ++ Raises ``RuntimeError`` on a bad manifest, a module that is not an ++ ``nn.Linear``, an unexpected checkpoint key, or any parameter / persistent ++ buffer still on ``meta``. Non-persistent buffers (rotary ``inv_freq``) may ++ stay on CPU; the caller moves the model afterwards. ++ """ ++ int8_dir = Path(int8_dir) ++ dev = torch.device(device) if not isinstance(device, torch.device) else device ++ manifest_path = int8_dir / MANIFEST_NAME ++ if not manifest_path.is_file(): ++ raise RuntimeError(f"missing int8 manifest: {manifest_path}") ++ manifest = json.loads(manifest_path.read_text(encoding="utf-8")) ++ if manifest.get("format") != "ming-int8-wo-v1": ++ raise RuntimeError( ++ f"unsupported int8 manifest format: {manifest.get('format')!r} ({manifest_path})" ++ ) ++ module_names = manifest.get("quantized_modules") ++ if not isinstance(module_names, list) or not all(isinstance(n, str) for n in module_names): ++ raise RuntimeError(f"{manifest_path} quantized_modules is not a list of strings") ++ ++ swapped = _swap_linears(model, module_names) ++ ++ index_path = int8_dir / INDEX_NAME ++ if not index_path.is_file(): ++ raise RuntimeError(f"missing index: {index_path}") ++ index = json.loads(index_path.read_text(encoding="utf-8")) ++ weight_map = index.get("weight_map") ++ if not isinstance(weight_map, dict) or not weight_map: ++ raise RuntimeError(f"{index_path} has no weight_map") ++ ++ shard_names: list[str] = [] ++ seen: set[str] = set() ++ for shard in weight_map.values(): ++ if shard not in seen: ++ seen.add(shard) ++ shard_names.append(shard) ++ ++ tensors_loaded = 0 ++ bytes_loaded = 0 ++ unexpected: list[str] = [] ++ for shard in shard_names: ++ rel = Path(shard) ++ if rel.is_absolute() or ".." in rel.parts: ++ raise RuntimeError(f"unsafe shard path in index: {shard}") ++ path = int8_dir / rel ++ if not path.is_file(): ++ raise RuntimeError(f"missing shard: {path}") ++ sd = load_file(str(path), device=str(dev)) ++ for tensor in sd.values(): ++ tensors_loaded += 1 ++ bytes_loaded += tensor.numel() * tensor.element_size() ++ incompatible = model.load_state_dict(sd, strict=False, assign=True) ++ unexpected.extend(incompatible.unexpected_keys) ++ del sd ++ ++ if unexpected: ++ listed = "\n".join(f" {key}" for key in unexpected) ++ raise RuntimeError( ++ f"unexpected keys in checkpoint (not present on the model):\n{listed}" ++ ) ++ ++ _assert_loaded(model, module_names, dev) ++ return { ++ "modules_swapped": swapped, ++ "tensors_loaded": tensors_loaded, ++ "bytes_loaded": bytes_loaded, ++ } ++ ++ ++def _swap_linears(model: nn.Module, module_names: list[str]) -> int: ++ for name in module_names: ++ try: ++ linear = model.get_submodule(name) ++ except AttributeError as exc: ++ raise RuntimeError(f"manifest module not found on model: {name}") from exc ++ if not isinstance(linear, nn.Linear): ++ raise RuntimeError( ++ f"{name} is {type(linear).__name__}, expected nn.Linear " ++ "(refusing to swap a router or other non-linear)" ++ ) ++ parent_name, _, leaf = name.rpartition(".") ++ if not leaf: ++ raise RuntimeError(f"cannot place shell for {name}") ++ parent = model.get_submodule(parent_name) if parent_name else model ++ has_bias = linear.bias is not None ++ bias_dtype = linear.bias.dtype if has_bias else torch.float32 ++ shell = Int8Linear.shell( ++ in_features=linear.in_features, ++ out_features=linear.out_features, ++ bias=has_bias, ++ bias_dtype=bias_dtype, ++ device="meta", ++ ) ++ setattr(parent, leaf, shell) ++ return len(module_names) ++ ++ ++def _assert_loaded(model: nn.Module, module_names: list[str], dev: torch.device) -> None: ++ offenders: list[str] = [] ++ for name, param in model.named_parameters(remove_duplicate=False): ++ if param is not None and param.device.type == "meta": ++ offenders.append(f"parameter {name} dtype={param.dtype} device={param.device}") ++ for mod_name, mod in model.named_modules(): ++ nonpersist = getattr(mod, "_non_persistent_buffers_set", set()) ++ for buf_name, buf in mod._buffers.items(): ++ if buf is None: ++ continue ++ full = f"{mod_name}.{buf_name}" if mod_name else buf_name ++ if buf.device.type != "meta": ++ # Non-persistent buffers (rotary inv_freq) are not in the ++ # checkpoint. accelerate leaves them on CPU; that is not an error. ++ continue ++ if buf_name in nonpersist: ++ offenders.append( ++ f"non-persistent buffer {full} dtype={buf.dtype} device={buf.device}" ++ ) ++ else: ++ offenders.append(f"buffer {full} dtype={buf.dtype} device={buf.device}") ++ if offenders: ++ listed = "\n".join(f" {line}" for line in offenders) ++ raise RuntimeError(f"tensors still on meta after load:\n{listed}") ++ ++ for name in module_names: ++ mod = model.get_submodule(name) ++ if not isinstance(mod, Int8Linear): ++ raise RuntimeError(f"{name} was not swapped to Int8Linear") ++ if mod.weight is None or mod.weight.dtype != torch.int8: ++ raise RuntimeError(f"{name}.weight is not int8 after load") ++ if mod.scale is None or mod.scale.dtype != torch.float32: ++ raise RuntimeError(f"{name}.scale is not float32 after load") ++ if mod.weight.device.type == "meta" or mod.scale.device.type == "meta": ++ raise RuntimeError(f"{name} still has meta tensors after load") ++ if mod.weight.device != dev or mod.scale.device != dev: ++ raise RuntimeError( ++ f"{name} loaded on weight={mod.weight.device} scale={mod.scale.device}, " ++ f"expected {dev}" ++ ) ++ if tuple(mod.scale.shape) != (mod.out_features,): ++ raise RuntimeError( ++ f"{name}.scale shape {tuple(mod.scale.shape)} != ({mod.out_features},)" ++ ) ++ if mod.bias is not None and mod.bias.device != dev: ++ raise RuntimeError(f"{name}.bias is on {mod.bias.device}, expected {dev}") +diff --git a/quant/quantize_stream.py b/quant/quantize_stream.py +new file mode 100644 +index 0000000..580c2f2 +--- /dev/null ++++ b/quant/quantize_stream.py +@@ -0,0 +1,556 @@ ++"""Stream a Ming MLLM directory to weight-only INT8 shards. ++ ++Never builds the model: it buffers at most one output shard (<= 5 GB) of tensors at a time. Measured on ++the real 34.0 GB checkpoint (AMD Strix Halo, 2026-09-23): 266 s wall, peak RSS 17.8 GiB. ++CLI: ``python quantize_stream.py SRC_MLLM_DIR DST_DIR [--exclude MODULE_REGEX]`` (matching modules stay BF16). ++""" ++ ++from __future__ import annotations ++ ++import json ++import math ++import re ++import os ++import shutil ++import sys ++from dataclasses import dataclass ++from pathlib import Path ++ ++import torch ++from safetensors import safe_open ++from safetensors.torch import save_file ++ ++try: # imported as the `quant` package ++ from .int8_linear import QUANT_RULE, quant_rule_leaf, quantize_weight ++except ImportError: # run as a script: python quant/quantize_stream.py SRC DST ++ from int8_linear import QUANT_RULE, quant_rule_leaf, quantize_weight ++ ++# Decimal GB, same unit Hugging Face uses for max_shard_size="5GB". ++MAX_SHARD_BYTES = 5 * 10**9 ++ ++_DTYPE_BYTES = { ++ "BOOL": 1, ++ "U8": 1, ++ "I8": 1, ++ "F8_E4M3": 1, ++ "F8_E5M2": 1, ++ "F8_E8M0": 1, ++ "U16": 2, ++ "I16": 2, ++ "F16": 2, ++ "BF16": 2, ++ "U32": 4, ++ "I32": 4, ++ "F32": 4, ++ "U64": 8, ++ "I64": 8, ++ "F64": 8, ++} ++ ++INDEX_NAME = "model.safetensors.index.json" ++MANIFEST_NAME = "int8_manifest.json" ++ ++ ++class QuantizeError(Exception): ++ """User-facing checkpoint error. main() prints it and returns 1.""" ++ ++ ++def _die(msg: str) -> None: ++ raise QuantizeError(msg) ++ ++ ++def _normalize_dtype(dtype_name) -> str: ++ text = str(dtype_name).upper() ++ if "." in text: ++ text = text.rsplit(".", 1)[-1] ++ aliases = { ++ "BFLOAT16": "BF16", ++ "FLOAT16": "F16", ++ "FLOAT32": "F32", ++ "FLOAT64": "F64", ++ "FLOAT8_E4M3FN": "F8_E4M3", ++ "FLOAT8_E5M2": "F8_E5M2", ++ "INT8": "I8", ++ "INT16": "I16", ++ "INT32": "I32", ++ "INT64": "I64", ++ "UINT8": "U8", ++ } ++ return aliases.get(text, text) ++ ++ ++def _dtype_nbytes(dtype_name: str) -> int: ++ try: ++ return _DTYPE_BYTES[dtype_name] ++ except KeyError: ++ _die(f"unsupported safetensors dtype {dtype_name!r}") ++ raise # unreachable; satisfies type checkers ++ ++ ++def _numel(shape: tuple[int, ...]) -> int: ++ n = 1 ++ for d in shape: ++ n *= int(d) ++ return n ++ ++ ++def _load_index(path: Path) -> dict: ++ if not path.is_file(): ++ _die(f"missing index: {path}") ++ ++ def _pairs(pairs): ++ keys = [k for k, _ in pairs] ++ dupes = sorted({k for k in keys if keys.count(k) > 1}) ++ if dupes: ++ _die(f"duplicate key(s) in {path}: {dupes}") ++ return dict(pairs) ++ ++ try: ++ raw = path.read_text(encoding="utf-8") ++ index = json.loads(raw, object_pairs_hook=_pairs) ++ except QuantizeError: ++ raise ++ except (OSError, json.JSONDecodeError) as exc: ++ _die(f"cannot read index {path}: {exc}") ++ if not isinstance(index, dict) or not isinstance(index.get("weight_map"), dict): ++ _die(f"index {path} has no weight_map object") ++ if not index["weight_map"]: ++ _die(f"index {path} weight_map is empty") ++ return index ++ ++ ++def _check_dst_clean(dst: Path) -> None: ++ if not dst.exists(): ++ return ++ if not dst.is_dir(): ++ _die(f"destination is not a directory: {dst}") ++ found = sorted(p.relative_to(dst).as_posix() for p in dst.rglob("*.safetensors")) ++ if found: ++ _die(f"destination already contains safetensors: {found}") ++ ++ ++def _reject_nested(src: Path, dst: Path) -> None: ++ src_r = src.resolve() ++ dst_r = dst.resolve() ++ if src_r == dst_r or src_r in dst_r.parents or dst_r in src_r.parents: ++ _die(f"SRC and DST must be distinct and not nested: {src} vs {dst}") ++ ++ ++def _shard_path(src: Path, shard_name: str) -> Path: ++ rel = Path(shard_name) ++ if rel.is_absolute() or ".." in rel.parts: ++ _die(f"unsafe shard path in index: {shard_name}") ++ path = src / rel ++ if not path.is_file(): ++ _die(f"index lists missing shard: {shard_name}") ++ return path ++ ++ ++@dataclass ++class Item: ++ src_shard: str ++ name: str ++ kind: str # "copy" or "quant" ++ shape: tuple[int, ...] ++ src_dtype: str ++ src_bytes: int ++ out_bytes: int ++ group: int = -1 ++ ++ ++def _scale_name(weight_name: str) -> str: ++ return weight_name[: -len("weight")] + "scale" ++ ++ ++def _plan(src: Path, index: dict, exclude: str | None = None) -> list[Item]: ++ """Metadata-only pass. Reads shapes and dtypes, not tensor bodies.""" ++ weight_map: dict[str, str] = index["weight_map"] ++ shard_order: list[str] = [] ++ seen_shards: set[str] = set() ++ for shard in weight_map.values(): ++ if shard not in seen_shards: ++ seen_shards.add(shard) ++ shard_order.append(shard) ++ ++ index_names_by_shard: dict[str, set[str]] = {s: set() for s in shard_order} ++ for name, shard in weight_map.items(): ++ if shard not in index_names_by_shard: ++ _die(f"weight_map value {shard!r} for {name} was not collected") ++ index_names_by_shard[shard].add(name) ++ ++ items: list[Item] = [] ++ seen_names: dict[str, str] = {} ++ for shard in shard_order: ++ path = _shard_path(src, shard) ++ with safe_open(str(path), framework="pt", device="cpu") as handle: ++ file_names = list(handle.keys()) ++ file_set = set(file_names) ++ if len(file_set) != len(file_names): ++ _die(f"shard {shard} header lists a tensor name twice") ++ missing = sorted(index_names_by_shard[shard] - file_set) ++ extra = sorted(file_set - index_names_by_shard[shard]) ++ if missing: ++ _die(f"index lists tensors missing from {shard}: {missing}") ++ if extra: ++ _die(f"{shard} contains tensors absent from the index: {extra}") ++ for name in file_names: ++ if name in seen_names: ++ _die( ++ f"tensor name appears twice: {name} " ++ f"({seen_names[name]} and {shard})" ++ ) ++ seen_names[name] = shard ++ sl = handle.get_slice(name) ++ if not hasattr(sl, "get_dtype") or not hasattr(sl, "get_shape"): ++ _die( ++ "safetensors safe_open slice is missing get_shape/get_dtype; " ++ "cannot plan shards without loading tensor bodies" ++ ) ++ shape = tuple(int(d) for d in sl.get_shape()) ++ dtype_name = _normalize_dtype(sl.get_dtype()) ++ src_bytes = _numel(shape) * _dtype_nbytes(dtype_name) ++ leaf = quant_rule_leaf(name) ++ if leaf is not None and exclude and re.search(exclude, name[: -len(".weight")]): ++ leaf = None # kept BF16 by --exclude ++ if leaf is not None and len(shape) != 2: ++ _die( ++ f"tensor {name} matches the quantize rule but is not 2-D " ++ f"(shape={list(shape)}, dtype={dtype_name})" ++ ) ++ if leaf is not None: ++ out_bytes = _numel(shape) * 1 + shape[0] * 4 # int8 weight + fp32 scale ++ items.append( ++ Item(shard, name, "quant", shape, dtype_name, src_bytes, out_bytes) ++ ) ++ else: ++ items.append( ++ Item(shard, name, "copy", shape, dtype_name, src_bytes, src_bytes) ++ ) ++ ++ index_names = set(weight_map) ++ planned = {it.name for it in items} ++ if planned != index_names: ++ _die( ++ "index / shard mismatch after scan: " ++ f"only_in_index={sorted(index_names - planned)[:8]} " ++ f"only_in_shards={sorted(planned - index_names)[:8]}" ++ ) ++ ++ produced = set(planned) ++ for it in items: ++ if it.kind != "quant": ++ continue ++ sname = _scale_name(it.name) ++ if sname in produced: ++ _die(f"scale name collides with an existing tensor: {sname}") ++ produced.add(sname) ++ return items ++ ++ ++def _assign_groups(items: list[Item], max_shard_bytes: int) -> list[list[Item]]: ++ if max_shard_bytes <= 0: ++ _die(f"max_shard_bytes must be positive, got {max_shard_bytes}") ++ groups: list[list[Item]] = [] ++ cur: list[Item] = [] ++ cur_bytes = 0 ++ for it in items: ++ if cur and cur_bytes + it.out_bytes > max_shard_bytes: ++ groups.append(cur) ++ cur = [] ++ cur_bytes = 0 ++ if cur_bytes == 0 and it.out_bytes > max_shard_bytes: ++ print( ++ f"warning: {it.name} contributes {it.out_bytes} bytes, " ++ f"over the {max_shard_bytes}-byte shard target; writing it alone", ++ file=sys.stderr, ++ flush=True, ++ ) ++ it.group = len(groups) ++ cur.append(it) ++ cur_bytes += it.out_bytes ++ if cur: ++ groups.append(cur) ++ return groups ++ ++ ++def _relative_frobenius(weight: torch.Tensor, q: torch.Tensor, scale: torch.Tensor) -> float: ++ w = weight.detach().to(dtype=torch.float64) ++ deq = q.detach().to(dtype=torch.float64) * scale.detach().to(dtype=torch.float64)[:, None] ++ denom = torch.linalg.matrix_norm(w, ord="fro") ++ numer = torch.linalg.matrix_norm(w - deq, ord="fro") ++ d = denom.item() ++ n = numer.item() ++ if d == 0.0: ++ return 0.0 if n == 0.0 else math.inf ++ return n / d ++ ++ ++def _percentile_linear(values: list[float], pct: float) -> float: ++ """NumPy-style linear percentile. Empty → 0.""" ++ if not values: ++ return 0.0 ++ ordered = sorted(values) ++ if len(ordered) == 1: ++ return ordered[0] ++ rank = (len(ordered) - 1) * (pct / 100.0) ++ lo = math.floor(rank) ++ hi = math.ceil(rank) ++ if lo == hi: ++ return ordered[lo] ++ w = rank - lo ++ return ordered[lo] * (1.0 - w) + ordered[hi] * w ++ ++ ++def _copy_sidecars(src: Path, dst: Path) -> list[str]: ++ copied: list[str] = [] ++ for dirpath, _dirnames, filenames in os.walk(src): ++ rel = Path(dirpath).relative_to(src) ++ out_dir = dst / rel ++ out_dir.mkdir(parents=True, exist_ok=True) ++ for filename in filenames: ++ if filename.endswith(".safetensors"): ++ continue ++ if filename == INDEX_NAME and rel == Path("."): ++ continue ++ src_file = Path(dirpath) / filename ++ dst_file = out_dir / filename ++ shutil.copy2(src_file, dst_file) ++ copied.append((rel / filename).as_posix()) ++ return copied ++ ++ ++def _write_shards( ++ src: Path, ++ dst: Path, ++ items: list[Item], ++ groups: list[list[Item]], ++) -> tuple[dict[str, str], int, int, list[tuple[str, float]], list[Path]]: ++ n_out = len(groups) ++ weight_map: dict[str, str] = {} ++ bytes_in = 0 ++ bytes_out = 0 ++ errors: list[tuple[str, float]] = [] ++ written: list[Path] = [] ++ ++ n_src = len({it.src_shard for it in items}) ++ src_seen = 0 ++ open_name: str | None = None ++ handle = None ++ buf: dict[str, torch.Tensor] = {} ++ buf_q = 0 ++ buf_c = 0 ++ current_group = 0 ++ ++ def flush() -> None: ++ nonlocal buf, buf_q, buf_c, current_group ++ if not buf: ++ return ++ fname = f"model-{current_group + 1:05d}-of-{n_out:05d}.safetensors" ++ path = dst / fname ++ for key, tensor in buf.items(): ++ if not tensor.is_contiguous(): ++ buf[key] = tensor.contiguous() ++ save_file(buf, str(path)) ++ shard_bytes = 0 ++ for key, tensor in buf.items(): ++ weight_map[key] = fname ++ shard_bytes += tensor.numel() * tensor.element_size() ++ written.append(path) ++ print( ++ f"wrote {fname}: tensors={len(buf)} quantized={buf_q} copied={buf_c} " ++ f"bytes={shard_bytes}", ++ flush=True, ++ ) ++ buf = {} ++ buf_q = 0 ++ buf_c = 0 ++ current_group += 1 ++ ++ try: ++ for it in items: ++ if it.src_shard != open_name: ++ if handle is not None: ++ handle.__exit__(None, None, None) ++ handle = None ++ path = _shard_path(src, it.src_shard) ++ handle = safe_open(str(path), framework="pt", device="cpu") ++ handle.__enter__() ++ open_name = it.src_shard ++ src_seen += 1 ++ n_here = sum(1 for x in items if x.src_shard == it.src_shard) ++ print( ++ f"reading source shard {src_seen}/{n_src} {it.src_shard} ({n_here} tensors)", ++ flush=True, ++ ) ++ assert handle is not None ++ tensor = handle.get_tensor(it.name) ++ got = tensor.numel() * tensor.element_size() ++ if got != it.src_bytes: ++ _die( ++ f"{it.name} byte size {got} != planned {it.src_bytes} " ++ f"(dtype={tensor.dtype}, shape={tuple(tensor.shape)})" ++ ) ++ bytes_in += got ++ if it.kind == "quant": ++ if not tensor.is_floating_point(): ++ _die( ++ f"{it.name} matches the quantize rule but dtype is {tensor.dtype}, " ++ "expected a floating dtype" ++ ) ++ if tuple(tensor.shape) != it.shape: ++ _die(f"{it.name} shape changed between passes: {tuple(tensor.shape)} vs {it.shape}") ++ q, scale = quantize_weight(tensor) ++ err = _relative_frobenius(tensor, q, scale) ++ if math.isnan(err) or math.isinf(err): ++ _die(f"non-finite relative error for {it.name}: {err}") ++ errors.append((it.name, err)) ++ del tensor ++ sname = _scale_name(it.name) ++ buf[it.name] = q ++ buf[sname] = scale ++ produced = q.numel() * q.element_size() + scale.numel() * scale.element_size() ++ if produced != it.out_bytes: ++ _die(f"{it.name} output bytes {produced} != planned {it.out_bytes}") ++ buf_q += 1 ++ else: ++ if not tensor.is_contiguous(): ++ tensor = tensor.contiguous() ++ buf[it.name] = tensor ++ buf_c += 1 ++ bytes_out += it.out_bytes ++ # Flush when this item closes its planned output shard. ++ group_items = groups[it.group] ++ if it is group_items[-1]: ++ flush() ++ finally: ++ if handle is not None: ++ handle.__exit__(None, None, None) ++ ++ if buf: ++ _die("internal error: output buffer not flushed") ++ if current_group != n_out: ++ _die(f"internal error: wrote {current_group} shards, planned {n_out}") ++ return weight_map, bytes_in, bytes_out, errors, written ++ ++ ++def _summary( ++ errors: list[tuple[str, float]], ++ n_quant: int, ++ n_copy: int, ++ bytes_in: int, ++ bytes_out: int, ++) -> dict: ++ vals = [e for _, e in errors] ++ if errors: ++ worst_name, worst_err = min( ++ errors, ++ key=lambda pair: (-pair[1], pair[0]), ++ ) ++ else: ++ worst_name, worst_err = None, 0.0 ++ mean = (sum(vals) / len(vals)) if vals else 0.0 ++ return { ++ "tensors_quantized": n_quant, ++ "tensors_copied": n_copy, ++ "bytes_in": bytes_in, ++ "bytes_out": bytes_out, ++ "mean_relative_error": mean, ++ "p99_relative_error": _percentile_linear(vals, 99.0), ++ "max_relative_error": worst_err if vals else 0.0, ++ "worst_tensor": worst_name, ++ } ++ ++ ++def run(src: Path, dst: Path, max_shard_bytes: int = MAX_SHARD_BYTES, exclude: str | None = None) -> dict: ++ src = src.resolve() ++ dst = dst.resolve() ++ if not src.is_dir(): ++ _die(f"SRC is not a directory: {src}") ++ _reject_nested(src, dst) ++ _check_dst_clean(dst) ++ index = _load_index(src / INDEX_NAME) ++ items = _plan(src, index, exclude) ++ groups = _assign_groups(items, max_shard_bytes) ++ dst.mkdir(parents=True, exist_ok=True) ++ ++ written: list[Path] = [] ++ try: ++ weight_map, bytes_in, bytes_out, errors, written = _write_shards(src, dst, items, groups) ++ copied = _copy_sidecars(src, dst) ++ n_quant = sum(1 for it in items if it.kind == "quant") ++ n_copy = sum(1 for it in items if it.kind == "copy") ++ measured = _summary(errors, n_quant, n_copy, bytes_in, bytes_out) ++ if measured["bytes_in"] != bytes_in or measured["bytes_out"] != bytes_out: ++ _die("internal error: summary byte counters diverged") ++ # Recompute the on-disk total from the tensors we recorded. weight_map ++ # values are what we just saved; bytes_out is that sum. ++ out_index = {"metadata": {"total_size": bytes_out}, "weight_map": weight_map} ++ (dst / INDEX_NAME).write_text( ++ json.dumps(out_index, indent=2) + "\n", encoding="utf-8" ++ ) ++ modules = sorted( ++ it.name[: -len(".weight")] for it in items if it.kind == "quant" ++ ) ++ manifest = { ++ "format": "ming-int8-wo-v1", ++ "scheme": "weight-only int8, per-output-channel symmetric, fp32 scales", ++ "rule": QUANT_RULE + (f" Additionally kept BF16: modules matching /{exclude}/." if exclude else ""), ++ "exclude": exclude, ++ "quantized_modules": modules, ++ "source_total_size": bytes_in, ++ "total_size": bytes_out, ++ "measured": measured, ++ } ++ (dst / MANIFEST_NAME).write_text( ++ json.dumps(manifest, indent=2, allow_nan=False) + "\n", encoding="utf-8" ++ ) ++ except Exception: ++ for path in written: ++ try: ++ path.unlink() ++ except OSError: ++ pass ++ raise ++ ++ print(f"copied {len(copied)} non-safetensors file(s)", flush=True) ++ m = measured ++ print( ++ "summary: " ++ f"quantized={m['tensors_quantized']} copied={m['tensors_copied']} " ++ f"bytes_in={m['bytes_in']} bytes_out={m['bytes_out']} " ++ f"mean_rel={m['mean_relative_error']:.8g} " ++ f"p99_rel={m['p99_relative_error']:.8g} " ++ f"max_rel={m['max_relative_error']:.8g} " ++ f"worst={m['worst_tensor']}", ++ flush=True, ++ ) ++ return manifest ++ ++ ++def main(argv: list[str] | None = None) -> int: ++ args = list(sys.argv[1:] if argv is None else argv) ++ exclude = None ++ if "--exclude" in args: ++ i = args.index("--exclude") ++ if i + 1 >= len(args): ++ print("--exclude needs a regex", file=sys.stderr) ++ return 2 ++ exclude = args[i + 1] ++ re.compile(exclude) ++ del args[i : i + 2] ++ if len(args) != 2: ++ print( ++ "usage: python quantize_stream.py SRC_MLLM_DIR DST_DIR [--exclude MODULE_REGEX]", ++ file=sys.stderr, ++ ) ++ return 2 ++ try: ++ run(Path(args[0]), Path(args[1]), max_shard_bytes=MAX_SHARD_BYTES, exclude=exclude) ++ except QuantizeError as exc: ++ print(f"error: {exc}", file=sys.stderr) ++ return 1 ++ return 0 ++ ++ ++if __name__ == "__main__": ++ sys.exit(main()) +diff --git a/quant/test_int8.py b/quant/test_int8.py +new file mode 100644 +index 0000000..03baa56 +--- /dev/null ++++ b/quant/test_int8.py +@@ -0,0 +1,665 @@ ++"""CPU tests for weight-only INT8 Ming MLLM quantize + load. ++ ++Run: HIP_VISIBLE_DEVICES=-1 python test_int8.py ++""" ++ ++from __future__ import annotations ++ ++import json ++import sys ++import tempfile ++import traceback ++from pathlib import Path ++ ++import torch ++import torch.nn.functional as F ++from safetensors.torch import load_file, save_file ++from torch import nn ++ ++import quantize_stream ++from int8_linear import Int8Linear, is_quantizable, quantize_weight ++from load_int8 import load_int8_mllm_ ++ ++# Tiny stand-in for Ming's MLLM names. Not the real model. ++HIDDEN = 32 ++INTER = 48 ++VOCAB = 64 ++N_EXPERTS = 2 ++ ++ ++class RMSNorm(nn.Module): ++ def __init__(self, dim: int, eps: float = 1e-6): ++ super().__init__() ++ self.weight = nn.Parameter(torch.ones(dim)) ++ self.eps = eps ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ var = x.float().pow(2).mean(dim=-1, keepdim=True) ++ y = x * torch.rsqrt(var + self.eps) ++ return (y * self.weight).to(dtype=x.dtype) ++ ++ ++class Attention(nn.Module): ++ def __init__(self, hidden: int): ++ super().__init__() ++ self.hidden = hidden ++ self.query_key_value = nn.Linear(hidden, hidden * 3, bias=True) ++ self.dense = nn.Linear(hidden, hidden, bias=False) ++ self.q_norm = RMSNorm(hidden) ++ self.k_norm = RMSNorm(hidden) ++ # Non-persistent, like BailingMoeV2RotaryEmbedding.inv_freq. ++ self.register_buffer( ++ "inv_freq", torch.arange(hidden // 2, dtype=torch.float32), persistent=False ++ ) ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ qkv = self.query_key_value(x) ++ h = self.hidden ++ q = self.q_norm(qkv[..., :h]) ++ k = self.k_norm(qkv[..., h : 2 * h]) ++ v = qkv[..., 2 * h :] ++ return self.dense(q + k + v) ++ ++ ++class DenseMLP(nn.Module): ++ def __init__(self, hidden: int, inter: int): ++ super().__init__() ++ self.gate_proj = nn.Linear(hidden, inter, bias=False) ++ self.up_proj = nn.Linear(hidden, inter, bias=True) ++ self.down_proj = nn.Linear(inter, hidden, bias=False) ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) ++ ++ ++class Expert(nn.Module): ++ def __init__(self, hidden: int, inter: int): ++ super().__init__() ++ self.gate_proj = nn.Linear(hidden, inter, bias=False) ++ self.up_proj = nn.Linear(hidden, inter, bias=True) ++ self.down_proj = nn.Linear(inter, hidden, bias=False) ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) ++ ++ ++class Router(nn.Module): ++ """Not an nn.Linear. Leaf name is gate / image_gate / audio_gate.""" ++ ++ def __init__(self, hidden: int, n_experts: int): ++ super().__init__() ++ self.weight = nn.Parameter(torch.empty(n_experts, hidden)) ++ self.expert_bias = nn.Parameter(torch.zeros(n_experts), requires_grad=False) ++ nn.init.kaiming_uniform_(self.weight, a=5**0.5) ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ return F.linear(x, self.weight, self.expert_bias) ++ ++ ++class MoeMLP(nn.Module): ++ def __init__(self, hidden: int, inter: int, n_experts: int): ++ super().__init__() ++ self.gate = Router(hidden, n_experts) ++ self.image_gate = Router(hidden, n_experts) ++ self.audio_gate = Router(hidden, n_experts) ++ self.experts = nn.ModuleList(Expert(hidden, inter) for _ in range(n_experts)) ++ self.shared_experts = Expert(hidden, inter) ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ scores = self.gate(x) + self.image_gate(x) + self.audio_gate(x) ++ weights = torch.softmax(scores, dim=-1) ++ mixed = self.shared_experts(x) ++ for i, expert in enumerate(self.experts): ++ mixed = mixed + expert(x) * weights[..., i : i + 1] ++ return mixed ++ ++ ++class DecoderLayer(nn.Module): ++ def __init__(self, hidden: int, mlp: nn.Module): ++ super().__init__() ++ self.input_layernorm = RMSNorm(hidden) ++ self.post_attention_layernorm = RMSNorm(hidden) ++ self.attention = Attention(hidden) ++ self.mlp = mlp ++ ++ def forward(self, x: torch.Tensor) -> torch.Tensor: ++ x = x + self.attention(self.input_layernorm(x)) ++ x = x + self.mlp(self.post_attention_layernorm(x)) ++ return x ++ ++ ++class TinyMing(nn.Module): ++ """Names match the real checkpoint: model.model.layers.*, model.lm_head, vision.*.""" ++ ++ def __init__(self): ++ super().__init__() ++ self.model = nn.Module() ++ self.model.model = nn.Module() ++ self.model.model.word_embeddings = nn.Embedding(VOCAB, HIDDEN) ++ self.model.model.layers = nn.ModuleList( ++ [ ++ DecoderLayer(HIDDEN, DenseMLP(HIDDEN, INTER)), ++ DecoderLayer(HIDDEN, MoeMLP(HIDDEN, INTER, N_EXPERTS)), ++ ] ++ ) ++ self.model.model.norm = RMSNorm(HIDDEN) ++ self.model.lm_head = nn.Linear(HIDDEN, VOCAB, bias=False) ++ block = nn.Module() ++ block.attn = nn.Module() ++ block.attn.qkv = nn.Linear(HIDDEN, HIDDEN, bias=False) ++ self.vision = nn.Module() ++ self.vision.blocks = nn.ModuleList([block]) ++ self.linear_proj = nn.ModuleList([nn.Linear(HIDDEN, HIDDEN, bias=True)]) ++ ++ def forward(self, input_ids: torch.Tensor) -> torch.Tensor: ++ h = self.model.model.word_embeddings(input_ids) ++ for layer in self.model.model.layers: ++ h = layer(h) ++ h = self.model.model.norm(h) ++ return self.model.lm_head(h) ++ ++ ++# Modules the rule must select for TinyMing. Hardcoded — not derived from is_quantizable. ++EXPECTED_QUANT_MODULES = [ ++ "model.model.layers.0.attention.dense", ++ "model.model.layers.0.attention.query_key_value", ++ "model.model.layers.0.mlp.down_proj", ++ "model.model.layers.0.mlp.gate_proj", ++ "model.model.layers.0.mlp.up_proj", ++ "model.model.layers.1.attention.dense", ++ "model.model.layers.1.attention.query_key_value", ++ "model.model.layers.1.mlp.experts.0.down_proj", ++ "model.model.layers.1.mlp.experts.0.gate_proj", ++ "model.model.layers.1.mlp.experts.0.up_proj", ++ "model.model.layers.1.mlp.experts.1.down_proj", ++ "model.model.layers.1.mlp.experts.1.gate_proj", ++ "model.model.layers.1.mlp.experts.1.up_proj", ++ "model.model.layers.1.mlp.shared_experts.down_proj", ++ "model.model.layers.1.mlp.shared_experts.gate_proj", ++ "model.model.layers.1.mlp.shared_experts.up_proj", ++] ++ ++MUST_NOT_QUANTIZE = [ ++ "model.model.layers.1.mlp.gate", ++ "model.model.layers.1.mlp.image_gate", ++ "model.model.layers.1.mlp.audio_gate", ++ "model.model.word_embeddings", ++ "model.model.norm", ++ "model.lm_head", ++ "vision.blocks.0.attn.qkv", ++ "linear_proj.0", ++ "model.model.layers.0.attention.q_norm", ++ "model.model.layers.0.input_layernorm", ++] ++ ++ ++def _move_parameters_to_meta(model: nn.Module) -> nn.Module: ++ """Parameters → meta, buffers stay where they are (CPU). Matches accelerate include_buffers=False.""" ++ for mod in model.modules(): ++ for name, param in list(mod._parameters.items()): ++ if param is None: ++ continue ++ mod._parameters[name] = nn.Parameter( ++ param.detach().to(device="meta"), ++ requires_grad=param.requires_grad, ++ ) ++ return model ++ ++ ++def _save_bf16_checkpoint(model: nn.Module, src: Path) -> None: ++ src.mkdir(parents=True, exist_ok=True) ++ sd = {k: v.detach().contiguous() for k, v in model.state_dict().items()} ++ if not sd: ++ raise AssertionError("empty state_dict") ++ for tensor in sd.values(): ++ if tensor.is_floating_point(): ++ assert tensor.dtype == torch.bfloat16, tensor.dtype ++ keys = list(sd) ++ mid = max(1, len(keys) // 2) ++ shards = { ++ "bf16-00001.safetensors": {k: sd[k] for k in keys[:mid]}, ++ "bf16-00002.safetensors": {k: sd[k] for k in keys[mid:]}, ++ } ++ weight_map = {} ++ total = 0 ++ for filename, tensors in shards.items(): ++ save_file(tensors, str(src / filename)) ++ for name, tensor in tensors.items(): ++ weight_map[name] = filename ++ total += tensor.numel() * tensor.element_size() ++ index = {"metadata": {"total_size": total}, "weight_map": weight_map} ++ (src / "model.safetensors.index.json").write_text( ++ json.dumps(index, indent=2) + "\n", encoding="utf-8" ++ ) ++ (src / "config.json").write_bytes(b'{"model_type":"tiny-ming","hidden":32}\n') ++ extra = src / "extra" ++ extra.mkdir() ++ (extra / "chat_template.jinja").write_text("{{ messages }}\n", encoding="utf-8") ++ ++ ++def _load_all(folder: Path) -> dict[str, torch.Tensor]: ++ index = json.loads((folder / "model.safetensors.index.json").read_text(encoding="utf-8")) ++ order: list[str] = [] ++ seen: set[str] = set() ++ for shard in index["weight_map"].values(): ++ if shard not in seen: ++ seen.add(shard) ++ order.append(shard) ++ sd: dict[str, torch.Tensor] = {} ++ for shard in order: ++ sd.update(load_file(str(folder / shard))) ++ return sd ++ ++ ++def _apply_int8_(model: nn.Module) -> None: ++ names = [] ++ for name, mod in model.named_modules(): ++ if isinstance(mod, nn.Linear) and is_quantizable( ++ f"{name}.weight", tuple(mod.weight.shape) ++ ): ++ names.append(name) ++ for name in names: ++ parent_name, _, leaf = name.rpartition(".") ++ parent = model.get_submodule(parent_name) if parent_name else model ++ setattr(parent, leaf, Int8Linear.from_linear(getattr(parent, leaf))) ++ ++ ++def _assert_no_meta(model: nn.Module) -> None: ++ for name, param in model.named_parameters(): ++ assert param.device.type != "meta", name ++ for mod_name, mod in model.named_modules(): ++ for buf_name, buf in mod._buffers.items(): ++ if buf is None: ++ continue ++ full = f"{mod_name}.{buf_name}" if mod_name else buf_name ++ assert buf.device.type != "meta", full ++ ++ ++def test_from_linear_roundtrip() -> None: ++ torch.manual_seed(0) ++ out_f, in_f = 5, 7 ++ lin = nn.Linear(in_f, out_f, bias=True) ++ scales = torch.tensor([0.5, 0.25, 0.125, 2.0, 4.0], dtype=torch.float32) ++ q = torch.randint(-127, 128, (out_f, in_f), dtype=torch.int8) ++ q[:, 0] = 127 ++ q[2, :] = 0 # all-zero row; must not NaN ++ weight = q.float() * scales[:, None] ++ with torch.no_grad(): ++ lin.weight.copy_(weight) ++ lin.bias.copy_(torch.tensor([0.1, -0.2, 0.3, -0.4, 0.5])) ++ mod = Int8Linear.from_linear(lin) ++ deq = mod.weight.float() * mod.scale[:, None] ++ for row in range(out_f): ++ if row == 2: ++ assert torch.equal(mod.weight[row], torch.zeros(in_f, dtype=torch.int8)) ++ assert float(mod.scale[row]) == 1.0 ++ assert torch.equal(deq[row], torch.zeros(in_f)) ++ else: ++ assert torch.equal(deq[row], weight[row]), (deq[row] - weight[row]).abs().max().item() ++ assert mod.bias is not None and torch.equal(mod.bias, lin.bias) ++ assert mod.bias.dtype == lin.bias.dtype ++ assert torch.isfinite(mod.scale).all() ++ ++ # Random weights: per-element error stays within half a bin (+ float slack). ++ lin_r = nn.Linear(13, 9, bias=False) ++ mod_r = Int8Linear.from_linear(lin_r) ++ w = lin_r.weight.detach().float() ++ deq_r = (mod_r.weight.double() * mod_r.scale.double()[:, None]).float() ++ err = (w.double() - deq_r.double()).abs() ++ half = mod_r.scale.double()[:, None] * 0.5 ++ slip = (err - half).max().item() ++ assert slip <= 1e-4, slip ++ assert torch.isfinite(mod_r.scale).all() ++ ++ # Entirely zero weight: finite forward, zero codes, scale 1. ++ lin_z = nn.Linear(4, 3, bias=True) ++ with torch.no_grad(): ++ lin_z.weight.zero_() ++ mod_z = Int8Linear.from_linear(lin_z) ++ assert torch.equal(mod_z.weight, torch.zeros_like(mod_z.weight)) ++ assert torch.equal(mod_z.scale, torch.ones(3)) ++ y = mod_z(torch.randn(8, 4)) ++ assert torch.isfinite(y).all() ++ assert torch.allclose(y, mod_z.bias.expand_as(y)) ++ ++ # Zero row contributes only its bias. ++ x = torch.randn(6, in_f) ++ y_mix = mod(x) ++ assert torch.isfinite(y_mix).all() ++ assert torch.allclose(y_mix[:, 2], mod.bias[2].expand(6)) ++ ++ # bf16 source linear: codes int8, scale fp32, bias stays bf16. ++ lin_b = nn.Linear(8, 4, bias=True).to(dtype=torch.bfloat16) ++ mod_b = Int8Linear.from_linear(lin_b) ++ assert mod_b.weight.dtype == torch.int8 ++ assert mod_b.scale.dtype == torch.float32 ++ assert mod_b.bias is not None and mod_b.bias.dtype == torch.bfloat16 ++ w_b = lin_b.weight.detach().float() ++ deq_b = mod_b.weight.float() * mod_b.scale[:, None] ++ err_b = (w_b.double() - deq_b.double()).abs() ++ half_b = mod_b.scale.double()[:, None] * 0.5 ++ assert (err_b - half_b).max().item() <= 1e-2, (err_b - half_b).max().item() ++ ++ ++def _assert_quant_dtypes(mod: Int8Linear, scale: torch.Tensor, weight: torch.Tensor, bias_dtype: torch.dtype) -> None: ++ assert mod.weight.dtype == torch.int8 ++ assert mod.scale.dtype == torch.float32 ++ assert torch.equal(mod.weight, weight) ++ assert torch.equal(mod.scale, scale) ++ assert mod.bias is not None and mod.bias.dtype == bias_dtype ++ ++ ++def test_dtype_cast_keeps_scale_fp32() -> None: ++ torch.manual_seed(1) ++ lin = nn.Linear(5, 3, bias=True) ++ fresh = Int8Linear.from_linear(lin) ++ scale = fresh.scale.detach().clone() ++ weight = fresh.weight.detach().clone() ++ bias = fresh.bias.detach().clone() ++ assert scale.dtype == torch.float32 and weight.dtype == torch.int8 and bias.dtype == torch.float32 ++ ++ # Each cast starts from fp32 so "bias follows the cast" is the single cast of the source bias. ++ mod = Int8Linear.from_linear(lin) ++ mod.bfloat16() ++ _assert_quant_dtypes(mod, scale, weight, torch.bfloat16) ++ assert torch.equal(mod.bias, bias.to(dtype=torch.bfloat16)) ++ ++ mod = Int8Linear.from_linear(lin) ++ mod.half() ++ _assert_quant_dtypes(mod, scale, weight, torch.float16) ++ assert torch.equal(mod.bias, bias.to(dtype=torch.float16)) ++ ++ mod = Int8Linear.from_linear(lin) ++ mod.to(torch.bfloat16) ++ _assert_quant_dtypes(mod, scale, weight, torch.bfloat16) ++ assert torch.equal(mod.bias, bias.to(dtype=torch.bfloat16)) ++ ++ mod = Int8Linear.from_linear(lin) ++ mod.to(dtype=torch.float16) ++ _assert_quant_dtypes(mod, scale, weight, torch.float16) ++ assert torch.equal(mod.bias, bias.to(dtype=torch.float16)) ++ ++ # A second cast applies to the bias's current dtype, not the original fp32 value. ++ mod = Int8Linear.from_linear(lin) ++ mod.to(torch.bfloat16) ++ mod.to(dtype=torch.float16) ++ _assert_quant_dtypes(mod, scale, weight, torch.float16) ++ assert torch.equal(mod.bias, bias.to(dtype=torch.bfloat16).to(dtype=torch.float16)) ++ ++ # What the caller actually does: parent.to(device=..., dtype=bf16). ++ parent = nn.Sequential(Int8Linear.from_linear(lin)) ++ parent.to(device="cpu", dtype=torch.bfloat16) ++ _assert_quant_dtypes(parent[0], scale, weight, torch.bfloat16) ++ assert torch.equal(parent[0].bias, bias.to(dtype=torch.bfloat16)) ++ ++ shell = Int8Linear.shell(4, 3, bias=True, bias_dtype=torch.bfloat16, device="meta") ++ assert shell.weight.dtype == torch.int8 and shell.weight.device.type == "meta" ++ assert shell.scale.dtype == torch.float32 and shell.scale.device.type == "meta" ++ assert shell.bias is not None ++ assert shell.bias.dtype == torch.bfloat16 and shell.bias.device.type == "meta" ++ shell_nb = Int8Linear.shell(4, 3, bias=False, bias_dtype=torch.float32, device="meta") ++ assert shell_nb.bias is None ++ ++ ++def test_forward_matches_reference() -> None: ++ torch.manual_seed(2) ++ for bias in (True, False): ++ lin = nn.Linear(6, 4, bias=bias) ++ # Bias is passed through unchanged, so it has to already match x's dtype ++ # (the caller does model.to(dtype=...) before the prefill). ++ modules = [ ++ (Int8Linear.from_linear(lin), torch.float32), ++ (Int8Linear.from_linear(lin).to(torch.bfloat16), torch.bfloat16), ++ (Int8Linear.from_linear(lin).to(dtype=torch.float16), torch.float16), ++ ] ++ for mod, dtype in modules: ++ if mod.bias is not None: ++ assert mod.bias.dtype == dtype ++ x = torch.randn(3, 5, 6, dtype=dtype) ++ ref_w = (mod.weight.float() * mod.scale[:, None]).to(dtype=x.dtype) ++ y = mod(x) ++ y_ref = F.linear(x, ref_w, mod.bias) ++ assert torch.equal(y, y_ref), (bias, dtype) ++ ++ ++def test_is_quantizable_rule() -> None: ++ false_cases = [ ++ ("model.model.layers.1.mlp.gate.weight", (256, 2048)), ++ ("model.model.layers.1.mlp.image_gate.weight", (256, 2048)), ++ ("model.model.layers.1.mlp.audio_gate.weight", (256, 2048)), ++ ("model.model.layers.1.mlp.gate.expert_bias", (256,)), ++ ("model.lm_head.weight", (151936, 2048)), ++ ("model.model.word_embeddings.weight", (151936, 2048)), ++ ("vision.blocks.0.attn.qkv.weight", (3072, 1280)), ++ ("model.model.layers.0.input_layernorm.weight", (2048,)), ++ ("model.model.layers.0.post_attention_layernorm.weight", (2048,)), ++ ("model.model.layers.0.attention.q_norm.weight", (128,)), ++ ("model.model.layers.0.attention.k_norm.weight", (128,)), ++ ("model.model.norm.weight", (2048,)), ++ ("linear_proj.0.weight", (2048, 2048)), ++ ("model.model.layers.0.attention.query_key_value.bias", (3072,)), ++ ("model.model.layers.0.mlp.experts.0.gate_proj.bias", (512,)), ++ # Right leaf, wrong rank: not quantizable (the stream must reject it). ++ ("model.model.layers.0.attention.query_key_value.weight", (3072,)), ++ ("model.model.layers.0.mlp.gate_proj.weight", (1024, 2048, 1)), ++ ] ++ true_cases = [ ++ ("model.model.layers.3.mlp.experts.3.gate_proj.weight", (512, 2048)), ++ ("model.model.layers.3.mlp.shared_experts.down_proj.weight", (2048, 512)), ++ ("model.model.layers.0.mlp.up_proj.weight", (512, 2048)), ++ ("layers.0.mlp.up_proj.weight", (512, 2048)), ++ ("model.model.layers.0.attention.query_key_value.weight", (3072, 2048)), ++ ("model.model.layers.0.attention.dense.weight", (2048, 2048)), ++ ("model.model.layers.0.mlp.gate_proj.weight", (512, 2048)), ++ ("model.model.layers.0.mlp.down_proj.weight", (2048, 512)), ++ ("model.model.layers.19.mlp.experts.255.up_proj.weight", (512, 2048)), ++ ] ++ for name, shape in false_cases: ++ assert is_quantizable(name, shape) is False, name ++ for name, shape in true_cases: ++ assert is_quantizable(name, shape) is True, name ++ ++ ++def _shard_groups(sd: dict[str, torch.Tensor]) -> set[str]: ++ """One copy-tensor, or one weight+scale pair, is one unsplittable group.""" ++ names = set(sd) ++ groups: set[str] = set() ++ for name in names: ++ if name.endswith(".scale") and name[: -len(".scale")] + ".weight" in names: ++ groups.add(name[: -len(".scale")]) ++ elif name.endswith(".weight") and name[: -len(".weight")] + ".scale" in names: ++ groups.add(name[: -len(".weight")]) ++ else: ++ groups.add(name) ++ return groups ++ ++ ++def test_end_to_end_stream_and_load() -> None: ++ assert quantize_stream.MAX_SHARD_BYTES == 5 * 10**9 ++ torch.manual_seed(3) ++ src_model = TinyMing().to(dtype=torch.bfloat16) ++ # Non-persistent rotary buffer is not part of the checkpoint. ++ assert "model.model.layers.0.attention.inv_freq" not in src_model.state_dict() ++ ++ with tempfile.TemporaryDirectory(prefix="ming-int8-") as tmp: ++ root = Path(tmp) ++ src = root / "src" ++ dst = root / "dst" ++ _save_bf16_checkpoint(src_model, src) ++ limit = 2048 ++ old = quantize_stream.MAX_SHARD_BYTES ++ quantize_stream.MAX_SHARD_BYTES = limit ++ try: ++ rc = quantize_stream.main([str(src), str(dst)]) ++ finally: ++ quantize_stream.MAX_SHARD_BYTES = old ++ assert rc == 0, rc ++ assert quantize_stream.MAX_SHARD_BYTES == 5 * 10**9 ++ ++ # Sidecars copied verbatim; original index replaced. ++ assert (dst / "config.json").read_bytes() == (src / "config.json").read_bytes() ++ assert (dst / "extra" / "chat_template.jinja").read_bytes() == ( ++ src / "extra" / "chat_template.jinja" ++ ).read_bytes() ++ assert not (dst / "bf16-00001.safetensors").exists() ++ ++ manifest = json.loads((dst / "int8_manifest.json").read_text(encoding="utf-8")) ++ assert manifest["format"] == "ming-int8-wo-v1" ++ assert manifest["scheme"] == ( ++ "weight-only int8, per-output-channel symmetric, fp32 scales" ++ ) ++ assert manifest["quantized_modules"] == sorted(EXPECTED_QUANT_MODULES) ++ for banned in MUST_NOT_QUANTIZE: ++ assert banned not in manifest["quantized_modules"], banned ++ ++ index = json.loads((dst / "model.safetensors.index.json").read_text(encoding="utf-8")) ++ assert index["metadata"]["total_size"] == manifest["total_size"] ++ measured = manifest["measured"] ++ assert measured["tensors_quantized"] == len(EXPECTED_QUANT_MODULES) ++ assert measured["bytes_in"] == manifest["source_total_size"] ++ assert measured["bytes_out"] == manifest["total_size"] ++ assert measured["bytes_out"] < measured["bytes_in"] ++ n_out_keys = measured["tensors_copied"] + 2 * measured["tensors_quantized"] ++ assert len(index["weight_map"]) == n_out_keys ++ ++ src_sd = _load_all(src) ++ dst_sd = _load_all(dst) ++ assert manifest["source_total_size"] == sum( ++ t.numel() * t.element_size() for t in src_sd.values() ++ ) ++ assert manifest["total_size"] == sum(t.numel() * t.element_size() for t in dst_sd.values()) ++ ++ shard_names = sorted({*index["weight_map"].values()}) ++ assert len(shard_names) >= 2, shard_names ++ for shard in shard_names: ++ shard_sd = load_file(str(dst / shard)) ++ total = sum(t.numel() * t.element_size() for t in shard_sd.values()) ++ if total > limit: ++ assert len(_shard_groups(shard_sd)) == 1, (shard, total, list(shard_sd)) ++ ++ errors = [] ++ for name, src_t in src_sd.items(): ++ if is_quantizable(name, tuple(src_t.shape)): ++ q = dst_sd[name] ++ scale_key = name[: -len("weight")] + "scale" ++ scale = dst_sd[scale_key] ++ assert q.dtype == torch.int8, name ++ assert scale.dtype == torch.float32, scale_key ++ q_ref, scale_ref = quantize_weight(src_t) ++ assert torch.equal(q, q_ref), name ++ assert torch.equal(scale, scale_ref), scale_key ++ errors.append((name, quantize_stream._relative_frobenius(src_t, q, scale))) ++ else: ++ assert name in dst_sd, name ++ assert dst_sd[name].dtype == src_t.dtype, (name, dst_sd[name].dtype, src_t.dtype) ++ assert torch.equal(dst_sd[name], src_t), name ++ # Router weights stayed BF16 and byte-identical (the gate vs gate_proj trap). ++ router = "model.model.layers.1.mlp.gate.weight" ++ assert dst_sd[router].dtype == torch.bfloat16 ++ assert torch.equal(dst_sd[router], src_sd[router]) ++ for suffix in ("image_gate.weight", "audio_gate.weight", "gate.expert_bias"): ++ key = f"model.model.layers.1.mlp.{suffix}" ++ assert torch.equal(dst_sd[key], src_sd[key]), key ++ ++ vals = [e for _, e in errors] ++ assert measured["max_relative_error"] == max(vals) ++ assert measured["mean_relative_error"] == sum(vals) / len(vals) ++ assert measured["worst_tensor"] in dict(errors) ++ assert measured["max_relative_error"] == dict(errors)[measured["worst_tensor"]] ++ assert 0.0 <= measured["mean_relative_error"] <= measured["p99_relative_error"] ++ assert measured["p99_relative_error"] <= measured["max_relative_error"] ++ assert measured["max_relative_error"] < 0.05, measured ++ ++ # Eager quant of the same BF16 bytes. ++ eager = TinyMing().to(dtype=torch.bfloat16) ++ incompatible = eager.load_state_dict(src_sd, strict=True) ++ assert not incompatible.missing_keys and not incompatible.unexpected_keys ++ _apply_int8_(eager) ++ ++ loaded = _move_parameters_to_meta(TinyMing()) ++ for layer in loaded.model.model.layers: ++ assert layer.attention.inv_freq.device.type == "cpu" ++ assert layer.attention.query_key_value.weight.device.type == "meta" ++ report = load_int8_mllm_(loaded, dst, "cpu") ++ assert report["modules_swapped"] == len(EXPECTED_QUANT_MODULES) ++ assert report["tensors_loaded"] == len(dst_sd) ++ assert report["bytes_loaded"] == manifest["total_size"] ++ _assert_no_meta(loaded) ++ for layer in loaded.model.model.layers: ++ assert layer.attention.inv_freq.device.type == "cpu" ++ assert layer.attention.inv_freq.dtype == torch.float32 ++ for name in EXPECTED_QUANT_MODULES: ++ mod = loaded.get_submodule(name) ++ assert isinstance(mod, Int8Linear), name ++ assert mod.weight.dtype == torch.int8 ++ assert mod.scale.dtype == torch.float32 ++ ++ eager.eval() ++ loaded.eval() ++ ids = torch.randint(0, VOCAB, (2, 6)) ++ with torch.no_grad(): ++ y_eager = eager(ids) ++ y_loaded = loaded(ids) ++ assert y_eager.dtype == y_loaded.dtype ++ assert torch.equal(y_eager, y_loaded), (y_eager - y_loaded).abs().max().item() ++ ++ # A second run into a non-empty safetensors dir must fail loudly. ++ print(" re-running into a non-empty dst (expect error on stderr)", flush=True) ++ rc_again = quantize_stream.main([str(src), str(dst)]) ++ assert rc_again == 1 ++ ++ ++def test_unknown_key_fails_loudly() -> None: ++ torch.manual_seed(4) ++ model = TinyMing().to(dtype=torch.bfloat16) ++ with tempfile.TemporaryDirectory(prefix="ming-int8-bad-") as tmp: ++ root = Path(tmp) ++ src = root / "src" ++ dst = root / "dst" ++ _save_bf16_checkpoint(model, src) ++ rc = quantize_stream.main([str(src), str(dst)]) ++ assert rc == 0, rc ++ shard = next(dst.glob("*.safetensors")) ++ sd = load_file(str(shard)) ++ sd["not.a.real.key"] = torch.zeros(4, dtype=torch.float32) ++ save_file(sd, str(shard)) ++ loaded = _move_parameters_to_meta(TinyMing()) ++ try: ++ load_int8_mllm_(loaded, dst, "cpu") ++ except RuntimeError as exc: ++ text = str(exc) ++ assert "unexpected" in text.lower(), text ++ assert "not.a.real.key" in text, text ++ print(f" caught RuntimeError: {text.splitlines()[0]}") ++ else: ++ raise AssertionError("load_int8_mllm_ returned instead of failing on an unknown key") ++ ++ ++def main() -> int: ++ import safetensors ++ ++ print(f"torch={torch.__version__} safetensors={safetensors.__version__}", flush=True) ++ tests = [ ++ test_from_linear_roundtrip, ++ test_dtype_cast_keeps_scale_fp32, ++ test_forward_matches_reference, ++ test_is_quantizable_rule, ++ test_end_to_end_stream_and_load, ++ test_unknown_key_fails_loudly, ++ ] ++ failed = 0 ++ for fn in tests: ++ try: ++ fn() ++ except Exception: ++ failed += 1 ++ print(f"FAIL {fn.__name__}", flush=True) ++ traceback.print_exc() ++ else: ++ print(f"PASS {fn.__name__}", flush=True) ++ print(f"{len(tests) - failed} passed, {failed} failed", flush=True) ++ return 1 if failed else 0 ++ ++ ++if __name__ == "__main__": ++ sys.exit(main()) +diff --git a/qwen2_5_vit.py b/qwen2_5_vit.py +index 3de0e73..efa8255 100644 +--- a/qwen2_5_vit.py ++++ b/qwen2_5_vit.py +@@ -36,7 +36,6 @@ from transformers.utils import ( + from typing import Union + + from transformers.configuration_utils import PretrainedConfig +-import transformer_engine.pytorch as te + + if is_flash_attn_2_available(): + from flash_attn import flash_attn_varlen_func +@@ -158,12 +157,28 @@ class Qwen2_5_VisionRotaryEmbedding(nn.Module): + new_inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim)) + self.inv_freq.copy_(new_inv_freq) + +-class Qwen2RMSNorm(te.RMSNorm): ++class Qwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ +- Qwen2RMSNorm is equivalent to T5LayerNorm ++ Qwen2RMSNorm is equivalent to T5LayerNorm. ++ ++ Replaces transformer_engine.pytorch.RMSNorm: ROCm has no transformer-engine. ++ te.RMSNorm defaults (zero_centered_gamma=False) are standard RMSNorm, and the ++ checkpoint stores this affine as `weight`. + """ +- super().__init__(hidden_size, eps=eps) ++ super().__init__() ++ self.weight = nn.Parameter(torch.ones(hidden_size)) ++ self.variance_epsilon = eps ++ ++ def forward(self, hidden_states): ++ input_dtype = hidden_states.dtype ++ hidden_states = hidden_states.to(torch.float32) ++ variance = hidden_states.pow(2).mean(-1, keepdim=True) ++ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) ++ return self.weight * hidden_states.to(input_dtype) ++ ++ def extra_repr(self): ++ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + class Qwen2_5_VLPatchMerger(nn.Module): + def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: +diff --git a/requirements-rocm.txt b/requirements-rocm.txt +new file mode 100644 +index 0000000..d8cda21 +--- /dev/null ++++ b/requirements-rocm.txt +@@ -0,0 +1,36 @@ ++# ROCm port of requirements.txt for AMD gfx1151 (ROCm 7.13). ++# Omitted on purpose — do not add them back: ++# torch, torchvision: the target interpreter already has a working ROCm ++# build (torch 2.10.0, torch.version.hip 7.13.99004). Reinstalling the ++# upstream CUDA pins would replace it. ++# transformer-engine: NVIDIA CUDA-only; ROCm has no TE. Qwen2RMSNorm in ++# qwen2_5_vit.py is pure PyTorch, and the unused TE import is gone. ++# Create the venv with system site packages so that ROCm torch is inherited: ++# python3 -m venv --system-site-packages .venv ++# .venv/bin/pip install -r requirements-rocm.txt ++# ++# Also omitted / relaxed versus upstream, because the ROCm interpreter is Python 3.13 ++# and its torch is built against numpy 2.x: ++# numpy upstream 1.23.1 has no Python 3.13 wheels, and downgrading would ++# break the inherited torch. Inherit the system numpy (validated 2.2.4). ++# Pillow upstream 10.4.0 has no Python 3.13 wheels. Inherit (validated 11.1.0). ++# safetensors inherit the system build (validated 0.8.0). ++# ++# Validated on halo (gfx1151, ROCm 7.13, Python 3.13.5) on 2026-09-22: ++# torch 2.10.0 (hip 7.13.99004) | numpy 2.2.4 | Pillow 11.1.0 | safetensors 0.8.0 ++# transformers 4.57.1 | diffusers 0.36.0 | accelerate 1.13.0 | tokenizers 0.22.2 ++# huggingface-hub 0.34.0 | peft 0.17.0 ++transformers==4.57.1 ++diffusers==0.36.0 ++accelerate==1.13.0 ++tokenizers==0.22.2 ++huggingface-hub==0.34.0 ++peft==0.17.0 ++requests==2.32.3 ++tqdm==4.67.1 ++typing-extensions==4.15.0 ++ ++# Optional FlashAttention 2 backend (validated: flash-attn==2.7.3). The CLI ++# default is eager attention; --attn-implementation flash_attention_2 needs ++# this package. ++# flash-attn==2.7.3 +diff --git a/tools/convert_connector.py b/tools/convert_connector.py +new file mode 100644 +index 0000000..8f3733c +--- /dev/null ++++ b/tools/convert_connector.py +@@ -0,0 +1,72 @@ ++#!/usr/bin/env python3 ++"""Store the connector component (Qwen2 1.5B, shipped as float32) as bfloat16. ++ ++infer.py loads the connector with torch_dtype=bfloat16, so the tensors it runs with are the ++fp32 values rounded to bf16 at load time. This does the same rounding once, offline, and proves ++every converted tensor equals `fp32_tensor.to(torch.bfloat16)` exactly — the runtime model is ++unchanged; only the download halves. ++ ++ usage: convert_connector.py SRC_CONNECTOR_DIR DST_CONNECTOR_DIR ++""" ++import json ++import shutil ++import sys ++from pathlib import Path ++ ++import torch ++from safetensors import safe_open ++from safetensors.torch import load_file, save_file ++ ++ ++def main(): ++ src, dst = Path(sys.argv[1]), Path(sys.argv[2]) ++ dst.mkdir(parents=True, exist_ok=True) ++ if any(dst.glob("*.safetensors")): ++ sys.exit(f"refusing: {dst} already contains safetensors") ++ index = json.loads((src / "model.safetensors.index.json").read_text()) ++ shards = sorted(set(index["weight_map"].values())) ++ ++ def converted(tensor): ++ return tensor.to(torch.bfloat16) if tensor.is_floating_point() else tensor ++ ++ out = {} ++ for shard in shards: ++ with safe_open(str(src / shard), "pt") as handle: ++ for key in handle.keys(): ++ if key in out: ++ sys.exit(f"duplicate tensor {key}") ++ out[key] = converted(handle.get_tensor(key)) ++ if set(out) != set(index["weight_map"]): ++ sys.exit("tensor set does not match the index weight_map") ++ target = dst / "model.safetensors" ++ save_file(out, str(target), metadata={"format": "pt"}) ++ del out ++ ++ back = load_file(str(target)) ++ checked = 0 ++ for shard in shards: ++ with safe_open(str(src / shard), "pt") as handle: ++ for key in handle.keys(): ++ reference = converted(handle.get_tensor(key)) ++ if back[key].dtype != reference.dtype or not torch.equal(back[key], reference): ++ sys.exit(f"MISMATCH {key}") ++ checked += 1 ++ if checked != len(back): ++ sys.exit(f"checked {checked} tensors but the output holds {len(back)}") ++ ++ for path in src.iterdir(): ++ if path.suffix == ".safetensors" or path.name == "model.safetensors.index.json": ++ continue ++ shutil.copy2(path, dst / path.name) ++ config = json.loads((dst / "config.json").read_text()) ++ key = "dtype" if "dtype" in config else "torch_dtype" ++ previous = config.get(key) ++ config[key] = "bfloat16" ++ (dst / "config.json").write_text(json.dumps(config, indent=2) + "\n") ++ dtypes = sorted({str(t.dtype) for t in back.values()}) ++ print(f"CONNECTOR_OK tensors={checked} exact=all dtypes={dtypes} bytes={target.stat().st_size} " ++ f"config.{key}: {previous} -> bfloat16") ++ ++ ++if __name__ == "__main__": ++ main() +diff --git a/tools/fidelity_compare.py b/tools/fidelity_compare.py +new file mode 100644 +index 0000000..99736d6 +--- /dev/null ++++ b/tools/fidelity_compare.py +@@ -0,0 +1,104 @@ ++#!/usr/bin/env python3 ++"""Compare two ming_bench.py output dirs (reference vs candidate), stem by stem. ++ ++Conditioning (what the DiT receives): cosine similarity over the whole tensor, the ++per-token cosine (mean and worst token), and relative L2 = |a - b| / |a|. ++Images: MAE, PSNR, windowed 7x7 SSIM on luminance, and alpha MAE for RGBA. ++ ++ usage: fidelity_compare.py [--json out.json] ++""" ++import json ++import sys ++from pathlib import Path ++ ++import numpy as np ++from PIL import Image ++from safetensors.numpy import load_file ++ ++ ++def load_image(path): ++ im = Image.open(path) ++ rgb = np.asarray(im.convert("RGB"), dtype=np.float64) ++ alpha = np.asarray(im.convert("RGBA"), dtype=np.float64)[..., 3] if im.mode in ("RGBA", "LA") else None ++ return rgb, alpha, im.size ++ ++ ++def box(x, k): ++ c = np.cumsum(np.cumsum(np.pad(x, ((1, 0), (1, 0))), 0), 1) ++ return (c[k:, k:] - c[:-k, k:] - c[k:, :-k] + c[:-k, :-k]) / (k * k) ++ ++ ++def ssim(a, b, k=7, L=255.0): ++ c1, c2 = (0.01 * L) ** 2, (0.03 * L) ** 2 ++ mu_a, mu_b = box(a, k), box(b, k) ++ va, vb = box(a * a, k) - mu_a ** 2, box(b * b, k) - mu_b ** 2 ++ cov = box(a * b, k) - mu_a * mu_b ++ s = ((2 * mu_a * mu_b + c1) * (2 * cov + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (va + vb + c2)) ++ return float(s.mean()) ++ ++ ++def image_metrics(ref_path, cand_path): ++ ra, aa, sa = load_image(ref_path) ++ rb, ab, sb = load_image(cand_path) ++ if sa != sb: ++ raise SystemExit(f"size mismatch {ref_path} {sa} vs {cand_path} {sb}") ++ lum = lambda x: 0.299 * x[..., 0] + 0.587 * x[..., 1] + 0.114 * x[..., 2] ++ mse = float(((ra - rb) ** 2).mean()) ++ out = { ++ "mae": round(float(np.abs(ra - rb).mean()), 3), ++ "psnr_db": None if mse == 0 else round(10 * np.log10(255.0 ** 2 / mse), 2), ++ "ssim_lum": round(ssim(lum(ra), lum(rb)), 4), ++ } ++ if aa is not None and ab is not None: ++ out["alpha_mae"] = round(float(np.abs(aa - ab).mean()), 3) ++ return out ++ ++ ++def cond_metrics(ref_path, cand_path): ++ ref, cand = load_file(str(ref_path)), load_file(str(cand_path)) ++ out = {} ++ for key in sorted(set(ref) & set(cand)): ++ a, b = ref[key].astype(np.float64), cand[key].astype(np.float64) ++ if a.shape != b.shape: ++ raise SystemExit(f"{key}: shape mismatch {a.shape} vs {b.shape}") ++ fa, fb = a.ravel(), b.ravel() ++ tok_a, tok_b = a.reshape(-1, a.shape[-1]), b.reshape(-1, b.shape[-1]) ++ tok_cos = (tok_a * tok_b).sum(-1) / (np.linalg.norm(tok_a, axis=-1) * np.linalg.norm(tok_b, axis=-1)) ++ out[key] = { ++ "shape": list(a.shape), ++ "cosine": round(float(fa @ fb / (np.linalg.norm(fa) * np.linalg.norm(fb))), 6), ++ "token_cos_mean": round(float(tok_cos.mean()), 6), ++ "token_cos_min": round(float(tok_cos.min()), 6), ++ "rel_l2": round(float(np.linalg.norm(fa - fb) / np.linalg.norm(fa)), 6), ++ } ++ missing = sorted(set(ref) ^ set(cand)) ++ if missing: ++ raise SystemExit(f"conditioning keys present on one side only: {missing}") ++ return out ++ ++ ++def main(): ++ ref_dir, cand_dir = Path(sys.argv[1]), Path(sys.argv[2]) ++ stems = sorted(p.stem for p in ref_dir.glob("*.png") if (cand_dir / p.name).exists()) ++ if not stems: ++ raise SystemExit(f"no common images between {ref_dir} and {cand_dir}") ++ rows = [] ++ for stem in stems: ++ row = {"stem": stem, "image": image_metrics(ref_dir / f"{stem}.png", cand_dir / f"{stem}.png")} ++ rc, cc = ref_dir / f"{stem}.cond.safetensors", cand_dir / f"{stem}.cond.safetensors" ++ if rc.exists() and cc.exists(): ++ row["cond"] = cond_metrics(rc, cc) ++ rows.append(row) ++ im = row["image"] ++ line = f"{stem:32s} SSIM {im['ssim_lum']:.4f} PSNR {im['psnr_db']} MAE {im['mae']:.2f}" ++ if "alpha_mae" in im: ++ line += f" aMAE {im['alpha_mae']:.2f}" ++ for key, c in row.get("cond", {}).items(): ++ line += f" | {key[:3]} cos {c['cosine']:.6f} tokmin {c['token_cos_min']:.4f} relL2 {c['rel_l2']:.4f}" ++ print(line) ++ if "--json" in sys.argv: ++ Path(sys.argv[sys.argv.index("--json") + 1]).write_text(json.dumps(rows, indent=2)) ++ ++ ++if __name__ == "__main__": ++ main() +diff --git a/tools/ming_bench.py b/tools/ming_bench.py +new file mode 100644 +index 0000000..c92bb3f +--- /dev/null ++++ b/tools/ming_bench.py +@@ -0,0 +1,127 @@ ++#!/usr/bin/env python3 ++"""Ming-Image speed + fidelity harness: one model load, N prompts. ++ ++Reuses infer.py's own loader and generation path unchanged. The only addition is a ++wrapper around model.diffusion_loss.sample that records the conditioning tensors the ++DiT receives (encoder_hidden_states / directvlm_hidden_states) and times the sampling ++stage (DiT steps + VAE decode) separately from the MLLM stage. ++ ++ usage: ming_bench.py --prompts a.json b.json --out DIR [--repeat-first N] -- ++ ++ are passed to infer.parse_args() as-is (e.g. --model, --resolution, ++ --steps, --seed, --device, --device-map none, --attn-implementation eager, --int8-mllm). ++ --repeat-first N re-runs the first prompt N more times at the same seed: the images ++ measure the platform's run-to-run noise floor and the timings are warm timings. ++""" ++import argparse ++import json ++import sys ++import time ++from pathlib import Path ++ ++ ++def main(): ++ ap = argparse.ArgumentParser() ++ ap.add_argument("--prompts", nargs="+", required=True) ++ ap.add_argument("--out", required=True) ++ ap.add_argument("--repeat-first", type=int, default=0) ++ own, rest = ap.parse_known_args() ++ if rest and rest[0] == "--": ++ rest = rest[1:] ++ sys.argv = [sys.argv[0], "--prompt", own.prompts[0]] + rest ++ ++ import torch ++ from safetensors.torch import save_file ++ import infer ++ ++ args = infer.parse_args() ++ model_directory = infer.resolve_model_directory( ++ args.model, revision=args.revision, cache_dir=args.cache_dir, ++ local_files_only=args.local_files_only, ++ ) ++ profile = infer.load_checkpoint_capabilities(model_directory) ++ resolution = infer.resolve_task_resolution(args.task, args.resolution) ++ sampling = profile.resolve_sampling_parameters(steps=args.steps, cfg=args.cfg) ++ dtype = infer._dtype(args.dtype) ++ out = Path(own.out) ++ out.mkdir(parents=True, exist_ok=True) ++ ++ def sync(): ++ if torch.cuda.is_available(): ++ torch.cuda.synchronize() ++ ++ sync() ++ t0 = time.perf_counter() ++ model, processor = infer.load_model_and_processor(model_directory, args) ++ sync() ++ load_s = time.perf_counter() - t0 ++ print(f"LOAD_S {load_s:.1f}", flush=True) ++ ++ captured = {} ++ original_sample = model.diffusion_loss.sample ++ ++ def recording_sample(*a, **kw): ++ for key in ("encoder_hidden_states", "directvlm_hidden_states"): ++ value = kw.get(key) ++ if isinstance(value, (list, tuple)): ++ value = torch.stack(list(value), dim=0) ++ if isinstance(value, torch.Tensor): ++ captured[key] = value.detach().float().cpu().contiguous() ++ sync() ++ ts = time.perf_counter() ++ result = original_sample(*a, **kw) ++ sync() ++ captured["_sample_s"] = time.perf_counter() - ts ++ return result ++ ++ model.diffusion_loss.sample = recording_sample ++ ++ runs = [(p, 0) for p in own.prompts] + [(own.prompts[0], i + 1) for i in range(own.repeat_first)] ++ results = [] ++ for prompt_path, rep in runs: ++ stem = Path(prompt_path).stem + (f"_rep{rep}" if rep else "") ++ prompt = infer._load_prompt(prompt_path) ++ captured.clear() ++ if torch.cuda.is_available(): ++ torch.cuda.reset_peak_memory_stats() ++ sync() ++ t1 = time.perf_counter() ++ images = infer.run_generation( ++ model, processor, profile, task=args.task, prompt=prompt, input_image=None, ++ resolution=resolution, sampling=sampling, seed=args.seed, num_layers=args.num_layers, ++ dtype=dtype, ++ ) ++ sync() ++ total_s = time.perf_counter() - t1 ++ if len(images) != 1: ++ raise RuntimeError(f"{stem}: expected 1 image, got {len(images)}") ++ image_path = out / f"{stem}.png" ++ images[0].save(image_path) ++ cond = {k: v for k, v in captured.items() if not k.startswith("_")} ++ if "encoder_hidden_states" not in cond: ++ raise RuntimeError(f"{stem}: conditioning was not captured") ++ save_file(cond, str(out / f"{stem}.cond.safetensors")) ++ sample_s = captured["_sample_s"] ++ row = { ++ "load_s": round(load_s, 1), ++ "prompt": str(prompt_path), "stem": stem, "seed": args.seed, "resolution": resolution, ++ "steps": sampling.steps, "cfg": sampling.cfg, "mode": images[0].mode, ++ "size": list(images[0].size), "total_s": round(total_s, 2), ++ "sample_s": round(sample_s, 2), "mllm_s": round(total_s - sample_s, 2), ++ "peak_alloc_gib": round(torch.cuda.max_memory_allocated() / 2**30, 2) ++ if torch.cuda.is_available() else None, ++ "cond_shapes": {k: list(v.shape) for k, v in cond.items()}, ++ } ++ results.append(row) ++ print("RUN " + json.dumps(row), flush=True) ++ with open(out / "runs.jsonl", "a") as fh: # accumulates across one-prompt-per-process runs ++ fh.write(json.dumps(row) + "\n") ++ ++ manifest = {"load_s": round(load_s, 1), "args": {k: str(v) for k, v in vars(args).items()}, ++ "runs": results} ++ (out / "manifest.json").write_text(json.dumps(manifest, indent=2)) ++ print("BENCH_DONE", out, flush=True) ++ ++ ++if __name__ == "__main__": ++ main() +diff --git a/tools/sdpa_layout.py b/tools/sdpa_layout.py +new file mode 100644 +index 0000000..cfb3720 +--- /dev/null ++++ b/tools/sdpa_layout.py +@@ -0,0 +1,49 @@ ++#!/usr/bin/env python3 ++"""Math SDPA with the DiT's real input layout: [B, L, H, D] permuted to [B, H, L, D] (non-contiguous, ++exactly what diffusers' native attention backend passes) vs the same tensors made contiguous. ++Speed and error vs an fp32 reference, masked, at the cabin prompt's real length. ++ ++ usage: sdpa_layout.py [L] ++""" ++import sys ++import time ++ ++import torch ++import torch.nn.functional as F ++from torch.nn.attention import SDPBackend, sdpa_kernel ++ ++L = int(sys.argv[1]) if len(sys.argv) > 1 else 5759 ++H, D, dev = 30, 128, "cuda" ++g = torch.Generator(device=dev).manual_seed(0) ++blhd = [torch.randn(1, L, H, D, device=dev, dtype=torch.bfloat16, generator=g) for _ in range(3)] ++q, k, v = (x.permute(0, 2, 1, 3) for x in blhd) # views, as diffusers passes them ++qc, kc, vc = (x.contiguous() for x in (q, k, v)) ++mask = torch.ones(1, 1, 1, L, dtype=torch.bool, device=dev) ++mask[..., L - 64:] = False ++with sdpa_kernel(SDPBackend.MATH): ++ ref = F.scaled_dot_product_attention(qc.float(), kc.float(), vc.float(), attn_mask=mask) ++ ++ ++def run(tag, a, b, c, bf16_reduction): ++ torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(bf16_reduction) ++ with sdpa_kernel(SDPBackend.MATH): ++ out = F.scaled_dot_product_attention(a, b, c, attn_mask=mask) ++ torch.cuda.synchronize() ++ t0 = time.perf_counter() ++ for _ in range(3): ++ out = F.scaled_dot_product_attention(a, b, c, attn_mask=mask) ++ torch.cuda.synchronize() ++ ms = (time.perf_counter() - t0) / 3 * 1000 ++ rel = ((out.float() - ref).norm() / ref.norm()).item() ++ exact = torch.equal(out, base) if base is not None else None ++ print(f" {tag:34s} {ms:8.2f} ms rel_l2 {rel:.3e} identical_to_default: {exact}") ++ return out ++ ++ ++base = None ++print(f"torch {torch.__version__} | L={L} | q strides {tuple(q.stride())} contiguous={q.is_contiguous()}") ++base = run("permuted views (DiT today), fp32", q, k, v, False) ++run("contiguous, fp32 (math unchanged)", qc, kc, vc, False) ++run("permuted views, bf16 reduction", q, k, v, True) ++run("contiguous, bf16 reduction", qc, kc, vc, True) ++torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(False) +diff --git a/tools/step_probe.py b/tools/step_probe.py +new file mode 100644 +index 0000000..dfd8af2 +--- /dev/null ++++ b/tools/step_probe.py +@@ -0,0 +1,99 @@ ++#!/usr/bin/env python3 ++"""Per-step timing of Ming-Image's DiT with allocator stats and a GPU clock/power sampler. ++ ++Diagnoses step time that grows within one generation. For every DiT call it records the ++synchronized wall time, the caching allocator's reserved/allocated bytes, how many device ++mallocs and malloc retries (fragmentation) have happened so far; a sampler thread reads the ++GPU sclk, power and temperature twice a second. ++ ++ usage: PYTHONPATH= step_probe.py --prompt P.json [--runs N] -- ++""" ++import argparse ++import glob ++import json ++import sys ++import threading ++import time ++ ++ ++def read_gpu(): ++ base = "/sys/class/drm/card0/device" ++ sclk = next((l.split(":")[1].strip().rstrip("*").strip() for l in open(f"{base}/pp_dpm_sclk") if "*" in l), "?") ++ hw = sorted(glob.glob(f"{base}/hwmon/hwmon*"))[0] ++ power = int(open(f"{hw}/power1_average").read()) / 1e6 ++ temp = int(open(f"{hw}/temp1_input").read()) / 1e3 ++ busy = int(open(f"{base}/gpu_busy_percent").read()) ++ return sclk, power, temp, busy ++ ++ ++def main(): ++ ap = argparse.ArgumentParser() ++ ap.add_argument("--prompt", required=True) ++ ap.add_argument("--runs", type=int, default=1) ++ own, rest = ap.parse_known_args() ++ if rest and rest[0] == "--": ++ rest = rest[1:] ++ sys.argv = [sys.argv[0], "--prompt", own.prompt] + rest ++ ++ import torch ++ import infer ++ ++ args = infer.parse_args() ++ model_directory = infer.resolve_model_directory(args.model, local_files_only=True) ++ caps = infer.load_checkpoint_capabilities(model_directory) ++ resolution = infer.resolve_task_resolution(args.task, args.resolution) ++ sampling = caps.resolve_sampling_parameters(steps=args.steps, cfg=args.cfg) ++ dtype = infer._dtype(args.dtype) ++ model, processor = infer.load_model_and_processor(model_directory, args) ++ prompt = infer._load_prompt(own.prompt) ++ ++ samples, stop = [], threading.Event() ++ ++ def sampler(): ++ t0 = time.perf_counter() ++ while not stop.is_set(): ++ samples.append((round(time.perf_counter() - t0, 1),) + read_gpu()) ++ time.sleep(0.5) ++ ++ dit = model.diffusion_loss.train_model ++ marks = {} ++ ++ def pre(_module, _args, _kwargs): ++ torch.cuda.synchronize() ++ marks["t"] = time.perf_counter() ++ ++ def post(_module, _args, _kwargs, _out): ++ torch.cuda.synchronize() ++ st = torch.cuda.memory_stats() ++ step_log.append({ ++ "step_s": round(time.perf_counter() - marks["t"], 2), ++ "reserved_gib": round(torch.cuda.memory_reserved() / 2**30, 2), ++ "allocated_gib": round(torch.cuda.memory_allocated() / 2**30, 2), ++ "device_mallocs": st.get("num_device_alloc", 0), ++ "device_frees": st.get("num_device_free", 0), ++ "alloc_retries": st.get("num_alloc_retries", 0), ++ }) ++ ++ dit.register_forward_pre_hook(pre, with_kwargs=True) ++ dit.register_forward_hook(post, with_kwargs=True) ++ thread = threading.Thread(target=sampler, daemon=True) ++ thread.start() ++ for run in range(own.runs): ++ step_log = [] ++ torch.cuda.synchronize() ++ t0 = time.perf_counter() ++ infer.run_generation(model, processor, caps, task=args.task, prompt=prompt, input_image=None, ++ resolution=resolution, sampling=sampling, seed=args.seed, ++ num_layers=args.num_layers, dtype=dtype) ++ torch.cuda.synchronize() ++ print(f"RUN {run} total_s {time.perf_counter() - t0:.1f}", flush=True) ++ for i, row in enumerate(step_log): ++ print("STEP " + json.dumps({"run": run, "i": i, **row}), flush=True) ++ stop.set() ++ thread.join() ++ for s in samples[:: max(1, len(samples) // 60)]: ++ print("GPU t=%6.1fs sclk=%s power=%.0fW temp=%.0fC busy=%d%%" % s) ++ ++ ++if __name__ == "__main__": ++ main() +diff --git a/tools/verify_package.py b/tools/verify_package.py +new file mode 100644 +index 0000000..cfb91f7 +--- /dev/null ++++ b/tools/verify_package.py +@@ -0,0 +1,110 @@ ++#!/usr/bin/env python3 ++"""Pre-upload verification of the INT8 package against the upstream download. Read-only on both trees ++except for writing SHA256SUMS into the package. Exits non-zero on the first failed check. ++ ++ usage: verify_package.py UPSTREAM_DIR PACKAGE_DIR ++""" ++import hashlib ++import json ++import os ++import sys ++from pathlib import Path ++ ++import torch ++from safetensors import safe_open ++ ++ ++def fail(msg): ++ sys.exit(f"VERIFY_FAIL {msg}") ++ ++ ++def shard_map(d): ++ index = json.loads((d / "model.safetensors.index.json").read_text()) ++ return index["weight_map"] ++ ++ ++def main(): ++ up, pkg = Path(sys.argv[1]), Path(sys.argv[2]) ++ ++ # 1. connector: bf16 file == upstream fp32 cast to bf16, tensor by tensor ++ up_map = shard_map(up / "connector") ++ with safe_open(str(pkg / "connector/model.safetensors"), "pt") as new: ++ if set(new.keys()) != set(up_map): ++ fail("connector tensor set differs from upstream") ++ n = 0 ++ for shard in sorted(set(up_map.values())): ++ with safe_open(str(up / "connector" / shard), "pt") as old: ++ for key in old.keys(): ++ ref = old.get_tensor(key) ++ ref = ref.to(torch.bfloat16) if ref.is_floating_point() else ref ++ got = new.get_tensor(key) ++ if got.dtype != ref.dtype or not torch.equal(got, ref): ++ fail(f"connector {key} != fp32->bf16") ++ n += 1 ++ print(f"OK connector: {n} tensors equal upstream fp32 -> bf16", flush=True) ++ ++ # 2. unchanged components: hardlink (same inode) or identical bytes ++ same = 0 ++ for comp in ("transformer", "vae", "mlp", "scheduler"): ++ for f in sorted((up / comp).rglob("*")): ++ if f.is_dir(): ++ continue ++ g = pkg / f.relative_to(up) ++ if not g.is_file(): ++ fail(f"missing {g}") ++ if os.stat(f).st_ino != os.stat(g).st_ino and f.read_bytes() != g.read_bytes(): ++ fail(f"{g} differs from upstream") ++ same += 1 ++ if (up / "LICENSE").read_bytes() != (pkg / "LICENSE").read_bytes(): ++ fail("LICENSE differs from upstream") ++ print(f"OK unchanged components: {same} files identical to upstream (+ LICENSE)", flush=True) ++ ++ # 3. mllm: copied tensors byte-identical, quantized ones present as int8 + fp32 scale ++ manifest = json.loads((pkg / "mllm/int8_manifest.json").read_text()) ++ quant = set(manifest["quantized_modules"]) ++ old_map, new_map = shard_map(up / "mllm"), shard_map(pkg / "mllm") ++ expect_new = {k for k in old_map if k[: -len(".weight")] not in quant or not k.endswith(".weight")} ++ expect_new |= {m + ".weight" for m in quant} | {m + ".scale" for m in quant} ++ if set(new_map) != expect_new: ++ fail(f"mllm index: {len(set(new_map) ^ expect_new)} names differ from the expected set") ++ handles = {} ++ ++ def tensor(tree, mapping, key): ++ path = str(tree / mapping[key]) ++ if path not in handles: ++ handles[path] = safe_open(path, "pt") ++ return handles[path].get_tensor(key) ++ ++ copied = quantized = 0 ++ for key in sorted(old_map): ++ module = key[: -len(".weight")] if key.endswith(".weight") else None ++ ref = tensor(up / "mllm", old_map, key) ++ if module in quant: ++ w, s = tensor(pkg / "mllm", new_map, key), tensor(pkg / "mllm", new_map, module + ".scale") ++ if w.dtype != torch.int8 or s.dtype != torch.float32 or w.shape != ref.shape or s.shape != (ref.shape[0],): ++ fail(f"{key}: int8/scale dtype or shape wrong") ++ quantized += 1 ++ else: ++ got = tensor(pkg / "mllm", new_map, key) ++ if got.dtype != ref.dtype or not torch.equal(got, ref): ++ fail(f"{key}: copied tensor differs from upstream") ++ copied += 1 ++ if len(handles) > 4: ++ handles.clear() ++ print(f"OK mllm: {copied} tensors byte-identical to upstream, {quantized} quantized (int8 + fp32 scale)", flush=True) ++ ++ # 4. sha256 of every file in the package ++ lines = [] ++ for f in sorted(p for p in pkg.rglob("*") if p.is_file() and p.name != "SHA256SUMS" and ".cache" not in p.parts): ++ h = hashlib.sha256() ++ with open(f, "rb") as fh: ++ for chunk in iter(lambda: fh.read(1 << 24), b""): ++ h.update(chunk) ++ lines.append(f"{h.hexdigest()} {f.relative_to(pkg).as_posix()}") ++ (pkg / "SHA256SUMS").write_text("\n".join(lines) + "\n") ++ print(f"OK sha256: {len(lines)} files -> SHA256SUMS", flush=True) ++ print("VERIFY_OK", flush=True) ++ ++ ++if __name__ == "__main__": ++ main()