diff --git a/.agents/skills/moe-models/SKILL.md b/.agents/skills/moe-models/SKILL.md index 76622ab96..aed597972 100644 --- a/.agents/skills/moe-models/SKILL.md +++ b/.agents/skills/moe-models/SKILL.md @@ -476,6 +476,25 @@ config.norm_topk_prob # Whether to normalize routing weights config.routed_scaling_factor # Post-normalization scale ``` +Current Transformers exposes NemotronH layer types as +`linear_attention` / `full_attention`; older configs use +`mamba` / `attention`. Normalize both vocabularies to Mobius +`mamba2` / `full_attention`, preserve `mlp` and `moe` distinctly, and reject +unknown values. Never map `mlp` to `moe` in parity fixtures. + +### Reduced-precision routing + +ONNX has no implicit mixed-float type promotion. NemotronH routing computes in +fp32, so keep the correction-bias initializer in fp32 and explicitly cast the +gate weight to fp32. Cast each expert output up, multiply and accumulate all +routed contributions in fp32, then cast the completed routed tensor back once. +Graph construction alone may miss this; execute fp16 and bf16 MoE paths. + +Official Nemotron 3.5 checkpoints also contain auxiliary `mtp.*` tensors. +The base `NemotronHForCausalLM` generation graph does not instantiate them and +marks them unexpected. Filter only that prefix and prove weight alignment still +populates every base-decoder initializer. + ### com.microsoft.MoE compatibility **Not compatible with NemotronH.** Three blockers: diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 91dfb2c8e..d0af55c4d 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -56,6 +56,12 @@ jobs: path: ~/.cache/huggingface key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }} + - name: Cache reduced Nemotron fixture + uses: actions/cache@v6 + with: + path: ${{ runner.temp }}/mobius-nemotron-cache + key: nemotron-reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1 + - name: Install PyTorch (CUDA) run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 @@ -82,11 +88,12 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} MOBIUS_TEST_DEVICE: cuda + MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors run: | AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then echo "Running all L4 golden comparison tests" - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m golden \ -v \ --timeout=300 \ @@ -98,7 +105,7 @@ jobs: # Convert JSON array to comma-separated list for --models MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))") if [ -n "$MODELS" ]; then - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m golden \ -v \ --models "$MODELS" \ diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index 23808927b..255e756ad 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -56,6 +56,12 @@ jobs: path: ~/.cache/huggingface key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }} + - name: Cache reduced Nemotron fixture + uses: actions/cache@v6 + with: + path: ${{ runner.temp }}/mobius-nemotron-cache + key: nemotron-reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1 + - name: Install PyTorch (CUDA) run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 @@ -64,17 +70,19 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' + pip install -r examples/olive/nemotron-3_5-lightning-30b/requirements.txt --index-url https://packagefeedproxy.microsoft.io/pypi/simple pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda - name: Run L5 generation E2E tests env: HF_TOKEN: ${{ secrets.HF_TOKEN }} MOBIUS_TEST_DEVICE: cuda + MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors run: | AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then echo "Running all L5 generation E2E tests" - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m generation \ -v \ --timeout=300 \ @@ -86,7 +94,7 @@ jobs: # Convert JSON array to comma-separated list for --models MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))") if [ -n "$MODELS" ]; then - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m generation \ -v \ --models "$MODELS" \ diff --git a/examples/olive/nemotron-3_5-lightning-30b/.gitignore b/examples/olive/nemotron-3_5-lightning-30b/.gitignore new file mode 100644 index 000000000..ab3bb1740 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/.gitignore @@ -0,0 +1,2 @@ +cache/ +output/ diff --git a/examples/olive/nemotron-3_5-lightning-30b/README.md b/examples/olive/nemotron-3_5-lightning-30b/README.md new file mode 100644 index 000000000..46ed93113 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/README.md @@ -0,0 +1,173 @@ +# Nemotron 3.5 Lightning: BF16 checkpoint + Olive + +This is **Option A** for +[`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16): +export the official BF16 checkpoint to supported FP16 ONNX, quantize the model +with Olive, assemble a direct ONNX Runtime package, then run cached generation. + +Every Hub access is pinned to revision +`d468880b6ad3c6e0d21377ce7242adaea4cc884d`. + +## Architecture and runtime contract + +The checkpoint is a real `nemotron_h` model, not an alias: + +- 52 base-decoder layers mixing Mamba2, sigmoid-routed MoE, and full GQA. +- 128 routed experts with top-6 selection and one shared expert. +- `mtp.*` contains 270 auxiliary multi-token-prediction tensors. This export + intentionally targets the base `NemotronHForCausalLM` decoder; its forward + graph does not instantiate MTP, and upstream marks those keys unexpected. + No base-decoder generation input, cache, logit, or weight depends on them. + +This model's graph mixes `conv_state` plus `ssm_state` with sparse +full-attention key/value caches. Mobius emits every field the current +`genai_config.json` schema can represent (semantic inputs, key/value and +convolution templates, and global cache-slot count) and leaves runtime +acceptance to ORT GenAI. The schema currently has no `ssm_state` template. The +validated recipe uses direct ONNX Runtime generation through `inference.py`; +downstream load/generation outcomes are informational and do not gate export. + +## Install + +From the repository root: + +```powershell +python -m pip install -e ".[transformers,testing]" ` + --index-url https://packagefeedproxy.microsoft.io/pypi/simple +python -m pip install -r examples\olive\nemotron-3_5-lightning-30b\requirements.txt ` + --index-url https://packagefeedproxy.microsoft.io/pypi/simple +``` + +Use an ONNX Runtime GPU build with CUDA 12 and cuDNN 9 for CUDA inference. + +## Full export, quantization, and smoke test + +```powershell +cd examples\olive\nemotron-3_5-lightning-30b +python optimize.py ` + --source-dir output\f16\cuda ` + --output-dir output\Q4_K_M\cuda ` + --ep cuda ` + --precision q4_k_m +``` + +The script performs four gated steps: + +1. Downloads the exact 14-shard BF16 checkpoint revision and exports FP16 ONNX. + BF16 execution is rejected explicitly because corrected reduced-real parity + reaches `0.8594` max logit error, above the `1e-2` reduced-precision gate. +2. Emits standard ONNX cache operations and applies the grouped-RMSNorm CUDA + workaround. CUDA still executes supported compute nodes; portable cache + operations avoid provider-dependent fused decode drift. +3. Runs Olive Q4 K-quant with an explicitly CPU-only target. It also suppresses + Olive 0.13's unrelated GPU-EP DLL auto-registration, so a missing TensorRT + installation cannot abort CPU weight-only quantization. +4. Reloads the assembled quantized package and generates four cached tokens. + +To reuse an existing source package: + +```powershell +python optimize.py --skip-export ` + --source-dir output\f16\cuda ` + --output-dir output\Q4_K_M\cuda ` + --ep cuda +``` + +`olive_q4.json` records the equivalent pass and provider-isolated target. Use +`optimize.py` rather than invoking the JSON directly when the installed ORT +wheel bundles unconfigured providers; the script contains the verified Olive +0.13 registration isolation. + +## Package layout + +```text +output/ +├── f16/cuda/ +│ ├── model.onnx +│ ├── model.onnx.data +│ ├── config.json +│ ├── generation_config.json +│ ├── tokenizer.json +│ ├── tokenizer_config.json +│ └── source_manifest.json +└── Q4_K_M/cuda/ + ├── model.onnx + ├── model.onnx.data + ├── config.json + ├── generation_config.json + ├── tokenizer.json + ├── tokenizer_config.json + └── source_manifest.json +``` + +This recipe intentionally uses direct ONNX Runtime, while +`mobius build --runtime ort-genai` remains allowed and emits the best current +schema metadata for downstream testing. + +## Direct generation and profiling + +```powershell +python inference.py ` + --model-dir output\Q4_K_M\cuda ` + --device cuda ` + --prompt "What is 84 * 3 / 2?" ` + --max-new-tokens 20 ` + --profile +``` + +The script initializes every cache from the saved graph, processes the prompt +token by token, carries Mamba and KV state independently, and fails if CUDA was +requested but not registered. + +## Reduced real-checkpoint validation + +The full checkpoint is 65.8 GB and cannot execute on the validation host's +8 GB RTX A1000. The reproducible reduced check range-downloads 236 MiB of real +weights while retaining production dimensions: + +- checkpoint layer 0: complete Mamba2 block; +- layer 1: router, shared expert, and four complete routed experts; +- layer 5: complete full-attention block; +- sliced real embedding and LM-head rows plus final norm. + +```powershell +python validate_reduced_checkpoint.py +``` + +The fixture is stored persistently under `~/.cache/mobius/` by default. Each +range request validates status, `Content-Range`, declared length, and payload +length, with three bounded attempts (1s then 2s backoff). The cache metadata +must match the pinned model, revision, and fixture schema; writes are atomic. +GPU CI restores the same revision/schema-keyed cache for L4 and L5. + +The supported matrix is intentionally limited to FP32/CPU and FP16/CUDA. +Reproduce the BF16 rejection evidence separately without creating a supported +package or weakening the production guard: + +```powershell +python validate_reduced_checkpoint.py --bf16-rejection-evidence +``` + +Validated results on ORT 1.28.0 / Olive 0.13.0: + +| Variant | Full-logit max abs | Generated IDs | Placement | +|---|---:|---|---| +| FP32 CPU | `9.54e-6` | `12, 13, 12, 12` | CPU | +| FP16 CUDA | `<= 0.0078125` (prefill + every cached step) | `12, 13, 12, 12` | portable ONNX graph on CUDA | +| BF16 CUDA | rejected (`0.8594`) | N/A | fails numerical gate | +| Olive Q4 | quantized | `12, 13, 12, 12` | portable ONNX graph on CUDA | + +The reduced package's portable weighted graph quantizes 17 matrix +multiplications to `com.microsoft::MatMulNBits` and reloads successfully for +multi-token generation. Record size/compression from the produced package; +it varies with external-data serialization and Olive version. + +## Evidence-based waivers + +- Full-checkpoint L4/L5 coherent-text generation: requires roughly 66 GB just + for checkpoint storage and substantially more than 8 GB accelerator memory. +- Full 30B Olive run: the recipe and reduced production-dimension pass are + validated; completing all 2,944 expert subgraphs requires a large-memory + host. +- Foundry Local was not available on this host. That downstream validation + remains informational and does not block Mobius export. diff --git a/examples/olive/nemotron-3_5-lightning-30b/inference.py b/examples/olive/nemotron-3_5-lightning-30b/inference.py new file mode 100644 index 000000000..272580f43 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/inference.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Direct ONNX Runtime generation for NemotronH hybrid-cache packages.""" + +from __future__ import annotations + +import argparse +import json +from collections import Counter +from collections.abc import Collection +from pathlib import Path +from typing import Any + +import numpy as np + +MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +REVISION = "d468880b6ad3c6e0d21377ce7242adaea4cc884d" +_BFLOAT16_ONNX_TYPE = 16 + + +def _numpy_dtype(ort_type: str): + if ort_type == "tensor(float)": + return np.float32 + if ort_type == "tensor(float16)": + return np.float16 + if ort_type == "tensor(bfloat16)": + import ml_dtypes + + return ml_dtypes.bfloat16 + raise TypeError(f"Unsupported state input type: {ort_type}") + + +def _concrete_state_shape(shape: list[Any]) -> tuple[int, ...]: + concrete: list[int] = [] + for dim in shape: + if isinstance(dim, int): + concrete.append(dim) + elif "batch" in str(dim): + concrete.append(1) + elif "past" in str(dim): + concrete.append(0) + else: + raise ValueError(f"Cannot resolve hybrid-cache dimension {dim!r}") + return tuple(concrete) + + +def _initial_states(session) -> dict[str, Any]: + import onnxruntime as ort + + states: dict[str, Any] = {} + for model_input in session.get_inputs(): + if not model_input.name.startswith("past_key_values."): + continue + shape = _concrete_state_shape(model_input.shape) + if model_input.type == "tensor(bfloat16)": + states[model_input.name] = ort.OrtValue.ortvalue_from_numpy_with_onnx_type( + np.zeros(shape, dtype=np.uint16), + _BFLOAT16_ONNX_TYPE, + ) + else: + states[model_input.name] = np.zeros( + shape, + dtype=_numpy_dtype(model_input.type), + ) + return states + + +def _update_states( + states: dict[str, Any], + output_names: list[str], + output_values: list[Any], +) -> None: + for name, value in zip(output_names, output_values): + if not name.startswith("present."): + continue + input_name = name.replace("present.", "past_key_values.", 1) + if input_name in states: + states[input_name] = value + + +def _token_feeds( + session, + token_ids: np.ndarray, + *, + total_length: int, + position_ids: np.ndarray, + states: dict[str, Any], +) -> dict[str, Any]: + available = {model_input.name for model_input in session.get_inputs()} + candidates = { + "input_ids": token_ids, + "attention_mask": np.ones((1, total_length), dtype=np.int64), + "position_ids": position_ids, + **states, + } + return {name: value for name, value in candidates.items() if name in available} + + +def _run_session(session, output_names: list[str], feeds: dict[str, Any]) -> list[Any]: + import onnxruntime as ort + + if not any(isinstance(value, ort.OrtValue) for value in feeds.values()): + return session.run(output_names, feeds) + ort_feeds = { + name: ( + value + if isinstance(value, ort.OrtValue) + else ort.OrtValue.ortvalue_from_numpy(value) + ) + for name, value in feeds.items() + } + return list(session.run_with_ort_values(output_names, ort_feeds)) + + +def _as_numpy(value: Any) -> np.ndarray: + import onnxruntime as ort + + if not isinstance(value, ort.OrtValue): + return value + if value.data_type() == "tensor(bfloat16)": + import torch + + return torch.from_dlpack(value).float().cpu().numpy() + return value.numpy() + + +def _logits_output_name(output_names: list[str]) -> str: + """Resolve the full-precision or Olive-renamed logits output.""" + for candidate in ("logits", "logits_Q4"): + if candidate in output_names: + return candidate + raise ValueError(f"Model has no logits output; found {output_names}") + + +def _create_session(model_path: Path, device: str, profile: bool): + if device == "cuda": + # Importing PyTorch first preloads its matching CUDA/cuDNN DLLs on Windows. + import torch # noqa: F401 + + import onnxruntime as ort + + if device == "cuda" and hasattr(ort, "preload_dlls"): + ort.preload_dlls() + options = ort.SessionOptions() + options.enable_profiling = profile + providers = ( + ["CUDAExecutionProvider", "CPUExecutionProvider"] + if device == "cuda" + else ["CPUExecutionProvider"] + ) + session = ort.InferenceSession( + str(model_path), + sess_options=options, + providers=providers, + ) + if device == "cuda" and session.get_providers()[0] != "CUDAExecutionProvider": + raise RuntimeError( + f"CUDAExecutionProvider was requested but providers are {session.get_providers()}" + ) + return session + + +def load_eos_token_ids(model_dir: str | Path) -> set[int]: + """Load scalar or list EOS IDs from the assembled package metadata.""" + model_dir = Path(model_dir) + eos_ids: set[int] = set() + for filename in ("generation_config.json", "config.json"): + path = model_dir / filename + if not path.is_file(): + continue + raw_eos = json.loads(path.read_text(encoding="utf-8")).get("eos_token_id") + if isinstance(raw_eos, int): + eos_ids.add(raw_eos) + elif isinstance(raw_eos, list): + eos_ids.update(value for value in raw_eos if isinstance(value, int)) + return eos_ids + + +def run_token_ids( + model_dir: str | Path, + input_ids: list[int], + *, + max_new_tokens: int, + device: str, + profile: bool = False, + eos_token_ids: Collection[int] | None = None, +) -> tuple[list[int], list[np.ndarray], str | None]: + """Run token-by-token hybrid-cache generation and return IDs plus logits.""" + model_path = Path(model_dir) / "model.onnx" + if not model_path.is_file(): + raise FileNotFoundError(f"Missing ONNX model: {model_path}") + if not input_ids: + raise ValueError("input_ids must not be empty") + + session = _create_session(model_path, device, profile) + states = _initial_states(session) + output_names = [output.name for output in session.get_outputs()] + logits_output_name = _logits_output_name(output_names) + generated: list[int] = [] + logits_by_step: list[np.ndarray] = [] + past_length = 0 + outputs: list[Any] | None = None + + for token_id in input_ids: + feeds = _token_feeds( + session, + np.array([[token_id]], dtype=np.int64), + total_length=past_length + 1, + position_ids=np.array([[past_length]], dtype=np.int64), + states=states, + ) + outputs = _run_session(session, output_names, feeds) + _update_states(states, output_names, outputs) + past_length += 1 + + assert outputs is not None + logits = _as_numpy(outputs[output_names.index(logits_output_name)])[0, -1].astype( + np.float32 + ) + eos_ids = set(eos_token_ids or ()) + for token_index in range(max_new_tokens): + logits_by_step.append(logits.copy()) + token_id = int(np.argmax(logits)) + generated.append(token_id) + if token_id in eos_ids or token_index + 1 == max_new_tokens: + break + feeds = _token_feeds( + session, + np.array([[token_id]], dtype=np.int64), + total_length=past_length + 1, + position_ids=np.array([[past_length]], dtype=np.int64), + states=states, + ) + outputs = _run_session(session, output_names, feeds) + _update_states(states, output_names, outputs) + past_length += 1 + logits = _as_numpy(outputs[output_names.index(logits_output_name)])[0, -1].astype( + np.float32 + ) + + profile_path = session.end_profiling() if profile else None + return generated, logits_by_step, profile_path + + +def summarize_profile(profile_path: str) -> dict[str, int]: + """Summarize actual node placement from an ORT profiling JSON file.""" + events = json.loads(Path(profile_path).read_text(encoding="utf-8")) + providers: Counter[str] = Counter() + for event in events: + args = event.get("args", {}) + provider = args.get("provider") + if event.get("cat") == "Node" and provider: + providers[str(provider)] += 1 + return dict(sorted(providers.items())) + + +def _tokenize_prompt(model_dir: Path, prompt: str, use_chat: bool): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + model_dir, + revision=None, + local_files_only=True, + ) + if use_chat: + ids = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=True, + add_generation_prompt=True, + ) + else: + ids = tokenizer.encode(prompt) + return tokenizer, list(ids) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", required=True) + parser.add_argument("--prompt", default="What is 84 * 3 / 2?") + parser.add_argument("--token-ids", nargs="+", type=int) + parser.add_argument("--max-new-tokens", type=int, default=20) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda") + parser.add_argument("--no-chat", action="store_true") + parser.add_argument("--profile", action="store_true") + args = parser.parse_args() + + model_dir = Path(args.model_dir) + tokenizer = None + if args.token_ids: + input_ids = args.token_ids + else: + tokenizer, input_ids = _tokenize_prompt(model_dir, args.prompt, not args.no_chat) + + generated, _logits, profile_path = run_token_ids( + model_dir, + input_ids, + max_new_tokens=args.max_new_tokens, + device=args.device, + profile=args.profile, + eos_token_ids=load_eos_token_ids(model_dir), + ) + if not generated: + raise RuntimeError("Generation produced no tokens") + + if tokenizer is None: + print("Generated token IDs:", generated) + else: + text = tokenizer.decode(generated, skip_special_tokens=True) + if not text.strip(): + raise RuntimeError("Generation produced only empty/special-token text") + print(text) + + if profile_path is not None: + print(f"ORT profile: {profile_path}") + print("Node placement:", summarize_profile(profile_path)) + + +if __name__ == "__main__": + main() diff --git a/examples/olive/nemotron-3_5-lightning-30b/olive_q4.json b/examples/olive/nemotron-3_5-lightning-30b/olive_q4.json new file mode 100644 index 000000000..e4d316cde --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/olive_q4.json @@ -0,0 +1,32 @@ +{ + "input_model": { + "type": "OnnxModel", + "model_path": "output/f16/cuda/model.onnx" + }, + "passes": { + "q4_k_m": { + "type": "OnnxKQuantQuantization", + "bits": 4, + "block_size": 32, + "save_as_external_data": true, + "all_tensors_to_one_file": true, + "external_data_name": "model.onnx.data", + "size_threshold": 1024 + } + }, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": [ + "CPUExecutionProvider" + ] + } + ] + } + }, + "no_artifacts": true, + "output_dir": "output/Q4_K_M/cuda" +} diff --git a/examples/olive/nemotron-3_5-lightning-30b/optimize.py b/examples/olive/nemotron-3_5-lightning-30b/optimize.py new file mode 100644 index 000000000..4ac40d470 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/optimize.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pinned BF16 export and Olive INT4 packaging for Nemotron 3.5 Lightning.""" + +from __future__ import annotations + +import argparse +import json +import shutil +import tempfile +from pathlib import Path + +from inference import MODEL_ID, REVISION, load_eos_token_ids, run_token_ids + +_METADATA_FILES = { + "added_tokens.json", + "chat_template.jinja", + "config.json", + "generation_config.json", + "merges.txt", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "vocab.json", +} + + +def _require_empty_output(path: Path) -> None: + if path.exists() and any(path.iterdir()): + raise FileExistsError(f"Output directory must be empty: {path}") + path.mkdir(parents=True, exist_ok=True) + + +def _save_pinned_metadata(output_dir: Path) -> None: + from transformers import AutoConfig, AutoTokenizer, GenerationConfig + + config = AutoConfig.from_pretrained( + MODEL_ID, + revision=REVISION, + trust_remote_code=False, + ) + config.save_pretrained(output_dir) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION) + tokenizer.save_pretrained(output_dir) + generation = GenerationConfig.from_pretrained(MODEL_ID, revision=REVISION) + generation.save_pretrained(output_dir) + + (output_dir / "source_manifest.json").write_text( + json.dumps( + { + "model_id": MODEL_ID, + "revision": REVISION, + "runtime": "onnxruntime-direct", + "ort_genai_supported": False, + }, + indent=2, + ), + encoding="utf-8", + ) + + +def export_checkpoint(output_dir: str | Path, *, ep: str) -> Path: + """Export the pinned BF16 checkpoint as a supported FP16 ONNX package.""" + from mobius import build + from mobius._flags import override_flags + + output = Path(output_dir) + _require_empty_output(output) + # CUDA fused cache kernels diverge during multi-step decode for this hybrid + # architecture. Keep a portable graph and let CUDA place supported standard + # ops; reduced-real validation enforces <=1e-2 at every cached step. + build_ep = "onnx-standard" if ep == "cuda" else ep + with override_flags(ort_cuda_grouped_rmsnorm_workaround=ep == "cuda"): + package = build( + MODEL_ID, + revision=REVISION, + dtype="f16", + load_weights=True, + trust_remote_code=False, + execution_provider=build_ep, + ) + package.save(output, external_data="onnx") + _save_pinned_metadata(output) + manifest_path = output / "source_manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest.update( + { + "source_dtype": "bf16", + "dtype": "f16", + "target_ep": ep, + "build_ep": build_ep, + } + ) + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return output + + +def _olive_config(source_model: Path, output_dir: Path, precision: str) -> dict: + if precision == "q4_k_m": + pass_config = { + "type": "OnnxKQuantQuantization", + "bits": 4, + "block_size": 32, + } + elif precision == "nf4": + pass_config = { + "type": "OnnxBnb4Quantization", + "precision": "nf4", + } + else: + raise ValueError(f"Unsupported quantization precision: {precision}") + + pass_config.update( + { + "save_as_external_data": True, + "all_tensors_to_one_file": True, + "external_data_name": "model.onnx.data", + "size_threshold": 1024, + } + ) + return { + "input_model": { + "type": "OnnxModel", + "model_path": str(source_model), + }, + "passes": {precision: pass_config}, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"], + } + ], + } + }, + "no_artifacts": True, + "output_dir": str(output_dir), + } + + +def _find_olive_model(output_dir: Path) -> Path: + candidates = list(output_dir.rglob("*.onnx")) + if len(candidates) != 1: + raise RuntimeError( + f"Expected exactly one Olive ONNX output under {output_dir}, got {candidates}" + ) + return candidates[0] + + +def _copy_olive_model(model_path: Path, destination: Path) -> None: + for child in model_path.parent.iterdir(): + if child.is_file(): + shutil.copy2(child, destination / child.name) + copied_model = destination / model_path.name + canonical_model = destination / "model.onnx" + if copied_model != canonical_model: + copied_model.replace(canonical_model) + + +def quantize_package( + source_dir: str | Path, + output_dir: str | Path, + *, + precision: str = "q4_k_m", +) -> Path: + """Quantize model.onnx with a CPU-isolated Olive workflow.""" + import olive.systems.local as olive_local + from olive.workflows import run as olive_run + + source = Path(source_dir) + source_model = source / "model.onnx" + if not source_model.is_file(): + raise FileNotFoundError(f"Missing source model: {source_model}") + output = Path(output_dir) + _require_empty_output(output) + + with tempfile.TemporaryDirectory(prefix="olive-nemotron-") as temp: + olive_output = Path(temp) / "output" + config = _olive_config(source_model, olive_output, precision) + config["cache_dir"] = str(Path(temp) / "cache") + config["clean_cache"] = True + # Olive 0.13 auto-registers every DLL bundled in a GPU ORT wheel, + # even for a CPU-only target. That makes an unrelated TensorRT DLL + # failure abort weight-only quantization. Suppress registration for + # this pass; the explicit workflow target remains CPU-only. + register_ep_libraries = olive_local.maybe_register_ep_libraries + olive_local.maybe_register_ep_libraries = lambda _paths: None + try: + olive_run(config) + finally: + olive_local.maybe_register_ep_libraries = register_ep_libraries + _copy_olive_model(_find_olive_model(olive_output), output) + + for name in _METADATA_FILES | {"source_manifest.json"}: + path = source / name + if path.is_file(): + shutil.copy2(path, output / name) + manifest_path = output / "source_manifest.json" + manifest = ( + json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest_path.is_file() + else {"model_id": MODEL_ID, "revision": REVISION} + ) + manifest.update({"quantization": precision, "olive_provider": "CPUExecutionProvider"}) + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return output + + +def smoke_test(model_dir: str | Path, *, device: str) -> list[int]: + """Load the assembled package and perform cached multi-token generation.""" + import numpy as np + + generated, logits, _profile = run_token_ids( + model_dir, + [1, 42, 17], + max_new_tokens=4, + device=device, + eos_token_ids=load_eos_token_ids(model_dir), + ) + if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): + raise RuntimeError(f"Quantized generation smoke test failed: {generated}") + return generated + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-dir", default="output/f16/cuda") + parser.add_argument("--output-dir", default="output/Q4_K_M/cuda") + parser.add_argument("--ep", choices=["cpu", "cuda"], default="cuda") + parser.add_argument("--precision", choices=["q4_k_m", "nf4"], default="q4_k_m") + parser.add_argument("--skip-export", action="store_true") + parser.add_argument("--skip-quantization", action="store_true") + parser.add_argument("--skip-smoke", action="store_true") + args = parser.parse_args() + + if not args.skip_export: + export_checkpoint(args.source_dir, ep=args.ep) + result_dir = Path(args.source_dir) + if not args.skip_quantization: + result_dir = quantize_package( + args.source_dir, + args.output_dir, + precision=args.precision, + ) + if not args.skip_smoke: + print("Generated token IDs:", smoke_test(result_dir, device=args.ep)) + + +if __name__ == "__main__": + main() diff --git a/examples/olive/nemotron-3_5-lightning-30b/requirements.txt b/examples/olive/nemotron-3_5-lightning-30b/requirements.txt new file mode 100644 index 000000000..b29382a5f --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/requirements.txt @@ -0,0 +1,3 @@ +olive-ai>=0.13.0 +requests>=2.25 +safetensors>=0.4 diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py new file mode 100644 index 000000000..839ae4a68 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -0,0 +1,664 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Validate a reduced model assembled from byte ranges of the pinned checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import math +import struct +import time +from collections import Counter +from pathlib import Path + +import numpy as np +import requests +import torch +from huggingface_hub import hf_hub_download +from inference import ( + MODEL_ID, + REVISION, + _as_numpy, + _create_session, + _initial_states, + _run_session, + _token_feeds, + load_eos_token_ids, + run_token_ids, + summarize_profile, +) +from optimize import quantize_package +from safetensors import safe_open +from safetensors.torch import load_file, save_file + +_VOCAB_SIZE = 256 +_NUM_EXPERTS = 4 +_LAYER_REMAP = {0: 0, 1: 1, 5: 2} +FIXTURE_SCHEMA_VERSION = 1 +_RANGE_ATTEMPTS = 3 +_DTYPES = { + "f32": (torch.float32, "FLOAT"), + "f16": (torch.float16, "FLOAT16"), +} + + +class _PinnedSafetensors: + """Read selected tensors via HTTP Range without downloading 65.8 GB.""" + + def __init__(self) -> None: + index_path = hf_hub_download( + MODEL_ID, + "model.safetensors.index.json", + revision=REVISION, + ) + index = json.loads(Path(index_path).read_text(encoding="utf-8")) + self.weight_map: dict[str, str] = index["weight_map"] + self._headers: dict[str, tuple[int, dict]] = {} + self._session = requests.Session() + + def _url(self, shard: str) -> str: + return f"https://huggingface.co/{MODEL_ID}/resolve/{REVISION}/{shard}" + + def _range(self, shard: str, start: int, end: int) -> bytes: + expected = end - start + 1 + expected_range_prefix = f"bytes {start}-{end}/" + last_error = "" + for attempt in range(_RANGE_ATTEMPTS): + try: + with self._session.get( + self._url(shard), + headers={"Range": f"bytes={start}-{end}"}, + timeout=180, + stream=True, + ) as response: + content_range = response.headers.get("Content-Range", "") + content_length = response.headers.get("Content-Length") + if response.status_code != 206: + last_error = f"status={response.status_code}" + elif not content_range.startswith(expected_range_prefix): + last_error = f"content-range={content_range!r}" + elif content_length is not None and ( + not content_length.isdecimal() or int(content_length) != expected + ): + last_error = f"content-length={content_length}, expected={expected}" + else: + payload = response.content + if len(payload) == expected: + return payload + last_error = f"bytes={len(payload)}, expected={expected}" + except requests.RequestException as error: + last_error = f"{type(error).__name__}: {error}" + + if attempt + 1 < _RANGE_ATTEMPTS: + time.sleep(2**attempt) + + raise RuntimeError( + f"Range fetch failed after {_RANGE_ATTEMPTS} attempts for " + f"{shard} bytes {start}-{end}: {last_error}" + ) + + def _header(self, shard: str) -> tuple[int, dict]: + if shard not in self._headers: + header_length = struct.unpack(" torch.Tensor: + shard = self.weight_map[name] + header_length, header = self._header(shard) + entry = header[name] + shape = list(entry["shape"]) + dtype_name = entry["dtype"] + dtype = {"BF16": torch.bfloat16, "F32": torch.float32}[dtype_name] + element_size = {"BF16": 2, "F32": 4}[dtype_name] + start, end = entry["data_offsets"] + if rows is not None: + if not shape or rows > shape[0]: + raise ValueError(f"Invalid row slice {rows} for {name}: {shape}") + row_elements = math.prod(shape[1:]) + end = start + rows * row_elements * element_size + shape[0] = rows + + data_start = 8 + header_length + payload = self._range(shard, data_start + start, data_start + end - 1) + tensor = torch.frombuffer(bytearray(payload), dtype=dtype).clone() + return tensor.reshape(shape) + + +def _source_to_target(name: str) -> str: + if name == "backbone.embeddings.weight": + return "model.embeddings.weight" + if name == "backbone.norm_f.weight": + return "model.norm_f.weight" + if not name.startswith("backbone.layers."): + return name + parts = name.split(".") + source_layer = int(parts[2]) + parts[2] = str(_LAYER_REMAP[source_layer]) + parts[0] = "model" + return ".".join(parts) + + +def default_reduced_cache_path() -> Path: + """Return the persistent, revision-and-schema-keyed fixture cache path.""" + return ( + Path.home() + / ".cache" + / "mobius" + / "nemotron-3_5-lightning" + / f"reduced-{REVISION}-schema-v{FIXTURE_SCHEMA_VERSION}.safetensors" + ) + + +def _build_reduced_state(cache_path: Path) -> dict[str, torch.Tensor]: + if cache_path.is_file(): + with safe_open(cache_path, framework="pt") as cached: + metadata = cached.metadata() or {} + expected_metadata = { + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": str(FIXTURE_SCHEMA_VERSION), + } + actual_metadata = {key: metadata.get(key) for key in expected_metadata} + if actual_metadata != expected_metadata: + raise ValueError( + "Reduced cache metadata mismatch: " + f"expected={expected_metadata}, actual={actual_metadata}. " + "Remove the stale cache file and retry." + ) + return load_file(cache_path) + + source = _PinnedSafetensors() + state: dict[str, torch.Tensor] = { + "model.embeddings.weight": source.tensor( + "backbone.embeddings.weight", + rows=_VOCAB_SIZE, + ), + "model.norm_f.weight": source.tensor("backbone.norm_f.weight"), + "lm_head.weight": source.tensor("lm_head.weight", rows=_VOCAB_SIZE), + } + + for source_layer in (0, 5): + prefix = f"backbone.layers.{source_layer}." + for name in sorted(source.weight_map): + if name.startswith(prefix): + state[_source_to_target(name)] = source.tensor(name) + + moe_prefix = "backbone.layers.1.mixer" + state["model.layers.1.norm.weight"] = source.tensor("backbone.layers.1.norm.weight") + state["model.layers.1.mixer.gate.weight"] = source.tensor( + f"{moe_prefix}.gate.weight", + rows=_NUM_EXPERTS, + ) + state["model.layers.1.mixer.gate.e_score_correction_bias"] = source.tensor( + f"{moe_prefix}.gate.e_score_correction_bias", + rows=_NUM_EXPERTS, + ) + for projection in ("up_proj", "down_proj"): + state[f"model.layers.1.mixer.experts.{projection}"] = torch.stack( + [ + source.tensor(f"{moe_prefix}.experts.{expert}.{projection}.weight") + for expert in range(_NUM_EXPERTS) + ] + ) + state[f"model.layers.1.mixer.shared_experts.{projection}.weight"] = source.tensor( + f"{moe_prefix}.shared_experts.{projection}.weight" + ) + + cache_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = cache_path.with_name(f"{cache_path.name}.tmp") + temporary_path.unlink(missing_ok=True) + save_file( + {name: tensor.contiguous() for name, tensor in state.items()}, + temporary_path, + metadata={ + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": str(FIXTURE_SCHEMA_VERSION), + }, + ) + temporary_path.replace(cache_path) + return state + + +def _hf_config(): + from transformers import NemotronHConfig + + return NemotronHConfig( + vocab_size=_VOCAB_SIZE, + hidden_size=2688, + layers_block_type=["linear_attention", "moe", "full_attention"], + num_attention_heads=32, + num_key_value_heads=2, + head_dim=128, + intermediate_size=1856, + mamba_num_heads=64, + mamba_head_dim=64, + ssm_state_size=128, + n_groups=8, + conv_kernel=4, + expand=2, + use_mamba_kernels=False, + moe_intermediate_size=1856, + moe_shared_expert_intermediate_size=3712, + n_routed_experts=_NUM_EXPERTS, + num_experts_per_tok=2, + routed_scaling_factor=2.5, + n_group=1, + topk_group=1, + norm_topk_prob=True, + layer_norm_epsilon=1e-5, + rescale_prenorm_residual=False, + max_position_embeddings=262144, + ) + + +def _hf_model( + state: dict[str, torch.Tensor], + *, + dtype: torch.dtype, + device: str, +): + from transformers import NemotronHForCausalLM + + model = NemotronHForCausalLM(_hf_config()).to(device=device, dtype=dtype) + target = model.state_dict() + if set(target) != set(state): + missing = sorted(set(target) - set(state)) + extra = sorted(set(state) - set(target)) + raise ValueError(f"Reduced state mismatch; missing={missing}, extra={extra}") + converted = { + name: tensor.to(device=device, dtype=target[name].dtype) + for name, tensor in state.items() + } + model.load_state_dict(converted, strict=True) + + # The production loader keeps this selection-only bias in fp32. + gate = model.model.layers[1].mixer.gate + gate.e_score_correction_bias = state[ + "model.layers.1.mixer.gate.e_score_correction_bias" + ].to(device=device, dtype=torch.float32) + return model.eval() + + +def _mobius_package( + state: dict[str, torch.Tensor], + *, + dtype_name: str, + ep: str, +): + import onnx_ir as ir + + from mobius import build_from_module + from mobius._configs import NemotronHConfig + from mobius._flags import override_flags + from mobius.models.nemotron_h import NemotronHCausalLMModel + + config = NemotronHConfig.from_transformers(_hf_config()) + config.dtype = getattr(ir.DataType, _DTYPES[dtype_name][1]) + module = NemotronHCausalLMModel(config) + # Keep CUDA execution portable: fused hybrid cache kernels show + # provider-dependent multi-step drift even when prefill matches. + build_ep = "onnx-standard" if ep == "cuda" else ep + with override_flags(ort_cuda_grouped_rmsnorm_workaround=ep == "cuda"): + package = build_from_module( + module, + config, + task="hybrid-text-generation", + execution_provider=build_ep, + trace_optimization=True, + ) + package.apply_weights(module.preprocess_weights(dict(state))) + unset = [ + name + for name, value in package["model"].graph.initializers.items() + if value.const_value is None + ] + if unset: + raise ValueError( + f"Weighted graph still has {len(unset)} unset parameters: {unset[:5]}" + ) + return package + + +def _bf16_rejection_evidence_package(state: dict[str, torch.Tensor]): + """Build a test-only BF16 graph without weakening the production guard.""" + import dataclasses + + import onnx_ir as ir + from onnxscript import OpBuilder, nn + + from mobius import build_from_module + from mobius._configs import NemotronHConfig + from mobius._flags import override_flags + from mobius.components import Linear + from mobius.models.nemotron_h import ( + NemotronHCausalLMModel, + _NemotronHTextModel, + ) + + config = NemotronHConfig.from_transformers(_hf_config()) + config.dtype = ir.DataType.BFLOAT16 + + # Prove this evidence path has not weakened or bypassed the production API. + try: + NemotronHCausalLMModel(config) + except ValueError as error: + if "BF16 execution is not numerically supported" not in str(error): + raise + else: + raise AssertionError("Production NemotronH BF16 guard did not reject the model") + + class _Bf16EvidenceCausalLM(nn.Module): + """Test-only wrapper around the production components.""" + + def __init__(self, evidence_config: NemotronHConfig): + super().__init__() + self.model = _NemotronHTextModel(evidence_config) + self.lm_head = Linear( + evidence_config.hidden_size, + evidence_config.vocab_size, + bias=False, + ) + + def forward( + self, + op: OpBuilder, + input_ids, + attention_mask, + position_ids, + past_key_values=None, + ): + hidden_states, present_key_values = self.model( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + return self.lm_head(op, hidden_states), present_key_values + + module = _Bf16EvidenceCausalLM(config) + with override_flags(ort_cuda_grouped_rmsnorm_workaround=True): + package = build_from_module( + module, + config, + task="hybrid-text-generation", + execution_provider="onnx-standard", + ) + + preprocessing_config = dataclasses.replace(config, dtype=ir.DataType.FLOAT) + preprocessor = NemotronHCausalLMModel(preprocessing_config) + package.apply_weights(preprocessor.preprocess_weights(dict(state))) + unset = [ + name + for name, value in package["model"].graph.initializers.items() + if value.const_value is None + ] + if unset: + raise ValueError(f"BF16 evidence graph has {len(unset)} unset parameters: {unset[:5]}") + return package + + +def _measure_bf16_rejection( + state: dict[str, torch.Tensor], + output_root: Path, +) -> float: + """Measure the rejected BF16 path on CUDA and return its maximum logit error.""" + if not torch.cuda.is_available(): + raise RuntimeError("BF16 rejection evidence requires CUDA") + + package = _bf16_rejection_evidence_package(state) + evidence_dir = output_root / "bf16-rejection-evidence" + evidence_dir.mkdir(parents=True, exist_ok=True) + package.save(evidence_dir, external_data="onnx") + + session = _create_session(evidence_dir / "model.onnx", "cuda", False) + prompt_ids = [1, 42, 17] + actual = _full_prefill(session, prompt_ids) + hf_model = _hf_model(state, dtype=torch.bfloat16, device="cuda") + expected = _hf_full_prefill(hf_model, prompt_ids, "cuda") + max_abs = float(np.max(np.abs(actual - expected))) + if not np.isfinite(max_abs): + raise AssertionError(f"BF16 rejection evidence is non-finite: {max_abs}") + print(f"BF16 rejection evidence: max_abs={max_abs:.6g} (limit=0.01)") + return max_abs + + +def _full_prefill(session, token_ids: list[int]) -> np.ndarray: + states = _initial_states(session) + output_names = [output.name for output in session.get_outputs()] + feeds = _token_feeds( + session, + np.array([token_ids], dtype=np.int64), + total_length=len(token_ids), + position_ids=np.arange(len(token_ids), dtype=np.int64)[None, :], + states=states, + ) + outputs = _run_session(session, output_names, feeds) + return _as_numpy(outputs[output_names.index("logits")]).astype(np.float32) + + +def _hf_full_prefill(model, token_ids: list[int], device: str) -> np.ndarray: + ids = torch.tensor([token_ids], dtype=torch.long, device=device) + with torch.no_grad(): + logits = model( + input_ids=ids, + attention_mask=torch.ones_like(ids), + position_ids=torch.arange(len(token_ids), device=device)[None, :], + use_cache=False, + ).logits + return logits.float().cpu().numpy() + + +def _hf_generate( + model, + token_ids: list[int], + device: str, + max_new_tokens: int, +) -> tuple[list[int], list[np.ndarray]]: + from transformers import DynamicCache + + cache = DynamicCache(config=model.config) + past_length = 0 + outputs = None + with torch.no_grad(): + for token_id in token_ids: + ids = torch.tensor([[token_id]], dtype=torch.long, device=device) + outputs = model( + input_ids=ids, + attention_mask=torch.ones( + (1, past_length + 1), dtype=torch.long, device=device + ), + position_ids=torch.tensor([[past_length]], dtype=torch.long, device=device), + past_key_values=cache, + use_cache=True, + ) + cache = outputs.past_key_values + past_length += 1 + + assert outputs is not None + generated: list[int] = [] + logits_by_step: list[np.ndarray] = [] + for _ in range(max_new_tokens): + logits_by_step.append(outputs.logits[0, -1].float().cpu().numpy()) + token_id = int(outputs.logits[0, -1].argmax()) + generated.append(token_id) + ids = torch.tensor([[token_id]], dtype=torch.long, device=device) + outputs = model( + input_ids=ids, + attention_mask=torch.ones( + (1, past_length + 1), dtype=torch.long, device=device + ), + position_ids=torch.tensor([[past_length]], dtype=torch.long, device=device), + past_key_values=cache, + use_cache=True, + ) + cache = outputs.past_key_values + past_length += 1 + return generated, logits_by_step + + +def _assert_logits_close( + actual: np.ndarray, + expected: np.ndarray, + *, + atol: float, + label: str, +) -> None: + max_abs = float(np.max(np.abs(actual - expected))) + cosine = float( + np.dot(actual.ravel(), expected.ravel()) + / (np.linalg.norm(actual) * np.linalg.norm(expected)) + ) + print(f"{label}: max_abs={max_abs:.6g}, cosine={cosine:.9f}") + if max_abs > atol: + raise AssertionError(f"{label}: max_abs={max_abs:.6g} exceeds {atol=}") + np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=atol) + + +def _graph_audit(model) -> dict[str, int]: + counts = Counter( + f"{node.domain or 'ai.onnx'}::{node.op_type}" for node in model.graph.all_nodes() + ) + return dict(sorted(counts.items())) + + +def _validate_variant( + state: dict[str, torch.Tensor], + output_root: Path, + *, + dtype_name: str, + device: str, +) -> Path: + torch_dtype = _DTYPES[dtype_name][0] + ep = "cuda" if device == "cuda" else "cpu" + package = _mobius_package(state, dtype_name=dtype_name, ep=ep) + variant_dir = output_root / f"{dtype_name}-{ep}" + variant_dir.mkdir(parents=True, exist_ok=True) + package.save(variant_dir, external_data="onnx") + _hf_config().save_pretrained(variant_dir) + + profile = device == "cuda" + session = _create_session(variant_dir / "model.onnx", device, False) + prompt_ids = [1, 42, 17] + actual = _full_prefill(session, prompt_ids) + hf_model = _hf_model(state, dtype=torch_dtype, device=device) + expected = _hf_full_prefill(hf_model, prompt_ids, device) + atol = 2e-3 if dtype_name == "f32" else 1e-2 + _assert_logits_close(actual, expected, atol=atol, label=f"{dtype_name}/{ep} prefill") + + generated, logits, profile_path = run_token_ids( + variant_dir, + prompt_ids, + max_new_tokens=4, + device=device, + profile=profile, + eos_token_ids=load_eos_token_ids(variant_dir), + ) + if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): + raise AssertionError(f"Invalid generation for {dtype_name}/{ep}: {generated}") + hf_generated, hf_step_logits = _hf_generate( + hf_model, + prompt_ids, + device, + max_new_tokens=4, + ) + if generated != hf_generated: + raise AssertionError( + f"{dtype_name}/{ep} generation mismatch: ONNX={generated}, HF={hf_generated}" + ) + for index, (actual_step, expected_step) in enumerate(zip(logits, hf_step_logits)): + _assert_logits_close( + actual_step, + expected_step, + atol=atol, + label=f"{dtype_name}/{ep} cached step {index}", + ) + print(f"{dtype_name}/{ep} generated IDs: {generated}") + print(f"{dtype_name}/{ep} weighted graph ops: {_graph_audit(package['model'])}") + + if profile_path is not None: + placement = summarize_profile(profile_path) + print(f"{dtype_name}/{ep} provider placement: {placement}") + if placement.get("CUDAExecutionProvider", 0) == 0: + raise AssertionError(f"No CUDA nodes found in profile: {placement}") + Path(profile_path).unlink(missing_ok=True) + del hf_model + if device == "cuda": + torch.cuda.empty_cache() + return variant_dir + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--cache", + default=default_reduced_cache_path(), + ) + parser.add_argument("--output-dir", default="output/reduced-validation") + parser.add_argument( + "--matrix", + nargs="+", + choices=["f32-cpu", "f16-cuda"], + default=["f32-cpu", "f16-cuda"], + ) + parser.add_argument( + "--bf16-rejection-evidence", + action="store_true", + help=( + "Measure the rejected BF16 CUDA path with a test-only component wrapper; " + "does not produce a supported package." + ), + ) + parser.add_argument("--skip-quantization", action="store_true") + args = parser.parse_args() + + state = _build_reduced_state(Path(args.cache)) + print(f"Loaded {len(state)} reduced real-weight tensors from revision {REVISION}") + output_root = Path(args.output_dir) + output_root.mkdir(parents=True, exist_ok=True) + if args.bf16_rejection_evidence: + max_abs = _measure_bf16_rejection(state, output_root) + if max_abs <= 1e-2: + raise AssertionError( + f"BF16 now meets the 1e-2 gate ({max_abs}); revisit the production rejection" + ) + return + + variants: dict[str, Path] = {} + for variant in args.matrix: + dtype_name, device = variant.split("-") + if device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA validation requested but unavailable: {variant}") + variants[variant] = _validate_variant( + state, + output_root, + dtype_name=dtype_name, + device=device, + ) + + if not args.skip_quantization: + source = variants.get("f16-cuda") + if source is None: + raise ValueError("Olive validation requires f16-cuda in --matrix") + quantized = quantize_package(source, output_root / "q4_k_m-cuda") + generated, logits, _profile = run_token_ids( + quantized, + [1, 42, 17], + max_new_tokens=4, + device="cuda", + profile=False, + eos_token_ids=load_eos_token_ids(quantized), + ) + if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): + raise AssertionError(f"Quantized generation failed: {generated}") + print(f"Olive Q4_K_M package loaded and generated IDs: {generated}") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index cb4bbfea7..76a711e31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,7 @@ markers = [ "arch_validation: L2 architecture validation tests that download real HF configs (no weights) and build full-size ONNX graphs (deselect with '-m \"not arch_validation\"')", "golden: L4 checkpoint-verified golden comparison tests (deselect with '-m \"not golden\"')", "generation: L5 generation end-to-end golden tests (deselect with '-m \"not generation\"')", + "quantization: quantized-package integration tests", ] [tool.mypy] diff --git a/scripts/detect_affected_models.py b/scripts/detect_affected_models.py index 55494f0cf..779b25317 100644 --- a/scripts/detect_affected_models.py +++ b/scripts/detect_affected_models.py @@ -53,6 +53,25 @@ "src/mobius/tasks/", ) +# Model-specific examples whose integration tests depend on files outside src/. +_MODEL_PATH_HINTS: tuple[tuple[str, tuple[str, ...]], ...] = ( + ("examples/olive/nemotron-3_5-lightning-30b/", ("nemotron_h",)), + ("testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml", ("nemotron_h",)), + ("testdata/golden/causal-lm/nemotron-3_5-lightning-30b-", ("nemotron_h",)), +) + + +def _model_type_hints(path: str) -> set[str]: + """Infer targeted model types from real-weight tests and example assets.""" + normalized = path.replace("\\", "/") + if normalized.startswith("tests/") and normalized.endswith("_real_weight_test.py"): + filename = normalized.rsplit("/", 1)[-1] + return {filename[: -len("_real_weight_test.py")]} + for prefix, model_types in _MODEL_PATH_HINTS: + if normalized.startswith(prefix): + return set(model_types) + return set() + def classify_file(path: str) -> str: """Classify a changed file path. @@ -435,6 +454,7 @@ def detect_affected_models( model_files: list[str] = [] traceable_files: list[str] = [] for path in changed_files: + affected.update(_model_type_hints(path)) category = classify_file(path) if category == "shared_infra": run_all = True @@ -462,7 +482,7 @@ def detect_affected_models( return {"affected": [], "run_all": True} if not model_files and not traceable_files: - return {"affected": [], "run_all": False} + return {"affected": sorted(affected), "run_all": False} # Build the registry map: source_module → [model_types] registry_map = _build_source_module_to_types() diff --git a/scripts/detect_affected_models_test.py b/scripts/detect_affected_models_test.py index a472f056f..1444621c2 100644 --- a/scripts/detect_affected_models_test.py +++ b/scripts/detect_affected_models_test.py @@ -26,6 +26,7 @@ _build_registry_class_to_types, _build_source_module_to_types, _find_reverse_dependents, + _model_type_hints, _parse_imports, classify_file, detect_affected_models, @@ -84,6 +85,16 @@ def test_windows_paths(self): assert classify_file("src\\mobius\\models\\falcon.py") == "model" +class TestModelTypeHints: + def test_real_weight_test_infers_model_type(self): + assert _model_type_hints("tests/nemotron_h_real_weight_test.py") == {"nemotron_h"} + + def test_model_example_uses_explicit_mapping(self): + assert _model_type_hints("examples/olive/nemotron-3_5-lightning-30b/optimize.py") == { + "nemotron_h" + } + + # ---------------------------------------------------------------- # AST registry parsing tests # ---------------------------------------------------------------- @@ -233,6 +244,16 @@ def test_test_file_no_affected(self): assert result["run_all"] is False assert result["affected"] == [] + def test_real_weight_test_targets_its_model(self): + result = detect_affected_models(["tests/nemotron_h_real_weight_test.py"]) + assert result == {"affected": ["nemotron_h"], "run_all": False} + + def test_model_example_targets_its_integration_tests(self): + result = detect_affected_models( + ["examples/olive/nemotron-3_5-lightning-30b/inference.py"] + ) + assert result == {"affected": ["nemotron_h"], "run_all": False} + def test_falcon_model_file(self): result = detect_affected_models(["src/mobius/models/falcon.py"]) assert result["run_all"] is False diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index faafabb8a..cd8771330 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -365,16 +365,6 @@ def _save_package( ) -> None: """Save a ModelPackage to disk, applying optimizations and runtime configs.""" runtime = getattr(args, "runtime", None) - if runtime == "ort-genai": - from mobius.integrations.ort_genai.auto_export import ( - _validate_ort_genai_compatibility, - ) - - try: - _validate_ort_genai_compatibility(pkg) - except ValueError as error: - raise SystemExit(f"Error: {error}") from error - components = (lambda name: name == component_filter) if component_filter else None for name, model in pkg.items(): if components is not None and not components(name): @@ -504,15 +494,6 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None: ) raise SystemExit(1) - if getattr(args, "runtime", None) == "ort-genai": - raise SystemExit( - "Error: mobius build-gguf does not yet support --runtime ort-genai. " - "The command cannot emit a valid genai_config.json until the selected " - "GGUF architecture's cache and tokenizer contracts have passed real " - "ORT GenAI generation. Use --runtime onnx-genai where supported, or " - "omit --runtime and run the ONNX model directly." - ) - mmproj_path = getattr(args, "mmproj", None) keep_quantized = not args.dequantize @@ -554,9 +535,9 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None: path = os.path.join(output_dir, "model.onnx") print(f"Saved {name} to {path}") - if getattr(args, "runtime", None) == "onnx-genai": + runtime = getattr(args, "runtime", None) + if runtime in ("ort-genai", "onnx-genai"): from mobius.integrations.gguf import write_gguf_tokenizer_json - from mobius.integrations.onnx_genai import write_onnx_genai_config # A GGUF checkpoint has no Hugging Face source directory, so the # tokenizer is reconstructed from the file's embedded ggml metadata @@ -564,9 +545,20 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None: tokenizer_path = write_gguf_tokenizer_json(gguf_path, output_dir) if tokenizer_path is not None: print(f" tokenizer: {tokenizer_path}") - artifacts = write_onnx_genai_config( - pkg, output_dir, config=getattr(pkg, "config", None), source=None - ) + if runtime == "onnx-genai": + from mobius.integrations.onnx_genai import write_onnx_genai_config + + artifacts = write_onnx_genai_config( + pkg, output_dir, config=getattr(pkg, "config", None), source=None + ) + else: + from mobius.integrations.ort_genai import write_ort_genai_config + + artifacts = write_ort_genai_config( + pkg, + output_dir, + ep=args.execution_provider, + ) for name, path in artifacts.items(): print(f" {name}: {path}") @@ -875,10 +867,9 @@ def main(argv: list[str] | None = None) -> None: default=None, help=( "Generate runtime-specific config files after building. " - "'onnx-genai' writes inference_metadata.yaml plus a tokenizer.json " - "reconstructed from the GGUF's embedded tokenizer metadata; " - "'ort-genai' is currently rejected until GGUF cache/tokenizer " - "contracts have runtime generation coverage." + "Both modes reconstruct tokenizer.json from GGUF metadata. " + "'onnx-genai' writes inference_metadata.yaml; 'ort-genai' writes " + "the best graph-derived genai_config.json metadata." ), ) gguf_parser.add_argument( diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 78dcb4e7e..9c31eb49b 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -40,6 +40,16 @@ def _resolve_dtype(config) -> ir.DataType | None: return None +def _resolve_explicit_dtype(value, *, field_name: str) -> ir.DataType: + """Normalize a required dtype string/torch/IR value.""" + if isinstance(value, ir.DataType): + return value + torch_dtype = getattr(torch, value, None) if isinstance(value, str) else value + if isinstance(torch_dtype, torch.dtype): + return tensor_adapters.from_torch_dtype(torch_dtype) + raise ValueError(f"Unsupported {field_name}: {value!r}") + + def _resolve_hidden_act(config, model_type: str) -> str | None: """Resolve the hidden activation function from HF config patterns. @@ -2643,6 +2653,7 @@ class NemotronHConfig(ArchitectureConfig): mamba_conv_bias: bool = True mamba_proj_bias: bool = False mamba_time_step_min: float = 0.001 + mamba_ssm_cache_dtype: ir.DataType = ir.DataType.FLOAT moe_latent_size: int | None = None @classmethod @@ -2660,15 +2671,29 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: "-": "mlp", "E": "moe", } - layers_block_type = [char_map.get(c, "mamba2") for c in pattern] + invalid_chars = sorted(set(pattern) - set(char_map)) + if invalid_chars: + raise ValueError( + "Unsupported NemotronH hybrid_override_pattern character(s): " + f"{invalid_chars}" + ) + layers_block_type = [char_map[c] for c in pattern] else: - # Convert HF names to mobius names + # Transformers 5.x uses ``linear_attention``/``full_attention``; + # older configs use ``mamba``/``attention``. Normalize both + # vocabularies to Mobius cache-layer names. type_map = { "mamba": "mamba2", + "linear_attention": "mamba2", "attention": "full_attention", + "full_attention": "full_attention", "moe": "moe", + "mlp": "mlp", } - layers_block_type = [type_map.get(t, t) for t in layers_block_type] + invalid_types = sorted(set(layers_block_type) - set(type_map)) + if invalid_types: + raise ValueError(f"Unsupported NemotronH layer type(s): {invalid_types}") + layers_block_type = [type_map[t] for t in layers_block_type] # Override num_hidden_layers based on actual pattern length n = len(layers_block_type) if layers_block_type else base.num_hidden_layers @@ -2717,6 +2742,10 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: mamba_conv_bias=getattr(config, "use_conv_bias", True), mamba_proj_bias=getattr(config, "mamba_proj_bias", False), mamba_time_step_min=getattr(config, "time_step_min", 0.001), + mamba_ssm_cache_dtype=_resolve_explicit_dtype( + getattr(config, "mamba_ssm_cache_dtype", "float32"), + field_name="mamba_ssm_cache_dtype", + ), moe_latent_size=getattr(config, "moe_latent_size", None), shared_expert_intermediate_size=shared_expert_intermediate_size, ) diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index f353b9084..a242485be 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -16,6 +16,7 @@ ArchitectureConfig, AudioConfig, MuseGlimmerConfig, + NemotronHConfig, QuantizationConfig, VisionConfig, _extract_audio_config, @@ -402,6 +403,60 @@ class FakeNemotronH: # rope_interleave stays at its inert False default. assert config.rope_interleave is False + +class TestNemotronHConfig: + @staticmethod + def _fake_config(layer_types: list[str]) -> SimpleNamespace: + return SimpleNamespace( + model_type="nemotron_h", + layers_block_type=layer_types, + num_hidden_layers=len(layer_types), + vocab_size=131072, + hidden_size=2688, + intermediate_size=1856, + num_attention_heads=32, + num_key_value_heads=2, + head_dim=128, + max_position_embeddings=262144, + pad_token_id=0, + layer_norm_epsilon=1e-5, + mamba_num_heads=64, + mamba_head_dim=64, + ssm_state_size=128, + n_groups=8, + conv_kernel=4, + expand=2, + n_routed_experts=128, + num_experts_per_tok=6, + moe_intermediate_size=1856, + moe_shared_expert_intermediate_size=3712, + routed_scaling_factor=2.5, + ) + + @pytest.mark.parametrize( + "hf_layer_types", + [ + ["mamba", "attention", "moe", "mlp"], + ["linear_attention", "full_attention", "moe", "mlp"], + ], + ids=["legacy-transformers", "current-transformers"], + ) + def test_normalizes_transformers_layer_type_vocabularies( + self, hf_layer_types: list[str] + ) -> None: + config = NemotronHConfig.from_transformers(self._fake_config(hf_layer_types)) + + assert config.layer_types == ["mamba2", "full_attention", "moe", "mlp"] + assert config.num_hidden_layers == 4 + assert config.num_local_experts == 128 + assert config.num_experts_per_tok == 6 + assert config.shared_expert_intermediate_size == 3712 + assert config.mamba_ssm_cache_dtype.name == "FLOAT" + + def test_rejects_unknown_layer_type(self) -> None: + with pytest.raises(ValueError, match="Unsupported NemotronH layer type"): + NemotronHConfig.from_transformers(self._fake_config(["linear_attention", "bogus"])) + def test_from_transformers_legacy_rotary_dim_enables_rope(self): """GPT-J / CodeGen-style legacy configs use ``rotary_dim``.""" diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index bbe9a68a3..698beed1c 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -919,7 +919,7 @@ def _create_default_registry() -> ModelRegistry: "command_r": "CohereForAI/c4ai-command-r-v01", "csm": "sesame/csm-1b", "evolla": "westlake-repl/Evolla-10B-hf", - "nemotron_h": "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", + "nemotron_h": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", "nemotron_parse": "nvidia/NVIDIA-Nemotron-Parse-2.0", "open-llama": "openlm-research/open_llama_3b", "persimmon": "adept/persimmon-8b-base", diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 82187b184..afb3495d4 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -22,6 +22,7 @@ __all__ = ["gguf_to_config"] +import contextlib import dataclasses import logging from typing import Any @@ -91,6 +92,18 @@ "ssm.conv_kernel": "linear_conv_kernel_dim", } +_TOKENIZER_KEY_MAP: dict[str, str] = { + "tokenizer.ggml.bos_token_id": "bos_token_id", + "tokenizer.ggml.padding_token_id": "pad_token_id", + # Non-standard but emitted by some multimodal converters. + "tokenizer.ggml.image_token_id": "image_token_id", +} +_STOP_TOKEN_KEYS = ( + "tokenizer.ggml.eos_token_id", + "tokenizer.ggml.eot_token_id", + "tokenizer.ggml.eom_token_id", +) + _ARCH_KEY_MAPS: dict[str, dict[str, str]] = { "deepseek4": { "attention.key_length": "head_dim", @@ -195,6 +208,42 @@ def _extract_config_fields( if isinstance(tokens, list): hf_fields["vocab_size"] = len(tokens) + vocab_size = hf_fields.get("vocab_size") + + def _valid_token_id(value: Any) -> int | None: + token_id = int(value) + if token_id < 0 or token_id == 0xFFFFFFFF: + return None + if vocab_size is not None and token_id >= int(vocab_size): + return None + return token_id + + for gguf_key, config_key in _TOKENIZER_KEY_MAP.items(): + if ( + gguf_key in metadata + and (token_id := _valid_token_id(metadata[gguf_key])) is not None + ): + hf_fields[config_key] = token_id + + stop_token_ids: list[int] = [] + for gguf_key in _STOP_TOKEN_KEYS: + if gguf_key not in metadata: + continue + token_id = _valid_token_id(metadata[gguf_key]) + if token_id is not None and token_id not in stop_token_ids: + stop_token_ids.append(token_id) + if stop_token_ids: + hf_fields["eos_token_id"] = ( + stop_token_ids[0] if len(stop_token_ids) == 1 else stop_token_ids + ) + + # Standard GGUF has no dedicated image-token key. Preserve the canonical + # HuggingFace placeholder when it is embedded in the tokenizer vocabulary. + tokens = metadata.get("tokenizer.ggml.tokens") + if "image_token_id" not in hf_fields and isinstance(tokens, list): + with contextlib.suppress(ValueError): + hf_fields["image_token_id"] = tokens.index("") + return hf_fields @@ -370,6 +419,13 @@ def gguf_to_config( if isinstance(swiglu_limit, (list, np.ndarray)): swiglu_limit = swiglu_limit[0] if len(swiglu_limit) else 0.0 + special_token_fields: dict[str, Any] = { + name: int(hf_fields[name]) + for name in ("bos_token_id", "pad_token_id", "image_token_id") + if hf_fields.get(name) is not None + } + if (eos_token_id := hf_fields.get("eos_token_id")) is not None: + special_token_fields["eos_token_id"] = eos_token_id config = ArchitectureConfig( hidden_size=hidden_size, intermediate_size=hf_fields.get("intermediate_size", 4 * hidden_size), @@ -435,6 +491,7 @@ def gguf_to_config( linear_key_head_dim=linear_key_head_dim, linear_value_head_dim=linear_value_head_dim, linear_conv_kernel_dim=(hf_fields.get("linear_conv_kernel_dim") or 4), + **special_token_fields, ) # Store model_type for registry lookup and tensor processor dispatch. diff --git a/src/mobius/integrations/gguf/_reader_test.py b/src/mobius/integrations/gguf/_reader_test.py index 0ec0e1089..c3cde77b5 100644 --- a/src/mobius/integrations/gguf/_reader_test.py +++ b/src/mobius/integrations/gguf/_reader_test.py @@ -559,6 +559,46 @@ def test_extract_config_fields_with_prefix(self): assert fields["hidden_size"] == 4096 assert fields["num_hidden_layers"] == 32 + def test_extract_config_fields_preserves_tokenizer_metadata(self): + fields = _extract_config_fields( + "llama", + { + "tokenizer.ggml.bos_token_id": 1, + "tokenizer.ggml.eos_token_id": 2, + "tokenizer.ggml.padding_token_id": 3, + "tokenizer.ggml.eot_token_id": 4, + "tokenizer.ggml.eom_token_id": 2, + "tokenizer.ggml.tokens": [ + "a", + "", + "b", + "c", + "d", + ], + }, + ) + + assert fields["bos_token_id"] == 1 + assert fields["eos_token_id"] == [2, 4] + assert fields["pad_token_id"] == 3 + assert fields["image_token_id"] == 1 + + def test_extract_config_fields_omits_invalid_token_sentinels(self): + fields = _extract_config_fields( + "llama", + { + "tokenizer.ggml.tokens": ["a"], + "tokenizer.ggml.bos_token_id": 0xFFFFFFFF, + "tokenizer.ggml.eos_token_id": 0xFFFFFFFF, + "tokenizer.ggml.eot_token_id": -1, + "tokenizer.ggml.padding_token_id": 999, + }, + ) + + assert "bos_token_id" not in fields + assert "eos_token_id" not in fields + assert "pad_token_id" not in fields + def test_infer_tie_embeddings_true(self, tied_gguf: Path): model = GGUFModel(tied_gguf) assert _infer_tie_embeddings(model) is True diff --git a/src/mobius/integrations/nemo/_genai_config.py b/src/mobius/integrations/nemo/_genai_config.py index 2b919cf3f..b284c72a5 100644 --- a/src/mobius/integrations/nemo/_genai_config.py +++ b/src/mobius/integrations/nemo/_genai_config.py @@ -269,8 +269,7 @@ def write_genai_bundle( Args: pkg: The :class:`~mobius._model_package.ModelPackage` produced by :func:`~mobius.integrations.nemo.build_from_nemo`. Must contain the - ``encoder_streaming``, ``decoder`` and ``joint`` graphs and be built - in float32 (the GenAI pipeline only supports float32 encoder I/O). + ``encoder_streaming``, ``decoder`` and ``joint`` graphs. archive: The source :class:`NeMoArchive` (for preprocessor parameters and the SentencePiece tokenizer). dest_dir: Output directory (created if needed). @@ -283,12 +282,6 @@ def write_genai_bundle( Returns: The resolved output directory path. """ - dtype = getattr(pkg.config, "dtype", None) - if dtype is not None and dtype != ir.DataType.FLOAT: - raise ValueError( - "ORT GenAI nemotron_speech only supports float32 encoder I/O; build " - f"the package with dtype='f32' (got {dtype})." - ) for key in ("encoder_streaming", "decoder", "joint"): if key not in pkg: raise KeyError(f"ModelPackage is missing required model {key!r}") diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 37271b90e..b281517a4 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -227,6 +227,71 @@ def _count_cache_layer_slots(model: ir.Model | None) -> int | None: return max(layer_indices) + 1 if layer_indices else None +_CACHE_INPUT_FIELDS = { + (None, "key"): "past_key_names", + (None, "value"): "past_value_names", + ("self", "key"): "past_key_names", + ("self", "value"): "past_value_names", + ("cross", "key"): "cross_past_key_names", + ("cross", "value"): "cross_past_value_names", + (None, "conv_state"): "past_conv_names", +} +_CACHE_OUTPUT_FIELDS = { + (None, "key"): "present_key_names", + (None, "value"): "present_value_names", + ("self", "key"): "present_key_names", + ("self", "value"): "present_value_names", + (None, "conv_state"): "present_conv_names", +} + + +def _cache_name_parts(name: str, prefix: str) -> tuple[str | None, str, str] | None: + """Return ``(scope, suffix, template)`` for indexed cache names.""" + parts = name.split(".") + if len(parts) not in (3, 4) or parts[0] != prefix or not parts[1].isdigit(): + return None + scope = parts[2] if len(parts) == 4 else None + suffix = parts[-1] + parts[1] = "%d" + return scope, suffix, ".".join(parts) + + +def _decoder_cache_templates(model: ir.Model) -> tuple[dict[str, str], dict[str, str]]: + """Map graph cache suffixes to every template the current config schema supports.""" + inputs: dict[str, str] = {} + outputs: dict[str, str] = {} + for value in model.graph.inputs: + if ( + value.name is None + or (parts := _cache_name_parts(value.name, "past_key_values")) is None + ): + continue + scope, suffix, template = parts + if (config_name := _CACHE_INPUT_FIELDS.get((scope, suffix))) is not None: + inputs[config_name] = template + for value in model.graph.outputs: + if value.name is None or (parts := _cache_name_parts(value.name, "present")) is None: + continue + scope, suffix, template = parts + if (config_name := _CACHE_OUTPUT_FIELDS.get((scope, suffix))) is not None: + outputs[config_name] = template + return inputs, outputs + + +def _decoder_output_mapping(model: ir.Model) -> dict[str, str] | None: + """Return semantic decoder outputs and graph-derived cache templates.""" + output_names = [value.name for value in model.graph.outputs if value.name is not None] + logits_name = next( + (name for name in output_names if name == "logits"), + next((name for name in output_names if name.startswith("logits_")), None), + ) + _cache_inputs, cache_outputs = _decoder_cache_templates(model) + outputs = dict(cache_outputs) + if logits_name is not None: + outputs["logits"] = logits_name + return outputs or None + + def _introspect_inputs(pkg: ModelPackage, key: str) -> dict[str, str] | None: """Return ``{name: name}`` identity mapping for a sub-model's inputs. @@ -991,45 +1056,23 @@ def _write_genai_config( # --- Discover decoder inputs from the ONNX graph --- decoder_key = "decoder" if "decoder" in pkg else "model" decoder_inputs = _introspect_inputs(pkg, decoder_key) - if decoder_inputs is not None: - # KV cache entries are template-based, not per-input - decoder_inputs["past_key_names"] = "past_key_values.%d.key" - decoder_inputs["past_value_names"] = "past_key_values.%d.value" + decoder_model = pkg.get(decoder_key) + decoder_outputs = None + if decoder_model is not None: + cache_inputs, _cache_outputs = _decoder_cache_templates(decoder_model) + if decoder_inputs is not None: + decoder_inputs.update(cache_inputs) + decoder_outputs = _decoder_output_mapping(decoder_model) # Derive decoder filename from the actual package key decoder_filename = ( f"{decoder_key}/model.onnx" if len(pkg) > 1 or decoder_key != "model" else "model.onnx" ) - # ORT GenAI's ``past_present_share_buffer`` mode requires the decoder - # graph to write the KV cache in place. Only ``com.microsoft. - # GroupQueryAttention`` does that; the standard ONNX ``Attention`` op - # concatenates ``past_key`` with the new ``K`` and returns a dynamic- - # shape ``present_key``, which is incompatible with the pre-allocated - # shared buffer. Introspect the graph: if there is at least one GQA - # node, the model supports shared-buffer mode; otherwise force it off - # regardless of the EP capability flag. - # - # ``com.microsoft.LinearAttention`` (linear/recurrent-attention layers, - # e.g. Qwen3.5's GatedDeltaNet) is a separate, *mandatory* case: its - # recurrent state requires ``past_present_share_buffer=True`` regardless - # of whether any other layer uses GQA (ORT GenAI raises "RecurrentState - # requires past_present_share_buffer=true" otherwise). - # - # Hybrid models mix LinearAttention layers with full-attention layers, - # which may lower to GQA *or* to the standard (non-GQA) ``Attention`` op - # depending on EP/dtype (e.g. the CPU EP only lowers to GQA for fp32; - # fp16 falls back to standard Attention -- see ``_execution_providers.py`` - # ``gqa_dtypes``). If a hybrid graph has LinearAttention but its - # full-attention layers are still standard (non-GQA) Attention, forcing - # ``past_present_share_buffer=True`` produces an unrunnable config: the - # recurrent state requires it, but standard Attention's dynamic-shape KV - # concat cannot honor a pre-allocated shared buffer, which fails at - # generation time with an ``attn_mask``/``total_sequence_length`` - # mismatch rather than at load time. Rather than silently emit a broken - # config, raise a clear error so the caller picks an EP/dtype combination - # (e.g. fp32 on CPU) that lowers full attention to GQA. - decoder_model = pkg.get(decoder_key) + # Derive shared-buffer metadata from the graph. GQA supports in-place KV, + # while LinearAttention requires a shared recurrent-state buffer. Mixed + # topologies are still emitted faithfully; downstream runtime acceptance + # must not become a Mobius export capability gate. supports_in_place_kv_cache: bool | None = None if decoder_model is not None: has_gqa = any( @@ -1040,26 +1083,6 @@ def _write_genai_config( node.op_type == "LinearAttention" and node.domain == "com.microsoft" for node in decoder_model.graph ) - has_standard_attention = any( - node.op_type == "Attention" and node.domain in ("", "ai.onnx") - for node in decoder_model.graph - ) - if has_recurrent_state and has_standard_attention: - # A GQA node elsewhere in the graph does NOT make a co-existing - # standard Attention node compatible with a shared buffer -- - # each op instance is independently (in)compatible, so this - # must reject on the mere presence of standard Attention, not - # only when GQA is completely absent (partial GQA fusion still - # leaves the unfused standard Attention layers broken). - raise ValueError( - "This decoder graph mixes com.microsoft.LinearAttention " - "(recurrent state, requires past_present_share_buffer=True) " - "with standard (non-GQA) Attention (incompatible with " - "past_present_share_buffer=True). This EP/dtype combination " - "cannot produce a runnable genai_config -- pick an EP/dtype " - "that lowers *all* full-attention layers to " - "GroupQueryAttention instead (e.g. fp32 on the CPU EP)." - ) supports_in_place_kv_cache = has_gqa or has_recurrent_state generator = GenaiConfigGenerator.from_config( @@ -1071,6 +1094,7 @@ def _write_genai_config( eos_token_id=eos_token_id, pad_token_id=pad_token_id, decoder_inputs=decoder_inputs, + decoder_outputs=decoder_outputs, decoder_filename=decoder_filename, supports_in_place_kv_cache=supports_in_place_kv_cache, num_cache_layer_slots=_count_cache_layer_slots(decoder_model), @@ -1221,30 +1245,6 @@ def _write_genai_config( return generator.write(output_dir) -def _validate_ort_genai_compatibility(pkg: ModelPackage) -> None: - """Reject packages whose required inputs cannot be supplied by ORT GenAI.""" - config = getattr(pkg, "config", None) - if getattr(config, "model_type", None) == "parakeet_ctc": - raise ValueError( - "ORT GenAI does not define a feature-input CTC ASR pipeline; " - "export Parakeet CTC as ONNX and run it directly with ONNX Runtime." - ) - if {"vision_encoder", "decoder"}.issubset(pkg) and "embedding" not in pkg: - model_type = getattr(config, "model_type", "unknown") - raise NotImplementedError( - "onnxruntime-genai does not support generic vision encoder-decoder " - f"packages such as {model_type!r}. Run the vision_encoder and decoder " - "ONNX sessions directly; emitting genai_config.json would create an " - "artifact that the runtime cannot load." - ) - if getattr(config, "model_type", None) == "mage_vl": - raise ValueError( - "ORT GenAI does not support Mage-VL's required patch_positions vision " - "input or its 1D decoder position_ids contract. Export without " - "--runtime ort-genai to save the runnable direct three-model ONNX package." - ) - - def write_ort_genai_config( pkg: ModelPackage, directory: str, @@ -1307,15 +1307,6 @@ def write_ort_genai_config( "This is set automatically when building with mobius.build(). " "Diffusion models (which have no config) are not supported." ) - _validate_ort_genai_compatibility(pkg) - - if getattr(config, "model_type", None) == "moonshine": - raise NotImplementedError( - "onnxruntime-genai does not support Moonshine's variable-length raw-waveform " - "encoder. Run the exported encoder and cached decoder directly with " - "ONNX Runtime." - ) - os.makedirs(directory, exist_ok=True) # Normalize EP: 'default' and 'onnx-standard' are portable-ONNX modes @@ -1375,6 +1366,10 @@ def write_ort_genai_config( # Gemma3 multimodal configs are unwrapped to the text sub-config # during build, but ORT GenAI needs the multimodal parent type. ort_model_type = "gemma3" + elif is_vlm and raw_type == "gemma4_text": + # GGUF multimodal builds retain the text checkpoint's model type + # after attaching the companion vision projector. + ort_model_type = "gemma4" elif is_vlm and raw_type == "gemma3n_text": # Same unwrapping for Gemma3n, whose parent type is "gemma3n". # Deliberately *not* aliased to "gemma3": the package threads @@ -1584,8 +1579,6 @@ def export_package( "Diffusion models (which have no config) are not supported — " "use ModelPackage.save() directly for those." ) - _validate_ort_genai_compatibility(pkg) - os.makedirs(output_dir, exist_ok=True) # 1. Save ONNX models + weights diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 222b83f8a..1731572b3 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -56,25 +56,6 @@ def _mock_model_with_outputs(names: list[str]) -> ir.Model: return _mock_model(outputs=names) -def test_moonshine_native_runtime_is_rejected(tmp_path): - from mobius._model_package import ModelPackage - - config = mock.MagicMock() - config.model_type = "moonshine" - package = ModelPackage( - {"encoder": _mock_model(), "decoder": _mock_model()}, - config=config, - ) - - with pytest.raises( - NotImplementedError, - match="variable-length raw-waveform encoder", - ) as error: - write_ort_genai_config(package, str(tmp_path)) - assert "onnx-genai" not in str(error.value) - assert "ONNX Runtime" in str(error.value) - - def _make_fake_llm_pkg(model_type: str = "qwen2"): """Build a minimal LLM-only ModelPackage with a fake config.""" import dataclasses @@ -1042,28 +1023,116 @@ def test_genai_config_json_is_written(self, tmp_path): assert "model" in data assert data["model"]["type"] == "qwen2" - def test_rejects_generic_vision_encoder_decoder_package(self, tmp_path): + def test_nemotron_h_mixed_cache_metadata_is_emitted(self, tmp_path): import dataclasses from mobius._model_package import ModelPackage @dataclasses.dataclass class FakeConfig: - model_type: str = "nemotron_parse" + model_type: str = "nemotron_h" + vocab_size: int = 256 + hidden_size: int = 64 + num_hidden_layers: int = 3 + num_attention_heads: int = 4 + num_key_value_heads: int = 2 + head_dim: int = 16 + max_position_embeddings: int = 128 + pad_token_id: int = 0 pkg = ModelPackage( { - "vision_encoder": _mock_model(), - "decoder": _mock_model(), + "model": _mock_model( + inputs=[ + "input_ids", + "attention_mask", + "past_key_values.0.conv_state", + "past_key_values.0.ssm_state", + "past_key_values.2.key", + "past_key_values.2.value", + ], + outputs=[ + "logits", + "present.0.conv_state", + "present.0.ssm_state", + "present.2.key", + "present.2.value", + ], + ) }, config=FakeConfig(), ) - with pytest.raises( - NotImplementedError, - match="does not support generic vision encoder-decoder", - ): - write_ort_genai_config(pkg, str(tmp_path)) - assert not (tmp_path / "genai_config.json").exists() + output_dir = tmp_path / "ort-genai" + + result = write_ort_genai_config(pkg, str(output_dir)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + decoder = generated["model"]["decoder"] + assert generated["model"]["type"] == "nemotron_h" + assert decoder["num_hidden_layers"] == 3 + assert decoder["inputs"] == { + "input_ids": "input_ids", + "attention_mask": "attention_mask", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value", + "past_conv_names": "past_key_values.%d.conv_state", + } + assert decoder["outputs"] == { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value", + "present_conv_names": "present.%d.conv_state", + } + + def test_olive_renamed_logits_output_is_emitted(self, tmp_path): + pkg = _make_fake_llm_pkg("qwen2") + pkg["model"] = _mock_model( + inputs=["input_ids", "past_key_values.0.key", "past_key_values.0.value"], + outputs=["logits_Q4", "present.0.key", "present.0.value"], + ) + + result = write_ort_genai_config(pkg, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + assert generated["model"]["decoder"]["outputs"]["logits"] == "logits_Q4" + + def test_nested_self_and_cross_cache_templates_are_emitted(self, tmp_path): + pkg = _make_fake_llm_pkg("decoder") + pkg["model"] = _mock_model( + inputs=[ + "input_ids", + "encoder_hidden_states", + "past_key_values.0.self.key", + "past_key_values.0.self.value", + "past_key_values.0.cross.key", + "past_key_values.0.cross.value", + ], + outputs=[ + "logits", + "present.0.self.key", + "present.0.self.value", + ], + ) + + result = write_ort_genai_config(pkg, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + decoder = json.load(config_file)["model"]["decoder"] + assert decoder["inputs"] == { + "input_ids": "input_ids", + "encoder_hidden_states": "encoder_hidden_states", + "past_key_names": "past_key_values.%d.self.key", + "past_value_names": "past_key_values.%d.self.value", + "cross_past_key_names": "past_key_values.%d.cross.key", + "cross_past_value_names": "past_key_values.%d.cross.value", + } + assert decoder["outputs"] == { + "logits": "logits", + "present_key_names": "present.%d.self.key", + "present_value_names": "present.%d.self.value", + } def test_processor_config_written_with_vision(self, tmp_path): """image_processor.json is written when pkg.config.vision is set.""" @@ -1164,48 +1233,6 @@ class FakeConfig: assert model["vision"]["patch_size"] == 16 assert model["vision"]["window_size"] == 64 - def test_mage_vl_is_rejected_before_writing_runtime_artifacts(self, tmp_path): - import dataclasses - - from mobius._model_package import ModelPackage - from mobius.integrations.ort_genai.auto_export import write_ort_genai_config - - @dataclasses.dataclass - class FakeVision: - image_size: int = 448 - patch_size: int = 16 - spatial_merge_size: int = 2 - - @dataclasses.dataclass - class FakeConfig: - model_type: str = "mage_vl" - vocab_size: int = 151936 - hidden_size: int = 2560 - num_hidden_layers: int = 1 - num_attention_heads: int = 32 - num_key_value_heads: int = 8 - head_dim: int = 128 - image_token_id: int = 151655 - temporal_patch_size: int = 1 - vision: FakeVision = dataclasses.field(default_factory=FakeVision) - - pkg = ModelPackage( - { - "decoder": _mock_model(), - "vision_encoder": _mock_model(), - "embedding": _mock_model(), - }, - config=FakeConfig(), - ) - - output_dir = tmp_path / "ort-genai" - with pytest.raises( - ValueError, - match=r"Mage-VL.*patch_positions.*1D decoder position_ids", - ): - write_ort_genai_config(pkg, str(output_dir)) - assert not output_dir.exists() - def test_processor_config_not_written_without_vision(self, tmp_path): """image_processor.json is NOT written when pkg.config has no vision attr.""" from mobius.integrations.ort_genai.auto_export import write_ort_genai_config @@ -1555,8 +1582,17 @@ class FakeConfig: # "gemma2" maps to "gemma" in _ORT_GENAI_MODEL_TYPE assert data["model"]["type"] == "gemma" - def test_config_mode_gemma3_text_vlm_uses_multimodal_model_type(self, tmp_path): - """Gemma3 VLM --config exports use ORT's multimodal gemma3 type.""" + @pytest.mark.parametrize( + ("text_model_type", "multimodal_model_type"), + [("gemma3_text", "gemma3"), ("gemma4_text", "gemma4")], + ) + def test_config_mode_text_vlm_uses_multimodal_model_type( + self, + tmp_path, + text_model_type, + multimodal_model_type, + ): + """Unwrapped text configs retain their multimodal runtime type.""" import dataclasses from mobius._model_package import ModelPackage @@ -1571,8 +1607,7 @@ class FakeVision: @dataclasses.dataclass class FakeConfig: - # build() stores the unwrapped text sub-config type on Gemma3 VLMs. - model_type: str = "gemma3_text" + model_type: str vocab_size: int = 262144 hidden_size: int = 64 num_hidden_layers: int = 2 @@ -1589,13 +1624,13 @@ class FakeConfig: "vision_encoder": _mock_model_with_inputs(["pixel_values"]), "embedding": _mock_model_with_inputs(["input_ids", "image_features"]), }, - config=FakeConfig(), + config=FakeConfig(model_type=text_model_type), ) result = write_ort_genai_config(pkg, str(tmp_path), hf_model_id=None) with open(result["genai_config"]) as f: data = json.load(f) - assert data["model"]["type"] == "gemma3" + assert data["model"]["type"] == multimodal_model_type def test_config_mode_gemma3n_text_vlm_uses_multimodal_model_type(self, tmp_path): """Gemma3n unwraps to "gemma3n_text" too, and must not alias to gemma3. @@ -1915,8 +1950,6 @@ def _make_pkg(): def test_writes_both_onnx_and_genai_config(self, tmp_path, monkeypatch): """export_package calls pkg.save AND writes genai_config.json.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() save_calls = [] @@ -1936,22 +1969,8 @@ def fake_save(self, directory, **kwargs): # ONNX path is in the manifest (single-component package) assert result["model"] == os.path.join(str(tmp_path), "model.onnx") - def test_mage_vl_is_rejected_before_saving_onnx(self, tmp_path): - pkg = self._make_pkg() - pkg.config.model_type = "mage_vl" - - with ( - mock.patch.object(pkg, "save") as save, - pytest.raises(ValueError, match=r"Mage-VL.*patch_positions"), - ): - export_package(pkg, str(tmp_path)) - - save.assert_not_called() - def test_propagates_save_kwargs(self, tmp_path, monkeypatch): """external_data and progress_bar are forwarded to pkg.save.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() save_calls = [] @@ -1972,8 +1991,6 @@ def fake_save(self, directory, **kwargs): def test_propagates_genai_config_kwargs(self, tmp_path, monkeypatch): """The ep and context_length kwargs reach the generated genai_config.json.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() monkeypatch.setattr(pkg.__class__, "save", lambda self, d, **kw: None) @@ -1999,7 +2016,6 @@ def test_preflights_missing_config(self, tmp_path, monkeypatch): no genai_config.json. """ from mobius._model_package import ModelPackage - from mobius.integrations.ort_genai.auto_export import export_package pkg = ModelPackage({"model": _mock_model()}, config=None) save_called = [] @@ -2017,8 +2033,6 @@ def fake_save(self, *a, **kw): def test_returns_manifest_with_all_artifacts(self, tmp_path, monkeypatch): """Returned manifest contains ONNX paths AND config artifacts.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() monkeypatch.setattr(pkg.__class__, "save", lambda self, d, **kw: None) @@ -2234,16 +2248,8 @@ def test_pixtral_config_filename_is_processor_config(self, tmp_path): assert data["model"]["image_token_id"] == 10 -class TestHybridAttentionShareBufferGuard: - """Tests for the LinearAttention/GQA past_present_share_buffer guard. - - See the comment above ``supports_in_place_kv_cache`` in - ``_write_genai_config``: recurrent-state layers (LinearAttention) - mandate ``past_present_share_buffer=True``, but standard (non-GQA) - Attention is incompatible with it. A hybrid graph with both, and no - GQA node to lower the standard Attention layers to, must raise a clear - build-time error rather than silently emit a broken config. - """ +class TestHybridAttentionShareBufferMetadata: + """Tests graph-derived shared-buffer metadata without runtime gating.""" @staticmethod def _make_pkg(node_op_types: list[tuple[str, str]]): @@ -2291,13 +2297,14 @@ def _write(self, pkg, tmp_path): has_speech=False, ) - def test_recurrent_state_with_standard_attention_and_no_gqa_raises(self, tmp_path): - """LinearAttention + standard Attention + no GQA is an unrunnable config.""" + def test_recurrent_state_with_standard_attention_is_emitted(self, tmp_path): pkg = self._make_pkg( [("LinearAttention", "com.microsoft"), ("Attention", "")], ) - with pytest.raises(ValueError, match="past_present_share_buffer"): - self._write(pkg, tmp_path) + path = self._write(pkg, tmp_path) + with open(path) as f: + data = json.load(f) + assert data["search"]["past_present_share_buffer"] is True def test_recurrent_state_with_gqa_does_not_raise(self, tmp_path): """LinearAttention + GQA (no standard Attention) is a valid hybrid config.""" @@ -2312,16 +2319,7 @@ def test_recurrent_state_with_gqa_does_not_raise(self, tmp_path): data = json.load(f) assert data["search"]["past_present_share_buffer"] is True - def test_recurrent_state_with_standard_attention_and_gqa_raises(self, tmp_path): - """Partial GQA fusion still leaves an incompatible standard Attention node. - - Regression test: the guard previously read - ``has_recurrent_state and has_standard_attention and not has_gqa``, so - a GQA node present *anywhere* in the graph would short-circuit the - check even though a separate, unfused standard Attention node - coexists. A GQA node on one layer doesn't make a standard Attention - node on another layer safe for ``past_present_share_buffer=True``. - """ + def test_recurrent_state_with_standard_attention_and_gqa_is_emitted(self, tmp_path): pkg = self._make_pkg( [ ("LinearAttention", "com.microsoft"), @@ -2329,8 +2327,10 @@ def test_recurrent_state_with_standard_attention_and_gqa_raises(self, tmp_path): ("Attention", ""), ], ) - with pytest.raises(ValueError, match="past_present_share_buffer"): - self._write(pkg, tmp_path) + path = self._write(pkg, tmp_path) + with open(path) as f: + data = json.load(f) + assert data["search"]["past_present_share_buffer"] is True def test_recurrent_state_only_does_not_raise(self, tmp_path): """LinearAttention with no full-attention layers at all is unaffected.""" @@ -2675,18 +2675,6 @@ def fake_export_package(pkg, output_dir, **kwargs): assert captured["execution_provider"] == "default" assert captured["text_only"] is False - def test_auto_export_rejects_mage_vl_before_saving(self, tmp_path): - pkg = _make_fake_llm_pkg("mage_vl") - - with ( - mock.patch("mobius._builder.build", return_value=pkg), - mock.patch.object(pkg, "save") as save, - pytest.raises(ValueError, match=r"Mage-VL.*patch_positions"), - ): - auto_export("microsoft/Mage-VL", str(tmp_path)) - - save.assert_not_called() - def test_auto_export_produces_genai_config(self, tmp_path): """Mock build() to return a tiny package, verify genai_config.""" import onnx_ir as ir diff --git a/src/mobius/integrations/ort_genai/genai_config.py b/src/mobius/integrations/ort_genai/genai_config.py index 532b92a58..a3f0027ed 100644 --- a/src/mobius/integrations/ort_genai/genai_config.py +++ b/src/mobius/integrations/ort_genai/genai_config.py @@ -156,6 +156,8 @@ class GenaiConfigGenerator: :func:`_default_decoder_inputs`. Must already include KV cache template entries (``past_key_names``, ``past_value_names``). + decoder_outputs: Explicit decoder output name mapping. When + provided, used instead of :func:`_default_decoder_outputs`. """ def __init__( @@ -174,6 +176,7 @@ def __init__( eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, decoder_inputs: dict[str, str] | None = None, + decoder_outputs: dict[str, str] | None = None, decoder_filename: str | None = None, supports_in_place_kv_cache: bool | None = None, decoder_graph_capture: bool | None = None, @@ -195,6 +198,8 @@ def __init__( # Explicit decoder inputs (from graph introspection); None -> use defaults self._decoder_inputs = decoder_inputs + # Explicit decoder outputs (from graph introspection); None -> use defaults + self._decoder_outputs = decoder_outputs # Explicit decoder filename; None -> use "model.onnx" self._decoder_filename = decoder_filename # Whether the exported decoder ONNX graph supports in-place KV-cache @@ -230,6 +235,7 @@ def from_config( eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, decoder_inputs: dict[str, str] | None = None, + decoder_outputs: dict[str, str] | None = None, decoder_filename: str | None = None, supports_in_place_kv_cache: bool | None = None, num_cache_layer_slots: int | None = None, @@ -278,6 +284,7 @@ def from_config( eos_token_id=eos_token_id, pad_token_id=pad, decoder_inputs=decoder_inputs, + decoder_outputs=decoder_outputs, decoder_filename=decoder_filename, supports_in_place_kv_cache=supports_in_place_kv_cache, layer_types=getattr(config, "layer_types", None), @@ -459,6 +466,10 @@ def generate(self) -> dict[str, Any]: decoder_inputs = dict(self._decoder_inputs) else: decoder_inputs = _default_decoder_inputs(is_vlm=is_multimodal) + if self._decoder_outputs is not None: + decoder_outputs = dict(self._decoder_outputs) + else: + decoder_outputs = _default_decoder_outputs() decoder_filename = "decoder/model.onnx" if is_multimodal else "model.onnx" decoder: dict[str, Any] = { "session_options": _make_session_options( @@ -469,7 +480,7 @@ def generate(self) -> dict[str, Any]: "head_size": self.head_dim, "hidden_size": self.hidden_size, "inputs": decoder_inputs, - "outputs": _default_decoder_outputs(), + "outputs": decoder_outputs, "num_attention_heads": self.num_attention_heads, "num_hidden_layers": self.num_hidden_layers, "num_key_value_heads": self.num_key_value_heads, diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index d56b4f87c..ff7cc7529 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -30,8 +30,8 @@ from __future__ import annotations import re -from typing import TYPE_CHECKING +import onnx_ir as ir import torch from onnxscript import OpBuilder, nn @@ -47,9 +47,6 @@ create_padding_mask, ) -if TYPE_CHECKING: - import onnx_ir as ir - # --------------------------------------------------------------------------- # Decoder layers # --------------------------------------------------------------------------- @@ -233,6 +230,7 @@ def __init__( self.weight = nn.Parameter([num_experts, hidden_size]) # Correction bias for expert selection (loaded from checkpoint) self.e_score_correction_bias = nn.Parameter([num_experts]) + self.e_score_correction_bias._keep_float32 = True def forward(self, op: OpBuilder, hidden_states: ir.Value): # Cast to float32 for numerical stability (eps=1e-20 underflows @@ -240,14 +238,17 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): # in NemotronHTopkRouter.forward and never casts back. hidden_states = op.Cast(hidden_states, to=1) # FLOAT32 - weight_t = op.Transpose(self.weight, perm=[1, 0]) + weight_t = op.Transpose(op.Cast(self.weight, to=1), perm=[1, 0]) router_logits = op.MatMul(hidden_states, weight_t) # Sigmoid probabilities (these become the final routing weights) probs = op.Sigmoid(router_logits) # Add correction bias for expert selection only - choice_scores = op.Add(probs, self.e_score_correction_bias) + choice_scores = op.Add( + probs, + op.Cast(self.e_score_correction_bias, to=1), + ) # Select top-k experts based on biased scores k = op.Constant(value_ints=[self.top_k]) @@ -374,12 +375,16 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): weighted = op.Mul(routing_weights, match_float) # Sum matched routing weights across top_k dim → per-token weight weight = op.ReduceSum(weighted, [-1], keepdims=True) - contribution = op.Mul(expert_output, weight) + # Match HF: accumulate routed expert contributions in the fp32 + # routing-weight dtype, then cast the completed routed result once. + contribution = op.Mul(op.Cast(expert_output, to=1), weight) if result is None: result = contribution else: result = op.Add(result, contribution) + result = op.CastLike(result, hidden_states) + # Optional latent projection back to hidden_size if self._has_latent: result = self.fc2_latent_proj(op, result) @@ -504,6 +509,9 @@ class NemotronHCausalLMModel(nn.Module): Uses ``HybridCausalLMTask`` with mixed ``"mamba2"``, ``"full_attention"``, and ``"mlp"`` layer types for the cache. + The exported task is the base decoder used by + ``NemotronHForCausalLM.forward``; auxiliary ``mtp.*`` training heads are + outside that generation graph and are intentionally not loaded. HuggingFace reference: ``NemotronHForCausalLM``. """ @@ -514,6 +522,12 @@ class NemotronHCausalLMModel(nn.Module): def __init__(self, config: NemotronHConfig): super().__init__() + if config.dtype == ir.DataType.BFLOAT16: + raise ValueError( + "NemotronH BF16 execution is not numerically supported: reduced real-weight " + "CUDA parity exceeds the 1e-2 logit tolerance. Build the BF16 checkpoint " + "with dtype='f16' instead." + ) self.config = config self.model = _NemotronHTextModel(config) self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) @@ -554,6 +568,7 @@ def preprocess_weights( - mlp: ``mixer.`` → ``mlp.`` - moe: ``mixer.`` → ``moe.`` 6. MoE stacked 3D expert tensors split into per-expert 2D weights + 7. Auxiliary ``mtp.*`` training heads omitted by NemotronHForCausalLM """ layer_types = self.config.layer_types or [] @@ -572,6 +587,11 @@ def preprocess_weights( new_state_dict: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): + # The official Nemotron 3.5 checkpoint includes multi-token + # prediction heads, but the trusted NemotronHForCausalLM forward + # does not instantiate them and marks ``mtp.*`` as unexpected. + if key.startswith("mtp."): + continue new_key = _rename_nemotron_h_weight(key, layer_types) # Split stacked 3D expert tensors into per-expert 2D weights. # HF stores experts.up_proj as (num_experts, inter, input) and diff --git a/src/mobius/models/parakeet_ctc_test.py b/src/mobius/models/parakeet_ctc_test.py index 929a68a5b..be1858ec1 100644 --- a/src/mobius/models/parakeet_ctc_test.py +++ b/src/mobius/models/parakeet_ctc_test.py @@ -168,11 +168,10 @@ def test_parakeet_synthetic_parity_with_padding(): np.testing.assert_allclose(actual, expected, atol=1e-5, rtol=1e-5) -def test_parakeet_rejects_unsupported_ort_genai_export(tmp_path): +def test_parakeet_emits_ort_genai_metadata(tmp_path): _, _, _, package = _build_tiny() - with pytest.raises( - ValueError, - match="does not define a feature-input CTC ASR pipeline", - ): - write_ort_genai_config(package, str(tmp_path)) + result = write_ort_genai_config(package, str(tmp_path)) + + assert "genai_config" in result + assert (tmp_path / "genai_config.json").is_file() diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index 68f994407..e2b09663d 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -310,7 +310,7 @@ def _make_hybrid_cache_inputs( ) ssm_state = builder.input( f"{prefix}.{i}.ssm_state", - dtype=dtype, + dtype=getattr(config, "mamba_ssm_cache_dtype", dtype), shape=[batch, mamba2_n_heads, mamba2_d_state, mamba2_d_head], ) pairs.append((conv_state, ssm_state)) diff --git a/testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml b/testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml new file mode 100644 index 000000000..a92660814 --- /dev/null +++ b/testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml @@ -0,0 +1,20 @@ +model_id: "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +model_type: "nemotron_h" +revision: "d468880b6ad3c6e0d21377ce7242adaea4cc884d" +task_type: "text-generation" +dtype: "float16" +trust_remote_code: false + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +skip_reason: "The pinned checkpoint is 65.8 GB with 30B total parameters; L4/L5 reference generation and weighted ONNX export exceed the storage and accelerator memory available to standard CI." +ci_skip_reason: "30B MoE checkpoint requires a large-memory CUDA host and about 66 GB of checkpoint storage." +notes: "NVIDIA Nemotron 3.5 Lightning 30B-A3B. Hybrid Mamba2 + sigmoid-routed MoE + full attention; 128 routed experts, top-6. The 270 auxiliary mtp.* checkpoint tensors are intentionally ignored by the trusted NemotronHForCausalLM inference model." diff --git a/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json new file mode 100644 index 000000000..bbe28a8bd --- /dev/null +++ b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json @@ -0,0 +1,43 @@ +{ + "model_id": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", + "revision": "d468880b6ad3c6e0d21377ce7242adaea4cc884d", + "fixture": "Production dimensions; complete checkpoint layers 0 (Mamba2), 1 (MoE with experts 0-3), and 5 (attention); embedding/lm_head rows 0-255.", + "input_ids": [ + 1, + 42, + 17 + ], + "top1_id": 12, + "top2_id": 13, + "top10_ids": [ + 12, + 13, + 14, + 1, + 11, + 17, + 16, + 10, + 170, + 201 + ], + "top10_logits": [ + "0x1.503e9a0000000p+1", + "0x1.1e09ce0000000p+1", + "0x1.0927420000000p+1", + "0x1.a5d06c0000000p-1", + "0x1.5fe3a60000000p-1", + "0x1.46611e0000000p-1", + "0x1.4523960000000p-1", + "0x1.3cc3800000000p-1", + "0x1.3bab820000000p-1", + "0x1.3b882c0000000p-1" + ], + "logits_summary": [ + "0x1.503e9a0000000p+1", + "-0x1.b120180000000p+0", + "0x1.3e00480000000p-1", + "0x1.e850aa0000000p-3" + ], + "full_logits_sha256": "7500c2ed86137e20dddf43a85b016746b8b38c4a932277c29d2efa926cc4a23d" +} diff --git a/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json new file mode 100644 index 000000000..51769b52f --- /dev/null +++ b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json @@ -0,0 +1,18 @@ +{ + "model_id": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", + "revision": "d468880b6ad3c6e0d21377ce7242adaea4cc884d", + "fixture": "Independent HuggingFace NemotronHForCausalLM greedy decode using the pinned reduced-real-weight fixture.", + "input_ids": [ + 1, + 42, + 17 + ], + "max_new_tokens": 4, + "do_sample": false, + "generated_tokens": [ + 12, + 13, + 12, + 12 + ] +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 0a12d5d94..16b2e9bcf 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -1303,6 +1303,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: "nemotron_h", { "hidden_act": "relu2", + "rms_norm_eps": 1e-5, "layer_types": ["mamba2", "mlp", "full_attention", "mamba2"], "_config_cls": NemotronHConfig, "num_hidden_layers": 4, @@ -1320,6 +1321,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: "nemotron_h", { "hidden_act": "relu2", + "rms_norm_eps": 1e-5, "layer_types": [ "mamba2", "moe", diff --git a/tests/arch_validation_test.py b/tests/arch_validation_test.py index d7317d47d..fc868fed1 100644 --- a/tests/arch_validation_test.py +++ b/tests/arch_validation_test.py @@ -26,20 +26,27 @@ from __future__ import annotations +import dataclasses import logging import pytest +from mobius._builder import resolve_dtype from mobius._config_resolver import ( _config_from_hf, _default_task_for_model, _try_load_config_json, ) from mobius._registry import registry +from mobius._testing.golden import discover_test_cases from mobius.tasks import get_task logger = logging.getLogger(__name__) +_TEST_CASES = discover_test_cases() +_PINNED_REVISIONS = {case.model_id: case.revision for case in _TEST_CASES} +_DECLARED_DTYPES = {case.model_id: case.dtype for case in _TEST_CASES} + # Build parametrized test cases from registry entries that have a test_model_id. # # We split known failures by which subset of tests they apply to: @@ -103,10 +110,39 @@ def _load_hf_config(model_id: str): """ import transformers + revision = _PINNED_REVISIONS.get(model_id) try: - return transformers.AutoConfig.from_pretrained(model_id, trust_remote_code=False) + return transformers.AutoConfig.from_pretrained( + model_id, + revision=revision, + trust_remote_code=False, + ) except (ValueError, OSError): - return _try_load_config_json(model_id) + return _try_load_config_json(model_id, revision=revision) + + +def test_load_hf_config_forwards_yaml_revision(monkeypatch): + model_id = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + calls = [] + + def _from_pretrained(received_model_id, **kwargs): + calls.append((received_model_id, kwargs)) + return object() + + import transformers + + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", _from_pretrained) + _load_hf_config(model_id) + + assert calls == [ + ( + model_id, + { + "revision": "d468880b6ad3c6e0d21377ce7242adaea4cc884d", + "trust_remote_code": False, + }, + ) + ] def _resolve_hf_config(hf_config): @@ -160,6 +196,11 @@ def _build_graph(model_type: str, model_id: str): parent_config=parent_config, module_class=registration.module_class, ) + if model_type == "nemotron_h": + config = dataclasses.replace( + config, + dtype=resolve_dtype(_DECLARED_DTYPES[model_id]), + ) module = registration.module_class(config) task_name = registration.task or _default_task_for_model(model_type) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index fd02237a0..04c462f9d 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -5160,6 +5160,61 @@ def test_nemotron_h_moe_preprocess_weights(self): for key in result: assert not key.startswith("backbone."), f"Unrenamed key: {key}" + def test_reduced_precision_keeps_router_bias_and_ssm_cache_float32(self): + import onnx_ir as ir + + from mobius import build_from_module + from mobius._configs import NemotronHConfig + from mobius.models.nemotron_h import NemotronHCausalLMModel + + config = NemotronHConfig( + vocab_size=TINY_VOCAB, + hidden_size=TINY_HIDDEN, + intermediate_size=TINY_INTERMEDIATE, + num_hidden_layers=2, + num_attention_heads=TINY_HEADS, + num_key_value_heads=TINY_KV_HEADS, + rms_norm_eps=1e-5, + layer_types=["mamba2", "moe"], + mamba_n_heads=TINY_KV_HEADS, + mamba_d_head=TINY_HEAD_DIM, + mamba_d_state=16, + mamba_n_groups=1, + mamba_d_conv=4, + hidden_act="relu2", + head_dim=TINY_HEAD_DIM, + num_local_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=TINY_INTERMEDIATE, + dtype=ir.DataType.FLOAT16, + ) + package = build_from_module( + NemotronHCausalLMModel(config), + config, + task="hybrid-text-generation", + execution_provider="cuda", + ) + model = package["model"] + + assert ( + model.graph.initializers["model.layers.1.moe.gate.e_score_correction_bias"].dtype + == ir.DataType.FLOAT + ) + inputs = {value.name: value for value in model.graph.inputs} + assert inputs["past_key_values.0.conv_state"].dtype == ir.DataType.FLOAT16 + assert inputs["past_key_values.0.ssm_state"].dtype == ir.DataType.FLOAT + + def test_bfloat16_is_rejected_with_actionable_error(self): + import onnx_ir as ir + + config = self._nemotron_h_config() + config.dtype = ir.DataType.BFLOAT16 + + with pytest.raises(ValueError, match=r"BF16.*dtype='f16'"): + from mobius.models.nemotron_h import NemotronHCausalLMModel + + NemotronHCausalLMModel(config) + # =========================================================================== # Hybrid SSM+Attention (Jamba) model tests diff --git a/tests/cli_test.py b/tests/cli_test.py index d5ce840df..ac75ebf78 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -479,34 +479,6 @@ def test_runtime_ort_genai_propagates_trust_remote_code(self): assert mock_export.call_args.kwargs["trust_remote_code"] is True - def test_runtime_ort_genai_rejects_mage_vl_before_saving(self): - with ( - tempfile.TemporaryDirectory() as tmpdir, - mock.patch("mobius._model_package.ModelPackage.save") as save, - mock.patch( - "mobius.integrations.ort_genai.write_ort_genai_config" - ) as config_writer, - pytest.raises( - SystemExit, - match=r"Mage-VL.*patch_positions.*1D decoder position_ids", - ), - ): - main( - [ - "build", - "--model", - "microsoft/Mage-VL", - tmpdir, - "--no-weights", - "--trust-remote-code", - "--runtime", - "ort-genai", - ] - ) - - save.assert_not_called() - config_writer.assert_not_called() - def test_runtime_onnx_genai_uses_native_vlm_emitter(self): pkg = mock.MagicMock() pkg.items.return_value = [] diff --git a/tests/gguf_test.py b/tests/gguf_test.py index 9f67f0916..6e5600a52 100644 --- a/tests/gguf_test.py +++ b/tests/gguf_test.py @@ -9,6 +9,7 @@ from __future__ import annotations +import json from unittest import mock import numpy as np @@ -52,6 +53,9 @@ def _create_tiny_gguf( writer.add_context_length(128) writer.add_feed_forward_length(ffn_size) writer.add_vocab_size(vocab) + writer.add_bos_token_id(1) + writer.add_eos_token_id(2) + writer.add_pad_token_id(3) # Tensors — unquantized random weights rng = np.random.default_rng(42) @@ -544,20 +548,25 @@ def test_contradictory_quantization_flags_error(self, tmp_path): ) assert exc_info.value.code == 2 - def test_ort_genai_runtime_is_rejected_before_artifacts(self, tmp_path): - """build-gguf must not silently ignore an ORT GenAI runtime request.""" + def test_ort_genai_runtime_writes_graph_derived_config(self, tmp_path): from mobius.__main__ import main + path = _create_tiny_gguf(tmp_path / "test.gguf") output_dir = tmp_path / "output" - with pytest.raises(SystemExit, match="does not yet support --runtime ort-genai"): - main( - [ - "build-gguf", - str(tmp_path / "not-downloaded.gguf"), - "--runtime", - "ort-genai", - "--output", - str(output_dir), - ] - ) - assert not output_dir.exists() + main( + [ + "build-gguf", + path, + "--runtime", + "ort-genai", + "--output", + str(output_dir), + ] + ) + + assert (output_dir / "model.onnx").is_file() + assert (output_dir / "genai_config.json").is_file() + config = json.loads((output_dir / "genai_config.json").read_text(encoding="utf-8")) + assert config["model"]["bos_token_id"] == 1 + assert config["model"]["eos_token_id"] == 2 + assert config["model"]["pad_token_id"] == 3 diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 4df4bf723..049c8878d 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -241,7 +241,6 @@ def _all_registered_with_test_id() -> dict[str, str]: "mctct": "Audio model — no test_model_id yet", "megatron-bert": "Encoder — no test_model_id yet", "modernbert-decoder": "Decoder variant — no test_model_id yet", - "nemotron_h": "No test_model_id — no suitable public checkpoint", "nllb-moe": "Seq2seq MoE — no test_model_id yet", "nllb_moe": "Seq2seq MoE — no test_model_id yet", "ovis2": "VL model — no test_model_id yet", diff --git a/tests/nemotron_h_real_weight_test.py b/tests/nemotron_h_real_weight_test.py new file mode 100644 index 000000000..d2dd232d9 --- /dev/null +++ b/tests/nemotron_h_real_weight_test.py @@ -0,0 +1,385 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pinned reduced-real-weight L4/L5 tests for Nemotron 3.5 Lightning.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +_ROOT = Path(__file__).parents[1] +_EXAMPLE_DIR = _ROOT / "examples" / "olive" / "nemotron-3_5-lightning-30b" +_VALIDATOR_PATH = _EXAMPLE_DIR / "validate_reduced_checkpoint.py" +_L4_PATH = ( + _ROOT / "testdata" / "golden" / "causal-lm" / "nemotron-3_5-lightning-30b-reduced.json" +) +_L5_PATH = ( + _ROOT + / "testdata" + / "golden" + / "causal-lm" + / "nemotron-3_5-lightning-30b-reduced_generation.json" +) + + +def _load_validator(): + sys.path.insert(0, str(_EXAMPLE_DIR)) + try: + spec = importlib.util.spec_from_file_location( + "nemotron_reduced_validator", _VALIDATOR_PATH + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + finally: + sys.path.pop(0) + + +def test_load_eos_token_ids_unions_generation_and_model_config(tmp_path): + validator = _load_validator() + (tmp_path / "generation_config.json").write_text( + json.dumps({"eos_token_id": [2, 11]}), + encoding="utf-8", + ) + (tmp_path / "config.json").write_text( + json.dumps({"eos_token_id": 2}), + encoding="utf-8", + ) + + assert validator.load_eos_token_ids(tmp_path) == {2, 11} + + +def test_run_token_ids_accepts_olive_logits_and_stops_on_eos(tmp_path, monkeypatch): + validator = _load_validator() + inference_globals = validator.run_token_ids.__globals__ + (tmp_path / "model.onnx").touch() + + class _Output: + name = "logits_Q4" + + class _Session: + @staticmethod + def get_inputs(): + return [] + + @staticmethod + def get_outputs(): + return [_Output()] + + calls = [] + + def _run_session(_session, _output_names, _feeds): + calls.append(1) + logits = np.zeros((1, 1, 16), dtype=np.float32) + logits[0, 0, 2] = 1.0 + return [logits] + + monkeypatch.setitem(inference_globals, "_create_session", lambda *_args: _Session()) + monkeypatch.setitem(inference_globals, "_initial_states", lambda _session: {}) + monkeypatch.setitem(inference_globals, "_run_session", _run_session) + + generated, logits, profile = validator.run_token_ids( + tmp_path, + [1], + max_new_tokens=4, + device="cpu", + eos_token_ids={2}, + ) + + assert generated == [2] + assert len(logits) == 1 + assert len(calls) == 1 + assert profile is None + + +def test_pinned_range_fetch_retries_and_validates_headers(monkeypatch): + validator = _load_validator() + + class _Response: + def __init__(self, status, headers, content=b""): + self.status_code = status + self.headers = headers + self.content = content + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + class _Session: + def __init__(self): + self.responses = [ + _Response(503, {}), + _Response(206, {"Content-Range": "bytes 1-4/10", "Content-Length": "4"}), + _Response( + 206, {"Content-Range": "bytes 0-3/10", "Content-Length": "4"}, b"data" + ), + ] + self.calls = 0 + + def get(self, *_args, **_kwargs): + self.calls += 1 + return self.responses.pop(0) + + reader = object.__new__(validator._PinnedSafetensors) + reader._session = _Session() + sleeps = [] + monkeypatch.setattr(validator.time, "sleep", sleeps.append) + + assert reader._range("model.safetensors", 0, 3) == b"data" + assert reader._session.calls == 3 + assert sleeps == [1, 2] + + +def test_pinned_range_fetch_fails_explicitly_after_retries(monkeypatch): + validator = _load_validator() + + class _Response: + def __init__(self): + self.status_code = 503 + self.headers = {} + self.content = b"" + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + class _Session: + calls = 0 + + def get(self, *_args, **_kwargs): + self.calls += 1 + return _Response() + + reader = object.__new__(validator._PinnedSafetensors) + reader._session = _Session() + monkeypatch.setattr(validator.time, "sleep", lambda _delay: None) + + with pytest.raises(RuntimeError, match=r"failed after 3 attempts.*status=503"): + reader._range("model.safetensors", 0, 3) + assert reader._session.calls == 3 + + +def test_reduced_cache_rejects_stale_fixture_schema(tmp_path): + validator = _load_validator() + from safetensors.torch import save_file + + cache_path = tmp_path / "stale.safetensors" + save_file( + {"placeholder": torch.zeros(1)}, + cache_path, + metadata={ + "model_id": validator.MODEL_ID, + "revision": validator.REVISION, + "fixture_schema": "0", + }, + ) + + with pytest.raises(ValueError, match=r"fixture_schema.*Remove the stale cache"): + validator._build_reduced_state(cache_path) + + +@pytest.fixture(scope="module") +def reduced_real_state(): + validator = _load_validator() + configured_cache = os.environ.get("MOBIUS_NEMOTRON_REDUCED_CACHE") + cache = ( + Path(configured_cache) if configured_cache else validator.default_reduced_cache_path() + ) + state = validator._build_reduced_state(cache) + return validator, state + + +@pytest.fixture(scope="module") +def reduced_real_outputs(reduced_real_state, tmp_path_factory): + validator, state = reduced_real_state + package = validator._mobius_package(state, dtype_name="f32", ep="cpu") + output_dir = tmp_path_factory.mktemp("nemotron-onnx") + package.save(output_dir, external_data="onnx") + session = validator._create_session(output_dir / "model.onnx", "cpu", False) + prompt_ids = [1, 42, 17] + onnx_logits = validator._full_prefill(session, prompt_ids) + onnx_tokens, _step_logits, _profile = validator.run_token_ids( + output_dir, + prompt_ids, + max_new_tokens=4, + device="cpu", + eos_token_ids=validator.load_eos_token_ids(output_dir), + ) + + hf_model = validator._hf_model(state, dtype=torch.float32, device="cpu") + hf_logits = validator._hf_full_prefill(hf_model, prompt_ids, "cpu") + hf_tokens, _hf_step_logits = validator._hf_generate(hf_model, prompt_ids, "cpu", 4) + return onnx_logits, onnx_tokens, hf_logits, hf_tokens + + +def _require_cuda() -> None: + if os.environ.get("MOBIUS_TEST_DEVICE") != "cuda": + pytest.skip("Set MOBIUS_TEST_DEVICE=cuda to run reduced-real CUDA coverage") + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is unavailable") + import onnxruntime as ort + + if hasattr(ort, "preload_dlls"): + ort.preload_dlls() + if "CUDAExecutionProvider" not in ort.get_available_providers(): + pytest.skip("ONNX Runtime CUDAExecutionProvider is unavailable") + + +@pytest.fixture(scope="module") +def reduced_real_fp16_cuda(reduced_real_state, tmp_path_factory): + validator, state = reduced_real_state + _require_cuda() + output_root = tmp_path_factory.mktemp("nemotron-fp16-cuda") + package_dir = validator._validate_variant( + state, + output_root, + dtype_name="f16", + device="cuda", + ) + return validator, state, package_dir + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.golden +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_reduced_real_l4(reduced_real_outputs, model_type): + del model_type + onnx_logits, _onnx_tokens, hf_logits, _hf_tokens = reduced_real_outputs + golden = json.loads(_L4_PATH.read_text(encoding="utf-8")) + + np.testing.assert_allclose(onnx_logits, hf_logits, rtol=1e-3, atol=2e-3) + last_logits = hf_logits[0, -1] + top10 = np.argsort(last_logits)[::-1][:10] + summary = [last_logits.max(), last_logits.min(), last_logits.mean(), last_logits.std()] + + assert top10.tolist() == golden["top10_ids"] + np.testing.assert_allclose( + last_logits[top10], + [float.fromhex(value) for value in golden["top10_logits"]], + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + summary, + [float.fromhex(value) for value in golden["logits_summary"]], + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.generation +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_reduced_real_l5(reduced_real_outputs, model_type): + del model_type + _onnx_logits, onnx_tokens, _hf_logits, hf_tokens = reduced_real_outputs + golden = json.loads(_L5_PATH.read_text(encoding="utf-8")) + + assert len(onnx_tokens) == golden["max_new_tokens"] + assert len(hf_tokens) == golden["max_new_tokens"] + assert hf_tokens == golden["generated_tokens"] + assert onnx_tokens == golden["generated_tokens"] + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.golden +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_reduced_real_fp16_cuda(reduced_real_fp16_cuda, model_type): + del model_type + _validator, _state, package_dir = reduced_real_fp16_cuda + import onnx_ir as ir + + assert (package_dir / "model.onnx").is_file() + assert (package_dir / "model.onnx.data").is_file() + model = ir.load(package_dir / "model.onnx") + op_types = {(node.domain, node.op_type) for node in model.graph.all_nodes()} + assert ("", "Attention") in op_types + assert ("", "Scan") in op_types + assert ("", "Conv") in op_types + assert not any( + domain == "com.microsoft" + and op_type in {"GroupQueryAttention", "LinearAttention", "CausalConvWithState"} + for domain, op_type in op_types + ) + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.golden +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_bf16_rejection_evidence( + reduced_real_state, + tmp_path, + model_type, +): + del model_type + validator, state = reduced_real_state + _require_cuda() + + max_abs = validator._measure_bf16_rejection(state, tmp_path) + + assert max_abs > 1e-2 + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.generation +@pytest.mark.quantization +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_olive_q4_final_package( + reduced_real_fp16_cuda, + tmp_path, + model_type, +): + del model_type + import onnx_ir as ir + + validator, _state, source_dir = reduced_real_fp16_cuda + quantized_dir = validator.quantize_package(source_dir, tmp_path / "q4_k_m-cuda") + + assert (quantized_dir / "model.onnx").is_file() + assert (quantized_dir / "model.onnx.data").is_file() + assert (quantized_dir / "config.json").is_file() + assert sum( + path.stat().st_size for path in quantized_dir.iterdir() if path.is_file() + ) < sum(path.stat().st_size for path in source_dir.iterdir() if path.is_file()) + + quantized_model = ir.load(quantized_dir / "model.onnx") + assert ( + sum( + node.domain == "com.microsoft" and node.op_type == "MatMulNBits" + for node in quantized_model.graph.all_nodes() + ) + == 17 + ) + + generated, logits, profile_path = validator.run_token_ids( + quantized_dir, + [1, 42, 17], + max_new_tokens=4, + device="cuda", + profile=True, + eos_token_ids=validator.load_eos_token_ids(quantized_dir), + ) + + assert generated == [12, 13, 12, 12] + assert all(np.isfinite(step).all() for step in logits) + assert profile_path is not None + assert validator.summarize_profile(profile_path).get("CUDAExecutionProvider", 0) > 0 + Path(profile_path).unlink(missing_ok=True) diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index c6b2784f6..d789b23ad 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -208,11 +208,6 @@ # DeepSeek MLA: deepseek_v2_0 uses group_limited_greedy routing which hits a # HF transformers 5.3.0 bug (DeepseekV2Moe missing num_experts attr). "deepseek_v2_0": "HF transformers 5.3.0 bug: DeepseekV2Moe missing num_experts attr", - # Additional divergences (newly registered models) - # NemotronH Mamba2 layers diverge (cos=0.65): LinearAttention gated-SSM - # recurrence on CPU produces different results than HF's naive Mamba2. - # Attention-only layers match perfectly (cos=0.9999). - "nemotron_h": "Mamba2 SSM recurrence diverges on CPU (LinearAttention vs HF naive)", } # Fields that are properties in HF configs and cannot be set directly, @@ -581,12 +576,18 @@ def _create_hf_config(model_type: str, config_overrides: dict): for lt in layer_types ] - # NemotronH uses layers_block_type with HF values {"mamba", "attention", "moe"}. - # Convert our internal layer_types names (mamba2, full_attention, mlp) to HF names. + # NemotronH uses layers_block_type with current HF values + # {"linear_attention", "full_attention", "moe", "mlp"}. + # Convert our internal layer_types names to that vocabulary. # Also translate mobius Mamba field names to HF NemotronHConfig field names. if hf_model_type in ("nemotron_h",) and "layer_types" in hf_kwargs: layer_types = hf_kwargs.pop("layer_types") - _nemotron_type_map = {"mamba2": "mamba", "full_attention": "attention", "mlp": "moe"} + _nemotron_type_map = { + "mamba2": "linear_attention", + "full_attention": "full_attention", + "moe": "moe", + "mlp": "mlp", + } hf_kwargs["layers_block_type"] = [_nemotron_type_map.get(lt, lt) for lt in layer_types] # Mobius NemotronHConfig → HF NemotronHConfig field name mapping _nemotron_field_map = { @@ -596,6 +597,7 @@ def _create_hf_config(model_type: str, config_overrides: dict): "mamba_n_groups": "n_groups", "mamba_d_conv": "conv_kernel", "mamba_expand": "expand", + "shared_expert_intermediate_size": "moe_shared_expert_intermediate_size", } for old_name, new_name in _nemotron_field_map.items(): if old_name in hf_kwargs: diff --git a/tests/weight_alignment_test.py b/tests/weight_alignment_test.py index 34e0f2dea..5cb6bbec4 100644 --- a/tests/weight_alignment_test.py +++ b/tests/weight_alignment_test.py @@ -159,6 +159,58 @@ def test_identity_state_dict_roundtrip(self, model_type: str, config_overrides: _assert_identity_roundtrip(model_type, config_overrides) +def test_nemotron_h_filters_only_auxiliary_mtp_weights() -> None: + """Official MTP tensors are dropped without losing any decoder parameter.""" + config_overrides = next( + overrides + for model_type, overrides, _ in ALL_CAUSAL_LM_CONFIGS + if model_type == "nemotron_h" and "moe" in overrides.get("layer_types", []) + ) + config = _base_config(**config_overrides) + module = registry.get("nemotron_h")(config) + pkg = get_task(_default_task_for_model("nemotron_h")).build(module, config) + parameter_names = _collect_parameter_names(pkg) + state_dict = _build_identity_state_dict(pkg, parameter_names) + + # The 3.5 checkpoint carries these auxiliary training heads, while the + # trusted NemotronHForCausalLM implementation intentionally ignores them. + state_dict["mtp.layers.0.eh_proj.weight"] = torch.ones(1) + state_dict["mtp.layers.1.mixer.experts.0.up_proj.weight"] = torch.ones(1) + + aligned = module.preprocess_weights(state_dict) + + assert not any(name.startswith("mtp.") for name in aligned) + assert parameter_names <= set(aligned) + + +def test_nemotron_h_maps_per_expert_checkpoint_weights() -> None: + """The official 3.5 per-expert safetensor names map to MoE initializers.""" + config_overrides = next( + overrides + for model_type, overrides, _ in ALL_CAUSAL_LM_CONFIGS + if model_type == "nemotron_h" and "moe" in overrides.get("layer_types", []) + ) + config = _base_config(**config_overrides) + module = registry.get("nemotron_h")(config) + moe_layer = config.layer_types.index("moe") + state_dict = { + f"backbone.layers.{moe_layer}.mixer.experts.0.up_proj.weight": torch.ones(1), + f"backbone.layers.{moe_layer}.mixer.experts.0.down_proj.weight": torch.ones(1), + f"backbone.layers.{moe_layer}.mixer.gate.weight": torch.ones(1), + f"backbone.layers.{moe_layer}.mixer.gate.e_score_correction_bias": torch.ones(1), + } + + aligned = module.preprocess_weights(state_dict) + + prefix = f"model.layers.{moe_layer}.moe" + assert set(aligned) == { + f"{prefix}.experts.0.up_proj.weight", + f"{prefix}.experts.0.down_proj.weight", + f"{prefix}.gate.weight", + f"{prefix}.gate.e_score_correction_bias", + } + + # --------------------------------------------------------------------------- # Encoder-only weight alignment # ---------------------------------------------------------------------------