diff --git a/conftest.py b/conftest.py index 55af14737..62399a924 100644 --- a/conftest.py +++ b/conftest.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Root conftest: ensure the local worktree src/ is on sys.path first.""" +"""Root conftest: resolve the worktree package and shared test configs locally.""" from __future__ import annotations @@ -14,6 +14,9 @@ _src = os.path.join(os.path.dirname(__file__), "src") if _src not in sys.path: sys.path.insert(0, _src) +_tests = os.path.join(os.path.dirname(__file__), "tests") +if _tests not in sys.path: + sys.path.insert(1, _tests) # Force reload of mobius modules if they were loaded from a different location # (e.g., from the main repo's editable install rather than this worktree). diff --git a/docs/getting-started.md b/docs/getting-started.md index 2c6fdf539..706097268 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -94,6 +94,392 @@ pkg = build("llava-hf/llava-1.5-7b-hf") pkg = build("facebook/wav2vec2-base") ``` +### Export Drama 3 token generation + +Drama 3 exports three graphs, not a standalone text-to-waveform service: + +```python +from mobius import build + +pkg = build( + "fishaudio/drama-3", + revision="b2b5ee998034b2975720906c815e7fb0e9f007bd", + dtype="f32", + keep_quantized=False, +) +pkg.save("output/drama-3", max_workers=1) +``` + +Equivalent CLI: + +```bash +mobius build --model fishaudio/drama-3 --revision b2b5ee998034b2975720906c815e7fb0e9f007bd --dtype f32 --dequantize --max-workers 1 --output output/drama-3 +``` + +The official model ID defaults to this pinned revision; explicit revisions and +local checkpoint directories are honored. The 35B checkpoint stores BF16 +weights and block-scaled FP8 experts. `--dequantize` is required: conversion +streams parameters lazily into dense FP32 external data; `max_workers=1` bounds +serialization to one parameter at a time. Plan for roughly 144 GB of dense FP32 +weights, plus the source checkpoint and runtime allocations. FP16 and BF16 +production exports remain disabled while real-weight qualification continues. +Test-only prototypes pass synthetic full-logit parity at `1e-2` and 20 exact +cached greedy frames. The FP16 +prototype implements FP32 residual streams, native FP16-input/FP32-output +`GemmFloat8` tower down projections, and logged FP32 saturation before backbone +weight casts. Checkpoint storage is unchanged; other adapters retain their +conversion behavior unless they opt into FP16 saturation. + +The pinned Fish Speech source declares `torch==2.8.0`. Run +`scripts/validate_drama3.py reference` in a separate environment with that release; +the validator checks the Torch version before loading the model. New references +record Torch, CUDA and Triton versions, and comparisons check that metadata when +both goldens include it. Reference collection checks embeddings, hidden states +and logits for finite values before writing a golden. +Earlier Torch 2.10 diagnostics are retained separately: +CUDA reduction order differs between these releases. The committed token golden +has been regenerated with Torch 2.8.0+cu126 / Triton 3.4.0; all 20 frame IDs are +unchanged. + +Fresh Torch 2.8 production-dimension, one-layer real-weight probes still exceed +the `atol=rtol=1e-2` gate: FP16 prefill-tower output has 5 violations among 12288 +values (maximum violating error `0.01248169`); BF16 cached tower output has +200 violations among 4096 values (maximum `0.125`). With identical source +hidden-state inputs, the BF16 tower matches bit-for-bit through prefill and three +cached calls. The remaining pipeline difference starts in the backbone and +propagates into the tower. + +Identical real BF16 expert-matmul operands isolate one cached discrepancy to +the CUDA library version: cuBLAS 12.6.4 matches the Torch 2.8+cu126 reference +exactly, while cuBLAS 13.6.0 changes three output values. Algorithm and workspace +controls reproduce the same distinction. A separate CUDA 12.6 runtime is being +built for model replay; merely matching the ORT version does not match the +loaded cuBLAS implementation. + +Recurrent state, nonlinear opmath and RoPE remain FP32. Masks are preserved +when padding is present; redundant causal masks are omitted on unpadded +reduced-precision attention to enable the source-compatible fused CUDA path. +ORT 1.28's native BF16 `CausalConvWithState` fallback also does not initialize +on CUDA; qualification prototypes use standard-ONNX expansion for this dtype. + +Native W8A8 execution is implemented in +[microsoft/onnxruntime#33644](https://github.com/microsoft/onnxruntime/pull/33644). +The exporter prepares typed E4M3FN expert banks and unchanged multiplicative +FP32 grids for `QMoE(activation_quantization="dynamic_fp8")`, without a +dequantize-to-half proxy. Byte-preserving save/load and opt-in FP32/FP16/BF16 +expert-only parity against pinned Triton pass on the custom runtime +(`MOBIUS_TEST_DRAMA3_NATIVE_FP8=1`). Production native-FP8 admission remains closed. +The earlier complete native-FP8/FP32 probe, using Torch 2.10, failed the prefill +hidden-state gate +(43436 of 61440 values, maximum absolute error `0.68518`), although all 40 +real expert layers pass when supplied identical source inputs and logits. +The first expert input differs by only `1.192e-7`, but changes one FP8 code; +those discontinuities amplify across subsequent layers. Qualification against +the declared Torch 2.8 environment is in progress. Complete native-FP8 source runs +in BF16 and FP16 each produce 20 finite frames. A private source-ordered BF16 +candidate on the CUDA 13 runtime fails the full prefill hidden-state gate: +22817 of 61440 values violate `atol=rtol=1e-2`, with maximum absolute error +`0.640625`. All 40 native experts execute on CUDA. Its first expert input differs +in nine BF16 values (maximum `0.000244140625`); later layers amplify the difference. +An independent source replay reproduces all 20 saved states bit-for-bit. +The candidate is limited to batch-one short prefill and explicitly bypasses a +failing symbolic-shape pass, so it is diagnostic, not a production export. +Normal graph construction now succeeds with the existing constraint-propagation +optimization in +[onnx/onnx-shape-inference#135](https://github.com/onnx/onnx-shape-inference/pull/135). +A normally built candidate reuses the unchanged external weights only after +checking computation graphs, generated constants and initializer contracts; +all 120 expert-boundary tensors reproduce the original prefill exactly. +FP16 tower inference additionally uses +[corrected `GemmFloat8` geometry and output typing](https://github.com/onnx/onnx-shape-inference/commit/d915e3dc6490abf07081e873ff43d4baef81de97); +the normal FP16 build completes without skipped shape passes and preserves all +120 original expert-boundary tensors. +On the CUDA 13 runtime, the complete native-FP8/FP16 candidate fails the same +prefill gate in 16862 of 61440 values, with maximum error `1.158203125`. +Its source replay reproduces all 20 reference states exactly. CUDA-library-aligned +BF16 and FP16 comparisons remain pending. Source-ordered reductions +and [a CUDA reciprocal-square-root fast path](https://github.com/microsoft/onnxruntime/pull/33662) +give bit-exact normalization on real first-layer inputs; source-layout `Gemm` +also matches all five isolated DeltaNet projections exactly. An artifact-only +cached recurrence matches source output and state bitwise for zero/nonzero +cache inputs. Remaining prefix work concerns softplus, CPU-folded learned +exponentials and chunked triangular solves. + +Stable `Log(Add(x, 1))` fusion is implemented for CUDA and native WebGPU in +[this independent ORT change](https://github.com/microsoft/onnxruntime/commit/3c0638c4c402900e42d1a9a0f6ee8d9ea7015460). +It keeps the exported graph in standard ONNX and activates at +`ORT_ENABLE_EXTENDED` or `ORT_ENABLE_ALL`, after provider assignment. +Use `scripts/validate_drama3.py compare --optimization-level extended` (with the +usual work directory and reference arguments) to qualify this path. Comparison +defaults to BASIC; the selected level applies to every component and is recorded +in session-start events and the validation report. +The complete FP32 pipeline also passes EXTENDED on the combined QMoE/XOR/Pow/Flash +runtime: 30-row prefill, all cached vectors, 20 exact frames and 40 actual CUDA +MoE executions, with maximum decoder/tower/hidden errors +`5.149841e-5` / `5.054474e-5` / `2.384186e-5`. +The portable FP32 source-softplus subgraph matches all 960 real reference values +exactly on CUDA while preserving the preceding `Exp` rounding. +Native WebGPU FP32 and feature-enabled FP16 were exercised through lavapipe +(software Vulkan); physical-WebGPU and browser/WASM execution remain untested. +Deprecated JSEP is excluded by ORT policy. + +The custom CUDA build needs `onnxruntime_USE_FP8_QMOE=ON`, SM89+, and dynamic +activation quantization. The runtime rejects more than +65,535 token/expert pairs per call (8,191 tokens at top-8); the graph records +that limit and never silently falls back to dense expert conversion. + +| Graph | Inputs | Outputs | +|---|---|---| +| `embedding` | `frames`: INT64 `[B,T,11]`, each row `[main_token, code0, ..., code9]` | `inputs_embeds`, `shared_input`: `[B,T,H]` | +| `decoder` | `inputs_embeds`, `attention_mask`, `position_ids`, per-layer KV/conv/recurrent state | `logits`, `hidden_states`, updated states | +| `audio_decoder` | `hidden_states`, `codes`: INT64 `[B,T,10]`, `semantic_mask`: BOOL `[B,T]`, `shared_input`, mask, positions, tower KV caches | `codebook_logits`: `[B,T,9,4096]`, updated tower KV caches | + +Use the checkpoint tokenizer and Fish Speech's `reference-bracket-v1` prompt +format. Initialize ordinary KV caches with sequence length zero; initialize +DeltaNet conv/recurrent state with zeros and the **actual batch size**. +Recurrent state is always FP32. Masks are INT64 `[B,past+T]` and positions are +INT64 `[B,T]`. The audio tower's cache persists across **time**, not just within +one frame. Serialized symbolic dimensions are component-qualified; initialize +KV caches using their sequence axis (axis 2), not a bare dimension-name match. +Padding must be removed before running the recurrent backbone; +padded rows would otherwise update its state. + +For teacher-forced prompt ingestion, the backbone sees frame `p`, while the +tower's `codes` and `semantic_mask` label frame `p+1`; `shared_input` comes +from frame `p`. `semantic_mask` tests the next main token against the inclusive +semantic band `[248099,252194]`. Code zero is that main token minus 248099 +(clipped to `[0,4095]`); EOS is 248046. + +For generation, sample the main token from decoder logits restricted to EOS +and the semantic band. Fill code zero, then run `audio_decoder` once per +remaining codebook, sampling code `i+1` from logit slice `i`. Each invocation +uses the **same previous-frame caches**, not the preceding invocation's +outputs; retain only the last invocation's present caches. Later code inputs +may initially be zero because logit slice `i` depends only on codes `0..i`. +This teacher-forced ABI favors a compact export over efficient fused sampling. + +For the faster generation ABI, explicitly select `task="fish-omni-generation"`. +It replaces `audio_decoder` with `audio_generator`, **not** a second copy of its +weights: + +```python +pkg = build( + "fishaudio/drama-3", + dtype="f32", + execution_provider="cuda", + task="fish-omni-generation", + keep_quantized=False, +) +pkg.save("output/drama-3-fast", max_workers=1) +``` + +The FP32 CUDA route uses packed expert banks and `com.microsoft::MoE`, executing +selected experts rather than evaluating all 256 experts per layer. Its raw-logit +routing and physical expert-byte layout match ORT's floating MoE ABI. Gate/up +rows are concatenated offline, with half-split SwiGLU fusion and no FC3 input: +the separate-FC3 runtime path would copy a complete 2 GiB expert bank on every +layer invocation. The vocabulary projection also uses output-major `Gemm`, +avoiding a 2 GiB runtime weight transpose above ORT's constant-folding limit. It still +stores dense FP32 weights: this is **not native FP8**, and it does not remove the +large model-memory requirement. Other floating routes gather and evaluate only +selected token/expert pairs, preserving their activation and routing rounding +boundaries. Prefill is processed in bounded 16-token chunks; short prompts and +single-token decode bypass padding/Scan. This avoids evaluating inactive experts +whose FP16 overflows could otherwise poison output through `Inf * 0`. + +`audio_generator` accepts the original tower inputs plus scalar BOOL `generate` +and `do_sample`, and per-row `temperature`, `top_k`, `top_p`, `seed`, `counter` +controls matching the sampling policies below. Set `generate=false` for +teacher-forced prompt prefill; no RNG/sorting runs in this mode. For generation, +use `T=1`, put the main semantic code in `codes[...,0]`, and set +`generate=true`. `do_sample=false` is exact greedy selection; `true` performs +the source sampling rule. The graph generates each residual code before the +next layer consumes it, traversing the tower **once per frame**, instead of +nine complete teacher-forced replays. With the released configuration this +reduces tower-block evaluations from 117 to 13 per frame; this is a work-count +reduction, **not a measured 9x latency claim**. + +Outputs include `frame_codes` INT64 `[B,1,10]`, `valid` BOOL `[B]`, +`next_counter` INT64 `[B]`, logits and the complete committed KV state. +Nonfinite active sampling/neural inputs or multi-frame generation produce +`valid=false`, `frame_codes=-1`, and an unchanged counter. Do not commit cache +outputs on an invalid result. Greedy generation and prefill do not advance RNG +state. The test-only reduced-precision prototypes exercise this one-pass task +and selected-only expert execution, without opening production admission or +claiming the FP32 CUDA route's fused-expert kernel. Source-correct FP32 RoPE/mask dispatch can retain standard +Attention rather than a GQA rewrite. No reduced-precision throughput claim is +made without measurement. + +Sampling math is exported in `policies/main_sampler.onnx` and +`policies/codebook_sampler.onnx` and restored in `pkg.policy_components` on load. +Both take last-step `logits` `[B,V]`, `temperature`/`top_p` FLOAT `[B]`, `top_k`, +`seed`, and `counter` INT64 `[B]`. The main sampler also takes the last 16 main +tokens as `history` INT64 `[B,16]`, restricts candidates to EOS and the semantic +band, and rerolls only repeated semantic tokens. It returns `semantic_code` +alongside `token`, `valid`, and `next_counter`. The codebook sampler returns the +latter three fields and is invoked after each conditioned codebook projection. +Check every `valid` result: invalid/nonfinite inputs produce `token=-1` and +leave the counter unchanged. + +These policies implement Fish's drop-crossing-token nucleus rule, temperature +floor `1e-5`, and exponential-race sampling. They use explicit, reproducible +counter RNG state, **not Torch/Philox seed-bit identity**. Main sampling always +consumes both first/retry draws; counters advance by twice the compact candidate +count for main tokens and by the codebook vocabulary size for residual tokens, +modulo `2**31-1`. The caller updates history and orchestrates the temporal loop; +no huge neural weights are duplicated in the sampling policies. + +Reference audio encoding and ORT GenAI generation integration remain external. +The waveform decoder can be +exported as a separate FP32 ONNX graph: + +```python +from mobius.tasks.fish_codec import load_drama3_codec_decoder + +codec = load_drama3_codec_decoder() # pinned codec.pth, or checkpoint_path=local_path +codec.save("output/drama-3-codec") +# For generated NumPy frames [B,T,11], feed: +codes = frames[..., 1:].transpose(0, 2, 1) # INT64 [B,10,T] +``` + +`waveform_decoder` produces FP32 `[B,1,T*2048]` at 44100 Hz. It is a full-sequence +decoder, not a streaming-cache codec. Keep it FP32 independently of the text/tower +dtype. The pinned real codec matches all 40960 samples from the 20-frame native +token golden on CPU: maximum absolute error `1.848e-6`, RMSE `1.184e-7`, +SNR `117.79 dB`. This verifies decoding, not reference encoding or speech quality. +Its residual-codebook upper clamping matches upstream, including tower IDs above +the codec's nine residual vocabularies of 1024. + +`--runtime onnx-genai` writes graph-derived +advisory component contracts, not an executable speech pipeline; +`--runtime ort-genai` is explicitly rejected. Fish Audio's checkpoint/source +license is separate from Mobius's MIT license. + +The original complete pinned FP32 export was validated on H200 with ORT 1.28 CUDA and +TF32 disabled: a 30-row prefill and 20 greedy frames match the independent +upstream FP32 pipeline. Maximum observed decoder/tower logit differences were +`3.624e-5` / `3.910e-5`. This is token-level evidence, not waveform or speech-quality +validation. The native reference uses an argmax sampler for determinism. + +The current pre-fused CUDA MoE plus single-pass tower route has also been +requalified against the complete independent Torch 2.8.0+cu126 reference: +all vectors for the 30-row prefill and 20 exact complete frames pass, using CPU +embedding and CUDA decoder/tower on one H200. Both prefill profiles contain +40 actual CUDA MoE executions. + +| Runtime | Maximum decoder-logit error | Maximum tower-logit error | Maximum hidden error | +|---|---:|---:|---:| +| Stock ORT 1.28 | `5.5313e-5` | `5.0306e-5` | `2.2888e-5` | +| Custom ORT 1.32, QMoE + XOR + Pow | `5.1498e-5` | `5.0545e-5` | `2.3842e-5` | + +The second correctness run uses the separately preserved combined-Pow wheel +(`da934e6b1791e3225c7f2e2a61a35a5ebd907d2946a8f0c0e02c7e455dd6395d`). +The performance measurements below retain their original two-patch runtime. + +Recorded single-H200 measurements of this dense-FP32 graph use ORT 1.28 +(`0368187f84`), TF32 disabled, two full-request warmups, the same 30-row prompt, +CPU embeddings, CUDA decoder/tower/codec, and no competing GPU processes: + +| Tower / policy | Complete frames | Steady frames/s | p50 / p95 frame ms | End-to-end RTF | +|---|---:|---:|---:|---:| +| Replay / greedy | 100 | 14.92 | 64.67 / 66.22 | 1.549 | +| Single pass / greedy | 100 | 32.79 | 28.58 / 28.71 | **0.770** | +| Single pass / sampled | 53 (native EOS) | 26.31 | 34.12 / 34.50 | 1.058 | + +The two greedy runs produce identical 100-frame hashes; the optimized +resident-state 20-frame prefix also matches the independent source golden. +Single-pass greedy is 2.20x faster in steady frame rate and 2.01x faster +end-to-end than replay in this controlled comparison, not 9x. Its prefill is +334 ms, first code frame 341 ms, and full waveform delivery 3.58 seconds for +4.64 seconds of audio. Model initialization is separate, approximately +101-105 seconds; GPU memory remains very large (about 142,161 MiB in the +single-pass greedy run). These are warmed single-request observations, not +universal throughput or streaming-audio claims. **The sampled run still misses +the RTF 1.0 target** and is not a performance pass. + +A separately qualified custom ORT 1.32 build combining +[microsoft/onnxruntime#33644](https://github.com/microsoft/onnxruntime/pull/33644) +and [microsoft/onnxruntime#33645](https://github.com/microsoft/onnxruntime/pull/33645), +with the codec's NC1D layout enabled, reaches **RTF 0.901** on the same sampled +53-frame case (29.28 steady frames/s, 30.28 / 30.57 ms p50 / p95, 348 ms first +code frame). It produces the identical sampled-frame hash and passes the +independent 20-frame/full-vector dense-FP32 reference on that runtime. +This run still uses **dense FP32 experts**, not the unqualified native-FP8 +pipeline. The combined build includes QMoE commit `e3181530` and XOR commit +`e90d897c`, CUDA 13.3 / SM90, FP8 QMoE enabled, and Flash/MEA attention disabled. +Its wheel SHA256 is +`f7e2f4cb9255e0556ca9cbc03c40189219880039eb76cdea23a000c0f2ae97dc`. +Initialization is separate (88.61 seconds), and waveform delivery is +full-sequence, not streaming. These cross-runtime observations do not isolate +the XOR kernel's speedup; separate kernel profiles confirm its CUDA placement. + +To reproduce on a sufficiently large CUDA host, select free devices and run: + +```bash +python scripts/validate_drama3.py download --work-dir drama3-validation +python scripts/validate_drama3.py export --work-dir drama3-validation +python scripts/validate_drama3.py reference --work-dir drama3-validation +python scripts/validate_drama3.py compare --work-dir drama3-validation +``` + +By default, reference loading spreads the backbone across GPUs 0-3 and places +the tower on GPU 4; comparison places embedding/audio/decoder on GPUs 0/1/2. +Override these with the script's device flags on a shared host. The script +retains full reference vectors and verifies the committed token-generation +golden, while the generic waveform golden runner remains out of scope. +To qualify the accelerated route instead, add +`--task fish-omni-generation --execution-provider cuda` to the export command +in a fresh output directory. Comparison automatically selects `audio_generator` +when present, uses one greedy conditioned tower call per frame, and checks its +full logits and frame IDs against the same independent reference. +`--audio-component audio_decoder` explicitly selects the replay graph when a +package contains both tower variants. Invalid generation or changed greedy RNG +state is an error, never a successful comparison. +Experimental native-FP8 references use `--native-fp8` and a separate work +directory; comparison requires its independent `--golden` until qualification. +The validator records and checks expert execution mode, so a native reference +cannot silently be compared as dense FP32. These flags do not bypass the +production native-export gate. + +For performance measurement, `scripts/benchmark_drama3.py` runs actual cached +requests with I/O binding: recurrent/KV state and neural intermediates remain on +the selected device, while only small policy outputs and the final waveform +return to the host. It supports both the original replay task and the one-pass +task, so their timings can be compared using the same ORT build and options. +For a memory-constrained single H200, `--embedding-device cpu` keeps the +roughly 2 GiB vocabulary lookup on the host while leaving decoder/tower caches +on CUDA. Only the small embedding activations cross devices; this reserves GPU +headroom for `--codec-device cuda`. Component placement is recorded in the +report. `--codec-device cpu` is also supported, with its actual decode time +included in end-to-end RTF. +Supply the pinned HF `config.json` and the independent reference stage's +`state-00.npz` prompt: + +```bash +python scripts/benchmark_drama3.py --package output/drama-3-fast --config drama3-config.json --prompt drama3-validation/reference/state-00.npz --codec output/drama-3-codec --device cuda --device-id 0 --embedding-device cpu --codec-device cuda --codec-conv1d-pad-to-nc1d --frames 100 --output drama3-perf.json +``` + +Add `--sample` for the ONNX stochastic policy. The recorded real-time sampled +case requires the custom runtime above; stock ORT 1.28 falls back to CPU for +integer XOR in the RNG. The codec layout option is explicit and should be +measured on the deployment GPU rather than assumed universally faster. + +The report separates initialization, prefill, first code frame, steady +frames/second, p50/p95 frame latency and waveform decoding. End-to-end RTF includes +prefill, generation and waveform decoding relative to actual audio duration. +The codec is full-sequence: **first-code-frame latency is not streaming +time-to-audio**. Omit `--codec` only for a token-pipeline diagnostic; such a run +does not attest full speech performance. + +Use an idle, sufficiently large GPU. Busy devices are rejected by default; +`--allow-contended` and `--profile` are explicitly unqualified diagnostics. +Performance sign-off requires at least 50 complete frames, no competing GPU +processes observed around the run, unprofiled timing and an independently +established quality gate. The optional `--max-rtf` target defaults to `1.0`; +qualified end-to-end runs exceeding it fail after writing their measurements. +The current optimizations have source-parity and execution-path coverage, but +**no fresh uncontended full-model H200 latency/RTF result is claimed**. + ### Build from a GGUF file Convert a GGUF model (e.g. from llama.cpp) to ONNX. Quantized target storage is diff --git a/pyproject.toml b/pyproject.toml index 0c262c3dc..8bba5b774 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ transformers = ["transformers", "safetensors"] gguf = ["gguf>=0.10.0"] testing = [ "accelerate", + "loguru", # Pinned Fish Speech reference implementation in Drama 3 parity tests. "diffusers", "onnxruntime-easy>=0.1.1", # onnxruntime>=1.28 is required for the fused Attention kernel to apply diff --git a/scripts/benchmark_drama3.py b/scripts/benchmark_drama3.py new file mode 100644 index 000000000..544178bee --- /dev/null +++ b/scripts/benchmark_drama3.py @@ -0,0 +1,671 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Benchmark Drama 3 requests with GPU-resident recurrent/KV state. + +Supports the teacher-forced baseline and single-pass generation task. Report +initialization, prefill, first frame, steady frame latency and optional waveform +RTF separately. A contended or short run is diagnostic, never performance sign-off. +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import io +import json +import os +import statistics +import subprocess +import tempfile +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +from onnxscript import GraphBuilder +from validate_drama3 import empty_cache_ports + +from mobius._testing.ort_inference import _numpy_to_ort_value, _ort_value_to_numpy + + +@dataclass(frozen=True) +class TensorView: + """Keep the owning OrtValue alive while binding a zero-copy contiguous slice.""" + + owner: ort.OrtValue + shape: tuple[int, ...] + offset: int = 0 + + +def sequence_view(value, start, length, *, squeeze_time=False): + shape = value.shape() + if len(shape) != 3 or shape[0] != 1 or not 0 <= start < start + length <= shape[1]: + raise ValueError(f"Invalid batch-one contiguous sequence slice: {shape}") + itemsize = {1: 4, 10: 2, 16: 2}[value.element_type()] + target = (1, shape[2]) if squeeze_time else (1, length, shape[2]) + if squeeze_time and length != 1: + raise ValueError("Only a length-one sequence can lose its time axis") + return TensorView(value, target, start * shape[2] * itemsize) + + +def codebook_view(value, index): + shape = value.shape() + if len(shape) != 4 or shape[:2] != [1, 1] or not 0 <= index < shape[2]: + raise ValueError(f"Invalid single-frame codebook view: {shape}, index={index}") + itemsize = {1: 4, 10: 2, 16: 2}[value.element_type()] + return TensorView(value, (1, shape[3]), index * shape[3] * itemsize) + + +class BoundSession: + """Bind caches/hidden states on-device; copy only small policy/audio outputs.""" + + def __init__( + self, + path, + *, + device, + device_id=0, + profile=False, + profile_directory=None, + cudnn_conv1d_pad_to_nc1d=False, + ): + options = ort.SessionOptions() + options.intra_op_num_threads = 8 + options.inter_op_num_threads = 1 + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + options.enable_mem_pattern = False + options.enable_profiling = profile + if profile: + profile_directory = profile_directory or Path(tempfile.gettempdir()) + profile_directory.mkdir(parents=True, exist_ok=True) + suffix = hashlib.sha256(str(path).encode()).hexdigest()[:12] + options.profile_file_prefix = str( + profile_directory / f"{path.parent.name}-{suffix}" + ) + providers = ( + [ + ( + "CUDAExecutionProvider", + { + "device_id": device_id, + "use_tf32": "0", + "arena_extend_strategy": "kSameAsRequested", + "cudnn_conv1d_pad_to_nc1d": str(int(cudnn_conv1d_pad_to_nc1d)), + }, + ), + "CPUExecutionProvider", + ] + if device == "cuda" + else ["CPUExecutionProvider"] + ) + self.session = ort.InferenceSession(str(path), options, providers=providers) + self.device, self.device_id = device, device_id if device == "cuda" else 0 + self.outputs = [port.name for port in self.session.get_outputs()] + + def run(self, feeds, *, fetches=None, cpu_outputs=()): + names = self.outputs if fetches is None else list(fetches) + if not set(names) <= set(self.outputs): + raise ValueError(f"Unknown output names: {set(names) - set(self.outputs)}") + binding = self.session.io_binding() + # This local mapping owns every raw-pointer view until execution completes. + owners = {} + for name, value in feeds.items(): + if isinstance(value, TensorView): + owners[name] = value.owner + owner_device = value.owner.device_name() + binding.bind_input( + name, + device_type=owner_device, + device_id=self.device_id if owner_device == "cuda" else 0, + element_type=value.owner.element_type(), + shape=value.shape, + buffer_ptr=value.owner.data_ptr() + value.offset, + ) + else: + owners[name] = ( + value + if isinstance(value, ort.OrtValue) + else _numpy_to_ort_value( + np.ascontiguousarray(value) if value.ndim else value + ) + ) + binding.bind_ortvalue_input(name, owners[name]) + for name in names: + target = "cpu" if name in cpu_outputs else self.device + binding.bind_output( + name, device_type=target, device_id=self.device_id if target == "cuda" else 0 + ) + self.session.run_with_iobinding(binding) + binding.synchronize_outputs() + values = binding.get_outputs() + if len(values) != len(names): + raise RuntimeError("ORT returned an unexpected bound-output count") + return dict(zip(names, values)) + + def empty_cache(self): + return empty_cache_ports(self.session.get_inputs()) + + +def next_cache(outputs, device): + result = { + name.replace("present.", "past_key_values.", 1): value + for name, value in outputs.items() + if name.startswith("present.") + } + if not result or any(value.device_name() != device for value in result.values()): + raise RuntimeError("Recurrent/KV state did not remain on the execution device") + return result + + +def gpu_snapshot(device_id): + text = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,name,uuid,memory.used,utilization.gpu", + "--format=csv,noheader,nounits", + ], + text=True, + ) + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + selected = visible.split(",")[device_id].strip() if visible else str(device_id) + processes = subprocess.check_output( + ["nvidia-smi", "--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader,nounits"], + text=True, + ) + for row in csv.reader(io.StringIO(text)): + if (selected.isdigit() and int(row[0]) == int(selected)) or ( + not selected.isdigit() and row[2].strip().startswith(selected) + ): + return { + "device_id": device_id, + "physical_index": int(row[0]), + "name": row[1].strip(), + "memory_used_mib": int(row[3]), + "utilization_percent": int(row[4]), + "foreign_pids": [ + int(process[1]) + for process in csv.reader(io.StringIO(processes)) + if process + and process[0].strip() == row[2].strip() + and int(process[1]) != os.getpid() + ], + } + raise ValueError(f"CUDA device {device_id} not found") + + +def require_idle(snapshot, allow_contended): + busy = ( + snapshot["utilization_percent"] > 5 + or snapshot["memory_used_mib"] > 1024 + or bool(snapshot.get("foreign_pids")) + ) + if busy and not allow_contended: + raise RuntimeError( + "GPU is occupied; refusing misleading performance measurements. " + "Use an idle device, or --allow-contended for an explicitly unqualified diagnostic." + ) + return not busy + + +def greedy_model(config, dtype): + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + logits = builder.input("logits", dtype, [1, config["text_config"]["vocab_size"]]) + candidates = builder.op.Constant( + value_ints=sorted( + [ + config["eos_token_id"], + *range(config["semantic_start_token_id"], config["semantic_end_token_id"] + 1), + ] + ) + ) + selected = builder.op.Gather(logits, candidates, axis=-1) + token = builder.op.Gather( + candidates, + builder.op.ArgMax( + builder.op.Cast(selected, to=ir.DataType.FLOAT), axis=-1, keepdims=0 + ), + ) + token.type, token.shape = ir.TensorType(ir.DataType.INT64), ir.Shape([1]) + builder.add_output(token, "token") + return ir.Model(graph, ir_version=12) + + +def benchmark(args): + if args.frames < 2 or args.warmup < 1: + raise ValueError("Use at least two measured frames and one warmup request") + if args.device == "cpu" and "cuda" in (args.embedding_device, args.codec_device): + raise ValueError("CUDA component placement requires --device cuda") + if args.codec_conv1d_pad_to_nc1d and ( + args.codec is None or (args.codec_device or args.device) != "cuda" + ): + raise ValueError("The cuDNN codec layout option requires a CUDA codec") + if not np.isfinite(args.max_rtf) or args.max_rtf <= 0 or args.device_id < 0: + raise ValueError("Use a finite positive RTF target and nonnegative device ID") + idle = True + before = None + if args.device == "cuda": + before = gpu_snapshot(args.device_id) + idle = require_idle(before, args.allow_contended) + ort.preload_dlls() + config = json.loads(args.config.read_text()) + with np.load(args.prompt) as saved: + prompt = saved["input_frames"] + count = config["audio_decoder_config"]["num_codebooks"] + if prompt.ndim != 2 or prompt.shape[1] != count + 1 or len(prompt) < 2: + raise ValueError("Prompt must contain at least two [main,code0,...] frame rows") + if len(prompt) + args.frames > config["text_config"]["max_position_embeddings"]: + raise ValueError("Benchmark exceeds the backbone position capacity") + fast = (args.package / "audio_generator" / "model.onnx").is_file() + audio_name = "audio_generator" if fast else "audio_decoder" + started = time.perf_counter() + kwargs = { + "device": args.device, + "device_id": args.device_id, + "profile": args.profile, + "profile_directory": args.output.parent / "drama3-profiles", + } + sessions = { + name: BoundSession( + args.package / name / "model.onnx", + **{ + **kwargs, + "device": (args.embedding_device or args.device) + if name == "embedding" + else args.device, + }, + ) + for name in ("embedding", "decoder", audio_name) + } + policy = { + "temperature": np.array([args.temperature], np.float32), + "top_k": np.array([args.top_k], np.int64), + "top_p": np.array([args.top_p], np.float32), + "seed": np.array([args.seed], np.int64), + } + controls = { + **policy, + "counter": np.zeros(1, np.int64), + "generate": np.array(False), + "do_sample": np.array(args.sample), + } + small_outputs = {"token", "frame_codes", "valid", "semantic_code", "next_counter"} + for name in ("main_sampler", "codebook_sampler"): + if args.sample: + sessions[name] = BoundSession(args.package / "policies" / f"{name}.onnx", **kwargs) + dtype = sessions["decoder"].session.get_outputs()[0].type + runtime_features = ( + sessions["decoder"] + .session.get_modelmeta() + .custom_metadata_map.get("mobius.required_runtime_features", "") + ) + element_type = { + "tensor(float)": ir.DataType.FLOAT, + "tensor(float16)": ir.DataType.FLOAT16, + "tensor(bfloat16)": ir.DataType.BFLOAT16, + }[dtype] + + def audio(feeds): + fetches = [ + name + for name in sessions[audio_name].outputs + if name != "codebook_logits" or not fast + ] + result = sessions[audio_name].run(feeds, fetches=fetches, cpu_outputs=small_outputs) + if fast and not _ort_value_to_numpy(result["valid"]).all(): + raise FloatingPointError( + "Single-pass tower rejected a frame; caches were not committed" + ) + return result + + with tempfile.TemporaryDirectory(prefix="drama3-benchmark-") as temporary: + greedy_path = Path(temporary) / "greedy.onnx" + ir.save(greedy_model(config, element_type), greedy_path) + sessions["greedy"] = BoundSession(greedy_path, **kwargs) + + def request(limit): + origin = time.perf_counter() + decoder_cache = sessions["decoder"].empty_cache() + tower_cache = sessions[audio_name].empty_cache() + counter = np.zeros(1, np.int64) + history = np.zeros((1, 16), np.int64) + positions = np.arange(len(prompt), dtype=np.int64)[None] + mask = np.ones_like(positions) + embedded = sessions["embedding"].run({"frames": prompt[None]}) + decoded = sessions["decoder"].run( + { + **decoder_cache, + "inputs_embeds": embedded["inputs_embeds"], + "position_ids": positions, + "attention_mask": mask, + } + ) + decoder_cache = next_cache(decoded, args.device) + prompt_semantic = (prompt[None, 1:, 0] >= config["semantic_start_token_id"]) & ( + prompt[None, 1:, 0] <= config["semantic_end_token_id"] + ) + prefilled = audio( + { + **tower_cache, + **(controls if fast else {}), + "hidden_states": sequence_view( + decoded["hidden_states"], 0, len(prompt) - 1 + ), + "shared_input": sequence_view( + embedded["shared_input"], 0, len(prompt) - 1 + ), + "codes": prompt[None, 1:, 1:], + "semantic_mask": prompt_semantic, + "position_ids": positions[:, :-1], + "attention_mask": mask[:, :-1], + } + ) + tower_cache = next_cache(prefilled, args.device) + prefill_seconds = time.perf_counter() - origin + durations, frames = [], [] + first_frame_seconds = None + frame = prompt[None, -1:] + for step in range(limit): + tick = time.perf_counter() + position = len(prompt) - 1 + step + if step: + embedded = sessions["embedding"].run({"frames": frame}) + positions = np.array([[position]], np.int64) + mask = np.ones((1, position + 1), np.int64) + decoded = sessions["decoder"].run( + { + **decoder_cache, + "inputs_embeds": embedded["inputs_embeds"], + "position_ids": positions, + "attention_mask": mask, + } + ) + decoder_cache = next_cache(decoded, args.device) + logits = sequence_view( + decoded["logits"], decoded["logits"].shape()[1] - 1, 1, squeeze_time=True + ) + selection = sessions["main_sampler" if args.sample else "greedy"].run( + { + "logits": logits, + **( + {**policy, "counter": counter, "history": history} + if args.sample + else {} + ), + }, + cpu_outputs=small_outputs, + ) + if args.sample: + if not _ort_value_to_numpy(selection["valid"]).all(): + raise FloatingPointError("Main sampling rejected invalid logits") + counter = _ort_value_to_numpy(selection["next_counter"]) + main = int(_ort_value_to_numpy(selection["token"])[0]) + if main == config["eos_token_id"]: + break + history = np.concatenate( + (history[:, 1:], np.array([[main]], np.int64)), axis=1 + ) + codes = np.zeros((1, 1, count), np.int64) + codes[0, 0, 0] = main - config["semantic_start_token_id"] + tower_feeds = { + **tower_cache, + "hidden_states": sequence_view( + decoded["hidden_states"], decoded["hidden_states"].shape()[1] - 1, 1 + ), + "shared_input": sequence_view( + embedded["shared_input"], embedded["shared_input"].shape()[1] - 1, 1 + ), + "semantic_mask": np.ones((1, 1), bool), + "position_ids": np.array([[position]], np.int64), + "attention_mask": np.ones((1, position + 1), np.int64), + } + for index in range(1 if fast else count - 1): + output = audio( + { + **tower_feeds, + "codes": codes, + **( + {**controls, "generate": np.array(True), "counter": counter} + if fast + else {} + ), + } + ) + if fast: + codes = _ort_value_to_numpy(output["frame_codes"]) + counter = _ort_value_to_numpy(output["next_counter"]) + elif args.sample: + sampled = sessions["codebook_sampler"].run( + { + **policy, + "counter": counter, + "logits": codebook_view(output["codebook_logits"], index), + }, + cpu_outputs=small_outputs, + ) + if not _ort_value_to_numpy(sampled["valid"]).all(): + raise FloatingPointError( + "Codebook sampling rejected invalid logits" + ) + codes[0, 0, index + 1] = _ort_value_to_numpy(sampled["token"])[0] + counter = _ort_value_to_numpy(sampled["next_counter"]) + else: + picked = sessions["greedy_codebook"].run( + {"logits": codebook_view(output["codebook_logits"], index)}, + cpu_outputs=small_outputs, + ) + codes[0, 0, index + 1] = _ort_value_to_numpy(picked["token"])[0] + tower_cache = next_cache(output, args.device) + frame = np.concatenate((np.array([[[main]]], np.int64), codes), axis=-1) + frames.append(frame[0, 0]) + durations.append(time.perf_counter() - tick) + if first_frame_seconds is None: + first_frame_seconds = time.perf_counter() - origin + elapsed = time.perf_counter() - origin + if not frames: + raise RuntimeError( + "No audio frames generated; no valid speech performance measurement" + ) + return np.stack(frames), durations, prefill_seconds, first_frame_seconds, elapsed + + if not fast and not args.sample: + vocabulary = config["audio_decoder_config"]["vocab_size"] + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + logits = builder.input("logits", element_type, [1, vocabulary]) + token = builder.op.ArgMax( + builder.op.Cast(logits, to=ir.DataType.FLOAT), axis=-1, keepdims=0 + ) + token.type, token.shape = ir.TensorType(ir.DataType.INT64), ir.Shape([1]) + builder.add_output(token, "token") + codebook_path = Path(temporary) / "greedy_codebook.onnx" + ir.save(ir.Model(graph, ir_version=12), codebook_path) + sessions["greedy_codebook"] = BoundSession(codebook_path, **kwargs) + if args.codec is not None: + path = args.codec / "waveform_decoder" / "model.onnx" + if not path.is_file(): + path = args.codec / "model.onnx" + sessions["waveform_decoder"] = BoundSession( + path, + **{ + **kwargs, + "device": args.codec_device or args.device, + "cudnn_conv1d_pad_to_nc1d": args.codec_conv1d_pad_to_nc1d, + }, + ) + metadata = sessions["waveform_decoder"].session.get_modelmeta().custom_metadata_map + if ( + metadata.get("sample_rate") != "44100" + or metadata.get("samples_per_frame") != "2048" + ): + raise ValueError( + "Waveform graph must declare the Drama 3 44100Hz/2048-sample frame contract" + ) + initialization = time.perf_counter() - started + + def decode_waveform(frames): + tick = time.perf_counter() + waveform = sessions["waveform_decoder"].run( + {"codes": frames[None, :, 1:].transpose(0, 2, 1)} + ) + samples = _ort_value_to_numpy(waveform["waveform"]) + seconds = time.perf_counter() - tick + if samples.shape != (1, 1, len(frames) * 2048) or not np.isfinite(samples).all(): + raise RuntimeError("Invalid waveform output; refusing a performance pass") + return seconds + + for _ in range(args.warmup): + warm_frames, *_ = request(args.frames) + if args.codec is not None: + decode_waveform(warm_frames) + after_warmup = gpu_snapshot(args.device_id) if args.device == "cuda" else None + frames, durations, prefill, first_frame, elapsed = request(args.frames) + waveform_seconds = decode_waveform(frames) if args.codec is not None else None + after_measurement = gpu_snapshot(args.device_id) if args.device == "cuda" else None + no_competitors = all( + not snapshot.get("foreign_pids") + for snapshot in (before, after_warmup, after_measurement) + if snapshot is not None + ) + steady = durations[1:] + audio_seconds = len(frames) * 2048 / 44100 + result = { + "ort_version": ort.__version__, + "ort_build": ort.get_build_info(), + "device": args.device, + "embedding_device": sessions["embedding"].device, + "codec_device": sessions["waveform_decoder"].device + if args.codec is not None + else None, + "codec_provider_options": sessions[ + "waveform_decoder" + ].session.get_provider_options() + if args.codec is not None + else None, + "provider_options": ( + { + "device_id": args.device_id, + "use_tf32": "0", + "arena_extend_strategy": "kSameAsRequested", + } + if args.device == "cuda" + else {} + ), + "graph_optimization_level": "ORT_ENABLE_ALL", + "gpu_before": before, + "gpu_after_warmup": after_warmup, + "gpu_after_measurement": after_measurement, + "dtype": dtype, + "expert_execution": "dynamic_fp8" + if "activation_quantization=dynamic_fp8" in runtime_features + else "dense", + "single_pass_tower": fast, + "sampling": args.sample, + "seed": args.seed, + "temperature": args.temperature, + "top_k": args.top_k, + "top_p": args.top_p, + "requested_frames": args.frames, + "warmup_requests": args.warmup, + "prompt_rows": len(prompt), + "prompt_sha256": hashlib.sha256(prompt.tobytes()).hexdigest(), + "tower_calls_per_generated_frame": 1 if fast else count - 1, + "graph_sha256": { + name: hashlib.sha256( + (args.package / name / "model.onnx").read_bytes() + ).hexdigest() + for name in ("embedding", "decoder", audio_name) + }, + "weights_hashed": False, + "gpu_resident_caches": args.device == "cuda", + "config_sha256": hashlib.sha256(args.config.read_bytes()).hexdigest(), + "initialization_seconds": initialization, + "prefill_seconds": prefill, + "first_frame_seconds": first_frame, + "generated_frames": len(frames), + "generated_frames_sha256": hashlib.sha256(frames.tobytes()).hexdigest(), + "steady_frames_per_second": len(steady) / sum(steady) if steady else None, + "p50_frame_ms": statistics.median(steady) * 1000 if steady else None, + "p95_frame_ms": float(np.percentile(steady, 95)) * 1000 if steady else None, + "token_pipeline_rtf": elapsed / audio_seconds, + "waveform_seconds": waveform_seconds, + "time_to_waveform_seconds": elapsed + waveform_seconds + if waveform_seconds is not None + else None, + "waveform_delivery": "Full sequence; first-frame latency is not streaming time-to-audio", + "end_to_end_rtf": (elapsed + waveform_seconds) / audio_seconds + if waveform_seconds is not None + else None, + "measurement_qualified": ( + args.device == "cuda" + and idle + and no_competitors + and len(frames) >= 50 + and not args.profile + and not args.allow_contended + ), + "quality_validation": "Must be established independently; this benchmark does not attest numerical parity", + } + if args.profile: + result["profiles"] = { + name: session.session.end_profiling() for name, session in sessions.items() + } + result["end_to_end_rtf_target"] = args.max_rtf + result["meets_end_to_end_rtf_target"] = ( + result["end_to_end_rtf"] <= args.max_rtf + if result["measurement_qualified"] and result["end_to_end_rtf"] is not None + else None + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + if result["meets_end_to_end_rtf_target"] is False: + raise RuntimeError(f"End-to-end RTF exceeds the requested {args.max_rtf} target") + return result + + +def parser(): + result = argparse.ArgumentParser(description=__doc__) + result.add_argument("--package", required=True, type=Path) + result.add_argument( + "--config", required=True, type=Path, help="Pinned Hugging Face config.json" + ) + result.add_argument( + "--prompt", required=True, type=Path, help="Independent state-00.npz with input_frames" + ) + result.add_argument("--output", required=True, type=Path) + result.add_argument("--codec", type=Path) + result.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + result.add_argument( + "--embedding-device", + choices=("cpu", "cuda"), + help="Override embedding placement to leave more GPU memory for the codec", + ) + result.add_argument("--codec-device", choices=("cpu", "cuda")) + result.add_argument( + "--codec-conv1d-pad-to-nc1d", + action="store_true", + help="Use cuDNN's alternative 1-D codec layout; measure on the target GPU", + ) + result.add_argument("--device-id", type=int, default=0) + result.add_argument("--frames", type=int, default=100) + result.add_argument("--warmup", type=int, default=2) + result.add_argument("--sample", action="store_true") + result.add_argument("--temperature", type=float, default=0.7) + result.add_argument("--top-k", type=int, default=30) + result.add_argument("--top-p", type=float, default=0.95) + result.add_argument("--seed", type=int, default=42) + result.add_argument( + "--profile", action="store_true", help="Diagnostic run; timings are unqualified" + ) + result.add_argument("--allow-contended", action="store_true", help="Diagnostic run only") + result.add_argument("--max-rtf", type=float, default=1.0) + return result + + +if __name__ == "__main__": + benchmark(parser().parse_args()) diff --git a/scripts/benchmark_drama3_test.py b/scripts/benchmark_drama3_test.py new file mode 100644 index 000000000..74fda11f8 --- /dev/null +++ b/scripts/benchmark_drama3_test.py @@ -0,0 +1,273 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Transport, contention guards and real-graph benchmark-driver checks.""" + +from __future__ import annotations + +import copy +import dataclasses +import json +from types import SimpleNamespace + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch +from _test_configs import FISH_OMNI_CONFIG +from benchmark_drama3 import ( + BoundSession, + benchmark, + codebook_view, + parser, + require_idle, + sequence_view, +) +from onnxscript import GraphBuilder +from safetensors.torch import save_file + +from mobius import build_from_module +from mobius._configs.fish_codec import FishCodecConfig +from mobius._configs.fish_omni import FishOmniConfig +from mobius._testing.ort_inference import _numpy_to_ort_value +from mobius.integrations.transformers._config_resolver import _dict_to_pretrained_config +from mobius.integrations.transformers._fish_omni_weights import ( + _source_name, + stream_fish_omni_safetensors_to_package, +) +from mobius.models import FishOmniModel +from mobius.models.fish_omni_test import reduced_precision_prototype_task +from mobius.tasks.fish_codec import build_fish_codec_decoder + + +def test_cpu_benchmark_rejects_untracked_cuda_component_placement(tmp_path): + args = parser().parse_args( + [ + "--package", + str(tmp_path), + "--config", + str(tmp_path / "config.json"), + "--prompt", + str(tmp_path / "prompt.npz"), + "--output", + str(tmp_path / "result.json"), + "--device", + "cpu", + "--embedding-device", + "cuda", + ] + ) + with pytest.raises(ValueError, match="CUDA component placement"): + benchmark(args) + + +def test_busy_gpu_cannot_be_a_performance_pass(): + busy = {"memory_used_mib": 130000, "utilization_percent": 100} + with pytest.raises(RuntimeError, match="refusing misleading"): + require_idle(busy, False) + assert require_idle(busy, True) is False + with pytest.raises(RuntimeError, match="refusing misleading"): + require_idle( + {"memory_used_mib": 10, "utilization_percent": 0, "foreign_pids": [123]}, False + ) + + +def test_scalar_boolean_keeps_its_rank(): + value = _numpy_to_ort_value(np.array(True)) + assert value.shape() == [] + + +def test_cpu_views_and_outputs_keep_cpu_device_zero_with_nonzero_cuda_id(): + value = _numpy_to_ort_value(np.zeros((1, 1, 4), np.float32)) + calls = [] + + def bind(kind, name, **kwargs): + calls.append((kind, name, kwargs["device_type"], kwargs["device_id"])) + + binding = SimpleNamespace( + bind_input=lambda name, **kwargs: bind("input", name, **kwargs), + bind_output=lambda name, **kwargs: bind("output", name, **kwargs), + synchronize_outputs=lambda: None, + get_outputs=lambda: [value], + ) + session = BoundSession.__new__(BoundSession) + session.session = SimpleNamespace( + io_binding=lambda: binding, run_with_iobinding=lambda binding: None + ) + session.outputs, session.device, session.device_id = ["output"], "cuda", 3 + actual = session.run({"x": sequence_view(value, 0, 1)}, cpu_outputs=("output",)) + assert actual["output"] is value + assert calls == [("input", "x", "cpu", 0), ("output", "output", "cpu", 0)] + + +def test_zero_copy_contiguous_views_bind_the_correct_bytes(tmp_path): + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + value = builder.input("x", ir.DataType.FLOAT, [1, 4]) + output = builder.op.Identity(value) + output.type, output.shape = value.type, value.shape + builder.add_output(output, "y") + path = tmp_path / "view.onnx" + ir.save(ir.Model(graph, ir_version=12), path) + session = BoundSession(path, device="cpu") + values = np.arange(12, dtype=np.float32).reshape(1, 3, 4) + owner = ort.OrtValue.ortvalue_from_numpy(values) + result = session.run({"x": sequence_view(owner, 2, 1, squeeze_time=True)}) + np.testing.assert_array_equal(result["y"].numpy(), values[:, 2]) + books = ort.OrtValue.ortvalue_from_numpy(values.reshape(1, 1, 3, 4)) + result = session.run({"x": codebook_view(books, 1)}) + np.testing.assert_array_equal(result["y"].numpy(), values[:, 1]) + with pytest.raises(ValueError, match="Invalid"): + sequence_view(owner, 2, 2) + + +@pytest.mark.parametrize("task", ["fish-omni", "fish-omni-generation"]) +@pytest.mark.parametrize("sample", [False, True]) +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="CUDA binding smoke test requires a GPU", + ), + ), + ], +) +def test_benchmark_runs_actual_cached_graphs_without_claiming_full_model_perf( + tmp_path, task, sample, device, dtype +): + if dtype != ir.DataType.FLOAT and device == "cpu": + pytest.skip("Reduced-precision benchmark routes require CUDA") + config = FishOmniConfig.from_transformers( + _dict_to_pretrained_config(copy.deepcopy(FISH_OMNI_CONFIG)) + ) + config = dataclasses.replace(config, dtype=dtype) + package = build_from_module( + FishOmniModel(config), + config, + task=reduced_precision_prototype_task(task) if dtype != ir.DataType.FLOAT else task, + execution_provider={ + ir.DataType.FLOAT: "default", + ir.DataType.FLOAT16: "cuda", + ir.DataType.BFLOAT16: "onnx-standard", + }[dtype], + ) + torch.manual_seed(42) + weights = {} + for model in package.values(): + for name, value in model.graph.initializers.items(): + if value.const_value is not None: + continue + tensor = torch.randn(tuple(value.shape)) * 0.1 + source = _source_name(name) + if ".packed_experts." in name: + prefix, projection = source.rsplit(".experts.", 1) + for expert, matrix in enumerate(tensor): + if projection == "gate_up_proj.weight": + gate, up = matrix.chunk(2, dim=0) + weights[f"{prefix}.experts.{expert}.gate_proj.weight"] = ( + gate.contiguous() + ) + weights[f"{prefix}.experts.{expert}.up_proj.weight"] = up.contiguous() + else: + weights[f"{prefix}.experts.{expert}.{projection}"] = ( + matrix.contiguous() + ) + else: + weights[source] = tensor + # Keep the request alive: for either sign of hidden[0], a semantic row wins. + weights["lm_head.weight"].zero_() + weights["lm_head.weight"][config.semantic_start_token_id, 0] = 100 + weights["lm_head.weight"][config.semantic_start_token_id + 1, 0] = -100 + save_file(weights, tmp_path / "model.safetensors") + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + model_dir = tmp_path / "onnx" + package.save(str(model_dir), max_workers=1, progress_bar=False) + source_config = tmp_path / "config.json" + source_config.write_text(json.dumps(FISH_OMNI_CONFIG)) + prompt = tmp_path / "prompt.npz" + np.savez(prompt, input_frames=np.array([[1, 0, 0, 0], [3, 0, 0, 0]], np.int64)) + mixed_placement = ( + task == "fish-omni-generation" + and device == "cuda" + and dtype == ir.DataType.FLOAT + and not sample + ) + with_codec = task == "fish-omni-generation" and (device == "cpu" or mixed_placement) + codec_path = tmp_path / "codec" + if with_codec: + codec = build_fish_codec_decoder( + FishCodecConfig( + latent_dim=16, + residual_codebooks=2, + transformer_layers=1, + transformer_heads=2, + transformer_intermediate_size=32, + max_frames=16, + decoder_dim=32, + ) + ) + rng = np.random.default_rng(42) + for value in codec["waveform_decoder"].graph.initializers.values(): + if value.const_value is None: + value.const_value = ir.tensor( + rng.normal(0, 0.02, tuple(value.shape)).astype(np.float32) + ) + codec.save(str(codec_path), max_workers=1, progress_bar=False) + args = parser().parse_args( + [ + "--package", + str(model_dir), + "--config", + str(source_config), + "--prompt", + str(prompt), + "--output", + str(tmp_path / "result.json"), + "--frames", + "4", + "--warmup", + "1", + "--device", + device, + "--top-k", + "1", + *(["--sample"] if sample else []), + *(["--allow-contended", "--profile"] if device == "cuda" else []), + *(["--codec", str(codec_path)] if with_codec else []), + *(["--embedding-device", "cpu"] if mixed_placement else []), + *(["--codec-conv1d-pad-to-nc1d"] if mixed_placement else []), + ] + ) + result = benchmark(args) + assert result["generated_frames"] == 4 + assert result["expert_execution"] == "dense" + assert result["single_pass_tower"] == (task == "fish-omni-generation") + assert result["measurement_qualified"] is False + assert result["meets_end_to_end_rtf_target"] is None + assert result["gpu_resident_caches"] == (device == "cuda") + assert result["embedding_device"] == ("cpu" if mixed_placement else device) + assert result["codec_device"] == (device if with_codec else None) + if mixed_placement: + assert ( + result["codec_provider_options"]["CUDAExecutionProvider"][ + "cudnn_conv1d_pad_to_nc1d" + ] + == "1" + ) + if with_codec: + assert result["end_to_end_rtf"] > 0 + assert result["waveform_seconds"] > 0 + assert result["time_to_waveform_seconds"] > result["first_frame_seconds"] + else: + assert result["end_to_end_rtf"] is None + assert result["steady_frames_per_second"] > 0 + assert result["p95_frame_ms"] > 0 diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py new file mode 100644 index 000000000..181f321d1 --- /dev/null +++ b/scripts/validate_drama3.py @@ -0,0 +1,768 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Validate full Drama 3 token graphs against pinned Fish Speech inference. + +Run download, export, reference, and compare as separate processes. The default +reference placement requires five CUDA devices; comparison uses three. Outputs +include full reference tensors, exact frame IDs, error bounds and an ORT CUDA +profile. The waveform codec is intentionally outside this validator. +""" + +from __future__ import annotations + +import argparse +import collections +import contextlib +import hashlib +import importlib.util +import json +import os +import pathlib +import sys +import time +import types +import urllib.request + +from mobius._configs.fish_omni import DRAMA3_REVISION, DRAMA3_SOURCE_REVISION + +ROOT: pathlib.Path +REVISION = DRAMA3_REVISION +SOURCE_REVISION = DRAMA3_SOURCE_REVISION +PROMPT = "Hello, this is a speech synthesis test." +REFERENCE_DEVICES: tuple[int, ...] = (0, 1, 2, 3) +REFERENCE_TOWER_DEVICE = 4 +EMBEDDING_DEVICE = 0 +AUDIO_DEVICE = 1 +DECODER_DEVICE = 2 +DTYPE = "f32" +EXPERT_EXECUTION = "dense" +DTYPE_NAMES = {"f32": "float32", "f16": "float16", "bf16": "bfloat16"} +ORT_OPTIMIZATION_LEVELS = { + "disabled": "ORT_DISABLE_ALL", + "basic": "ORT_ENABLE_BASIC", + "extended": "ORT_ENABLE_EXTENDED", + "all": "ORT_ENABLE_ALL", +} +GOLDEN = ( + pathlib.Path(__file__).resolve().parents[1] + / "testdata" + / "golden" + / "speech" + / "drama-3_tokens_generation.json" +) + + +def event(kind, **details): + print(json.dumps({"event": kind, **details}), flush=True) + + +def empty_cache_ports(ports): + import ml_dtypes + import numpy as np + + caches = {} + for port in ports: + if not port.name.startswith("past_key_values."): + continue + shape = list(port.shape) + shape[0] = 1 + if port.name.endswith((".key", ".value")): + assert len(shape) == 4 + shape[2] = 0 + assert all(isinstance(dim, int) for dim in shape), (port.name, shape) + type_name = getattr(port, "type", "tensor(float)") + dtypes = { + "tensor(float)": np.float32, + "tensor(float16)": np.float16, + "tensor(bfloat16)": ml_dtypes.bfloat16, + } + if type_name not in dtypes: + raise ValueError(f"Unsupported cache dtype {type_name}: {port.name}") + caches[port.name] = np.zeros(shape, dtype=dtypes[type_name]) + return caches + + +def download(): + from huggingface_hub import snapshot_download + + path = snapshot_download( + "fishaudio/drama-3", + revision=REVISION, + cache_dir=str(ROOT / "hf-cache"), + allow_patterns=[ + "config.json", + "model*.safetensors", + "model.safetensors.index.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "chat_template.jinja", + "LICENSE", + ], + max_workers=8, + ) + (ROOT / "checkpoint-path.txt").write_text(path) + event("checkpoint_ready", revision=REVISION) + + +def checkpoint_path(): + path = pathlib.Path((ROOT / "checkpoint-path.txt").read_text().strip()).resolve() + if not path.is_dir() or path.name != REVISION: + raise ValueError(f"Expected the immutable checkpoint snapshot {REVISION}, got {path}") + return path + + +def load_upstream(): + prefix = "fish_speech.models.text2semantic" + for name in ("fish_speech", "fish_speech.models", prefix, f"{prefix}.fish_omni"): + package = types.ModuleType(name) + package.__path__ = [] + sys.modules[name] = package + result = {} + source_dir = ROOT / "upstream" / SOURCE_REVISION + source_dir.mkdir(parents=True, exist_ok=True) + for name in ( + "errors", + "sampling", + "fish_omni.fp8_moe", + "fish_omni.config", + "fish_omni.layers", + "fish_omni.moe", + "fish_omni.backbone", + "fish_omni.tower", + "fish_omni.model", + "fish_omni.generate", + ): + path = source_dir / (name + ".py") + if not path.exists(): + url = ( + "https://raw.githubusercontent.com/fishaudio/fish-speech/" + f"{SOURCE_REVISION}/fish_speech/models/text2semantic/" + f"{name.replace('.', '/')}.py" + ) + with urllib.request.urlopen(url, timeout=60) as response: + path.write_bytes(response.read()) + spec = importlib.util.spec_from_file_location(f"{prefix}.{name}", path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load the pinned reference module {name!r}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + result[name.rsplit(".", 1)[-1]] = module + return result + + +def export(checkpoint, *, task="fish-omni", execution_provider=None, native_fp8=False): + from mobius import build + + start = time.monotonic() + event("export_start") + package = build( + str(checkpoint), + revision=REVISION, + dtype=DTYPE, + keep_quantized=native_fp8, + task=task, + execution_provider=execution_provider + or ( + "cuda" + if native_fp8 + else {"f32": "default", "f16": "cuda", "bf16": "onnx-standard"}[DTYPE] + ), + ) + event("graph_and_bindings_ready", seconds=time.monotonic() - start) + package.save(str(ROOT / "onnx"), max_workers=1) + event("export_complete", seconds=time.monotonic() - start) + (ROOT / "export-result.json").write_text( + json.dumps( + { + "revision": REVISION, + "dtype": DTYPE, + "expert_execution": "dynamic_fp8" if native_fp8 else "dense", + "seconds": time.monotonic() - start, + "weight_loading_report": package.weight_loading_report, + "files": { + str(path.relative_to(ROOT / "onnx")): path.stat().st_size + for path in (ROOT / "onnx").rglob("*") + if path.is_file() + }, + }, + indent=2, + ) + ) + + +def validate_reference_torch_version(version: str) -> None: + if version.split("+", 1)[0] != "2.8.0": + raise RuntimeError( + "The pinned Fish Speech reference requires torch==2.8.0; " + f"found {version}. Use a separate reference environment." + ) + + +def reference_array(tensor, *, name: str): + import numpy as np + + value = tensor.detach().float().cpu().numpy() + if not np.isfinite(value).all(): + raise FloatingPointError( + f"Pinned Fish Speech reference produced nonfinite {name} in {DTYPE} execution" + ) + return value + + +def reference(checkpoint, *, native_fp8=False): + os.environ["NVIDIA_TF32_OVERRIDE"] = "0" + os.environ["FISH_DRAMA3_FP8_EXPERTS"] = "1" if native_fp8 else "0" + os.environ["FISH_DRAMA3_CUDA_GRAPH"] = "0" + import numpy as np + import torch + + validate_reference_torch_version(torch.__version__) + torch.set_num_threads(8) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + upstream = load_upstream() + start = time.monotonic() + event( + "reference_load_start", + devices=REFERENCE_DEVICES, + tower_device=REFERENCE_TOWER_DEVICE, + ) + model = upstream["model"].FishOmniModel.from_pretrained( + checkpoint, + devices=[f"cuda:{device}" for device in REFERENCE_DEVICES], + tower_device=f"cuda:{REFERENCE_TOWER_DEVICE}", + dtype={"f32": torch.float32, "f16": torch.float16, "bf16": torch.bfloat16}[DTYPE], + ) + event("reference_loaded", seconds=time.monotonic() - start) + backbone, tower = model.backbone, model.audio_decoder + if any(layer.mlp.experts.fp8 != native_fp8 for layer in backbone.model.layers): + raise AssertionError("Pinned source did not activate the requested expert execution") + prompt = upstream["generate"].build_drama3_prompt(PROMPT, model.tokenizer, model.config) + states = [] + active_heads: list[np.ndarray] = [] + original_embed = tower.embed_backbone + original_forward = backbone.forward + original_prefill = tower.prefill + original_decode = tower.decode_step + + def embed(frames, text): + output = original_embed(frames, text) + unseen = tower.embedding_sum( + frames.to(tower.device), + 1 + tower.config.vq_backbone_codebooks, + tower.config.num_codebooks, + ) + unseen = torch.where( + tower.is_semantic(frames[:, 0].to(tower.device))[:, None], unseen, 0.0 + ) + states.append( + { + "input_frames": frames.detach().cpu().numpy(), + "inputs_embeds": reference_array(output, name="embedding")[None], + "shared_input": reference_array(unseen, name="shared_input")[None], + } + ) + return output + + def forward(*args, **kwargs): + hidden = original_forward(*args, **kwargs) + states[-1]["hidden_states"] = reference_array(hidden, name="decoder.hidden_states") + states[-1]["logits"] = reference_array( + backbone.compute_logits(hidden), name="decoder.logits" + ) + states[-1]["position_ids"] = kwargs["input_pos"].cpu().numpy()[None] + return hidden + + def prefill(*args, **kwargs): + active_heads.clear() + output = original_prefill(*args, **kwargs) + states[-1]["prefill_tower_logits"] = np.stack(active_heads, axis=-2)[None] + return output + + def decode(*args, **kwargs): + active_heads.clear() + output = original_decode(*args, **kwargs) + states[-1]["tower_logits"] = np.stack(active_heads, axis=-2)[None] + states[-1]["sampled_frame"] = output.detach().cpu().numpy() + return output + + def capture_head(head): + def capture(_module, _inputs, output): + active_heads.append( + reference_array(head(output), name=f"tower.head_{len(active_heads)}.logits") + ) + + return capture + + handles = [ + layer.register_forward_hook(capture_head(head)) + for layer, head in zip(tower.residual_layers, tower.codebook_heads) + ] + tower.embed_backbone = embed + backbone.forward = forward + tower.prefill = prefill + tower.decode_step = decode + upstream["tower"].sample_top_k_top_p = lambda logits, *a, **kw: logits.argmax(-1) + with torch.inference_mode(): + generated, finish = upstream["generate"].generate_drama3( + model, + prompt, + max_new_tokens=20, + temperature=1.0, + top_k=30, + top_p=1.0, + ) + for handle in handles: + handle.remove() + output_dir = ROOT / "reference" + output_dir.mkdir(exist_ok=True) + for index, state in enumerate(states): + np.savez(output_dir / f"state-{index:02}.npz", **state) + logits = states[0]["logits"] + top = np.argsort(logits[0, -1])[-10:][::-1] + golden = { + "model_id": "fishaudio/drama-3", + "revision": REVISION, + "source_revision": SOURCE_REVISION, + "dtype": DTYPE_NAMES[DTYPE], + "expert_execution": "dynamic_fp8" if native_fp8 else "dense", + "reference_environment": { + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "triton": getattr(upstream["fp8_moe"].triton, "__version__", None), + }, + "prompt": PROMPT, + "prompt_rows": len(prompt), + "prefill_shape": list(logits.shape), + "prefill_last_top10": top.tolist(), + "prefill_last_top10_logits": logits[0, -1, top].tolist(), + "prefill_logits_sha256": hashlib.sha256(logits.tobytes()).hexdigest(), + "max_new_tokens": 20, + "do_sample": False, + "finish_reason": finish, + "generated_frames": generated.tolist(), + "attempted_frames": len(states), + "seconds": time.monotonic() - start, + } + (output_dir / "golden.json").write_text(json.dumps(golden, indent=2)) + event( + "reference_complete", + **{ + k: golden[k] + for k in ("prompt_rows", "finish_reason", "attempted_frames", "seconds") + }, + ) + + +def validate_reference_golden(golden): + import numpy as np + + committed = json.loads(GOLDEN.read_text()) + assert golden.get("expert_execution", "dense") == EXPERT_EXECUTION, "expert_execution" + assert committed.get("expert_execution", "dense") == EXPERT_EXECUTION, "expert_execution" + if "reference_environment" in golden and "reference_environment" in committed: + assert golden["reference_environment"] == committed["reference_environment"], ( + "reference_environment" + ) + assert golden["dtype"] == DTYPE_NAMES[DTYPE], "dtype" + for field in ( + "model_id", + "revision", + "source_revision", + "dtype", + "prompt", + "max_new_tokens", + "do_sample", + "finish_reason", + ): + assert golden[field] == committed[field], field + assert len(golden["generated_frames"]) == len(committed["generated_frames"]) + np.testing.assert_array_equal(golden["generated_frames"], committed["generated_frames"]) + assert golden["prefill_shape"] == committed["prefill_shape"] + np.testing.assert_array_equal( + golden["prefill_last_top10"], committed["prefill_last_top10"] + ) + np.testing.assert_allclose( + golden["prefill_last_top10_logits"], + committed["prefill_last_top10_logits"], + atol=1e-2 if DTYPE != "f32" else 1e-4, + rtol=1e-2 if DTYPE != "f32" else 1e-4, + ) + + +def run_session(session, feeds): + """Bridge NumPy/BF16 feeds and outputs without changing their tensor types.""" + import numpy as np + + from mobius._testing.ort_inference import _numpy_to_ort_value, _ort_value_to_numpy + + converted = { + name: _numpy_to_ort_value(np.ascontiguousarray(value) if value.ndim else value) + for name, value in feeds.items() + } + cuda_options = session.get_provider_options().get("CUDAExecutionProvider") + device_context: contextlib.AbstractContextManager[object] = contextlib.nullcontext() + if cuda_options is not None: + import torch + + # Cross-component calls can leave another GPU current on this thread; + # ORT's stream events must be created/recorded on the session's device. + device_context = torch.cuda.device(int(cuda_options["device_id"])) + with device_context: + values = session(**converted) + ports = session.get_outputs() + assert len(ports) == len(values) + return {port.name: _ort_value_to_numpy(value) for port, value in zip(ports, values)} + + +def run_tower_session(session, feeds, *, generate=False): + """Run teacher-forced or single-pass greedy tower against the same reference.""" + import numpy as np + + names = {port.name for port in session.get_inputs()} + if "generate" in names: + feeds = { + **feeds, + "generate": np.array(generate), + "do_sample": np.array(False), + "temperature": np.array([0.7], np.float32), + "top_k": np.array([30], np.int64), + "top_p": np.array([0.95], np.float32), + "seed": np.array([42], np.int64), + "counter": np.zeros(1, np.int64), + } + result = run_session(session, feeds) + if "generate" in names: + if not result["valid"].all(): + raise FloatingPointError("Single-pass tower rejected reference inputs") + if result["next_counter"].any(): + raise AssertionError("Greedy single-pass tower advanced RNG state") + return result + + +def compare(*, audio_model_name=None, optimization_level="basic"): + os.environ["NVIDIA_TF32_OVERRIDE"] = "0" + import numpy as np + import onnxruntime as ort + import onnxruntime_easy as ort_easy + + if optimization_level not in ORT_OPTIMIZATION_LEVELS: + raise ValueError(f"Unknown ORT optimization level: {optimization_level!r}") + optimization_name = ORT_OPTIMIZATION_LEVELS[optimization_level] + optimization = getattr(ort.GraphOptimizationLevel, optimization_name) + golden = json.loads((ROOT / "reference" / "golden.json").read_text()) + validate_reference_golden(golden) + config = json.loads((checkpoint_path() / "config.json").read_text()) + ort.preload_dlls() + sessions = {} + if audio_model_name is None: + audio_model_name = ( + "audio_generator" + if (ROOT / "onnx" / "audio_generator" / "model.onnx").is_file() + else "audio_decoder" + ) + start = time.monotonic() + for name, device in ( + ("embedding", EMBEDDING_DEVICE), + ("audio_decoder", AUDIO_DEVICE), + ("decoder", DECODER_DEVICE), + ): + options = ort.SessionOptions() + options.intra_op_num_threads = 8 + options.inter_op_num_threads = 1 + options.graph_optimization_level = optimization + options.enable_mem_pattern = False + options.add_session_config_entry("session.disable_prepacking", "1") + if name == "decoder": + options.enable_profiling = True + options.profile_file_prefix = str(ROOT / "decoder-profile") + event( + "ort_session_start", + component=name, + device=device, + ort_graph_optimization=optimization_name, + ) + component = audio_model_name if name == "audio_decoder" else name + sessions[name] = ort_easy.EasySession( + str(ROOT / "onnx" / component / "model.onnx"), + sess_options=options, + providers=[ + ( + "CUDAExecutionProvider", + { + "device_id": device, + "use_tf32": "0", + "arena_extend_strategy": "kSameAsRequested", + }, + ), + "CPUExecutionProvider", + ], + ) + event("ort_session_ready", component=name, seconds=time.monotonic() - start) + + def run(name, feeds, *, generate=False): + if name == "audio_decoder": + return run_tower_session(sessions[name], feeds, generate=generate) + return run_session(sessions[name], feeds) + + expected_type = { + "f32": "tensor(float)", + "f16": "tensor(float16)", + "bf16": "tensor(bfloat16)", + }[DTYPE] + assert sessions["decoder"].get_outputs()[0].type == expected_type + runtime_features = ( + sessions["decoder"] + .get_modelmeta() + .custom_metadata_map.get("mobius.required_runtime_features", "") + ) + actual_experts = ( + "dynamic_fp8" if "activation_quantization=dynamic_fp8" in runtime_features else "dense" + ) + assert actual_experts == EXPERT_EXECUTION, "Graph and reference expert execution differ" + + def empty_cache(name): + return empty_cache_ports(sessions[name].get_inputs()) + + def next_cache(output): + return { + key.replace("present.", "past_key_values.", 1): value + for key, value in output.items() + if key.startswith("present.") + } + + metrics = {} + + def check(name, actual, expected): + assert actual.shape == expected.shape, (name, actual.shape, expected.shape) + actual = actual.astype(np.float32) + assert np.isfinite(actual).all(), name + difference = np.abs(actual - expected) + maximum = float(difference.max()) if difference.size else 0.0 + metrics[name] = { + "shape": list(actual.shape), + "max_abs_diff": maximum, + "mean_abs_diff": float(difference.mean()) if difference.size else 0.0, + } + event("comparison", name=name, **metrics[name]) + tolerance = 1e-2 if DTYPE != "f32" else 1e-3 + np.testing.assert_allclose( + actual, expected, atol=tolerance, rtol=tolerance, err_msg=name + ) + + count = config["audio_decoder_config"]["num_codebooks"] + vocab = config["audio_decoder_config"]["vocab_size"] + semantic_start = config["semantic_start_token_id"] + semantic_end = config["semantic_end_token_id"] + eos = config["eos_token_id"] + candidates = np.array([eos, *range(semantic_start, semantic_end + 1)]) + decoder_cache, tower_cache = empty_cache("decoder"), empty_cache("audio_decoder") + generated = [] + current_frames: np.ndarray | None = None + profile_summary = None + for step in range(golden["attempted_frames"]): + with np.load(ROOT / "reference" / f"state-{step:02}.npz") as state: + if step == 0: + current_frames = state["input_frames"] + else: + assert current_frames is not None + np.testing.assert_array_equal(current_frames, state["input_frames"]) + assert current_frames is not None + embedded = run("embedding", {"frames": current_frames[None]}) + check(f"state_{step}.embedding", embedded["inputs_embeds"], state["inputs_embeds"]) + check( + f"state_{step}.shared_input", embedded["shared_input"], state["shared_input"] + ) + positions = state["position_ids"] + mask = np.ones((1, int(positions[0, -1]) + 1), dtype=np.int64) + decoded = run( + "decoder", + { + **decoder_cache, + "inputs_embeds": embedded["inputs_embeds"], + "position_ids": positions, + "attention_mask": mask, + }, + ) + if step == 0: + profile_path = pathlib.Path(sessions["decoder"].end_profiling()) + counts: collections.Counter[tuple[str, str]] = collections.Counter() + for entry in json.loads(profile_path.read_text()): + args = entry.get("args", {}) + if "provider" in args: + counts[(args["provider"], args.get("op_name", "unknown"))] += 1 + profile_summary = { + f"{provider}/{op}": n for (provider, op), n in counts.items() + } + assert ( + sum( + counts[("CUDAExecutionProvider", op)] + for op in ("MatMul", "FusedMatMul", "Gemm") + ) + > 0 + ) + if EXPERT_EXECUTION == "dynamic_fp8": + assert ( + counts[("CUDAExecutionProvider", "QMoE")] + == config["text_config"]["num_hidden_layers"] + ), "Native FP8 experts did not execute on CUDA" + (ROOT / "profile-summary.json").write_text( + json.dumps(profile_summary, indent=2) + ) + event("cuda_profile_verified", counts=profile_summary) + check(f"state_{step}.hidden", decoded["hidden_states"], state["hidden_states"]) + check(f"state_{step}.logits", decoded["logits"], state["logits"]) + decoder_cache = next_cache(decoded) + if step == 0: + prompt = current_frames + prefilled = run( + "audio_decoder", + { + **tower_cache, + "hidden_states": decoded["hidden_states"][:, :-1], + "shared_input": embedded["shared_input"][:, :-1], + "codes": prompt[None, 1:, 1:], + "semantic_mask": ( + (prompt[None, 1:, 0] >= semantic_start) + & (prompt[None, 1:, 0] <= semantic_end) + ), + "position_ids": positions[:, :-1], + "attention_mask": mask[:, :-1], + }, + ) + check( + "prefill.tower", + prefilled["codebook_logits"], + state["prefill_tower_logits"], + ) + tower_cache = next_cache(prefilled) + main = int( + candidates[decoded["logits"][0, -1, candidates].astype(np.float32).argmax()] + ) + codes = np.zeros((1, 1, count), dtype=np.int64) + codes[0, 0, 0] = np.clip(main - semantic_start, 0, vocab - 1) + semantic = np.array([[semantic_start <= main <= semantic_end]]) + single_pass = audio_model_name == "audio_generator" + for index in range(1 if single_pass else count - 1): + tower_output = run( + "audio_decoder", + { + **tower_cache, + "hidden_states": decoded["hidden_states"][:, -1:], + "shared_input": embedded["shared_input"][:, -1:], + "codes": codes, + "semantic_mask": semantic, + "position_ids": positions[:, -1:], + "attention_mask": mask, + }, + generate=True, + ) + if single_pass: + codes = tower_output["frame_codes"] + else: + codes[0, 0, index + 1] = ( + tower_output["codebook_logits"][0, 0, index] + .astype(np.float32) + .argmax() + ) + check( + f"state_{step}.tower", tower_output["codebook_logits"], state["tower_logits"] + ) + current_frames = np.concatenate((np.array([[[main]]]), codes), axis=-1)[0] + np.testing.assert_array_equal(current_frames, state["sampled_frame"]) + tower_cache = next_cache(tower_output) + event("frame_matches", step=step, frame=current_frames.tolist()) + if main == eos: + break + generated.extend(current_frames.tolist()) + assert len(generated) == len(golden["generated_frames"]) + np.testing.assert_array_equal(generated, golden["generated_frames"]) + report = { + "status": "passed", + "revision": REVISION, + "source_revision": SOURCE_REVISION, + "dtype": DTYPE_NAMES[DTYPE], + "expert_execution": EXPERT_EXECUTION, + "reference_environment": golden.get("reference_environment"), + "audio_component": audio_model_name, + "onnxruntime_version": ort.__version__, + "provider": "CUDAExecutionProvider", + "ort_graph_optimization": optimization_name, + "atol": 1e-2 if DTYPE != "f32" else 1e-3, + "rtol": 1e-2 if DTYPE != "f32" else 1e-3, + "prompt_rows": golden["prompt_rows"], + "generated_frames": len(generated), + "attempted_frames": golden["attempted_frames"], + "seconds": time.monotonic() - start, + "metrics": metrics, + "profile_summary": profile_summary, + } + (ROOT / "validation-result.json").write_text(json.dumps(report, indent=2)) + event( + "REAL_CHECKPOINT_VALIDATION_PASSED", + **{key: report[key] for key in ("generated_frames", "attempted_frames", "seconds")}, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("stage", choices=("download", "export", "reference", "compare")) + parser.add_argument("--work-dir", required=True, type=pathlib.Path) + parser.add_argument("--dtype", choices=("f32", "f16", "bf16"), default="f32") + parser.add_argument("--native-fp8", action="store_true") + parser.add_argument( + "--task", choices=("fish-omni", "fish-omni-generation"), default="fish-omni" + ) + parser.add_argument("--execution-provider", choices=("default", "cuda", "onnx-standard")) + parser.add_argument("--audio-component", choices=("audio_decoder", "audio_generator")) + parser.add_argument( + "--optimization-level", + choices=tuple(ORT_OPTIMIZATION_LEVELS), + default="basic", + help="ORT comparison optimization level (default: basic)", + ) + parser.add_argument("--golden", type=pathlib.Path) + parser.add_argument("--reference-devices", default="0,1,2,3") + parser.add_argument("--reference-tower-device", default=4, type=int) + parser.add_argument("--embedding-device", default=0, type=int) + parser.add_argument("--audio-device", default=1, type=int) + parser.add_argument("--decoder-device", default=2, type=int) + args = parser.parse_args() + ROOT = args.work_dir.resolve() + ROOT.mkdir(parents=True, exist_ok=True) + REFERENCE_DEVICES = tuple(int(device) for device in args.reference_devices.split(",")) + REFERENCE_TOWER_DEVICE = args.reference_tower_device + EMBEDDING_DEVICE = args.embedding_device + AUDIO_DEVICE = args.audio_device + DECODER_DEVICE = args.decoder_device + DTYPE = args.dtype + EXPERT_EXECUTION = "dynamic_fp8" if args.native_fp8 else "dense" + if args.golden is not None: + GOLDEN = args.golden.resolve() + elif DTYPE != "f32" or args.native_fp8: + suffix = f"fp8_{DTYPE}" if args.native_fp8 else DTYPE + GOLDEN = GOLDEN.with_name(f"drama-3_{suffix}_tokens_generation.json") + if args.stage == "compare" and not GOLDEN.is_file(): + parser.error( + "Native-FP8/reduced-precision comparison requires --golden pointing to the independent " + "reference stage's golden.json until a matching golden is committed" + ) + if args.stage == "download": + download() + elif args.stage == "export": + export( + checkpoint_path(), + task=args.task, + execution_provider=args.execution_provider, + native_fp8=args.native_fp8, + ) + elif args.stage == "reference": + reference(checkpoint_path(), native_fp8=args.native_fp8) + else: + compare( + audio_model_name=args.audio_component, + optimization_level=args.optimization_level, + ) diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py new file mode 100644 index 000000000..a3220c169 --- /dev/null +++ b/scripts/validate_drama3_test.py @@ -0,0 +1,459 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Network-free checks for the full-checkpoint validation driver.""" + +from __future__ import annotations + +import copy +import json +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock + +import ml_dtypes +import numpy as np +import pytest +import validate_drama3 +from validate_drama3 import ( + GOLDEN, + REVISION, + SOURCE_REVISION, + empty_cache_ports, + event, + validate_reference_golden, +) + + +def test_logging_accepts_a_metric_name(capsys): + event("comparison", name="decoder.logits") + assert json.loads(capsys.readouterr().out) == { + "event": "comparison", + "name": "decoder.logits", + } + + +@pytest.mark.parametrize( + "level, expected", + [ + (None, "ORT_ENABLE_BASIC"), + ("disabled", "ORT_DISABLE_ALL"), + ("basic", "ORT_ENABLE_BASIC"), + ("extended", "ORT_ENABLE_EXTENDED"), + ("all", "ORT_ENABLE_ALL"), + ], +) +def test_comparison_sets_every_session_and_reports_optimization_level( + tmp_path, monkeypatch, capsys, level, expected +): + import onnxruntime as ort + import onnxruntime_easy as ort_easy + + (tmp_path / "reference").mkdir() + (tmp_path / "reference" / "golden.json").write_text( + json.dumps({"attempted_frames": 0, "generated_frames": [], "prompt_rows": 0}) + ) + (tmp_path / "config.json").write_text( + json.dumps( + { + "audio_decoder_config": {"num_codebooks": 1, "vocab_size": 2}, + "semantic_start_token_id": 4, + "semantic_end_token_id": 5, + "eos_token_id": 3, + } + ) + ) + monkeypatch.setattr(validate_drama3, "ROOT", tmp_path, raising=False) + monkeypatch.setattr(validate_drama3, "checkpoint_path", lambda: tmp_path) + monkeypatch.setattr(validate_drama3, "validate_reference_golden", lambda golden: None) + monkeypatch.setattr(ort, "preload_dlls", lambda: None) + options = [] + + def session(path, *, sess_options, providers): + options.append(sess_options.graph_optimization_level) + return SimpleNamespace( + get_inputs=list, + get_outputs=lambda: [SimpleNamespace(type="tensor(float)")], + get_modelmeta=lambda: SimpleNamespace(custom_metadata_map={}), + ) + + monkeypatch.setattr(ort_easy, "EasySession", session) + kwargs = {} if level is None else {"optimization_level": level} + validate_drama3.compare(**kwargs) + assert options == [getattr(ort.GraphOptimizationLevel, expected)] * 3 + report = json.loads((tmp_path / "validation-result.json").read_text()) + assert report["ort_graph_optimization"] == expected + starts = [ + json.loads(line) + for line in capsys.readouterr().out.splitlines() + if json.loads(line)["event"] == "ort_session_start" + ] + assert [entry["component"] for entry in starts] == [ + "embedding", + "audio_decoder", + "decoder", + ] + assert all(entry["ort_graph_optimization"] == expected for entry in starts) + + +def test_comparison_rejects_unknown_optimization_level_before_loading(monkeypatch): + monkeypatch.delattr(validate_drama3, "ROOT", raising=False) + with pytest.raises(ValueError, match="Unknown ORT optimization level"): + validate_drama3.compare(optimization_level="invalid") + + +def test_empty_cache_uses_the_abi_not_symbolic_dimension_names(): + ports = [ + SimpleNamespace( + name="past_key_values.3.key", + shape=["component.decoder.batch", 2, "component.decoder.past_sequence_len", 256], + ), + SimpleNamespace( + name="past_key_values.0.conv_state", + shape=["component.decoder.batch", 8192, 3], + ), + SimpleNamespace( + name="past_key_values.0.recurrent_state", + shape=["component.decoder.batch", 32, 128, 128], + ), + ] + cache = empty_cache_ports(ports) + assert cache["past_key_values.3.key"].shape == (1, 2, 0, 256) + assert cache["past_key_values.0.conv_state"].shape == (1, 8192, 3) + assert cache["past_key_values.0.recurrent_state"].shape == (1, 32, 128, 128) + + +def test_unknown_nonbatch_cache_dimension_is_an_error(): + with pytest.raises(AssertionError): + empty_cache_ports( + [ + SimpleNamespace( + name="past_key_values.0.conv_state", shape=["batch", "unknown", 3] + ) + ] + ) + + +def test_empty_cache_preserves_bfloat16_and_float32_state_types(): + cache = empty_cache_ports( + [ + SimpleNamespace( + name="past_key_values.0.key", + shape=["batch", 2, "past", 8], + type="tensor(bfloat16)", + ), + SimpleNamespace( + name="past_key_values.0.recurrent_state", + shape=["batch", 2, 8, 8], + type="tensor(float)", + ), + ] + ) + assert cache["past_key_values.0.key"].dtype == ml_dtypes.bfloat16 + assert cache["past_key_values.0.recurrent_state"].dtype == np.float32 + + +def test_unknown_cache_dtype_is_an_error(): + with pytest.raises(ValueError, match="Unsupported cache dtype"): + empty_cache_ports( + [ + SimpleNamespace( + name="past_key_values.0.conv_state", + shape=["batch", 8, 3], + type="tensor(double)", + ) + ] + ) + + +@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) +def test_runtime_bridge_preserves_types_and_output_names(tmp_path, dtype): + import onnx_ir as ir + import onnxruntime as ort + import onnxruntime_easy as ort_easy + from onnxscript import GraphBuilder + + element_type = ir.DataType.FLOAT if dtype == "float32" else ir.DataType.BFLOAT16 + array_type = np.float32 if dtype == "float32" else ml_dtypes.bfloat16 + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + x = builder.input("x", element_type, [2]) + result = builder.op.Identity(x) + result.type, result.shape = x.type, x.shape + builder.add_output(result, "result") + path = tmp_path / "bridge.onnx" + ir.save(ir.Model(graph, ir_version=12), path) + session = ort_easy.EasySession( + str(path), sess_options=ort.SessionOptions(), providers=["CPUExecutionProvider"] + ) + values = np.array([0.75, -1.5], dtype=array_type) + output = validate_drama3.run_session(session, {"x": values}) + assert set(output) == {"result"} + assert output["result"].dtype == values.dtype + np.testing.assert_array_equal(output["result"], values) + + +def test_runtime_bridge_selects_the_session_cuda_device(monkeypatch): + import torch + + active = [] + events = [] + + @contextmanager + def device(index): + events.append(("enter", index)) + active.append(index) + try: + yield + finally: + assert active.pop() == index + events.append(("exit", index)) + + class Session: + def get_provider_options(self): + return {"CUDAExecutionProvider": {"device_id": "3"}} + + def __call__(self, **feeds): + assert active == [3] + return [feeds["x"]] + + def get_outputs(self): + return [SimpleNamespace(name="result")] + + monkeypatch.setattr(torch.cuda, "device", device) + values = np.array([0.75, -1.5], dtype=np.float32) + result = validate_drama3.run_session(Session(), {"x": values}) + np.testing.assert_array_equal(result["result"], values) + assert events == [("enter", 3), ("exit", 3)] + + +@pytest.mark.parametrize("generate", [False, True]) +def test_tower_bridge_selects_teacher_or_single_pass_without_rng(monkeypatch, generate): + session = SimpleNamespace( + get_inputs=lambda: [SimpleNamespace(name="generate")], + ) + output = {"valid": np.array([True]), "next_counter": np.zeros(1, np.int64)} + hidden = np.zeros((1, 1, 8), np.float32) + + def run(actual, feeds): + assert actual is session + assert feeds["hidden_states"] is hidden + assert feeds["generate"].shape == () + assert bool(feeds["generate"]) is generate + assert not bool(feeds["do_sample"]) + np.testing.assert_array_equal(feeds["counter"], [0]) + return output + + monkeypatch.setattr(validate_drama3, "run_session", run) + assert ( + validate_drama3.run_tower_session( + session, {"hidden_states": hidden}, generate=generate + ) + is output + ) + + +@pytest.mark.parametrize("invalid", ["valid", "counter"]) +def test_tower_bridge_rejects_invalid_generation_or_greedy_rng_advance(monkeypatch, invalid): + session = SimpleNamespace(get_inputs=lambda: [SimpleNamespace(name="generate")]) + monkeypatch.setattr( + validate_drama3, + "run_session", + lambda *args: { + "valid": np.array([invalid != "valid"]), + "next_counter": np.array([int(invalid == "counter")], np.int64), + }, + ) + with pytest.raises( + FloatingPointError if invalid == "valid" else AssertionError, + match=r"rejected|advanced", + ): + validate_drama3.run_tower_session(session, {}, generate=True) + + +def test_tower_bridge_preserves_legacy_teacher_feeds(monkeypatch): + session = SimpleNamespace(get_inputs=lambda: [SimpleNamespace(name="hidden_states")]) + feeds = {"hidden_states": np.zeros((1, 3, 8), np.float32)} + output = {"codebook_logits": np.zeros((1, 3, 2, 16), np.float32)} + + def run(actual, actual_feeds): + assert actual is session and actual_feeds is feeds + return output + + monkeypatch.setattr(validate_drama3, "run_session", run) + assert validate_drama3.run_tower_session(session, feeds) is output + + +def test_full_export_forwards_explicit_task_and_execution_provider(tmp_path, monkeypatch): + import mobius + + calls = {} + + class Package: + def __init__(self): + self.weight_loading_report = {} + + def save(self, path, *, max_workers): + calls["save"] = (path, max_workers) + + def build(model, **kwargs): + calls["build"] = (model, kwargs) + return Package() + + monkeypatch.setattr(mobius, "build", build) + monkeypatch.setattr(validate_drama3, "ROOT", tmp_path, raising=False) + monkeypatch.setattr(validate_drama3, "DTYPE", "f32") + validate_drama3.export( + tmp_path / "checkpoint", task="fish-omni-generation", execution_provider="cuda" + ) + model, options = calls["build"] + assert model == str(tmp_path / "checkpoint") + assert options["revision"] == REVISION + assert options["task"] == "fish-omni-generation" + assert options["execution_provider"] == "cuda" + assert options["keep_quantized"] is False + assert calls["save"] == (str(tmp_path / "onnx"), 1) + + +def test_committed_token_golden_is_pinned_and_complete(): + golden = json.loads(GOLDEN.read_text()) + validate_reference_golden(golden) + assert golden["revision"] == REVISION + assert golden["source_revision"] == SOURCE_REVISION + assert golden["expert_execution"] == "dense" + assert golden["reference_environment"] == { + "torch": "2.8.0+cu126", + "torch_cuda": "12.6", + "triton": "3.4.0", + } + assert golden["max_new_tokens"] == len(golden["generated_frames"]) == 20 + for frame in golden["generated_frames"]: + assert len(frame) == 11 + assert 248099 <= frame[0] <= 252194 + assert frame[1] == frame[0] - 248099 + assert all(0 <= code < 4096 for code in frame[1:]) + + +@pytest.mark.parametrize("change", ["revision", "length", "frame"]) +def test_reference_mismatch_fails_before_model_loading(change): + golden = json.loads(GOLDEN.read_text()) + if change == "revision": + golden["revision"] = "not-the-pinned-checkpoint" + elif change == "length": + golden["generated_frames"].pop() + else: + golden["generated_frames"][0][2] += 1 + with pytest.raises(AssertionError): + validate_reference_golden(golden) + + +def test_native_fp8_reference_cannot_be_compared_as_dense(): + golden = json.loads(GOLDEN.read_text()) + golden["expert_execution"] = "dynamic_fp8" + with pytest.raises(AssertionError, match="expert_execution"): + validate_reference_golden(golden) + + +@pytest.mark.parametrize("version", ["2.8.0", "2.8.0+cu126", "2.8.0+cu128"]) +def test_reference_accepts_the_pinned_torch_release(version): + validate_drama3.validate_reference_torch_version(version) + + +@pytest.mark.parametrize("version", ["2.7.1", "2.8.0a0", "2.10.0+cu128"]) +def test_reference_rejects_a_different_torch_release(version): + with pytest.raises(RuntimeError, match=r"requires torch==2\.8\.0"): + validate_drama3.validate_reference_torch_version(version) + + +@pytest.mark.parametrize("value", [np.nan, np.inf, -np.inf]) +def test_reference_rejects_nonfinite_outputs(value, monkeypatch): + tensor = Mock() + tensor.detach.return_value.float.return_value.cpu.return_value.numpy.return_value = ( + np.array([value], np.float32) + ) + monkeypatch.setattr(validate_drama3, "DTYPE", "f16") + with pytest.raises(FloatingPointError, match=r"nonfinite decoder\.logits in f16"): + validate_drama3.reference_array(tensor, name="decoder.logits") + + +@pytest.mark.parametrize("value", [[], [-1.0, 0.0, 1.0]]) +def test_reference_preserves_finite_outputs(value): + expected = np.array(value, np.float32) + tensor = Mock() + tensor.detach.return_value.float.return_value.cpu.return_value.numpy.return_value = ( + expected + ) + assert validate_drama3.reference_array(tensor, name="decoder.logits") is expected + + +def test_reference_environment_changes_are_not_silently_compared(tmp_path, monkeypatch): + committed = json.loads(GOLDEN.read_text()) + committed["reference_environment"] = {"torch": "2.8.0+cu126"} + path = tmp_path / "golden.json" + path.write_text(json.dumps(committed)) + monkeypatch.setattr(validate_drama3, "GOLDEN", path) + actual = copy.deepcopy(committed) + actual["reference_environment"]["torch"] = "2.10.0+cu128" + with pytest.raises(AssertionError, match="reference_environment"): + validate_reference_golden(actual) + + +@pytest.mark.parametrize( + ("record_committed", "record_actual"), [(True, True), (True, False), (False, True)] +) +def test_reference_environment_preserves_matching_and_legacy_goldens( + tmp_path, monkeypatch, record_committed, record_actual +): + committed = json.loads(GOLDEN.read_text()) + committed.pop("reference_environment", None) + actual = copy.deepcopy(committed) + environment = {"torch": "2.8.0+cu126", "torch_cuda": "12.6", "triton": "3.4.0"} + if record_committed: + committed["reference_environment"] = environment + if record_actual: + actual["reference_environment"] = environment + path = tmp_path / "golden.json" + path.write_text(json.dumps(committed)) + monkeypatch.setattr(validate_drama3, "GOLDEN", path) + validate_reference_golden(actual) + + +@pytest.mark.parametrize("native_fp8", [False, True]) +def test_export_preserves_the_requested_expert_execution(tmp_path, monkeypatch, native_fp8): + import mobius + + calls = {} + package = SimpleNamespace( + weight_loading_report={}, + save=lambda *args, **kwargs: None, + ) + + def build(*args, **kwargs): + calls.update(kwargs) + return package + + monkeypatch.setattr(mobius, "build", build) + monkeypatch.setattr(validate_drama3, "ROOT", tmp_path, raising=False) + monkeypatch.setattr(validate_drama3, "DTYPE", "f32") + validate_drama3.export(tmp_path, native_fp8=native_fp8) + assert calls["keep_quantized"] is native_fp8 + assert calls["execution_provider"] == ("cuda" if native_fp8 else "default") + report = json.loads((tmp_path / "export-result.json").read_text()) + assert report["expert_execution"] == ("dynamic_fp8" if native_fp8 else "dense") + + +def test_stale_checkpoint_path_is_not_relabelled_as_the_pinned_revision(tmp_path, monkeypatch): + wrong = tmp_path / "other-revision" + wrong.mkdir() + (tmp_path / "checkpoint-path.txt").write_text(str(wrong)) + monkeypatch.setattr(validate_drama3, "ROOT", tmp_path, raising=False) + with pytest.raises(ValueError, match="immutable checkpoint snapshot"): + validate_drama3.checkpoint_path() + + +def test_pinned_checkpoint_path_accepts_trailing_newline(tmp_path, monkeypatch): + snapshot = tmp_path / REVISION + snapshot.mkdir() + (tmp_path / "checkpoint-path.txt").write_text(f"{snapshot}\n") + monkeypatch.setattr(validate_drama3, "ROOT", tmp_path, raising=False) + assert validate_drama3.checkpoint_path() == snapshot diff --git a/src/mobius/_configs/fish_codec.py b/src/mobius/_configs/fish_codec.py new file mode 100644 index 000000000..e8cc13149 --- /dev/null +++ b/src/mobius/_configs/fish_codec.py @@ -0,0 +1,79 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pinned Drama 3 waveform codec configuration, independent of the token tower.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +DRAMA3_CODEC_REPO = "fishaudio/drama-3" +DRAMA3_CODEC_REVISION = "b2b5ee998034b2975720906c815e7fb0e9f007bd" +FISH_CODEC_SOURCE_REVISION = "541566932abbaebcb0a40e562a1f846738383f38" + + +@dataclass(frozen=True) +class FishCodecConfig: + """Decode-only configuration of upstream ``modded_dac_vq.yaml``. + + The encoder and pre-quantization transformer are deliberately excluded. + Upstream constructs, but never calls, the configured decoder transformer. + ``max_frames`` bounds the precomputed, BF16-rounded RoPE table. + """ + + sample_rate: int = 44100 + latent_dim: int = 1024 + residual_codebooks: int = 9 + semantic_codebook_size: int = 4096 + residual_codebook_size: int = 1024 + codebook_dim: int = 8 + transformer_layers: int = 8 + transformer_heads: int = 16 + transformer_intermediate_size: int = 3072 + transformer_window: int = 128 + norm_eps: float = 1e-5 + rope_base: float = 10000.0 + max_frames: int = 327680 + quantizer_rates: tuple[int, ...] = (2, 2) + decoder_dim: int = 1536 + decoder_rates: tuple[int, ...] = (8, 8, 4, 2) + + def __post_init__(self): + positive = ( + self.sample_rate, + self.latent_dim, + self.residual_codebooks, + self.semantic_codebook_size, + self.residual_codebook_size, + self.codebook_dim, + self.transformer_heads, + self.transformer_intermediate_size, + self.transformer_window, + self.max_frames, + self.decoder_dim, + *self.quantizer_rates, + *self.decoder_rates, + ) + if any(value <= 0 for value in positive): + raise ValueError("Codec dimensions, rates and window must be positive") + if self.transformer_layers < 0: + raise ValueError("transformer_layers must be nonnegative") + if self.latent_dim % self.transformer_heads or self.head_dim % 2: + raise ValueError("latent_dim must divide into heads with even head_dim") + if not self.decoder_rates or self.decoder_dim % (2 ** len(self.decoder_rates)): + raise ValueError("decoder_dim must be divisible by 2**len(decoder_rates)") + if self.norm_eps <= 0 or self.rope_base <= 0: + raise ValueError("norm_eps and rope_base must be positive") + + @property + def head_dim(self) -> int: + return self.latent_dim // self.transformer_heads + + @property + def num_codebooks(self) -> int: + return self.residual_codebooks + 1 + + @property + def samples_per_frame(self) -> int: + return math.prod(self.quantizer_rates) * math.prod(self.decoder_rates) diff --git a/src/mobius/_configs/fish_omni.py b/src/mobius/_configs/fish_omni.py new file mode 100644 index 000000000..348a79d3a --- /dev/null +++ b/src/mobius/_configs/fish_omni.py @@ -0,0 +1,223 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Configuration for Drama 3's hybrid backbone and temporal codebook tower.""" + +from __future__ import annotations + +import copy +import dataclasses +from collections.abc import Mapping +from typing import Any + +import onnx_ir as ir + +from mobius._configs import ArchitectureConfig, _as_attribute_config + +DRAMA3_REVISION = "b2b5ee998034b2975720906c815e7fb0e9f007bd" +DRAMA3_SOURCE_REVISION = "541566932abbaebcb0a40e562a1f846738383f38" + + +@dataclasses.dataclass +class FishOmniConfig(ArchitectureConfig): + """Drama 3 token generation; the waveform codec is exported separately.""" + + tower: ArchitectureConfig | None = None + num_codebooks: int = 10 + shared_enrichment_layers: int = 4 + vq_backbone_codebooks: int = 1 + semantic_start_token_id: int = 248099 + semantic_end_token_id: int = 252194 + fp8_experts: bool = False + native_fp8_experts: bool = False + + @classmethod + def from_transformers(cls, config, parent_config=None) -> FishOmniConfig: + # HF config classes and raw JSON expose architecture-specific dynamic fields. + parent: Any = _as_attribute_config(parent_config or config) + text: Any = _as_attribute_config(getattr(parent, "text_config", config)) + audio: Any = _as_attribute_config(parent.audio_decoder_config) + if parent.model_type != "fish_omni": + raise ValueError("FishOmniConfig requires model_type='fish_omni'") + if text.model_type not in {"fish_omni_text", "qwen3_5_moe_text"}: + raise ValueError("Unsupported Drama 3 text_config.model_type") + if getattr(parent, "audio_encoder_config", None) is not None: + raise ValueError("Drama 3 audio_encoder_config is not supported") + if getattr(parent, "speech_template_version", None) != "reference-bracket-v1": + raise ValueError("Drama 3 requires speech_template_version='reference-bracket-v1'") + text_flags = { + "rotary_interleaved": False, + "use_qk_norm": True, + "attn_output_gate": True, + "layernorm_zero_centered_gamma": True, + "shared_expert_gate": True, + "tie_word_embeddings": False, + "attention_bias": False, + "hidden_act": "silu", + } + audio_flags = { + "decoder_type": "codebook_mtp_shared", + "rotary_interleaved": True, + "attention_output_gate": True, + "layernorm_zero_centered_gamma": True, + "use_moe": False, + "mtp_use_swa": False, + "attention_qkv_bias": False, + "attention_o_bias": False, + "dropout": 0.0, + } + for source, flags in ((text, text_flags), (audio, audio_flags)): + for name, required in flags.items(): + if getattr(source, name, required) != required: + raise ValueError(f"Unsupported Drama 3 {name}; expected {required!r}") + if any(kind != "moe" for kind in getattr(text, "mlp_layer_types", None) or []): + raise ValueError("Drama 3 requires MoE in every backbone layer") + text = copy.copy(text) + rope = text.rope_parameters + if not isinstance(rope, Mapping): + rope = vars(rope) + text.rope_parameters = dict(rope) + if rope.get("rope_type", "default") != "default": + raise ValueError("Drama 3 requires default RoPE") + base = ArchitectureConfig.from_transformers(text) + if len(base.layer_types or []) != base.num_hidden_layers or not set( + base.layer_types or [] + ) <= {"linear_attention", "full_attention"}: + raise ValueError("Drama 3 requires one valid attention type per backbone layer") + if ( + base.num_experts_per_tok is None + or base.num_local_experts is None + or not 0 < base.num_experts_per_tok <= base.num_local_experts + ): + raise ValueError("Invalid Drama 3 expert counts") + dim = audio.dim + n_shared = audio.shared_enrichment_layers + n_codebooks = audio.num_codebooks + n_backbone = audio.vq_backbone_codebooks + vocab = audio.vocab_size + if dim != base.hidden_size or getattr(audio, "text_dim", dim) != dim: + raise ValueError("Drama 3 tower and backbone hidden dimensions must match") + if audio.n_layer != n_shared or n_shared < 1: + raise ValueError("Drama 3 n_layer must equal positive shared_enrichment_layers") + if not 0 <= n_backbone < n_codebooks or n_codebooks < 2: + raise ValueError("Invalid Drama 3 codebook counts") + placement = getattr(audio, "vq_feature_placement", None) + if placement not in {"backbone_partial", "mtp_shared_pre"} or ( + placement == "mtp_shared_pre" and n_backbone != 0 + ): + raise ValueError("Unsupported Drama 3 vq_feature_placement") + start = parent.semantic_start_token_id + end = parent.semantic_end_token_id + eos = getattr(parent, "eos_token_id", None) + if ( + not 0 <= start <= end < base.vocab_size + or vocab != end - start + 1 + or eos is None + or not 0 <= eos < base.vocab_size + or start <= eos <= end + ): + raise ValueError("Invalid Drama 3 semantic band or EOS token") + quant = getattr(parent, "quantization_config", None) + if quant is not None and not isinstance(quant, Mapping): + quant = vars(quant) + if quant is not None and ( + quant.get("quant_method") != "bsgemm-moe" + or quant.get("weight_block_size") != [128, 128] + or quant.get("activation_scheme") != "dynamic" + ): + raise ValueError("Drama 3 expects dynamic bsgemm-moe FP8 with 128x128 blocks") + tower = ArchitectureConfig( + hidden_size=dim, + intermediate_size=audio.intermediate_size, + num_attention_heads=audio.n_head, + num_key_value_heads=audio.n_local_heads, + head_dim=audio.head_dim, + num_hidden_layers=n_shared + n_codebooks - 1, + # The tower attends over time, not over the 11 codebook positions + # named by the stale max_seq_len field in the released config. + max_position_embeddings=base.max_position_embeddings, + vocab_size=vocab, + rms_norm_eps=audio.norm_eps, + rope_type="default", + rope_theta=audio.rope_base, + partial_rotary_factor=audio.rotary_percent, + rope_interleave=True, + hidden_act="silu", + ) + for heads, kv_heads, head_dim, fraction in ( + ( + base.num_attention_heads, + base.num_key_value_heads, + base.head_dim, + base.partial_rotary_factor, + ), + ( + tower.num_attention_heads, + tower.num_key_value_heads, + tower.head_dim, + tower.partial_rotary_factor, + ), + ): + if fraction is None: + raise ValueError("Drama 3 requires a partial rotary factor") + width = int(head_dim * fraction) + if kv_heads < 1 or heads % kv_heads or not 0 < width <= head_dim or width % 2: + raise ValueError("Invalid Drama 3 attention or rotary geometry") + fields = {field.name: getattr(base, field.name) for field in dataclasses.fields(base)} + fields.update( + model_type="fish_omni", + hidden_act="silu", + attn_qk_norm=True, + attn_qkv_bias=False, + attn_o_bias=False, + mamba_ssm_dtype=ir.DataType.FLOAT, + norm_topk_prob=getattr(text, "norm_topk_prob", True), + eos_token_id=eos, + # bsgemm-moe is block-FP8, not the generic affine INT4 format. + quantization=None, + component_quantization={}, + ) + return cls( + **fields, + tower=tower, + num_codebooks=n_codebooks, + shared_enrichment_layers=n_shared, + vq_backbone_codebooks=n_backbone, + semantic_start_token_id=start, + semantic_end_token_id=end, + fp8_experts=quant is not None, + ) + + def validate(self) -> None: + """Reject unsafe precision instead of overflowing the released checkpoint.""" + super().validate() + if self.quantization is not None or self.component_quantization: + raise ValueError( + "Drama 3 affine component quantization is unsupported; " + "use a floating graph with explicit FP8 checkpoint dequantization." + ) + if self.dtype not in (ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16): + raise ValueError("Drama 3 requires floating FP32, FP16 or BF16 compute") + if self.native_fp8_experts: + if not self.fp8_experts or self.dtype not in ( + ir.DataType.FLOAT, + ir.DataType.FLOAT16, + ir.DataType.BFLOAT16, + ): + raise ValueError( + "Drama 3 native FP8 requires FP8 source weights and floating compute; " + "use keep_quantized=False (--dequantize) for FP32." + ) + if ( + self.hidden_size % 128 + or not self.moe_intermediate_size + or self.moe_intermediate_size % 128 + or not self.norm_topk_prob + ): + raise ValueError( + "Drama 3 native FP8 requires 128-aligned experts and normalized routing; " + "use --dequantize for the dense FP32 route." + ) + if self.tower is None: + raise ValueError("Drama 3 requires an audio tower configuration") + self.tower.validate() diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 930991813..49884080b 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -54,6 +54,7 @@ WhisperConfig, XverseConfig, ) +from mobius._configs.fish_omni import DRAMA3_REVISION from mobius.models import ( ApertusCausalLMModel, ArceeCausalLMModel, @@ -187,6 +188,7 @@ MPTCausalLMModel, ) from mobius.models.falcon_h1 import FalconH1ForCausalLM +from mobius.models.fish_omni import FishOmniModel from mobius.models.fun_asr import FunASRForConditionalGeneration from mobius.models.gemma3n import Gemma3nCausalLMModel, Gemma3nMultiModalModel from mobius.models.gguf_embeddings import GemmaEmbeddingGGUFModel, LlamaEmbedGGUFModel @@ -914,6 +916,13 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: ), "sensevoice_small": ModelRegistration(SenseVoiceSmallModel, task="audio-ctc"), "qwen3_tts": ModelRegistration(Qwen3TTSForConditionalGeneration), + "fish_omni": ModelRegistration( + FishOmniModel, + test_model_id="fishaudio/drama-3", + test_revision=DRAMA3_REVISION, + family="fish", + variant="temporal-codebook-mtp", + ), "qwen3_tts_tokenizer_12hz": ModelRegistration(Qwen3TTSTokenizerV2Model, task="codec"), "vibevoice": ModelRegistration( VibeVoiceForConditionalGeneration, diff --git a/src/mobius/_testing/ort_inference.py b/src/mobius/_testing/ort_inference.py index 439e24f30..1950384d1 100644 --- a/src/mobius/_testing/ort_inference.py +++ b/src/mobius/_testing/ort_inference.py @@ -78,7 +78,9 @@ def _numpy_to_ort_value(value: np.ndarray) -> ort.OrtValue: if value.dtype == np.bool_: # ORT supports tensor(bool) natively but DLPack does not, so # bypass ort_easy's DLPack-first path. - return ort.OrtValue.ortvalue_from_numpy(np.ascontiguousarray(value)) + return ort.OrtValue.ortvalue_from_numpy( + np.ascontiguousarray(value) if value.ndim else value + ) try: import ml_dtypes except ImportError: # pragma: no cover diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index d9cee1d06..d3c43a191 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -121,6 +121,8 @@ "SequenceMambaBlock", "SequenceSelectiveScan", "SiLU", + "Float32SiLU", + "float32_silu", "Siglip2NaFlexVisionEmbeddings", "Siglip2NaFlexVisionModel", "SigmoidTopKGate", @@ -165,7 +167,7 @@ "siglip2_naflex_attention_mask", ] -from mobius.components._activations import SiLU, get_activation +from mobius.components._activations import Float32SiLU, SiLU, float32_silu, get_activation from mobius.components._attention import ( Attention, FusedQKVAttention, diff --git a/src/mobius/components/_activations.py b/src/mobius/components/_activations.py index 2bb8ad8a7..4a2db66fc 100644 --- a/src/mobius/components/_activations.py +++ b/src/mobius/components/_activations.py @@ -4,13 +4,11 @@ from __future__ import annotations from collections import OrderedDict -from typing import TYPE_CHECKING +from collections.abc import Callable +import onnx_ir as ir from onnxscript import OpBuilder, nn -if TYPE_CHECKING: - import onnx_ir as ir - def silu(op: OpBuilder, x): """SiLU (Swish) activation.""" @@ -66,7 +64,7 @@ def linear(op: OpBuilder, x): # Mapping from activation string names to functions. # Each function takes (op: OpBuilder, x) and returns the activated tensor. -ACT2FN: OrderedDict[str, callable] = OrderedDict( +ACT2FN: OrderedDict[str, Callable[[OpBuilder, ir.Value], ir.Value]] = OrderedDict( { "gegelu": gelu, "gelu": gelu, @@ -120,3 +118,16 @@ class SiLU(nn.Module): def forward(self, op: OpBuilder, x: ir.Value): return silu(op, x) + + +class Float32SiLU(SiLU): + """SiLU with FP32 opmath and a single cast to the activation dtype.""" + + def forward(self, op: OpBuilder, x: ir.Value): + return op.CastLike(float32_silu(op, x), x) + + +def float32_silu(op: OpBuilder, x: ir.Value): + """Torch CUDA SiLU opmath: FP32 division before any output narrowing.""" + wide = op.Cast(x, to=ir.DataType.FLOAT) + return op.Div(wide, op.Add(1.0, op.Exp(op.Neg(wide)))) diff --git a/src/mobius/components/_activations_test.py b/src/mobius/components/_activations_test.py index cfe2e39be..8b164a967 100644 --- a/src/mobius/components/_activations_test.py +++ b/src/mobius/components/_activations_test.py @@ -5,9 +5,16 @@ from __future__ import annotations +import ml_dtypes +import numpy as np +import onnx_ir as ir import pytest +import torch +from onnxscript import GraphBuilder from mobius._testing import count_op_type, create_test_builder, create_test_input +from mobius._testing.ort_inference import OnnxModelSession +from mobius.components import Float32SiLU from mobius.components._activations import ( ACT2FN, gelu, @@ -21,6 +28,32 @@ class TestActivations: + @pytest.mark.parametrize( + ("dtype", "array_dtype", "torch_dtype"), + [ + (ir.DataType.FLOAT16, np.float16, torch.float16), + (ir.DataType.BFLOAT16, ml_dtypes.bfloat16, torch.bfloat16), + ], + ) + def test_float32_silu_rounds_only_at_output(self, dtype, array_dtype, torch_dtype): + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + x = builder.input("x", dtype, [10]) + output = Float32SiLU()(builder.op, x) + output.type, output.shape = x.type, x.shape + builder.add_output(output, "output") + division = next(node for node in graph if node.op_type == "Div") + assert division.inputs[0].producer().attributes["to"].as_int() == ir.DataType.FLOAT + values = np.array([-1000, -8, -3, -0.75, -0.1, 0, 0.1, 0.75, 3, 8], np.float32) + inputs = values.astype(array_dtype) + reference = torch.nn.functional.silu(torch.tensor(values).to(torch_dtype)) + session = OnnxModelSession(ir.Model(graph, ir_version=12)) + try: + actual = session.run({"x": inputs})["output"] + finally: + session.close() + np.testing.assert_array_equal(actual.astype(np.float32), reference.float().numpy()) + def _run_activation(self, act_fn): builder, op, _graph = create_test_builder() x = create_test_input(builder, "x", [2, 3, 64]) diff --git a/src/mobius/components/_gated_deltanet.py b/src/mobius/components/_gated_deltanet.py index 55ea823c7..51ef741cc 100644 --- a/src/mobius/components/_gated_deltanet.py +++ b/src/mobius/components/_gated_deltanet.py @@ -227,6 +227,8 @@ def forward( # Decomposed form of op.LpNormalization(x, axis=-1, p=2). # TODO: Use op.LpNormalization directly once ORT >=1.25 supports it. q_4d = op.Reshape(query, qk_4d_shape) # (B, T, num_k_heads, head_k_dim) + if self._stash_type != self._dtype: + q_4d = op.Cast(q_4d, to=self._stash_type) q_norm_squared = op.ReduceSumSquare(q_4d, [-1], keepdims=1) if self._qk_l2norm_eps: q_norm_squared = op.Add( @@ -237,6 +239,8 @@ def forward( query = op.Reshape(op.Div(q_4d, q_l2), qk_3d_shape) # (B, T, key_dim) k_4d = op.Reshape(key, qk_4d_shape) # (B, T, num_k_heads, head_k_dim) + if self._stash_type != self._dtype: + k_4d = op.Cast(k_4d, to=self._stash_type) k_norm_squared = op.ReduceSumSquare(k_4d, [-1], keepdims=1) if self._qk_l2norm_eps: k_norm_squared = op.Add( @@ -248,7 +252,8 @@ def forward( # === Compute gating parameters === # beta: (B, T, num_v_heads) - beta = op.Sigmoid(b) + beta_input = op.Cast(b, to=self._stash_type) if self._stash_type != self._dtype else b + beta = op.Sigmoid(beta_input) # Compute decay: g = -exp(A_log) * softplus(a + dt_bias). # Mirrors HF: -A_log.float().exp() * softplus(a.float() + dt_bias). if self._dtype == ir.DataType.FLOAT16 or self._stash_type != self._dtype: diff --git a/src/mobius/components/_gated_deltanet_test.py b/src/mobius/components/_gated_deltanet_test.py index 93039de17..53503ecb3 100644 --- a/src/mobius/components/_gated_deltanet_test.py +++ b/src/mobius/components/_gated_deltanet_test.py @@ -45,6 +45,36 @@ def test_gated_deltanet_propagates_activation_and_state_dtypes(dtype, state_dtyp assert output.dtype == dtype +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_gated_deltanet_normalizes_and_gates_in_state_precision(dtype): + config = make_config( + dtype=dtype, + mamba_ssm_dtype=ir.DataType.FLOAT, + linear_num_value_heads=4, + linear_num_key_heads=2, + linear_key_head_dim=16, + linear_value_head_dim=16, + linear_conv_kernel_dim=4, + ) + module = GatedDeltaNet(config) + _cast_module_dtype(module, dtype) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 2, 64], dtype) + conv = create_test_input(builder, "conv_state", [1, 128, 3], dtype) + state = create_test_input(builder, "state", [1, 4, 16, 16], ir.DataType.FLOAT) + module(op, hidden, conv, state) + reductions = [node for node in graph if node.op_type == "ReduceSumSquare"] + assert len(reductions) == 2 + for node in reductions: + cast = node.inputs[0].producer() + assert cast.op_type == "Cast" + assert cast.attributes["to"].as_int() == ir.DataType.FLOAT + beta = next(node for node in graph if node.op_type == "Sigmoid") + cast = beta.inputs[0].producer() + assert cast.op_type == "Cast" + assert cast.attributes["to"].as_int() == ir.DataType.FLOAT + + class TestGatedDeltaNet: """Tests for the GatedDeltaNet linear attention module.""" diff --git a/src/mobius/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 4e286d40b..e0a779255 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -107,8 +107,9 @@ class BaseRope(nn.Module): dtype: Model compute dtype. When not FP32, a Cast is inserted after gathering cos/sin embeddings so that the ``RotaryEmbedding`` op receives inputs with matching types. - The cache itself stays FP32 for precision; only the gathered - per-position embeddings are cast. + Caches are initialized in FP32, but the exporter normally casts + parameters to the model dtype. Architectures needing FP32 rotation + explicitly preserve their cache parameters. """ def __init__( @@ -140,8 +141,8 @@ def _cast_embeddings( """Cast cos/sin embeddings to model dtype if needed. The RotaryEmbedding op requires all tensor inputs to share the - same type ``T``. Since caches are stored in FP32, a Cast is - needed when the model uses FP16 or BF16. + same type ``T``. An explicit Cast also handles architectures that + preserve FP32 cache parameters in an FP16 or BF16 model. """ if self._dtype != ir.DataType.FLOAT: cos = op.Cast(cos, to=self._dtype) diff --git a/src/mobius/functions/causal_conv.py b/src/mobius/functions/causal_conv.py index 80502d99c..cd5763181 100644 --- a/src/mobius/functions/causal_conv.py +++ b/src/mobius/functions/causal_conv.py @@ -16,6 +16,7 @@ from onnxscript._internal.builder import build_function from mobius._constants import OPSET_VERSION +from mobius.components import float32_silu DOMAIN = "com.microsoft" @@ -128,7 +129,10 @@ def body(op, input_val, weight_val, bias_val, conv_state_val): # Step 5: Apply activation. if activation in ("silu", "swish"): - output = op.Swish(conv_out) + output = op.CastLike( + float32_silu(op, conv_out), + conv_out, + ) elif activation == "none": output = conv_out else: diff --git a/src/mobius/functions/causal_conv_test.py b/src/mobius/functions/causal_conv_test.py index f7e3b5422..a070bde12 100644 --- a/src/mobius/functions/causal_conv_test.py +++ b/src/mobius/functions/causal_conv_test.py @@ -53,10 +53,12 @@ def test_activations(self, activation: str): ) op_types = [node.op_type for node in func.graph] if activation in ("silu", "swish"): - assert "Swish" in op_types + assert "Div" in op_types + assert "Exp" in op_types assert "Sigmoid" not in op_types else: - assert "Swish" not in op_types + assert "Div" not in op_types + assert "Exp" not in op_types assert "Sigmoid" not in op_types def test_invalid_ndim_raises(self): diff --git a/src/mobius/generation/_fish_omni.py b/src/mobius/generation/_fish_omni.py new file mode 100644 index 000000000..e6e12a62a --- /dev/null +++ b/src/mobius/generation/_fish_omni.py @@ -0,0 +1,270 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Drama 3's drop-crossing nucleus sampler and semantic-only repetition reroll. + +Policies use explicit counter RNG state, not Torch/Philox seed-bit identity. +They perform sampling in ONNX; the caller still orchestrates temporal decoding. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import onnx_ir as ir +from onnxscript import GraphBuilder, OpBuilder + +from mobius._configs.fish_omni import FishOmniConfig +from mobius.generation._policy_components import ( + _RNG_MODULUS, + PolicyComponent, + _component, + _counter_uniform, + _make_graph, + _set_public_shape, +) + + +class FishAudioSamplingInputs(NamedTuple): + """Tensor controls for the single-pass temporal tower task.""" + + generate: ir.Value + do_sample: ir.Value + temperature: ir.Value + top_k: ir.Value + top_p: ir.Value + seed: ir.Value + counter: ir.Value + + +def sample_codebook_or_greedy( + op: OpBuilder, + logits: ir.Value, + sampling: FishAudioSamplingInputs, + counter: ir.Value, + *, + index: int, + vocabulary: int, +) -> tuple[ir.Value, ir.Value, ir.Value]: + """Skip RNG/sorting entirely for greedy generation and prompt ingestion.""" + active = op.And(sampling.generate, sampling.do_sample) + + def stochastic(branch_op): + noises, _ = _noise(branch_op, logits, sampling.seed, counter, 1) + return _race( + branch_op, + logits, + sampling.temperature, + sampling.top_k, + sampling.top_p, + noises[0], + ) + + def greedy(branch_op): + float_logits = branch_op.Cast(logits, to=ir.DataType.FLOAT) + return ( + branch_op.ArgMax(float_logits, axis=-1, keepdims=False), + _all_rows(branch_op, _finite(branch_op, float_logits)), + ) + + branches = [] + for suffix, body in (("random", stochastic), ("greedy", greedy)): + name = f"codebook_{index}_{suffix}" + op.builder.push_module(name) + branch = op.builder.subgraph( + body, + inputs=[], + outputs=[ir.Value(name=f"{name}_token"), ir.Value(name=f"{name}_valid")], + name=name, + ) + op.builder.pop_module() + branches.append(branch) + token, valid = op.If(active, then_branch=branches[0], else_branch=branches[1], _outputs=2) + advanced = op.Mod( + op.Add(op.Mod(counter, _RNG_MODULUS, fmod=0), vocabulary), _RNG_MODULUS, fmod=0 + ) + return token, op.Where(active, advanced, counter), valid + + +def _finite(op: OpBuilder, value: ir.Value) -> ir.Value: + return op.Not(op.Or(op.IsNaN(value), op.IsInf(value))) + + +def _all_rows(op: OpBuilder, value: ir.Value) -> ir.Value: + return op.Equal( + op.ReduceMin(op.Cast(value, to=ir.DataType.INT64), [-1], keepdims=False), 1 + ) + + +def _race( + op: OpBuilder, + logits: ir.Value, + temperature: ir.Value, + top_k: ir.Value, + top_p: ir.Value, + noise: ir.Value, +) -> tuple[ir.Value, ir.Value]: + """Fish Speech's sorted-probability exponential race with explicit noise.""" + scaled = op.Div( + op.Cast(logits, to=ir.DataType.FLOAT), + op.Unsqueeze(op.Max(temperature, 1e-5), [-1]), + ) + width = op.Shape(scaled, start=1, end=2) + values, ids = op.TopK(scaled, width, axis=-1, sorted=1, _outputs=2) + ranks = op.Range(0, op.Squeeze(width, [0]), 1) + cumulative = op.CumSum(op.Softmax(values, axis=-1), -1) + dropped = op.And( + op.Greater(ranks, 0), + op.Or( + op.Greater(cumulative, op.Unsqueeze(top_p, [-1])), + op.GreaterOrEqual(ranks, op.Unsqueeze(top_k, [-1])), + ), + ) + probabilities = op.Softmax(op.Where(dropped, float("-inf"), values), axis=-1) + choice = op.ArgMax(op.Div(probabilities, noise), axis=-1, keepdims=False) + token = op.Squeeze(op.GatherElements(ids, op.Unsqueeze(choice, [-1]), axis=-1), [-1]) + valid = op.And( + _all_rows(op, _finite(op, scaled)), + op.And( + op.And(_finite(op, temperature), _finite(op, top_p)), + _all_rows(op, op.And(_finite(op, noise), op.Greater(noise, 0.0))), + ), + ) + return token, valid + + +def _noise( + op: OpBuilder, logits: ir.Value, seed: ir.Value, counter: ir.Value, draws: int +) -> tuple[list[ir.Value], ir.Value]: + width = op.Squeeze(op.Shape(logits, start=1, end=2), [0]) + ranks = op.Range(0, width, 1) + # Bound integer multiplication in the shared RNG, including arbitrary + # signed seeds/counters. Its counter stream is periodic modulo 2**31-1. + seed = op.Mod(seed, _RNG_MODULUS, fmod=0) + counter = op.Mod(counter, _RNG_MODULUS, fmod=0) + noises = [] + for draw in range(draws): + offsets = op.Mod( + op.Add(op.Unsqueeze(op.Add(counter, op.Mul(draw, width)), [-1]), ranks), + _RNG_MODULUS, + fmod=0, + ) + uniform = _counter_uniform(op, op.Unsqueeze(seed, [-1]), offsets) + # Float32 conversion can round a uniform to 1; keep -log(U) positive. + uniform = op.Clip(uniform, 1.1754943508222875e-38, 1.0 - 2**-24) + noises.append(op.Neg(op.Log(uniform))) + next_counter = op.Mod(op.Add(counter, op.Mul(draws, width)), _RNG_MODULUS, fmod=0) + return noises, next_counter + + +def _inputs( + builder: GraphBuilder, dtype: ir.DataType, width: int +) -> tuple[ir.Value, ir.Value, ir.Value, ir.Value, ir.Value, ir.Value]: + return ( + builder.input("logits", dtype, ["batch", width]), + builder.input("temperature", ir.DataType.FLOAT, ["batch"]), + builder.input("top_k", ir.DataType.INT64, ["batch"]), + builder.input("top_p", ir.DataType.FLOAT, ["batch"]), + builder.input("seed", ir.DataType.INT64, ["batch"]), + builder.input("counter", ir.DataType.INT64, ["batch"]), + ) + + +def _output(builder: GraphBuilder, value: ir.Value, name: str, dtype: ir.DataType) -> None: + value.type = ir.TensorType(dtype) + _set_public_shape(value, ["batch"]) + builder.add_output(value, name) + + +_CONTRACT = { + "sampling": "sorted-exponential-race", + "top_p": "drop-crossing-token-except-rank-zero", + "temperature_floor": 1e-5, + "rng": "mobius.counter-uniform-exponential@1", + "counter_modulus": _RNG_MODULUS, + "torch_seed_identity": False, + "invalid_input": "token=-1, valid=false, counter unchanged; caller must stop", +} + + +def build_fish_codebook_sampler(config: FishOmniConfig) -> PolicyComponent: + """Sample one residual code; invoke after each conditioned tower projection.""" + config.validate() + assert config.tower is not None + graph, builder = _make_graph("fish_codebook_sampler") + op = builder.op + logits, temperature, top_k, top_p, seed, counter = _inputs( + builder, config.dtype, config.tower.vocab_size + ) + noises, advanced = _noise(op, logits, seed, counter, 1) + token, valid = _race(op, logits, temperature, top_k, top_p, noises[0]) + _output(builder, op.Where(valid, token, -1), "token", ir.DataType.INT64) + _output(builder, valid, "valid", ir.DataType.BOOL) + _output(builder, op.Where(valid, advanced, counter), "next_counter", ir.DataType.INT64) + return _component("mobius.policy.fish-codebook-sampler@1", graph, dict(_CONTRACT)) + + +def build_fish_main_sampler(config: FishOmniConfig) -> PolicyComponent: + """Draw main/retry every step and reroll only repeated semantic main tokens.""" + config.validate() + assert config.tower is not None + assert config.eos_token_id is not None + graph, builder = _make_graph("fish_main_sampler") + op = builder.op + logits, temperature, top_k, top_p, seed, counter = _inputs( + builder, config.dtype, config.vocab_size + ) + history = builder.input("history", ir.DataType.INT64, ["batch", 16]) + candidates = op.Constant( + value_ints=sorted( + [ + config.eos_token_id, + *range(config.semantic_start_token_id, config.semantic_end_token_id + 1), + ] + ) + ) + selected_logits = op.Gather(logits, candidates, axis=-1) + noises, advanced = _noise(op, selected_logits, seed, counter, 2) + first, first_valid = _race(op, selected_logits, temperature, top_k, top_p, noises[0]) + retry, retry_valid = _race( + op, + selected_logits, + op.Expand(op.Constant(value_float=1.0), op.Shape(temperature)), + op.Expand(op.Constant(value_int=30), op.Shape(top_k)), + op.Expand(op.Constant(value_float=0.95), op.Shape(top_p)), + noises[1], + ) + first, retry = op.Gather(candidates, first), op.Gather(candidates, retry) + repeated = op.Greater( + op.ReduceSum( + op.Cast(op.Equal(op.Unsqueeze(first, [-1]), history), to=ir.DataType.INT64), + [-1], + keepdims=False, + ), + 0, + ) + semantic = op.And( + op.GreaterOrEqual(first, config.semantic_start_token_id), + op.LessOrEqual(first, config.semantic_end_token_id), + ) + main = op.Where(op.And(repeated, semantic), retry, first) + valid = op.And( + op.And(first_valid, retry_valid), + _all_rows(op, _finite(op, op.Cast(logits, to=ir.DataType.FLOAT))), + ) + _output(builder, op.Where(valid, main, -1), "token", ir.DataType.INT64) + _output( + builder, + op.Where( + valid, + op.Clip( + op.Sub(main, config.semantic_start_token_id), 0, config.tower.vocab_size - 1 + ), + -1, + ), + "semantic_code", + ir.DataType.INT64, + ) + _output(builder, valid, "valid", ir.DataType.BOOL) + _output(builder, op.Where(valid, advanced, counter), "next_counter", ir.DataType.INT64) + return _component("mobius.policy.fish-main-sampler@1", graph, dict(_CONTRACT)) diff --git a/src/mobius/generation/_fish_omni_test.py b/src/mobius/generation/_fish_omni_test.py new file mode 100644 index 000000000..32366b763 --- /dev/null +++ b/src/mobius/generation/_fish_omni_test.py @@ -0,0 +1,197 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Network-free source-equivalent sampling math and explicit RNG-state tests.""" + +from __future__ import annotations + +import copy + +import numpy as np +import onnx_ir as ir +import pytest +import torch +from _test_configs import FISH_OMNI_CONFIG +from onnxscript import GraphBuilder + +from mobius._configs.fish_omni import FishOmniConfig +from mobius._testing.ort_inference import OnnxModelSession +from mobius.generation._fish_omni import ( + _race, + build_fish_codebook_sampler, + build_fish_main_sampler, +) +from mobius.integrations.transformers._config_resolver import _dict_to_pretrained_config + + +def _config(): + return FishOmniConfig.from_transformers( + _dict_to_pretrained_config(copy.deepcopy(FISH_OMNI_CONFIG)) + ) + + +def _reference(logits, temperature, top_k, top_p, noise): + values, ids = torch.sort( + torch.tensor(logits) / torch.tensor(temperature).clamp_min(1e-5)[:, None], + descending=True, + ) + ranks = torch.arange(values.shape[-1]) + dropped = (values.softmax(-1).cumsum(-1) > torch.tensor(top_p)[:, None]) | ( + ranks >= torch.tensor(top_k)[:, None] + ) + dropped[:, 0] = False + probabilities = values.masked_fill(dropped, -torch.inf).softmax(-1) + chosen = (probabilities / torch.tensor(noise)).argmax(-1, keepdim=True) + return ids.gather(-1, chosen).squeeze(-1).numpy() + + +@pytest.mark.parametrize("top_k,top_p", [([3, 3], [0.7, 1.0]), ([0, 1], [-1.0, 0.0])]) +def test_exponential_race_drops_crossing_token(top_k, top_p): + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + logits = builder.input("logits", ir.DataType.FLOAT, [2, 3]) + noise = builder.input("noise", ir.DataType.FLOAT, [2, 3]) + token, valid = _race( + builder.op, + logits, + builder.op.Constant(value_floats=[1.0, 0.7]), + builder.op.Constant(value_ints=top_k), + builder.op.Constant(value_floats=top_p), + noise, + ) + for name, value, dtype in ( + ("token", token, ir.DataType.INT64), + ("valid", valid, ir.DataType.BOOL), + ): + value.type, value.shape = ir.TensorType(dtype), ir.Shape([2]) + builder.add_output(value, name) + scores = np.log(np.array([[0.45, 0.35, 0.2], [0.1, 0.65, 0.25]], dtype=np.float32)) + draws = np.array([[100, 1e-6, 1], [1, 0.01, 1]], dtype=np.float32) + session = OnnxModelSession(ir.Model(graph, ir_version=12)) + try: + result = session.run({"logits": scores, "noise": draws}) + finally: + session.close() + np.testing.assert_array_equal( + result["token"], _reference(scores, [1.0, 0.7], top_k, top_p, draws) + ) + assert result["token"][0] == 0 + assert result["valid"].all() + + +def _noise(seed, counter, width): + modulus = 2147483647 + result = [] + for rank in range(width): + state = ( + int(seed) % modulus + 1 + ((int(counter) + rank) % modulus) * 2654435761 + ) % modulus + for _ in range(3): + state = (state * 48271) % modulus + state ^= state // 2048 + state = (state * 48271) % modulus + uniform = np.float32(state + 1) / np.float32(modulus + 2) + uniform = np.clip(uniform, np.float32(1.1754943508222875e-38), np.float32(1 - 2**-24)) + result.append(-np.log(uniform)) + return np.array(result, dtype=np.float32) + + +def test_main_sampler_rerolls_semantics_but_not_eos(): + config = _config() + component = build_fish_main_sampler(config) + candidates = sorted( + [ + config.eos_token_id, + *range(config.semantic_start_token_id, config.semantic_end_token_id + 1), + ] + ) + width = len(candidates) + logits = np.random.default_rng(42).normal(size=(2, config.vocab_size)).astype(np.float32) + logits[0, config.eos_token_id] = -100 + logits[1, config.eos_token_id] = 100 + seed, counter = np.array([42, 2026], np.int64), np.array([0, 17], np.int64) + noise = np.stack([_noise(s, c, width) for s, c in zip(seed, counter)]) + retry_noise = np.stack([_noise(s, c + width, width) for s, c in zip(seed, counter)]) + params = { + "temperature": np.array([0.8, 0.8], np.float32), + "top_k": np.array([5, 5], np.int64), + "top_p": np.array([0.9, 0.9], np.float32), + } + first = np.array(candidates)[ + _reference( + logits[:, candidates], + params["temperature"], + params["top_k"], + params["top_p"], + noise, + ) + ] + retry = np.array(candidates)[ + _reference(logits[:, candidates], [1.0, 1.0], [30, 30], [0.95, 0.95], retry_noise) + ] + history = np.full((2, 16), -1, np.int64) + history[:, -1] = first + session = OnnxModelSession(component.model) + try: + feeds = { + **params, + "logits": logits, + "seed": seed, + "counter": counter, + "history": history, + } + result = session.run(feeds) + repeated = session.run(feeds) + finally: + session.close() + np.testing.assert_array_equal(result["token"], [retry[0], config.eos_token_id]) + np.testing.assert_array_equal(result["next_counter"], counter + 2 * width) + assert result["valid"].all() + for name in result: + np.testing.assert_array_equal(result[name], repeated[name]) + + +@pytest.mark.parametrize("invalid", ["logits", "temperature", "top_p"]) +def test_codebook_sampler_reports_invalid_inputs_without_advancing_rng(invalid): + component = build_fish_codebook_sampler(_config()) + feeds = { + "logits": np.arange(16, dtype=np.float32)[None], + "temperature": np.array([1.0], np.float32), + "top_k": np.array([16], np.int64), + "top_p": np.array([1.0], np.float32), + "seed": np.array([42], np.int64), + "counter": np.array([17], np.int64), + } + feeds[invalid].flat[0] = np.nan + session = OnnxModelSession(component.model) + try: + result = session.run(feeds) + finally: + session.close() + assert result["token"].tolist() == [-1] + assert result["valid"].tolist() == [False] + assert result["next_counter"].tolist() == [17] + + +def test_main_sampler_rejects_nonfinite_noncandidate_logits(): + config = _config() + logits = np.zeros((1, config.vocab_size), dtype=np.float32) + logits[0, 0] = np.nan + session = OnnxModelSession(build_fish_main_sampler(config).model) + try: + result = session.run( + { + "logits": logits, + "temperature": np.array([1.0], np.float32), + "top_k": np.array([16], np.int64), + "top_p": np.array([1.0], np.float32), + "seed": np.array([42], np.int64), + "counter": np.array([17], np.int64), + "history": np.full((1, 16), -1, np.int64), + } + ) + finally: + session.close() + assert result["token"].tolist() == [-1] + assert result["valid"].tolist() == [False] + assert result["next_counter"].tolist() == [17] diff --git a/src/mobius/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index 47aebecca..67fe00073 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -96,17 +96,23 @@ @dataclasses.dataclass(frozen=True) class StreamingWeightSource: - """One checkpoint tensor bound to one dense ONNX initializer.""" + """One checkpoint tensor, reconstructed or explicitly byte-preserved as FP8.""" source_name: str - mode: Literal["direct", "fp8_scalar", "fp8_block_128"] = "direct" + mode: Literal["direct", "fp8_scalar", "fp8_block_128", "fp8_native"] = "direct" scale_name: str | None = None expected_scale: float | None = None + # Existing checkpoint adapters reconstruct through BF16. Architectures + # whose reference loader reconstructs directly in the compute dtype can + # opt in without changing the existing adapters' rounding semantics. + reconstruction_dtype: torch.dtype = torch.bfloat16 + # Opt-in source-loader saturation; other precisions/adapters are unchanged. + saturate_float16: bool = False @dataclasses.dataclass(frozen=True) class StreamingExpertBankSource: - """Per-expert source matrices packed into one dense rank-3 initializer.""" + """Per-expert source matrices packed into one typed rank-3 initializer.""" experts: tuple[tuple[StreamingWeightSource, ...], ...] @@ -700,6 +706,29 @@ def _load_indexed_tensor( return tensor +def _cast_preprocessed_source( + tensor: torch.Tensor, source: StreamingWeightSource, target_dtype: torch.dtype +) -> torch.Tensor: + if source.saturate_float16 and target_dtype == torch.float16: + # Compare the scalar maximum in Python: BF16 cannot represent 65504, + # so a BF16 tensor comparison/clamp would round that bound to 65536. + maximum = float(tensor.abs().max().item()) if tensor.numel() else 0.0 + if not math.isfinite(maximum): + raise ValueError(f"Nonfinite streaming weight '{source.source_name}'") + limit = torch.finfo(torch.float16).max + if maximum > limit: + wide = tensor.to(torch.float32) + overflow = int((wide.abs() > limit).sum().item()) + logger.warning( + "Clamping %d values outside the FP16 range in '%s'", + overflow, + source.source_name, + ) + # Never modify caller-owned/memory-mapped checkpoint storage. + tensor = wide.clamp(-limit, limit) + return tensor if tensor.dtype == target_dtype else tensor.to(target_dtype) + + def _materialize_preprocessed_source( source: StreamingWeightSource, key_index: Mapping[str, tuple[str, list[int], str]], @@ -708,6 +737,12 @@ def _materialize_preprocessed_source( """Read and reconstruct one classified source tensor.""" source_dtype = key_index[source.source_name][2] tensor = _load_indexed_tensor(source.source_name, key_index) + if source.mode == "fp8_native": + if not source_dtype.startswith("F8") or tensor.dtype != target_dtype: + raise ValueError( + f"Native FP8 source '{source.source_name}' must preserve its FP8 dtype" + ) + return tensor if source.mode == "fp8_block_128": assert source.scale_name is not None scale = _load_indexed_tensor(source.scale_name, key_index) @@ -715,9 +750,9 @@ def _materialize_preprocessed_source( tensor, scale, name=source.source_name, - target_dtype=torch.bfloat16, + target_dtype=source.reconstruction_dtype, ) - return tensor if target_dtype == torch.bfloat16 else tensor.to(target_dtype) + return _cast_preprocessed_source(tensor, source, target_dtype) if source.mode == "fp8_scalar": if not source_dtype.startswith("F8"): raise ValueError( @@ -734,14 +769,14 @@ def _materialize_preprocessed_source( ) tensor = tensor.to(torch.bfloat16) tensor.mul_(scale.to(torch.bfloat16)) - return tensor if target_dtype == torch.bfloat16 else tensor.to(target_dtype) + return _cast_preprocessed_source(tensor, source, target_dtype) if source.mode == "direct": if source_dtype.startswith("F8"): raise ValueError( f"Streaming source '{source.source_name}' is FP8 but was not explicitly " "classified as scaled FP8" ) - return tensor if tensor.dtype == target_dtype else tensor.to(target_dtype) + return _cast_preprocessed_source(tensor, source, target_dtype) raise AssertionError(f"Unknown streaming weight mode: {source.mode}") @@ -847,7 +882,7 @@ def _assign_lazy_expert_bank( key_index: Mapping[str, tuple[str, list[int], str]], target_name: str, ) -> None: - """Pack independently scaled expert matrices into one dense rank-3 bank.""" + """Pack reconstructed or byte-preserved expert matrices into one rank-3 bank.""" onnx_dtype = initializer.dtype assert onnx_dtype is not None assert initializer.shape is not None @@ -893,12 +928,13 @@ def stream_preprocessed_safetensors_to_model( revision: str | None = None, _resolved_paths: list[str] | None = None, ) -> dict[str, object]: - """Stream a fully classified transformed checkpoint into a dense ONNX graph. + """Stream a fully classified checkpoint into typed ONNX initializers. The planner must classify every source tensor as a target, a consumed scale, a validated deterministic constant, or an explicitly ignored sidecar tensor. Any unclassified key, missing target, malformed FP8 grid, or changed constant - fails before serialization. The resulting package is dense; this path never + fails before serialization. FP8 bytes require an explicit native classification; + direct sources cannot silently reinterpret FP8 storage. This path never claims native FP8 preservation. """ paths = ( @@ -926,10 +962,24 @@ def stream_preprocessed_safetensors_to_model( def validate_source( source: StreamingWeightSource, + target_dtype: ir.DataType, ) -> tuple[list[int], int, int]: if source.source_name not in key_index: raise ValueError(f"Streaming source '{source.source_name}' does not exist") _source_path, source_shape, source_dtype = key_index[source.source_name] + if source.mode == "fp8_native": + target_torch_dtype = tensor_adapters.to_torch_dtype(target_dtype) + if ( + source_dtype not in {"F8_E4M3", "F8_E5M2"} + or _TORCH_TO_SAFETENSORS_DTYPE.get(target_torch_dtype) != source_dtype + or source.scale_name is not None + or source.expected_scale is not None + or source.saturate_float16 + ): + raise ValueError( + f"Native FP8 source '{source.source_name}' requires identical FP8 " + "storage dtype and independently bound scales; no conversion is allowed" + ) if source.mode in {"fp8_block_128", "fp8_scalar"}: if source_dtype not in {"F8_E4M3", "F8_E5M2"}: raise ValueError( @@ -1042,19 +1092,45 @@ def validate_source( source_bytes + target_bytes + source.scratch_bytes, ) elif isinstance(source, StreamingWeightSource): - source_shape, source_bytes, scale_bytes = validate_source(source) + source_shape, source_bytes, scale_bytes = validate_source( + source, initializer.dtype + ) if expected_shape != source_shape: raise ValueError( f"Weight shape mismatch for '{target_name}': model expects " f"{expected_shape}, checkpoint source '{source.source_name}' has " f"{source_shape}" ) - bf16_bytes = math.prod(source_shape) * 2 - cast_bytes = target_bytes if initializer.dtype != ir.DataType.BFLOAT16 else 0 + reconstruction_dtype = ( + source.reconstruction_dtype + if source.mode == "fp8_block_128" + else tensor_adapters.to_torch_dtype(initializer.dtype) + if source.mode == "fp8_native" + else torch.bfloat16 + ) + reconstruction_bytes = ( + math.prod(source_shape) + * torch.empty((), dtype=reconstruction_dtype).element_size() + ) + cast_bytes = ( + target_bytes + if tensor_adapters.to_torch_dtype(initializer.dtype) != reconstruction_dtype + else 0 + ) + # Conservative bound: two FP32 temporaries plus the overflow mask. + saturation_bytes = ( + math.prod(source_shape) * 9 + if source.saturate_float16 and initializer.dtype == ir.DataType.FLOAT16 + else 0 + ) largest_source_tensor_bytes = max(largest_source_tensor_bytes, source_bytes) largest_reconstruction_working_set_bytes = max( largest_reconstruction_working_set_bytes, - source_bytes + bf16_bytes + cast_bytes + scale_bytes, + source_bytes + + reconstruction_bytes + + cast_bytes + + scale_bytes + + saturation_bytes, ) else: if len(expected_shape) != 3 or len(source.experts) != expected_shape[0]: @@ -1066,7 +1142,9 @@ def validate_source( for expert_index, projections in enumerate(source.experts): rows = 0 for projection in projections: - source_shape, source_bytes, scale_bytes = validate_source(projection) + source_shape, source_bytes, scale_bytes = validate_source( + projection, initializer.dtype + ) if len(source_shape) != 2 or source_shape[1] != expected_shape[2]: raise ValueError( f"Expert source '{projection.source_name}' has shape " @@ -1076,10 +1154,27 @@ def validate_source( dense_projection_bytes = ( math.prod(source_shape) * initializer.dtype.bitwidth + 7 ) // 8 - bf16_projection_bytes = math.prod(source_shape) * 2 + reconstruction_dtype = ( + projection.reconstruction_dtype + if projection.mode == "fp8_block_128" + else tensor_adapters.to_torch_dtype(initializer.dtype) + if projection.mode == "fp8_native" + else torch.bfloat16 + ) + reconstruction_bytes = ( + math.prod(source_shape) + * torch.empty((), dtype=reconstruction_dtype).element_size() + ) cast_projection_bytes = ( dense_projection_bytes - if initializer.dtype != ir.DataType.BFLOAT16 + if tensor_adapters.to_torch_dtype(initializer.dtype) + != reconstruction_dtype + else 0 + ) + saturation_bytes = ( + math.prod(source_shape) * 9 + if projection.saturate_float16 + and initializer.dtype == ir.DataType.FLOAT16 else 0 ) largest_source_tensor_bytes = max( @@ -1089,9 +1184,10 @@ def validate_source( max_transient_bytes = max( max_transient_bytes, source_bytes - + bf16_projection_bytes + + reconstruction_bytes + cast_projection_bytes - + scale_bytes, + + scale_bytes + + saturation_bytes, ) if rows != expected_shape[1]: raise ValueError( diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index a858548d8..892ccece7 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -1056,6 +1056,17 @@ def write_onnx_genai_config( "runtime configuration is claimed." ) return _write_advisory_component_contract(pkg, output_dir, warning=warning) + if "fish_omni" in config_types: + return _write_advisory_component_contract( + pkg, + output_dir, + warning=( + "Drama 3 exports token graphs, not a complete speech runtime. " + "Next-frame codebook sampling, temporal audio-decoder cache orchestration, " + "prompt formatting and waveform codec execution are host-owned. " + "Only graph-derived advisory component contracts are exported." + ), + ) if _looks_like_vibevoice_streaming_asr(pkg): if kv_native_dtype is not None: raise ValueError( diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index fa8b235e6..33c1b2a39 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -2034,6 +2034,13 @@ def write_ort_genai_config( "KV caches, DPM-Solver sampling, and prefilled voice-prompt caches; it " "cannot be represented by an ONNX Runtime GenAI configuration." ) + if getattr(config, "model_type", None) == "fish_omni": + raise ValueError( + "Drama 3 requires orchestration of next-frame codebook sampling and " + "temporal audio caches across multiple graphs; an ONNX Runtime GenAI " + "speech pipeline configuration is not implemented. Exported sampling " + "and codec graphs do not provide that runtime orchestration." + ) os.makedirs(directory, exist_ok=True) if set(pkg) in ({"audio_encoder"}, {"speaker_encoder"}) and getattr( pkg, "gguf_projector_type", None diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index d4d7f27f9..2b9f3c378 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -446,6 +446,10 @@ def build_transformers_model( from mobius.models.reuse import REUSE_REVISION detection_revision = REUSE_REVISION + if model_id == "fishaudio/drama-3" and revision is None: + from mobius._configs.fish_omni import DRAMA3_REVISION + + revision = detection_revision = DRAMA3_REVISION hf_config, loaded_from_raw_json = _load_transformers_config( config_model_id, @@ -644,6 +648,21 @@ def build_transformers_model( if task is None: task = _default_task_for_model(model_type) + if model_type == "fish_omni": + from mobius.tasks import FishOmniTask + + if task not in ("fish-omni", "fish-omni-generation") and not isinstance( + task, FishOmniTask + ): + raise ValueError( + "Drama 3 requires task='fish-omni' or 'fish-omni-generation'; " + "it is not a generic causal LM" + ) + if output_layer_indices is not None or fp8_kv_cache or kv_cache_scales is not None: + raise ValueError("Drama 3 does not support hidden-layer capture or FP8 KV caches") + if load_weights and config.fp8_experts and keep_quantized: + config = dataclasses.replace(config, native_fp8_experts=True) + from mobius.tasks import get_task resolved_task = get_task(task) @@ -677,6 +696,7 @@ def build_transformers_model( for name, model in package.items(): model.graph.name = f"{graph_source_name}/{name}" if model_type in _QWEN4_MODEL_TYPES | { + "fish_omni", "vibevoice", "vibevoice_streaming", "vibevoice_asr", @@ -693,7 +713,15 @@ def build_transformers_model( if load_weights: _reject_unsupported_affine_qwen4(model_type, config) - if is_gptoss_mxfp4_source and keep_quantized: + if model_type == "fish_omni": + from mobius.integrations.transformers._fish_omni_weights import ( + stream_fish_omni_safetensors_to_package, + ) + + stream_fish_omni_safetensors_to_package( + package, model_id, config, revision=revision + ) + elif is_gptoss_mxfp4_source and keep_quantized: from mobius.integrations.transformers._gptoss_weights import ( stream_gptoss_mxfp4_safetensors_to_package, ) diff --git a/src/mobius/integrations/transformers/_fish_omni_weights.py b/src/mobius/integrations/transformers/_fish_omni_weights.py new file mode 100644 index 000000000..0a81cf36f --- /dev/null +++ b/src/mobius/integrations/transformers/_fish_omni_weights.py @@ -0,0 +1,178 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Fail-closed bounded-memory loading for Drama 3 BF16/block-FP8 checkpoints.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import onnx_ir as ir +import torch + +from mobius._configs.fish_omni import FishOmniConfig +from mobius._model_package import ModelPackage +from mobius.integrations._weight_loading import ( + StreamingExpertBankSource, + StreamingWeightPlan, + StreamingWeightSource, + stream_preprocessed_safetensors_to_model, +) + + +def _classify_source(source, index, config, *, expert, saturate_float16=False): + if source not in index: + raise ValueError(f"Missing Drama 3 checkpoint tensor {source!r}") + source_dtype = index[source][2] + if expert and source.endswith(".weight_scale_inv"): + if not config.native_fp8_experts or source_dtype != "F32": + raise ValueError(f"Native Drama 3 FP8 scales must be FP32: {source}") + return StreamingWeightSource(source) + if config.fp8_experts and expert: + if source_dtype != "F8_E4M3": + raise ValueError(f"Drama 3 FP8 expert {source!r} must have dtype F8_E4M3") + if config.native_fp8_experts: + return StreamingWeightSource(source, mode="fp8_native") + return StreamingWeightSource( + source, + mode="fp8_block_128", + scale_name=source.removesuffix(".weight") + ".weight_scale_inv", + reconstruction_dtype=torch.float32, + saturate_float16=saturate_float16, + ) + if source_dtype not in {"F32", "BF16", "F16"}: + raise ValueError(f"Unexpected Drama 3 weight dtype {source_dtype}: {source}") + return StreamingWeightSource(source, saturate_float16=saturate_float16) + + +def _source_name(target: str) -> str: + if target == "embedding.embed_tokens.weight": + return "model.language_model.embed_tokens.weight" + if target == "embedding.codebook_embeddings.weight": + return "audio_decoder.codebook_embeddings.weight" + if target.startswith("decoder.model."): + source = "model.language_model." + target.removeprefix("decoder.model.") + source = source.replace(".mlp.packed_experts.", ".mlp.experts.") + return source.replace(".linear_attn.conv1d_weight", ".linear_attn.conv1d.weight") + if target == "decoder.lm_head.weight": + return "lm_head.weight" + for local, source in ( + ("audio_decoder.shared_layers.", "audio_decoder_mtp_shared.layers."), + ("audio_decoder.residual_layers.", "audio_decoder_mtp.layers."), + ("audio_decoder.codebook_heads.", "audio_decoder_mtp.heads."), + ): + if target.startswith(local): + return source + target.removeprefix(local) + raise ValueError(f"Unrecognized Drama 3 graph parameter {target!r}") + + +def build_fish_omni_streaming_plan( + index: Mapping[str, tuple[str, list[int], str]], + initializers: Mapping[str, ir.Value], + config: FishOmniConfig, +) -> StreamingWeightPlan: + """Map every learned parameter and classify each independently scaled expert.""" + targets: dict[str, StreamingWeightSource | StreamingExpertBankSource] = {} + for name, value in initializers.items(): + if value.const_value is not None: + continue + source = _source_name(name) + saturate = value.dtype == ir.DataType.FLOAT16 and ( + source.startswith("model.language_model.") or source == "lm_head.weight" + ) + if any( + source.endswith(f".experts.{projection}.{suffix}") + for projection in ("gate_proj", "up_proj", "down_proj", "gate_up_proj") + for suffix in ("weight", "weight_scale_inv") + ): + assert config.num_local_experts is not None + prefix, projection = source.rsplit(".experts.", 1) + projections = ( + ("gate_proj.weight", "up_proj.weight") + if projection == "gate_up_proj.weight" + else (projection,) + ) + targets[name] = StreamingExpertBankSource( + tuple( + tuple( + _classify_source( + f"{prefix}.experts.{expert}.{part}", + index, + config, + expert=True, + saturate_float16=saturate, + ) + for part in projections + ) + for expert in range(config.num_local_experts) + ) + ) + else: + targets[name] = _classify_source( + source, + index, + config, + expert=".mlp.experts." in source, + saturate_float16=saturate, + ) + return StreamingWeightPlan( + targets=targets, + report={ + "architecture": "fish_omni", + "codec_included": False, + "native_fp8": config.native_fp8_experts, + "output_weight_format": "mixed-block-fp8-128" + if config.native_fp8_experts + else "dense", + "activation_quantization": "dynamic_fp8" if config.native_fp8_experts else "none", + }, + ) + + +def stream_fish_omni_safetensors_to_package( + package: ModelPackage, + model_id: str, + config: FishOmniConfig, + *, + revision: str | None, +) -> None: + """Validate the entire package before binding any lazy weight sources. + + The aggregate graph is only a binding inventory and is never executed or + serialized. One generic planner checks all three components transactionally, including unexpected + source keys, rather than ignoring another component's unchecked tensors. + """ + parameters: dict[str, ir.Value] = {} + for model in package.values(): + for name, value in model.graph.initializers.items(): + if value.const_value is not None: + continue + if name in parameters: + raise ValueError(f"Duplicate Drama 3 parameter {name!r}") + parameters[name] = value + inventory = ir.Model( + ir.Graph( + [], + [], + nodes=[], + initializers=[ + ir.Value(name=name, shape=value.shape, type=value.type) + for name, value in parameters.items() + ], + opset_imports={"": 24}, + ), + ir_version=12, + ) + report = stream_preprocessed_safetensors_to_model( + inventory, + model_id, + lambda index, initializers: build_fish_omni_streaming_plan( + index, initializers, config + ), + revision=revision, + ) + for name, value in parameters.items(): + value.const_value = inventory.graph.initializers[name].const_value + package.weight_loading_report = report + for model in package.values(): + model.metadata_props.update(inventory.metadata_props) diff --git a/src/mobius/integrations/transformers/_fish_omni_weights_test.py b/src/mobius/integrations/transformers/_fish_omni_weights_test.py new file mode 100644 index 000000000..e82531d3c --- /dev/null +++ b/src/mobius/integrations/transformers/_fish_omni_weights_test.py @@ -0,0 +1,405 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import copy +import dataclasses +import json + +import numpy as np +import onnx_ir as ir +import pytest +import torch +from _test_configs import FISH_OMNI_CONFIG +from safetensors.torch import load_file, save_file + +from mobius import build_from_module +from mobius.__main__ import main +from mobius._configs.fish_omni import DRAMA3_REVISION +from mobius._model_package import ModelPackage +from mobius.integrations.transformers import _builder +from mobius.integrations.transformers._config_resolver import _dict_to_pretrained_config +from mobius.integrations.transformers._fish_omni_weights import ( + _source_name, + build_fish_omni_streaming_plan, + stream_fish_omni_safetensors_to_package, +) +from mobius.models import FishOmniModel +from mobius.models.fish_omni_test import ( + native_fp8_prototype_config, + reduced_precision_prototype_task, + tiny_config, +) + + +def _checkpoint(package): + torch.manual_seed(19) + weights = {} + for model in package.values(): + for name, value in model.graph.initializers.items(): + if value.const_value is not None: + continue + tensor = torch.randn(tuple(value.shape), dtype=torch.float32) * 0.02 + source = _source_name(name) + if ".packed_experts." in name: + prefix, projection = source.rsplit(".experts.", 1) + for expert, matrix in enumerate(tensor): + if projection == "gate_up_proj.weight": + gate, up = matrix.chunk(2, dim=0) + weights[f"{prefix}.experts.{expert}.gate_proj.weight"] = ( + gate.contiguous() + ) + weights[f"{prefix}.experts.{expert}.up_proj.weight"] = up.contiguous() + else: + weights[f"{prefix}.experts.{expert}.{projection}"] = ( + matrix.contiguous() + ) + else: + weights[source] = tensor + return weights + + +@pytest.mark.parametrize("failure", [None, "scale_dtype", "scale_shape", "target_dtype"]) +def test_native_fp8_streaming_preserves_original_code_bytes_and_fp32_grids(tmp_path, failure): + config = native_fp8_prototype_config() + package = build_from_module( + FishOmniModel(config), + config, + task=reduced_precision_prototype_task("fish-omni"), + execution_provider="cuda", + ) + weights = _checkpoint(package) + for name, value in list(weights.items()): + if ".mlp.experts." not in name: + continue + if name.endswith(".weight"): + # Every finite E4M3FN byte, including negative zero and both extrema. + raw = torch.arange(value.numel()).remainder(254) + raw = torch.where(raw >= 127, raw + 1, raw) + weights[name] = raw.to(torch.uint8).reshape(value.shape).view(torch.float8_e4m3fn) + else: + weights[name] = value.abs() + 0.01 + scale = next(name for name in weights if name.endswith(".weight_scale_inv")) + if failure == "scale_dtype": + weights[scale] = weights[scale].bfloat16() + elif failure == "scale_shape": + weights[scale] = torch.ones((1, 2)) + elif failure == "target_dtype": + target = next( + value + for name, value in package["decoder"].graph.initializers.items() + if ".packed_experts." in name and name.endswith(".weight") + ) + target.type = ir.TensorType(ir.DataType.FLOAT) + save_file(weights, tmp_path / "model.safetensors") + if failure: + with pytest.raises((ValueError, TypeError), match=r"FP32|shape|identical FP8"): + stream_fish_omni_safetensors_to_package( + package, str(tmp_path), config, revision=None + ) + assert all( + value.const_value is None + for name, value in package["decoder"].graph.initializers.items() + if ".packed_experts." in name + ) + return + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + assert package.weight_loading_report["native_fp8"] is True + assert package.weight_loading_report["activation_quantization"] == "dynamic_fp8" + path = tmp_path / "decoder.onnx" + ir.save(package["decoder"], path, external_data="decoder.data") + saved = ir.load(path) + for name, value in saved.graph.initializers.items(): + if ".packed_experts." not in name: + continue + prefix, projection = _source_name(name).rsplit(".experts.", 1) + source = [ + weights[f"{prefix}.experts.{i}.{projection}"] + for i in range(config.num_local_experts) + ] + if name.endswith(".weight"): + assert value.dtype == ir.DataType.FLOAT8E4M3FN + expected = torch.stack([tensor.view(torch.uint8) for tensor in source]).numpy() + np.testing.assert_array_equal(value.const_value.numpy().view(np.uint8), expected) + else: + assert value.dtype == ir.DataType.FLOAT + np.testing.assert_array_equal( + value.const_value.numpy(), torch.stack(source).numpy() + ) + + +def test_fish_omni_streams_all_components(tmp_path): + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task="fish-omni") + weights = _checkpoint(package) + save_file(weights, tmp_path / "model.safetensors") + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + for model in package.values(): + for name, value in model.graph.initializers.items(): + assert value.const_value is not None, name + if ".packed_experts." in name: + prefix, projection = _source_name(name).rsplit(".experts.", 1) + projections = ( + ("gate_proj.weight", "up_proj.weight") + if projection == "gate_up_proj.weight" + else (projection,) + ) + expected = torch.stack( + [ + torch.cat( + [weights[f"{prefix}.experts.{i}.{part}"] for part in projections] + ) + for i in range(value.shape[0]) + ] + ) + torch.testing.assert_close( + torch.from_numpy(value.const_value.numpy()), expected + ) + elif name.startswith( + ( + "embedding.", + "decoder.model.layers.", + "decoder.lm_head.", + "audio_decoder.shared_layers.", + "audio_decoder.residual_layers.", + "audio_decoder.codebook_heads.", + ) + ): + torch.testing.assert_close( + torch.from_numpy(value.const_value.numpy()), weights[_source_name(name)] + ) + assert package.weight_loading_report["native_fp8"] is False + assert package.weight_loading_report["codec_included"] is False + + +@pytest.mark.parametrize("failure", ["unexpected", "missing", "shape"]) +def test_fish_omni_rejects_bad_checkpoint_transactionally(tmp_path, failure): + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task="fish-omni") + weights = _checkpoint(package) + if failure == "unexpected": + weights["unknown.weight"] = torch.ones(1) + elif failure == "missing": + weights.pop("lm_head.weight") + else: + weights["lm_head.weight"] = torch.ones(1) + save_file(weights, tmp_path / "model.safetensors") + unbound = [ + value + for model in package.values() + for value in model.graph.initializers.values() + if value.const_value is None + ] + with pytest.raises(ValueError): + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + assert all(value.const_value is None for value in unbound) + + +def test_fish_omni_fp8_classification(): + name = "decoder.model.layers.0.mlp.experts.0.gate_proj.weight" + value = ir.Value( + name=name, shape=ir.Shape([128, 128]), type=ir.TensorType(ir.DataType.FLOAT) + ) + source = _source_name(name) + index = {source: ("unused", [128, 128], "F8_E4M3")} + plan = build_fish_omni_streaming_plan( + index, {name: value}, dataclasses.replace(tiny_config(), fp8_experts=True) + ) + assert plan.targets[name].mode == "fp8_block_128" + assert plan.targets[name].reconstruction_dtype == torch.float32 + index[source] = ("unused", [128, 128], "BF16") + with pytest.raises(ValueError, match="F8_E4M3"): + build_fish_omni_streaming_plan( + index, {name: value}, dataclasses.replace(tiny_config(), fp8_experts=True) + ) + + +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.BFLOAT16]) +def test_fish_omni_fp8_scales_are_reconstructed_in_fp32(tmp_path, dtype): + config = dataclasses.replace(tiny_config(), fp8_experts=True, dtype=dtype) + name = "decoder.model.layers.0.mlp.experts.0.gate_proj.weight" + value = ir.Value(name=name, shape=ir.Shape([128, 128]), type=ir.TensorType(dtype)) + model = ir.Model(ir.Graph([], [], nodes=[], initializers=[value]), ir_version=12) + package = ModelPackage({"decoder": model}, config=config) + source = _source_name(name) + weight = torch.linspace(-1, 1, 128 * 128).reshape(128, 128).to(torch.float8_e4m3fn) + scale = torch.tensor([[0.123456]], dtype=torch.float32) + save_file( + { + source: weight, + source.removesuffix(".weight") + ".weight_scale_inv": scale, + }, + tmp_path / "model.safetensors", + ) + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + expected = weight.float() * scale + if dtype == ir.DataType.BFLOAT16: + expected = expected.to(torch.bfloat16).float() + actual = torch.tensor(value.const_value.numpy().astype("float32")) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_fish_omni_cuda_banks_preserve_fp8_reconstruction_and_expert_order(tmp_path): + config = dataclasses.replace(tiny_config(), fp8_experts=True) + name = "decoder.model.layers.0.mlp.packed_experts.gate_proj.weight" + count = config.num_local_experts + value = ir.Value( + name=name, shape=ir.Shape([count, 128, 128]), type=ir.TensorType(ir.DataType.FLOAT) + ) + package = ModelPackage( + {"decoder": ir.Model(ir.Graph([], [], nodes=[], initializers=[value]), ir_version=12)}, + config=config, + ) + tensors = {} + expected = [] + for expert in range(count): + source = f"model.language_model.layers.0.mlp.experts.{expert}.gate_proj.weight" + weight = torch.linspace(-1, 1, 128 * 128).reshape(128, 128).to(torch.float8_e4m3fn) + scale = torch.tensor([[0.123456 * (expert + 1)]], dtype=torch.float32) + tensors[source] = weight + tensors[source.removesuffix(".weight") + ".weight_scale_inv"] = scale + expected.append(weight.float() * scale) + save_file(tensors, tmp_path / "model.safetensors") + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + torch.testing.assert_close( + torch.from_numpy(value.const_value.numpy()), torch.stack(expected), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("revision", [None, "explicit-revision"]) +def test_fish_omni_revision_reaches_first_config_read(monkeypatch, revision): + seen = [] + + def load(model_id, **kwargs): + seen.append(kwargs["revision"]) + return _dict_to_pretrained_config(copy.deepcopy(FISH_OMNI_CONFIG)), True + + monkeypatch.setattr(_builder, "_load_transformers_config", load) + package = _builder.build_transformers_model( + "fishaudio/drama-3", revision=revision, load_weights=False + ) + expected = revision or DRAMA3_REVISION + assert seen == [expected] + assert all( + m.metadata_props["mobius.source_revision"] == expected for m in package.values() + ) + + +def test_fish_omni_requires_explicit_dequantization_before_weight_io(monkeypatch): + raw = copy.deepcopy(FISH_OMNI_CONFIG) + raw["quantization_config"] = { + "quant_method": "bsgemm-moe", + "weight_block_size": [128, 128], + "activation_scheme": "dynamic", + } + monkeypatch.setattr( + _builder, + "_load_transformers_config", + lambda *args, **kwargs: (_dict_to_pretrained_config(raw), True), + ) + with pytest.raises(ValueError, match="--dequantize"): + _builder.build_transformers_model("fishaudio/drama-3") + + +@pytest.mark.parametrize( + "task,ep", + [ + ("fish-omni", "default"), + ("fish-omni-generation", "default"), + ("fish-omni-generation", "cuda"), + ], +) +def test_fish_omni_cli_persists_three_weighted_models(tmp_path, task, ep): + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task="fish-omni") + (tmp_path / "config.json").write_text(json.dumps(FISH_OMNI_CONFIG), encoding="utf-8") + save_file(_checkpoint(package), tmp_path / "model.safetensors") + output = tmp_path / "output" + main( + [ + "build", + "--model", + str(tmp_path), + "--dtype", + "f32", + "--task", + task, + "--execution-provider", + ep, + "--dequantize", + "--output", + str(output), + ] + ) + audio = "audio_generator" if task == "fish-omni-generation" else "audio_decoder" + for component in ("embedding", "decoder", audio): + model = ir.load(output / component / "model.onnx") + assert all( + value.const_value is not None for value in model.graph.initializers.values() + ) + report = json.loads((output / "weight-loading-report.json").read_text()) + assert report["native_fp8"] is False + + +@pytest.mark.parametrize("source_dtype", [torch.float32, torch.bfloat16, torch.float8_e4m3fn]) +def test_fp16_source_saturation_is_finite_logged_and_does_not_modify_checkpoint( + tmp_path, caplog, source_dtype +): + fp8 = source_dtype == torch.float8_e4m3fn + config = dataclasses.replace(tiny_config(), dtype=ir.DataType.FLOAT16, fp8_experts=fp8) + name = "decoder.model.layers.0.mlp.experts.0.gate_proj.weight" + source = _source_name(name) + value = ir.Value( + name=name, shape=ir.Shape([128, 128]), type=ir.TensorType(ir.DataType.FLOAT16) + ) + package = ModelPackage( + {"decoder": ir.Model(ir.Graph([], [], nodes=[], initializers=[value]), ir_version=12)}, + config=config, + ) + if fp8: + weights = torch.linspace(-1, 1, 128 * 128).reshape(128, 128).to(source_dtype) + scale = torch.tensor([[100000.0]], dtype=torch.float32) + checkpoint = { + source: weights, + source.removesuffix(".weight") + ".weight_scale_inv": scale, + } + reconstructed = weights.float() * scale + else: + weights = torch.linspace(-100000, 100000, 128 * 128).reshape(128, 128).to(source_dtype) + weights[0, 0] = 65536 + checkpoint = {source: weights} + reconstructed = weights.float() + path = tmp_path / "model.safetensors" + save_file(checkpoint, path) + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + actual = torch.from_numpy(value.const_value.numpy()) + torch.testing.assert_close( + actual, reconstructed.clamp(-65504, 65504).half(), rtol=0, atol=0 + ) + assert torch.isfinite(actual).all() + assert "Clamping" in caplog.text + restored = load_file(path)[source] + assert torch.equal(restored.view(torch.uint8), weights.view(torch.uint8)) + expected_bound = weights.numel() * (weights.element_size() + 8 + 2) + assert ( + package.weight_loading_report["largest_reconstruction_working_set_bytes"] + >= expected_bound + ) + + +def test_fp16_saturation_rejects_nonfinite_source_weights(tmp_path): + config = dataclasses.replace(tiny_config(), dtype=ir.DataType.FLOAT16) + name = "decoder.lm_head.weight" + source = _source_name(name) + value = ir.Value( + name=name, shape=ir.Shape([1, 2]), type=ir.TensorType(ir.DataType.FLOAT16) + ) + package = ModelPackage( + {"decoder": ir.Model(ir.Graph([], [], nodes=[], initializers=[value]), ir_version=12)}, + config=config, + ) + save_file({source: torch.tensor([[float("nan"), 1.0]])}, tmp_path / "model.safetensors") + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + with pytest.raises(ValueError, match="Nonfinite streaming weight"): + value.const_value.numpy() diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 638ae99db..b9e74efc5 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -4,6 +4,7 @@ from __future__ import annotations __all__ = [ + "FishOmniModel", "ApertusCausalLMModel", "ArcticGGUFCausalLMModel", "ArceeCausalLMModel", @@ -271,6 +272,7 @@ MPTCausalLMModel, ) from mobius.models.falcon_h1 import FalconH1ForCausalLM +from mobius.models.fish_omni import FishOmniModel from mobius.models.flux_sd3 import FluxTransformer2DModel, SD3Transformer2DModel from mobius.models.fun_asr import FunASRForConditionalGeneration from mobius.models.gemma import Gemma2CausalLMModel, GemmaCausalLMModel diff --git a/src/mobius/models/fish_codec.py b/src/mobius/models/fish_codec.py new file mode 100644 index 000000000..7987342dc --- /dev/null +++ b/src/mobius/models/fish_codec.py @@ -0,0 +1,289 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Declarative Fish Speech DAC ``from_indices`` waveform decoder. + +Replicates the decode path of pinned Fish Speech ``modded_dac.DAC`` and +``rvq.DownsampleResidualVectorQuantize``. Input is INT64 [B,10,T] code IDs +(without the main text token); output is FP32 [B,1,T*2048] at 44100 Hz. +This is a full-sequence decoder, not a streaming convolution/KV-cache API. +""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs.fish_codec import FishCodecConfig +from mobius.components import ( + CausalConv1d, + CausalTransConv1d, + Conv1d, + ConvNeXtBlock, + Embedding, + Linear, + RMSNorm, +) + + +class _Snake1d(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter([1, channels, 1]) + + def forward(self, op: OpBuilder, x: ir.Value): + # DAC Snake has one direct alpha parameter, unlike log-space SnakeBeta. + sine = op.Sin(op.Mul(x, self.alpha)) + return op.Add(x, op.Div(op.Mul(sine, sine), op.Add(self.alpha, 1e-9))) + + +class _ResidualUnit(nn.Module): + def __init__(self, channels: int, dilation: int): + super().__init__() + self.block = nn.Sequential( + _Snake1d(channels), + CausalConv1d(channels, channels, 7, dilation=dilation), + _Snake1d(channels), + CausalConv1d(channels, channels, 1), + ) + + def forward(self, op: OpBuilder, x: ir.Value): + return op.Add(x, self.block(op, x)) + + +class _DecoderBlock(nn.Module): + def __init__(self, channels: int, rate: int): + super().__init__() + self.block = nn.Sequential( + _Snake1d(channels), + CausalTransConv1d(channels, channels // 2, 2 * rate, rate), + *[_ResidualUnit(channels // 2, dilation) for dilation in (1, 3, 9)], + ) + + def forward(self, op: OpBuilder, x: ir.Value): + return self.block(op, x) + + +class _WaveformDecoder(nn.Module): + def __init__(self, config: FishCodecConfig): + super().__init__() + channels = config.decoder_dim + layers: list[nn.Module] = [CausalConv1d(config.latent_dim, channels, 7)] + for rate in config.decoder_rates: + layers.append(_DecoderBlock(channels, rate)) + channels //= 2 + layers.extend([_Snake1d(channels), CausalConv1d(channels, 1, 7)]) + self.model = nn.Sequential(*layers) + + def forward(self, op: OpBuilder, x: ir.Value): + return op.Tanh(self.model(op, x)) + + +class _Codebook(nn.Module): + def __init__(self, config: FishCodecConfig, size: int): + super().__init__() + self.codebook = Embedding(size, config.codebook_dim) + # DAC out_proj is a weight-normalized 1x1 convolution, not a linear. + self.out_proj = Conv1d(config.codebook_dim, config.latent_dim, 1) + + def forward(self, op: OpBuilder, codes: ir.Value): + embedded = op.Transpose(self.codebook(op, codes), perm=[0, 2, 1]) + return self.out_proj(op, embedded) + + +class _RVQ(nn.Module): + def __init__(self, config: FishCodecConfig, count: int, size: int): + super().__init__() + self.quantizers = nn.ModuleList([_Codebook(config, size) for _ in range(count)]) + self.size = size + + def forward(self, op: OpBuilder, codes: ir.Value): + # The token tower has 4096 outputs for every book; upstream clamps only + # the upper bound before embedding lookup (residual books have 1024 rows). + codes = op.Min(codes, op.Constant(value_int=self.size - 1)) + values = [ + book(op, op.Gather(codes, i, axis=1)) for i, book in enumerate(self.quantizers) + ] + result = values[0] + for value in values[1:]: + result = op.Add(result, value) + return result + + +class _Scale(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.gamma = nn.Parameter([channels]) + + def forward(self, op: OpBuilder, x: ir.Value): + return op.Mul(x, self.gamma) + + +class _Attention(nn.Module): + def __init__(self, config: FishCodecConfig): + super().__init__() + self.config = config + self.wqkv = Linear(config.latent_dim, 3 * config.latent_dim, bias=False) + self.wo = Linear(config.latent_dim, config.latent_dim, bias=False) + + def forward(self, op: OpBuilder, x: ir.Value, mask: ir.Value, cos, sin): + cfg = self.config + q, k, v = op.Split(self.wqkv(op, x), [cfg.latent_dim] * 3, axis=-1, _outputs=3) + + def rotate(value): + pairs = op.Reshape(value, [0, 0, cfg.transformer_heads, cfg.head_dim // 2, 2]) + even, odd = op.Gather(pairs, 0, axis=-1), op.Gather(pairs, 1, axis=-1) + real = op.Sub(op.Mul(even, cos), op.Mul(odd, sin)) + imag = op.Add(op.Mul(odd, cos), op.Mul(even, sin)) + return op.Reshape( + op.Concat(op.Unsqueeze(real, [-1]), op.Unsqueeze(imag, [-1]), axis=-1), + [0, 0, cfg.latent_dim], + ) + + q, k = rotate(q), rotate(k) + attended = op.Attention( + q, + k, + v, + mask, + q_num_heads=cfg.transformer_heads, + kv_num_heads=cfg.transformer_heads, + ) + return self.wo(op, attended) + + +class _FeedForward(nn.Module): + def __init__(self, config: FishCodecConfig): + super().__init__() + self.w1 = Linear(config.latent_dim, config.transformer_intermediate_size, bias=False) + self.w3 = Linear(config.latent_dim, config.transformer_intermediate_size, bias=False) + self.w2 = Linear(config.transformer_intermediate_size, config.latent_dim, bias=False) + + def forward(self, op: OpBuilder, x: ir.Value): + gate = self.w1(op, x) + return self.w2(op, op.Mul(op.Mul(gate, op.Sigmoid(gate)), self.w3(op, x))) + + +class _TransformerBlock(nn.Module): + def __init__(self, config: FishCodecConfig): + super().__init__() + self.attention = _Attention(config) + self.feed_forward = _FeedForward(config) + self.attention_norm = RMSNorm(config.latent_dim, eps=config.norm_eps) + self.ffn_norm = RMSNorm(config.latent_dim, eps=config.norm_eps) + self.attention_layer_scale = _Scale(config.latent_dim) + self.ffn_layer_scale = _Scale(config.latent_dim) + + def forward(self, op: OpBuilder, x: ir.Value, mask: ir.Value, cos, sin): + attended = self.attention(op, self.attention_norm(op, x), mask, cos, sin) + x = op.Add(x, self.attention_layer_scale(op, attended)) + return op.Add(x, self.ffn_layer_scale(op, self.feed_forward(op, self.ffn_norm(op, x)))) + + +class _PostTransformer(nn.Module): + def __init__(self, config: FishCodecConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList( + [_TransformerBlock(config) for _ in range(config.transformer_layers)] + ) + self.norm = RMSNorm(config.latent_dim, eps=config.norm_eps) + # Upstream precompute_freqs_cis rounds BOTH sin and cos to BF16 even + # when all weights and activations are FP32. Keep those exact values. + frequencies = 1.0 / ( + config.rope_base ** (torch.arange(0, config.head_dim, 2).float() / config.head_dim) + ) + angles = torch.outer(torch.arange(config.max_frames).float(), frequencies) + polar = torch.polar(torch.ones_like(angles), angles) + self._rope = ( + torch.stack((polar.real, polar.imag), dim=-1).to(torch.bfloat16).float().numpy() + ) + + def forward(self, op: OpBuilder, x: ir.Value): + x = op.Transpose(x, perm=[0, 2, 1]) # (B,T,D) + length = op.Squeeze(op.Shape(x, start=1, end=2), [0]) + positions = op.Range(op.Constant(value_int=0), length, op.Constant(value_int=1)) + rows, columns = op.Unsqueeze(positions, [1]), op.Unsqueeze(positions, [0]) + mask = op.And( + op.LessOrEqual(columns, rows), + op.Greater(columns, op.Sub(rows, self.config.transformer_window)), + ) + mask = op.Unsqueeze(mask, [0, 1]) + rope = op.Gather(op.Constant(value=ir.tensor(self._rope)), positions) + cos = op.Unsqueeze(op.Gather(rope, 0, axis=-1), [0, 2]) + sin = op.Unsqueeze(op.Gather(rope, 1, axis=-1), [0, 2]) + for layer in self.layers: + x = layer(op, x, mask, cos, sin) + return op.Transpose(self.norm(op, x), perm=[0, 2, 1]) + + +class _Quantizer(nn.Module): + def __init__(self, config: FishCodecConfig): + super().__init__() + self.semantic_quantizer = _RVQ(config, 1, config.semantic_codebook_size) + self.quantizer = _RVQ(config, config.residual_codebooks, config.residual_codebook_size) + self.post_module = _PostTransformer(config) + self.upsample = nn.Sequential( + *[ + nn.Sequential( + CausalTransConv1d(config.latent_dim, config.latent_dim, rate, rate), + ConvNeXtBlock(config.latent_dim), + ) + for rate in reversed(config.quantizer_rates) + ] + ) + + def forward(self, op: OpBuilder, codes: ir.Value): + semantic = self.semantic_quantizer(op, op.Slice(codes, [0], [1], [1])) + residual = self.quantizer(op, op.Slice(codes, [1], [np.iinfo(np.int64).max], [1])) + return self.upsample(op, self.post_module(op, op.Add(semantic, residual))) + + +class FishCodecDecoder(nn.Module): + """Fish Speech DAC decode-only model; use ``build_fish_codec_decoder`` to export.""" + + category = "Audio" + default_task = "fish-codec-decode" + + def __init__(self, config: FishCodecConfig): + super().__init__() + self.config = config + self.quantizer = _Quantizer(config) + self.decoder = _WaveformDecoder(config) + + def forward(self, op: OpBuilder, codes: ir.Value): + return self.decoder(op, self.quantizer(op, codes)) + + def preprocess_weights(self, state_dict: dict[str, torch.Tensor]): + """Select decode weights and materialize both upstream weight-norm formats. + + At most the decode subset plus one convolution normalization temporary + is materialized; the checkpoint can remain memory-mapped on CPU. + """ + expected = dict(self.named_parameters()) + result: dict[str, torch.Tensor] = {} + for name, parameter in expected.items(): + source = name + if source in state_dict: + tensor = state_dict[source] + elif source.endswith(".weight"): + prefix = source.removesuffix(".weight") + modern = prefix + ".parametrizations.weight.original" + if modern + "0" in state_dict and modern + "1" in state_dict: + g, v = state_dict[modern + "0"], state_dict[modern + "1"] + elif prefix + ".weight_g" in state_dict and prefix + ".weight_v" in state_dict: + g, v = state_dict[prefix + ".weight_g"], state_dict[prefix + ".weight_v"] + else: + raise ValueError(f"Missing codec weight {source!r}") + tensor = torch._weight_norm(v.float(), g.float(), 0) + else: + raise ValueError(f"Missing codec weight {source!r}") + if tuple(tensor.shape) != tuple(parameter.shape): + raise ValueError( + f"Codec weight {source!r}: expected {tuple(parameter.shape)}, " + f"got {tuple(tensor.shape)}" + ) + result[name] = tensor.float() + return result diff --git a/src/mobius/models/fish_codec_test.py b/src/mobius/models/fish_codec_test.py new file mode 100644 index 000000000..b8416b2ab --- /dev/null +++ b/src/mobius/models/fish_codec_test.py @@ -0,0 +1,121 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import dataclasses +from unittest.mock import patch + +import onnx_ir as ir +import pytest +import torch + +from mobius._configs.fish_codec import ( + DRAMA3_CODEC_REVISION, + FishCodecConfig, +) +from mobius.models.fish_codec import FishCodecDecoder +from mobius.tasks.fish_codec import build_fish_codec_decoder, load_drama3_codec_decoder + + +def _tiny(): + return FishCodecConfig( + latent_dim=16, + transformer_layers=1, + transformer_heads=2, + transformer_intermediate_size=32, + max_frames=16, + decoder_dim=32, + decoder_rates=(2, 2), + ) + + +def test_native_codec_contract(): + config = FishCodecConfig() + assert config.num_codebooks == 10 + assert config.samples_per_frame == 2048 + assert config.sample_rate == 44100 + model = build_fish_codec_decoder(_tiny())["waveform_decoder"] + assert list(model.graph.inputs[0].shape) == ["batch", 10, "frames"] + assert model.graph.inputs[0].dtype == ir.DataType.INT64 + assert model.graph.outputs[0].name == "waveform" + assert model.metadata_props["codec_revision"] == DRAMA3_CODEC_REVISION + assert any(node.op_type == "Attention" for node in model.graph) + assert not any(node.domain == "com.microsoft" for node in model.graph) + + +def test_missing_and_mismatched_weights_fail_closed(): + module = FishCodecDecoder(_tiny()) + with pytest.raises(ValueError, match="Missing codec weight"): + module.preprocess_weights({}) + name = next(iter(dict(module.named_parameters()))) + with pytest.raises(ValueError, match=r"expected .*got"): + module.preprocess_weights({name: torch.zeros(1)}) + + +def test_preprocess_both_weight_norm_formats(): + module = FishCodecDecoder(_tiny()) + state = {name: torch.ones(tuple(param.shape)) for name, param in module.named_parameters()} + for index, name in enumerate( + [ + "quantizer.semantic_quantizer.quantizers.0.out_proj.weight", + "decoder.model.0.conv.weight", + ] + ): + v = state.pop(name) + prefix = name.removesuffix(".weight") + g = torch.full((v.shape[0], 1, 1), 0.5) + if index == 0: + state[prefix + ".weight_g"] = g + state[prefix + ".weight_v"] = v + else: + state[prefix + ".parametrizations.weight.original0"] = g + state[prefix + ".parametrizations.weight.original1"] = v + converted = module.preprocess_weights(state) + assert converted.keys() == dict(module.named_parameters()).keys() + for name in ( + "quantizer.semantic_quantizer.quantizers.0.out_proj.weight", + "decoder.model.0.conv.weight", + ): + tensor = converted[name] + torch.testing.assert_close( + torch.linalg.vector_norm(tensor, dim=(1, 2)), torch.full((tensor.shape[0],), 0.5) + ) + + +def test_loader_pins_revision_and_safe_checkpoint_options(): + with ( + patch("mobius.tasks.fish_codec.hf_hub_download", return_value="codec.pth") as download, + patch("mobius.tasks.fish_codec.torch.load", return_value={}) as load, + patch("mobius.tasks.fish_codec.build_fish_codec_decoder") as build, + ): + load_drama3_codec_decoder(cache_dir="cache") + assert download.call_args.kwargs["revision"] == DRAMA3_CODEC_REVISION + load.assert_called_once_with( + "codec.pth", map_location="cpu", mmap=True, weights_only=True + ) + build.assert_called_once_with(state_dict={}) + + +@pytest.mark.parametrize("value", [None, [], {"weight": "not a tensor"}]) +def test_loader_rejects_non_tensor_checkpoints(value): + with ( + patch("mobius.tasks.fish_codec.torch.load", return_value=value), + pytest.raises(ValueError, match="tensor state dictionary"), + ): + load_drama3_codec_decoder("caller-owned.pth") + + +@pytest.mark.parametrize( + "overrides", + [ + {"latent_dim": 15}, + {"transformer_heads": 0}, + {"decoder_dim": 15}, + {"transformer_window": 0}, + {"decoder_rates": ()}, + ], +) +def test_invalid_config(overrides): + with pytest.raises(ValueError): + dataclasses.replace(_tiny(), **overrides) diff --git a/src/mobius/models/fish_omni.py b/src/mobius/models/fish_omni.py new file mode 100644 index 000000000..6f9929125 --- /dev/null +++ b/src/mobius/models/fish_omni.py @@ -0,0 +1,964 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Drama 3 frame embeddings, hybrid text decoder, and temporal codebook decoder. + +Replicates Fish Speech's FishOmniBackbone and FishOmniAudioDecoder. Graphs +consume token/code frames rather than waveforms. The task attaches ONNX sampling +policies; the waveform codec is exported independently by mobius.tasks.fish_codec. +The caller still formats prompts and orchestrates temporal decoding. +""" + +from __future__ import annotations + +import dataclasses +import math +from typing import ClassVar + +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._build_context import ep_capabilities +from mobius._configs import ArchitectureConfig +from mobius._configs.fish_omni import FishOmniConfig +from mobius.components import ( + Embedding, + Float32SiLU, + Linear, + OffsetRMSNorm, + PostGatedRMSNorm, + Qwen35Attention, + SoftmaxTopKGate, + apply_rotary_pos_emb, + create_attention_bias, + float32_silu, + initialize_rope, +) +from mobius.components._attention import _apply_attention +from mobius.components._scan_utils import create_body_graph, rename_subgraph_values +from mobius.generation._fish_omni import ( + FishAudioSamplingInputs, + _all_rows, + _finite, + sample_codebook_or_greedy, +) +from mobius.models.qwen35 import ( + Qwen35MoEBlock, + Qwen35MoEDecoderLayer, + Qwen35MoETextModel, +) + + +def _float32_rope(config: ArchitectureConfig) -> nn.Module: + rope = initialize_rope(dataclasses.replace(config, dtype=ir.DataType.FLOAT)) + assert rope is not None + for parameter in rope.parameters(): + parameter._keep_float32 = True # type: ignore[attr-defined] + return rope + + +def _sigmoid_float32(op, value): + # Avoid runtime sigmoid approximations losing small negative-tail values + # before the source's single FP16/BF16 rounding boundary. + value_f32 = op.Cast(value, to=ir.DataType.FLOAT) + sigmoid = op.Div(1.0, op.Add(1.0, op.Exp(op.Neg(value_f32)))) + return op.CastLike(sigmoid, value) + + +def _apply_fish_attention( + op, query, key, value, bias, past_key, past_value, *, dtype, **kwargs +): + if dtype == ir.DataType.FLOAT or bias is None: + return _apply_attention(op, query, key, value, bias, past_key, past_value, **kwargs) + + # The last causal query can see every key. An all-zero last bias row + # therefore proves there is no padding; Attention supplies causality itself. + # Omitting that redundant mask enables the source's fused BF16 attention + # instead of rounding softmax probabilities in ORT's unfused fallback. + last_row = op.Cast(op.Gather(bias, -1, axis=-2), to=ir.DataType.FLOAT) + unpadded = op.Equal( + op.ReduceMax(op.Abs(last_row), keepdims=False), + op.CastLike(0.0, last_row), + ) + + def unmasked(branch_op): + return _apply_attention( + branch_op, query, key, value, None, past_key, past_value, **kwargs + ) + + def masked(branch_op): + return _apply_attention( + branch_op, query, key, value, bias, past_key, past_value, **kwargs + ) + + branches = [] + for name, body in (("unpadded_attention", unmasked), ("padded_attention", masked)): + op.builder.push_module(name) + branch = op.builder.subgraph( + body, + inputs=[], + outputs=[ + ir.Value(name=f"{name}_{suffix}") for suffix in ("output", "key", "value") + ], + name=name, + ) + op.builder.pop_module() + branches.append(branch) + return op.If(unpadded, then_branch=branches[0], else_branch=branches[1], _outputs=3) + + +def _rotate_float32(op, value, positions, *, heads, rotary_dim, interleaved): + rotated = apply_rotary_pos_emb( + op, + op.Cast(value, to=ir.DataType.FLOAT), + positions, + num_heads=heads, + rotary_embedding_dim=rotary_dim, + interleaved=interleaved, + ) + return op.CastLike(rotated, value) + + +class _FishGatedNorm(PostGatedRMSNorm): + """Normalize in FP32, round activation/scale, then apply the FP32 SiLU gate.""" + + def forward(self, op, hidden_states, gate): + hidden_f32 = op.Cast(hidden_states, to=ir.DataType.FLOAT) + unit_scale = op.Expand(op.Constant(value_float=1.0), [self.hidden_size]) + normalized = op.RMSNormalization( + hidden_f32, + unit_scale, + epsilon=self.variance_epsilon, + axis=-1, + ) + rounded = op.CastLike(normalized, hidden_states) + scaled = op.Mul(rounded, self.weight) + gated = op.Mul( + op.Cast(scaled, to=ir.DataType.FLOAT), + float32_silu(op, gate), + ) + return op.CastLike(gated, hidden_states) + + +class _FishAttention(Qwen35Attention): + """Gated GQA with FP32 partial rotation before rounding Q/K back.""" + + def __init__(self, config): + super().__init__(config) + self._compute_dtype = config.dtype + + def forward( + self, op, hidden_states, attention_bias, position_embeddings, past_key_value=None + ): + q_gate = op.Reshape( + self.q_proj(op, hidden_states), + [0, 0, self.num_attention_heads, self.head_dim * 2], + ) + query, gate = op.Split(q_gate, num_outputs=2, axis=-1, _outputs=2) + gate = op.Reshape(gate, [0, 0, -1]) + key = op.Reshape(self.k_proj(op, hidden_states), [0, 0, -1, self.head_dim]) + query = op.Reshape(self.q_norm(op, query), [0, 0, -1]) + key = op.Reshape(self.k_norm(op, key), [0, 0, -1]) + query = _rotate_float32( + op, + query, + position_embeddings, + heads=self.num_attention_heads, + rotary_dim=self.rotary_embedding_dim, + interleaved=self._rope_interleave, + ) + key = _rotate_float32( + op, + key, + position_embeddings, + heads=self.num_key_value_heads, + rotary_dim=self.rotary_embedding_dim, + interleaved=self._rope_interleave, + ) + attended, pk, pv = _apply_fish_attention( + op, + query, + key, + self.v_proj(op, hidden_states), + attention_bias, + past_key_value[0] if past_key_value is not None else None, + past_key_value[1] if past_key_value is not None else None, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + scale=self.scaling, + dtype=self._compute_dtype, + ) + return self.o_proj(op, op.Mul(attended, _sigmoid_float32(op, gate))), (pk, pv) + + +class _FishMoEBlock(Qwen35MoEBlock): + """Accumulate weighted expert outputs in FP32 before returning model dtype.""" + + def __init__(self, config, gate=None): + if gate is None: + assert config.num_local_experts is not None + assert config.num_experts_per_tok is not None + gate = _FishRouter( + config.hidden_size, + config.num_local_experts, + config.num_experts_per_tok, + norm_topk_prob=config.norm_topk_prob, + ) + super().__init__(config, gate=gate) + self.packed_experts = ( + _FishFP8Experts(config) if config.native_fp8_experts else _FishExperts(config) + ) + self.shared_expert.act_fn = Float32SiLU() + + def forward(self, op, hidden_states): + routing, selected, logits = self.gate(op, hidden_states) + routed = self.packed_experts(op, hidden_states, routing, selected, logits) + return self._add_shared_expert(op, hidden_states, routed) + + def _add_shared_expert(self, op, hidden_states, routed): + shared = self.shared_expert(op, hidden_states) + gate = _sigmoid_float32(op, self.shared_expert_gate(op, hidden_states)) + return op.Add(routed, op.Mul(shared, gate)) + + +class _FishExperts(nn.Module): + """Selected-only expert execution with bounded portable prefill scratch.""" + + def __init__(self, config): + super().__init__() + assert config.num_local_experts is not None + assert config.moe_intermediate_size is not None + assert config.num_experts_per_tok is not None + self.num_experts = config.num_local_experts + self.hidden_size = config.hidden_size + self.intermediate_size = config.moe_intermediate_size + self.top_k = config.num_experts_per_tok + self.normalize = int(config.norm_topk_prob) + self.dtype = config.dtype + self.act_fn = Float32SiLU() + self.gate_up_proj = _FishExpertBank( + [self.num_experts, 2 * self.intermediate_size, self.hidden_size] + ) + self.down_proj = _FishExpertBank( + [self.num_experts, self.hidden_size, self.intermediate_size] + ) + + def forward(self, op, hidden_states, weights, selected, router_logits): + # Materialize banks in the root module scope, never inside If/Scan. + caps = ep_capabilities() + if self.dtype == ir.DataType.FLOAT and caps.name == "cuda" and caps.supports_fused_moe: + return self._fused( + op, + hidden_states, + router_logits, + self.gate_up_proj(op), + self.down_proj(op), + ) + gate_up, down = self.gate_up_proj(op), self.down_proj(op) + hidden = op.Reshape(hidden_states, [-1, self.hidden_size]) + weights = op.Reshape(weights, [-1, self.top_k]) + selected = op.Reshape(selected, [-1, self.top_k]) + length = op.Shape(hidden, start=0, end=1) + count = op.Squeeze(length, [0]) + chunk_size = 16 + + def direct(branch_op): + return self._selected(branch_op, hidden, weights, selected, gate_up, down) + + def chunked(branch_op): + padding = branch_op.Mod(branch_op.Neg(count), chunk_size, fmod=0) + rows = branch_op.Unsqueeze(padding, [0]) + + def pad(value, width, zero): + tail = branch_op.Expand(zero, branch_op.Concat(rows, [width], axis=0)) + return branch_op.Reshape( + branch_op.Concat(value, tail, axis=0), [-1, chunk_size, width] + ) + + h_chunks = pad(hidden, self.hidden_size, branch_op.CastLike(0.0, hidden)) + w_chunks = pad(weights, self.top_k, branch_op.CastLike(0.0, weights)) + i_chunks = pad(selected, self.top_k, branch_op.Constant(value_int=0)) + assert gate_up.name is not None + prefix = gate_up.name + "_selected_chunks_" + h = ir.Value( + name=prefix + "hidden", + type=ir.TensorType(self.dtype), + shape=ir.Shape([chunk_size, self.hidden_size]), + ) + w = ir.Value( + name=prefix + "weights", + type=ir.TensorType(self.dtype), + shape=ir.Shape([chunk_size, self.top_k]), + ) + ids = ir.Value( + name=prefix + "ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape([chunk_size, self.top_k]), + ) + body, builder = create_body_graph([], [h, w, ids], name=prefix + "body") + output = self._selected(builder.op, h, w, ids, gate_up, down) + output.type = ir.TensorType(self.dtype) + output.shape = ir.Shape([chunk_size, self.hidden_size]) + output.name = prefix + "output" + body.outputs.append(output) + rename_subgraph_values(body, prefix) + scanned = branch_op.Scan( + h_chunks, w_chunks, i_chunks, body=body, num_scan_inputs=3 + ) + return branch_op.Slice( + branch_op.Reshape(scanned, [-1, self.hidden_size]), [0], length, [0] + ) + + branches = [] + for name, body in (("selected_direct", direct), ("selected_chunks", chunked)): + op.builder.push_module(name) + branch = op.builder.subgraph( + body, + inputs=[], + outputs=[ir.Value(name=name + "_output", type=ir.TensorType(self.dtype))], + name=name, + ) + op.builder.pop_module() + branches.append(branch) + result = op.If( + op.LessOrEqual(count, chunk_size), + then_branch=branches[0], + else_branch=branches[1], + ) + return op.Reshape(result, op.Shape(hidden_states)) + + def _selected(self, op, hidden, weights, selected, gate_up_bank, down_bank): + ids = op.Reshape(selected, [-1]) + pairs = op.Range(0, op.Size(ids), 1) + inputs = op.Unsqueeze(op.Gather(hidden, op.Div(pairs, self.top_k), axis=0), [-1]) + projected = op.Squeeze(op.MatMul(op.Gather(gate_up_bank, ids, axis=0), inputs), [-1]) + gate, up = op.Split( + projected, [self.intermediate_size, self.intermediate_size], axis=-1, _outputs=2 + ) + activated = op.Mul(self.act_fn(op, gate), up) + projected = op.Squeeze( + op.MatMul(op.Gather(down_bank, ids, axis=0), op.Unsqueeze(activated, [-1])), + [-1], + ) + # Both coefficient rounding and expert-output rounding precede FP32 weighting. + contribution = op.Mul( + op.Cast(projected, to=ir.DataType.FLOAT), + op.Cast(op.Reshape(weights, [-1, 1]), to=ir.DataType.FLOAT), + ) + summed = op.ReduceSum( + op.Reshape(contribution, [-1, self.top_k, self.hidden_size]), [1], keepdims=0 + ) + return op.CastLike(summed, hidden) + + def _fused(self, op, hidden_states, router_logits, gate_up_bank, down_bank): + # Legacy ABI shapes consume output-major bytes, not transposed storage. + # Prepack [gate | up] offline: passing FC3 would copy the complete + # 2 GiB expert bank on every layer invocation in the released model. + gate_up = op.Reshape( + gate_up_bank, [self.num_experts, self.hidden_size, 2 * self.intermediate_size] + ) + down = op.Reshape( + down_bank, [self.num_experts, self.intermediate_size, self.hidden_size] + ) + output = op.MoE( + op.Reshape(hidden_states, [-1, self.hidden_size]), + op.Reshape(router_logits, [-1, self.num_experts]), + gate_up, + None, + down, + None, + activation_type="swiglu", + activation_alpha=1.0, + activation_beta=0.0, + swiglu_limit=float("inf"), + k=self.top_k, + normalize_routing_weights=self.normalize, + swiglu_fusion=2, + _domain="com.microsoft", + ) + output.type = ir.TensorType(ir.DataType.FLOAT) + return op.Reshape(output, op.Shape(hidden_states)) + + +class _FishFP8ExpertBank(nn.Module): + def __init__(self, experts, rows, columns): + super().__init__() + self.weight = nn.Parameter([experts, rows, columns], dtype=ir.DataType.FLOAT8E4M3FN) + self.weight_scale_inv = nn.Parameter( + [experts, rows // 128, columns // 128], dtype=ir.DataType.FLOAT + ) + self.weight_scale_inv._keep_float32 = True # type: ignore[attr-defined] + + def forward(self, op): + return op.Identity(self.weight), op.Identity(self.weight_scale_inv) + + +class _FishFP8Experts(nn.Module): + """Source dynamic W8A8 with unchanged E4M3 bytes and multiplicative FP32 grids.""" + + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + self.dtype = config.dtype + self.gate_proj = _FishFP8ExpertBank( + config.num_local_experts, config.moe_intermediate_size, config.hidden_size + ) + self.up_proj = _FishFP8ExpertBank( + config.num_local_experts, config.moe_intermediate_size, config.hidden_size + ) + self.down_proj = _FishFP8ExpertBank( + config.num_local_experts, config.hidden_size, config.moe_intermediate_size + ) + + def forward(self, op, hidden_states, weights, selected, router_logits): + if ep_capabilities().name != "cuda": + raise ValueError("Drama 3 native dynamic FP8 requires execution_provider='cuda'") + gate, gate_scale = self.gate_proj(op) + up, up_scale = self.up_proj(op) + down, down_scale = self.down_proj(op) + # QMoE owns softmax/top-k and the source's coefficient narrowing. + # Unsupported devices/route counts must error in the opt-in runtime, + # never silently fall back to weight-only FP8 or dense expert scratch. + output = op.QMoE( + op.Reshape(hidden_states, [-1, self.hidden_size]), + op.Reshape(router_logits, [-1, self.num_experts]), + gate, + gate_scale, + None, + down, + down_scale, + None, + up, + up_scale, + quant_type="fp8", + expert_weight_bits=8, + block_size=128, + activation_quantization="dynamic_fp8", + activation_type="silu", + swiglu_fusion=0, + k=self.top_k, + normalize_routing_weights=1, + _domain="com.microsoft", + ) + output.type = ir.TensorType(self.dtype) + return op.Reshape(output, op.Shape(hidden_states)) + + +class _FishExpertBank(nn.Module): + def __init__(self, shape): + super().__init__() + self.weight = nn.Parameter(shape) + + def forward(self, op): + return op.Identity(self.weight) + + +class _FishRouter(SoftmaxTopKGate): + def forward(self, op: OpBuilder, hidden_states: ir.Value): + logits = op.MatMul(hidden_states, op.Transpose(self.weight, perm=[1, 0])) + probabilities = op.Softmax(op.Cast(logits, to=ir.DataType.FLOAT), axis=-1) + weights, indices = op.TopK( + probabilities, op.Constant(value_ints=[self.top_k]), axis=-1, _outputs=2 + ) + if self.norm_topk_prob: + weights = op.Div(weights, op.ReduceSum(weights, [-1], keepdims=True)) + return op.CastLike(weights, hidden_states), indices, logits + + +class _FishDecoderLayer(Qwen35MoEDecoderLayer): + def __init__(self, config, layer_idx): + super().__init__(config, layer_idx) + self._compute_dtype = config.dtype + self.mlp = _FishMoEBlock( + config, + gate=_FishRouter( + config.hidden_size, + config.num_local_experts, + config.num_experts_per_tok, + norm_topk_prob=config.norm_topk_prob, + ), + ) + if self.layer_type == "linear_attention": + self.linear_attn.norm = _FishGatedNorm( + config.linear_value_head_dim, + eps=config.rms_norm_eps, + ) + else: + self.self_attn = _FishAttention(config) + + def forward(self, op, hidden_states, attention_bias, position_embeddings, past_key_value): + if self._compute_dtype != ir.DataType.FLOAT16: + return super().forward( + op, hidden_states, attention_bias, position_embeddings, past_key_value + ) + # Native FP16 uses FP32 residuals, but both mixers consume rounded FP16 norms. + residual = op.Cast(hidden_states, to=ir.DataType.FLOAT) + normed = op.Cast(self.input_layernorm(op, residual), to=ir.DataType.FLOAT16) + if self.layer_type == "linear_attention": + attended, conv, state = self.linear_attn( + op, normed, past_key_value[0], past_key_value[1] + ) + present = (conv, state) + else: + attended, present = self.self_attn( + op, normed, attention_bias, position_embeddings, past_key_value + ) + residual = op.Add(residual, op.Cast(attended, to=ir.DataType.FLOAT)) + normed = op.Cast(self.post_attention_layernorm(op, residual), to=ir.DataType.FLOAT16) + output = self.mlp(op, normed) + return op.Add(residual, op.Cast(output, to=ir.DataType.FLOAT)), present + + +class _FishTextModel(Qwen35MoETextModel): + def __init__(self, config): + super().__init__(config) + self.layers = nn.ModuleList( + [_FishDecoderLayer(config, index) for index in range(config.num_hidden_layers)] + ) + self.rotary_emb = _float32_rope(config) + + +class FishOmniBackbone(nn.Module): + """Hybrid Qwen3.5-MoE decoder producing both logits and normalized hidden states.""" + + def __init__(self, config: FishOmniConfig): + super().__init__() + self.model = _FishTextModel(config) + # The released head exceeds ORT's 1 GiB constant-folding limit; + # retain output-major storage instead of transposing 2 GiB every step. + self.lm_head = _OutputMajorLinear(config.hidden_size, config.vocab_size) + self._compute_dtype = config.dtype + + def forward(self, op, inputs_embeds, attention_mask, position_ids, past_key_values): + if self._compute_dtype == ir.DataType.FLOAT16: + inputs_embeds = op.Cast(inputs_embeds, to=ir.DataType.FLOAT) + hidden, present = self.model( + op, + input_ids=None, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + if self._compute_dtype == ir.DataType.FLOAT16: + hidden = op.Cast(hidden, to=ir.DataType.FLOAT16) + return self.lm_head(op, hidden), hidden, present + + +class FishOmniEmbedding(nn.Module): + """Mix frame codes into text embeddings and expose unseen codes to the shared tower.""" + + def __init__(self, config: FishOmniConfig): + super().__init__() + assert config.tower is not None + self.config = config + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + self.codebook_embeddings = Embedding( + config.num_codebooks * config.tower.vocab_size, config.hidden_size + ) + + def forward(self, op: OpBuilder, frames: ir.Value): + cfg = self.config + assert cfg.tower is not None + tokens = op.Gather(frames, 0, axis=-1) + semantic = op.And( + op.GreaterOrEqual(tokens, cfg.semantic_start_token_id), + op.LessOrEqual(tokens, cfg.semantic_end_token_id), + ) + semantic = op.Unsqueeze(semantic, [-1]) + text = self.embed_tokens(op, tokens) # (B, T, H) + codebook_vocab = cfg.tower.vocab_size + + def code_sum(first: int, end: int): + if first == end: + return op.Mul(text, op.CastLike(0.0, text)) + codes = op.Slice(frames, [first + 1], [end + 1], [-1]) + offsets = op.Constant(value_ints=[i * codebook_vocab for i in range(first, end)]) + embeds = self.codebook_embeddings(op, op.Add(codes, offsets)) + return op.ReduceSum(embeds, [-2], keepdims=False) + + count = cfg.vq_backbone_codebooks + mixed = op.Mul( + op.Cast(op.Add(text, code_sum(1, count + 1)), to=ir.DataType.FLOAT), + 1.0 / math.sqrt(1 + count), + ) + mixed = op.CastLike(mixed, text) + unseen = code_sum(count + 1, cfg.num_codebooks) + zero = op.CastLike(0.0, text) + return op.Where(semantic, mixed, text), op.Where(semantic, unseen, zero) + + +class _TowerAttention(nn.Module): + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + assert config.partial_rotary_factor is not None + self.rotary_dim = int(config.head_dim * config.partial_rotary_factor) + self.q_width = config.num_attention_heads * config.head_dim + self.kv_width = config.num_key_value_heads * config.head_dim + self.wqkv = Linear(config.hidden_size, 2 * (self.q_width + self.kv_width), bias=False) + self.wo = Linear(self.q_width, config.hidden_size, bias=False) + # Released configs say attention_qk_norm=False, but the trained tower + # always uses zero-centered per-head Q/K norms (Fish Speech TowerAttention). + self.q_norm = OffsetRMSNorm(config.head_dim, eps=config.rms_norm_eps) + self.k_norm = OffsetRMSNorm(config.head_dim, eps=config.rms_norm_eps) + + def forward(self, op, hidden, bias, positions, past): + cfg = self.config + q, gate, k, v = op.Split( + self.wqkv(op, hidden), + [self.q_width, self.q_width, self.kv_width, self.kv_width], + axis=-1, + _outputs=4, + ) + q = self.q_norm(op, op.Reshape(q, [0, 0, -1, cfg.head_dim])) + k = self.k_norm(op, op.Reshape(k, [0, 0, -1, cfg.head_dim])) + q, k = op.Reshape(q, [0, 0, -1]), op.Reshape(k, [0, 0, -1]) + q = _rotate_float32( + op, + q, + positions, + heads=cfg.num_attention_heads, + rotary_dim=self.rotary_dim, + interleaved=True, + ) + k = _rotate_float32( + op, + k, + positions, + heads=cfg.num_key_value_heads, + rotary_dim=self.rotary_dim, + interleaved=True, + ) + attended, pk, pv = _apply_fish_attention( + op, + q, + k, + v, + bias, + past[0], + past[1], + num_attention_heads=cfg.num_attention_heads, + num_key_value_heads=cfg.num_key_value_heads, + scale=cfg.head_dim**-0.5, + dtype=cfg.dtype, + ) + return self.wo(op, op.Mul(attended, _sigmoid_float32(op, gate))), (pk, pv) + + +class _OutputMajorLinear(Linear): + """Keep the source's physical [out, in] weight layout in a 2-D GEMM.""" + + def __init__(self, in_features, out_features): + super().__init__(in_features, out_features, bias=False) + self.in_features, self.out_features = in_features, out_features + + def forward(self, op, hidden): + projected = self._project(op, op.Reshape(hidden, [-1, self.in_features])) + projected.shape = ir.Shape([None, self.out_features]) + shape = op.Concat( + op.Shape(hidden, end=-1), op.Constant(value_ints=[self.out_features]), axis=0 + ) + return op.Reshape(projected, shape) + + def _project(self, op, hidden): + # A folded Transpose + MatMul changes CUDA's accumulation order; + # a single BF16 rounding difference can amplify in the next tower layer. + return op.Gemm(hidden, self.weight, transB=1) + + +class _FP32OutputLinear(_OutputMajorLinear): + """Native FP16-input GEMM with FP32 output, without promoting the weight bank.""" + + def _project(self, op, hidden): + projected = op.GemmFloat8( + hidden, + self.weight, + transA=0, + transB=1, + alpha=1.0, + beta=0.0, + dtype=int(ir.DataType.FLOAT), + _domain="com.microsoft", + ) + projected.type = ir.TensorType(ir.DataType.FLOAT) + return projected + + +class _TowerMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.w1 = Linear(config.hidden_size, config.intermediate_size, bias=False) + self.w3 = Linear(config.hidden_size, config.intermediate_size, bias=False) + if config.dtype == ir.DataType.FLOAT16: + self.w2 = _FP32OutputLinear(config.intermediate_size, config.hidden_size) + elif config.dtype == ir.DataType.BFLOAT16: + self.w2 = _OutputMajorLinear(config.intermediate_size, config.hidden_size) + else: + self.w2 = Linear(config.intermediate_size, config.hidden_size, bias=False) + self.act_fn = Float32SiLU() + + def forward(self, op, hidden): + gate = self.w1(op, hidden) + return self.w2(op, op.Mul(self.act_fn(op, gate), self.w3(op, hidden))) + + +class _TowerTransformer(nn.Module): + def __init__(self, config): + super().__init__() + self.attention = _TowerAttention(config) + self.attention_norm = OffsetRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = OffsetRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.feed_forward = _TowerMLP(config) + self._compute_dtype = config.dtype + + def forward(self, op, hidden, bias, positions, past): + if self._compute_dtype == ir.DataType.FLOAT16: + hidden = op.Cast(hidden, to=ir.DataType.FLOAT) + normed = op.Cast(self.attention_norm(op, hidden), to=ir.DataType.FLOAT16) + attended, present = self.attention(op, normed, bias, positions, past) + residual = op.Add(hidden, op.Cast(attended, to=ir.DataType.FLOAT)) + normed = op.Cast(self.ffn_norm(op, residual), to=ir.DataType.FLOAT16) + return op.Add(residual, self.feed_forward(op, normed)), present + attended, present = self.attention( + op, self.attention_norm(op, hidden), bias, positions, past + ) + hidden = op.Add(hidden, attended) + return op.Add(hidden, self.feed_forward(op, self.ffn_norm(op, hidden))), present + + +class _TowerLayer(nn.Module): + def __init__(self, config, *, shared: bool): + super().__init__() + self.shared = shared + self.vocab_size = config.vocab_size + if shared: + self.semantic_embedding = Embedding(config.vocab_size, config.hidden_size) + else: + self.codebook_embedding = Embedding(config.vocab_size, config.hidden_size) + self.enorm = OffsetRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = OffsetRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = Linear(2 * config.hidden_size, config.hidden_size, bias=False) + self.transformer_layer = _TowerTransformer(config) + self.final_layernorm = OffsetRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, op, hidden, codes, semantic, bias, positions, past, shared_input=None): + table = self.semantic_embedding if self.shared else self.codebook_embedding + embedded = table(op, op.Clip(codes, 0, self.vocab_size - 1)) + embedded = op.Where(op.Unsqueeze(semantic, [-1]), embedded, op.CastLike(0.0, hidden)) + if shared_input is not None: + embedded = op.Add(embedded, shared_input) + mixed = self.eh_proj( + op, op.Concat(self.enorm(op, embedded), self.hnorm(op, hidden), axis=-1) + ) + output, present = self.transformer_layer(op, mixed, bias, positions, past) + return op.CastLike(self.final_layernorm(op, output), mixed), present + + +class FishOmniAudioDecoder(nn.Module): + """Teacher-forced temporal codebook tower; caches persist across audio frames.""" + + def __init__(self, config: FishOmniConfig): + super().__init__() + assert config.tower is not None + self.config = config + tower_config = dataclasses.replace(config.tower, dtype=config.dtype) + self.shared_layers = nn.ModuleList( + [ + _TowerLayer(tower_config, shared=True) + for _ in range(config.shared_enrichment_layers) + ] + ) + self.residual_layers = nn.ModuleList( + [_TowerLayer(tower_config, shared=False) for _ in range(config.num_codebooks - 1)] + ) + self.codebook_heads = nn.ModuleList( + [ + Linear(config.hidden_size, config.tower.vocab_size, bias=False) + for _ in range(config.num_codebooks - 1) + ] + ) + self.rotary_emb = _float32_rope(config.tower) + + def forward( + self, + op, + hidden_states, + codes, + semantic_mask, + shared_input, + attention_mask, + position_ids, + past_key_values, + ): + hidden, bias, positions, present = self._shared_forward( + op, + hidden_states, + codes, + semantic_mask, + shared_input, + attention_mask, + position_ids, + past_key_values, + ) + logits = [] + offset = len(self.shared_layers) + for i, (layer, head) in enumerate(zip(self.residual_layers, self.codebook_heads)): + hidden, kv = layer( + op, + hidden, + op.Gather(codes, i, axis=-1), + semantic_mask, + bias, + positions, + past_key_values[offset + i], + ) + present.append(kv) + logits.append(op.Unsqueeze(head(op, hidden), [-2])) + return op.Concat(*logits, axis=-2), present + + def _shared_forward( + self, + op, + hidden_states, + codes, + semantic_mask, + shared_input, + attention_mask, + position_ids, + past_key_values, + ): + positions = self.rotary_emb(op, position_ids) + bias = create_attention_bias( + op, input_ids=semantic_mask, attention_mask=attention_mask, dtype=self.config.dtype + ) + hidden = hidden_states + present = [] + for i, layer in enumerate(self.shared_layers): + hidden, kv = layer( + op, + hidden, + op.Gather(codes, 0, axis=-1), + semantic_mask, + bias, + positions, + past_key_values[i], + shared_input if i == 0 else None, + ) + present.append(kv) + return hidden, bias, positions, present + + def forward_sampled( + self, + op, + hidden_states, + codes, + semantic_mask, + shared_input, + attention_mask, + position_ids, + past_key_values, + sampling: FishAudioSamplingInputs, + ): + """Prefill teacher-forced rows or sample one frame in one tower traversal.""" + hidden, bias, positions, present = self._shared_forward( + op, + hidden_states, + codes, + semantic_mask, + shared_input, + attention_mask, + position_ids, + past_key_values, + ) + last_codes = op.Gather(codes, -1, axis=1) + sampled = [op.Gather(last_codes, 0, axis=-1)] + valid = op.Expand( + op.Or( + op.Not(sampling.generate), + op.Equal(op.Squeeze(op.Shape(hidden_states, start=1, end=2), [0]), 1), + ), + op.Shape(sampling.counter), + ) + counter = sampling.counter + logits = [] + offset = len(self.shared_layers) + for i, (layer, head) in enumerate(zip(self.residual_layers, self.codebook_heads)): + incoming = op.Where( + sampling.generate, + op.Unsqueeze(sampled[-1], [-1]), + op.Gather(codes, i, axis=-1), + ) + hidden, kv = layer( + op, + hidden, + incoming, + semantic_mask, + bias, + positions, + past_key_values[offset + i], + ) + present.append(kv) + scores = head(op, hidden) + logits.append(op.Unsqueeze(scores, [-2])) + last_scores = op.Gather(scores, -1, axis=1) + assert self.config.tower is not None + token, counter, sample_valid = sample_codebook_or_greedy( + op, + last_scores, + sampling, + counter, + index=i, + vocabulary=self.config.tower.vocab_size, + ) + finite = _all_rows( + op, + _finite(op, op.Reshape(op.Cast(scores, to=ir.DataType.FLOAT), [0, -1])), + ) + valid = op.And(valid, op.And(finite, sample_valid)) + sampled.append(token) + frame = op.Unsqueeze( + op.Concat(*[op.Unsqueeze(code, [-1]) for code in sampled], axis=-1), [1] + ) + frame = op.Where(sampling.generate, frame, op.Unsqueeze(last_codes, [1])) + frame = op.Where(op.Unsqueeze(valid, [-1, -2]), frame, -1) + return ( + op.Concat(*logits, axis=-2), + present, + frame, + op.Where(valid, counter, sampling.counter), + valid, + ) + + +class FishOmniModel(nn.Module): + """Fish Audio Drama 3 token-generation pipeline with temporal codebook MTP. + + Exports frame embedding, hybrid MoE decoder, and teacher-forced audio decoder + graphs with ONNX sampling policies. Reference-bracket prompt formatting and + temporal orchestration remain caller-owned; the waveform codec is exported + separately. FP8 experts require explicit streaming dequantization. + """ + + default_task: str = "fish-omni" + category: str = "Audio" + config_class = FishOmniConfig + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ("model.language_model", "lm_head"), + "embedding": ( + "model.language_model.embed_tokens", + "audio_decoder.codebook_embeddings", + ), + "audio_decoder": ("audio_decoder_mtp_shared", "audio_decoder_mtp"), + "audio_generator": ("audio_decoder_mtp_shared", "audio_decoder_mtp"), + } + + def __init__(self, config: FishOmniConfig): + super().__init__() + config.validate() + self.config = config + self.decoder = FishOmniBackbone(config) + self.embedding = FishOmniEmbedding(config) + self.audio_decoder = FishOmniAudioDecoder(config) diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py new file mode 100644 index 000000000..d1b613a19 --- /dev/null +++ b/src/mobius/models/fish_omni_test.py @@ -0,0 +1,737 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import copy +import dataclasses +import json + +import ml_dtypes +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch +from _test_configs import FISH_OMNI_CONFIG +from onnxscript import GraphBuilder, nn + +from mobius import build_from_module +from mobius._build_context import ep_capabilities +from mobius._configs.fish_omni import FishOmniConfig +from mobius._model_package import ModelPackage +from mobius._registry import registry +from mobius._testing.ort_inference import OnnxModelSession +from mobius.integrations.transformers._config_resolver import ( + _config_from_hf, + _dict_to_pretrained_config, +) +from mobius.models import FishOmniModel +from mobius.models.fish_omni import ( + _apply_fish_attention, + _FishMoEBlock, + _FP32OutputLinear, + _OutputMajorLinear, + _sigmoid_float32, +) +from mobius.tasks import FishOmniGenerationTask, FishOmniTask, ModelTask +from mobius.tasks._base import _make_graph, _make_model + + +def tiny_config(): + return FishOmniConfig.from_transformers( + _dict_to_pretrained_config(copy.deepcopy(FISH_OMNI_CONFIG)) + ) + + +class _ReducedPrecisionQualificationOnly: + def _validate_precision(self, config): + if not config.native_fp8_experts and config.dtype not in ( + ir.DataType.FLOAT16, + ir.DataType.BFLOAT16, + ): + return super()._validate_precision(config) + expected = ( + "cuda" + if config.dtype == ir.DataType.FLOAT16 or config.native_fp8_experts + else "onnx-standard" + ) + if ep_capabilities().name != expected: + raise ValueError("Qualification probe has the wrong execution provider") + + def build(self, module, config): + package = super().build(module, config) + for model in package.values(): + model.metadata_props["mobius.precision_validation"] = ( + "unqualified-native-fp8-prototype" + if config.native_fp8_experts + else "unqualified-reduced-precision-prototype" + ) + return package + + +class _TeacherProbe(_ReducedPrecisionQualificationOnly, FishOmniTask): + pass + + +class _GenerationProbe(_ReducedPrecisionQualificationOnly, FishOmniGenerationTask): + pass + + +def reduced_precision_prototype_task(task): + """Test only: probe prepared math without opening production admission.""" + return _GenerationProbe() if task == "fish-omni-generation" else _TeacherProbe() + + +def native_fp8_prototype_config(dtype=ir.DataType.FLOAT16): + config = tiny_config() + return dataclasses.replace( + config, + dtype=dtype, + hidden_size=128, + moe_intermediate_size=128, + fp8_experts=True, + native_fp8_experts=True, + tower=dataclasses.replace(config.tower, hidden_size=128), + ) + + +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +def test_native_fp8_prototype_preserves_storage_and_dynamic_runtime_contract(dtype): + config = native_fp8_prototype_config(dtype) + package = build_from_module( + FishOmniModel(config), + config, + task=_TeacherProbe(), + execution_provider="cuda", + ) + decoder = package["decoder"] + nodes = [node for node in decoder.graph if node.op_type == "QMoE"] + assert len(nodes) == config.num_hidden_layers + for node in nodes: + assert node.domain == "com.microsoft" + assert node.attributes["activation_quantization"].as_string() == "dynamic_fp8" + assert node.attributes["block_size"].as_int() == 128 + assert node.attributes["expert_weight_bits"].as_int() == 8 + assert node.attributes["swiglu_fusion"].as_int() == 0 + assert node.inputs[4] is None and node.inputs[7] is None + assert all(node.inputs[i].dtype == ir.DataType.FLOAT8E4M3FN for i in (2, 5, 8)) + assert all(node.inputs[i].dtype == ir.DataType.FLOAT for i in (3, 6, 9)) + assert node.inputs[1].dtype == dtype + banks = { + name: value + for name, value in decoder.graph.initializers.items() + if ".packed_experts." in name + } + assert len(banks) == config.num_hidden_layers * 6 + assert all( + value.dtype + == ( + ir.DataType.FLOAT + if name.endswith(".weight_scale_inv") + else ir.DataType.FLOAT8E4M3FN + ) + for name, value in banks.items() + ) + assert "dynamic_fp8" in decoder.metadata_props["mobius.required_runtime_features"] + assert decoder.metadata_props["mobius.precision_validation"].startswith("unqualified") + with pytest.raises(ValueError, match="disabled pending real-weight"): + build_from_module( + FishOmniModel(config), config, task="fish-omni", execution_provider="cuda" + ) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize( + "dtype,torch_dtype,numpy_dtype", + [ + (ir.DataType.FLOAT16, torch.float16, np.float16), + (ir.DataType.BFLOAT16, torch.bfloat16, ml_dtypes.bfloat16), + ], +) +def test_fish_sigmoid_preserves_native_negative_tail_rounding( + device, dtype, torch_dtype, numpy_dtype +): + if device == "cuda" and ( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers() + ): + pytest.skip("CUDA sigmoid parity requires Torch and ORT CUDA") + values = np.array( + [ + -1000.0, + -20.0, + -15.4453125, + -12.828125, + -12.46875, + -9.5234375, + -9.1328125, + -8.859375, + -8.3984375, + -8.3671875, + -7.6015625, + -7.08984375, + -5.54296875, + -3.875, + -2.443359375, + 0.0, + 20.0, + 1000.0, + ], + dtype=numpy_dtype, + ) + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + x = builder.input("x", dtype, [len(values)]) + y = _sigmoid_float32(builder.op, x) + y.type, y.shape = ir.TensorType(dtype), ir.Shape([len(values)]) + builder.add_output(y, "y") + assert not any(node.op_type == "Sigmoid" for node in graph) + reference = torch.tensor(values.astype(np.float32), device=device, dtype=torch_dtype) + expected = reference.sigmoid().float().cpu().numpy() + session = OnnxModelSession(ir.Model(graph, ir_version=12), device=device) + try: + actual = session.run({"x": values})["y"].astype(np.float32) + finally: + session.close() + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +def test_fish_omni_build(dtype): + config = dataclasses.replace(tiny_config(), dtype=dtype) + package = build_from_module( + FishOmniModel(config), + config, + task=reduced_precision_prototype_task("fish-omni") + if dtype != ir.DataType.FLOAT + else "fish-omni", + execution_provider={ + ir.DataType.FLOAT: "default", + ir.DataType.FLOAT16: "cuda", + ir.DataType.BFLOAT16: "onnx-standard", + }[dtype], + ) + assert set(package) == {"decoder", "embedding", "audio_decoder"} + assert set(package.policy_components) == {"main_sampler", "codebook_sampler"} + assert {v.name for v in package["embedding"].graph.outputs} == { + "inputs_embeds", + "shared_input", + } + decoder = package["decoder"] + head = decoder.graph.initializers["decoder.lm_head.weight"] + consumers = tuple(head.consumers()) + assert consumers + assert all(node.op_type == "Gemm" for node in consumers) + assert "past_key_values.0.recurrent_state" in {v.name for v in decoder.graph.inputs} + recurrent = next(v for v in decoder.graph.inputs if v.name.endswith("recurrent_state")) + assert recurrent.dtype == ir.DataType.FLOAT + assert {v.name for v in decoder.graph.outputs} >= {"logits", "hidden_states"} + tower = package["audio_decoder"] + if dtype == ir.DataType.BFLOAT16: + assert sum(node.op_type == "If" for node in tower.graph) == 3 + assert sum(node.op_type == "Gemm" for node in tower.graph) == 3 + assert tuple(tower.graph.outputs[0].shape)[-2:] == (2, 16) + assert sum(v.name.startswith("past_key_values.") for v in tower.graph.inputs) == 6 + assert any("q_norm.weight" in key for key in tower.graph.initializers) + assert all( + v.const_value is not None or v.dtype == dtype + for m in package.values() + for v in m.graph.initializers.values() + ) + rotary_caches = [ + value + for model in package.values() + for name, value in model.graph.initializers.items() + if name.endswith((".cos_cache", ".sin_cache")) + ] + assert rotary_caches + assert all(value.dtype == ir.DataType.FLOAT for value in rotary_caches) + + +def test_fish_omni_registry_config_resolution(): + parent = _dict_to_pretrained_config(copy.deepcopy(FISH_OMNI_CONFIG)) + resolved = _config_from_hf( + parent.text_config, parent_config=parent, module_class=registry.get("fish_omni") + ) + assert isinstance(resolved, FishOmniConfig) + assert resolved.partial_rotary_factor == resolved.tower.partial_rotary_factor + assert resolved.tower.partial_rotary_factor == pytest.approx(0.25) + assert resolved.tower.rope_interleave + assert resolved.tower.max_position_embeddings == resolved.max_position_embeddings + assert resolved.quantization is None + + +def test_fish_omni_cuda_build_uses_sparse_fused_expert_banks(): + config = tiny_config() + package = build_from_module( + FishOmniModel(config), config, task="fish-omni", execution_provider="cuda" + ) + graph = package["decoder"].graph + assert sum(node.op_type == "MoE" for node in graph) == config.num_hidden_layers + assert not any(node.op_type == "TopK" for node in graph) + banks = [value for name, value in graph.initializers.items() if ".packed_experts." in name] + assert len(banks) == config.num_hidden_layers * 2 + assert all(len(value.shape) == 3 for value in banks) + for node in graph: + if node.op_type == "MoE": + assert len(node.inputs) <= 6 or node.inputs[6] is None + assert node.attributes["activation_type"].as_string() == "swiglu" + assert node.attributes["swiglu_fusion"].as_int() == 2 + assert not any(".experts.0." in name for name in graph.initializers) + + +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +def test_fish_single_pass_task_reuses_tower_weights_and_layer_count(dtype): + config = dataclasses.replace(tiny_config(), dtype=dtype) + ep = { + ir.DataType.FLOAT: "default", + ir.DataType.FLOAT16: "cuda", + ir.DataType.BFLOAT16: "onnx-standard", + }[dtype] + teacher = build_from_module( + FishOmniModel(config), + config, + task=reduced_precision_prototype_task("fish-omni") + if dtype != ir.DataType.FLOAT + else "fish-omni", + execution_provider=ep, + )["audio_decoder"] + generated = build_from_module( + FishOmniModel(config), + config, + task=reduced_precision_prototype_task("fish-omni-generation") + if dtype != ir.DataType.FLOAT + else "fish-omni-generation", + execution_provider=ep, + )["audio_generator"] + assert { + name for name, value in teacher.graph.initializers.items() if value.const_value is None + } == { + name + for name, value in generated.graph.initializers.items() + if value.const_value is None + } + assert {value.name for value in generated.graph.outputs} >= { + "frame_codes", + "valid", + "next_counter", + "codebook_logits", + } + assert ( + sum(value.name.startswith("past_key_values.") for value in generated.graph.inputs) == 6 + ) + assert sum( + node.op_type == "Attention" + for node in ir.traversal.RecursiveGraphIterator(teacher.graph) + ) == sum( + node.op_type == "Attention" + for node in ir.traversal.RecursiveGraphIterator(generated.graph) + ) + + +@pytest.mark.parametrize("dtype", [ir.DataType.DOUBLE]) +def test_fish_omni_rejects_unqualified_precision(dtype): + with pytest.raises(ValueError, match="requires floating"): + FishOmniModel(dataclasses.replace(tiny_config(), dtype=dtype)) + + +@pytest.mark.parametrize( + "dtype,ep", + [(ir.DataType.BFLOAT16, "onnx-standard"), (ir.DataType.FLOAT16, "cuda")], +) +def test_fish_omni_reduced_precision_rejects_unqualified_ep(dtype, ep): + config = dataclasses.replace(tiny_config(), dtype=dtype) + with pytest.raises(ValueError, match=f"requires execution_provider='{ep}'"): + build_from_module(FishOmniModel(config), config, task="fish-omni") + + +@pytest.mark.parametrize( + "dtype,ep", + [(ir.DataType.FLOAT16, "cuda"), (ir.DataType.BFLOAT16, "onnx-standard")], +) +def test_reduced_precision_production_export_stays_closed_after_real_weight_failure(dtype, ep): + config = dataclasses.replace(tiny_config(), dtype=dtype) + with pytest.raises(ValueError, match="disabled pending real-weight numerical"): + build_from_module( + FishOmniModel(config), config, task="fish-omni", execution_provider=ep + ) + + +@pytest.mark.parametrize( + ("section", "key", "value"), + [ + ("text_config", "attn_output_gate", False), + ("text_config", "layer_types", ["full_attention"]), + ("text_config", "tie_word_embeddings", True), + ("audio_decoder_config", "vq_feature_placement", "other"), + ("audio_decoder_config", "dim", 64), + ("audio_decoder_config", "rotary_interleaved", False), + ("audio_decoder_config", "vocab_size", 17), + ], +) +def test_fish_omni_rejects_incompatible_configs(section, key, value): + raw = copy.deepcopy(FISH_OMNI_CONFIG) + raw[section][key] = value + with pytest.raises(ValueError): + FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)) + + +@pytest.mark.parametrize("task", ["fish-omni", "fish-omni-generation"]) +def test_fish_omni_ort_genai_metadata_rejects_incomplete_runtime(tmp_path, task): + from mobius.integrations.ort_genai import write_ort_genai_config + + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task=task) + output = tmp_path / "ort-genai" + with pytest.raises(ValueError, match="next-frame codebook sampling"): + write_ort_genai_config(package, str(output), hf_model_id="must-not-download") + assert not output.exists() + + +@pytest.mark.parametrize("task", ["fish-omni", "fish-omni-generation"]) +def test_fish_omni_onnx_genai_metadata_is_graph_derived_and_advisory(tmp_path, task): + from mobius.integrations.onnx_genai import write_onnx_genai_config + + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task=task) + artifacts = write_onnx_genai_config(package, str(tmp_path)) + metadata = json.loads((tmp_path / "inference_metadata.yaml").read_text()) + assert set(artifacts) == {"inference_metadata", "runtime_compatibility"} + assert metadata["runtime_validation_status"] == "unsupported-by-tested-runtime" + assert set(metadata["components"]) == set(package) + for name, model in package.items(): + component = metadata["components"][name] + assert component["filename"] == f"{name}/model.onnx" + assert component["inputs"] == [value.name for value in model.graph.inputs] + assert component["outputs"] == [value.name for value in model.graph.outputs] + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="BF16 Attention requires CUDA", +) +@pytest.mark.parametrize("padded", [False, True]) +def test_fish_attention_preserves_padding_and_cache(padded): + graph = ir.Graph([], [], nodes=[], opset_imports={"": 24}) + builder = GraphBuilder(graph) + dtype = ir.DataType.BFLOAT16 + inputs = { + name: builder.input(name, dtype, shape) + for name, shape in ( + ("query", [2, 2, 32]), + ("key", [2, 2, 16]), + ("value", [2, 2, 16]), + ("past_key", [2, 2, 2, 8]), + ("past_value", [2, 2, 2, 8]), + ) + } + bias = builder.input("bias", dtype, [2, 1, 2, 4]) + outputs = _apply_fish_attention( + builder.op, + inputs["query"], + inputs["key"], + inputs["value"], + bias, + inputs["past_key"], + inputs["past_value"], + dtype=dtype, + num_attention_heads=4, + num_key_value_heads=2, + scale=8**-0.5, + ) + for name, value, shape in zip( + ("output", "present_key", "present_value"), + outputs, + ([2, 2, 32], [2, 2, 4, 8], [2, 2, 4, 8]), + ): + value.type, value.shape = ir.TensorType(dtype), ir.Shape(shape) + builder.add_output(value, name) + rng = np.random.default_rng(42) + feeds = { + name: rng.normal(0, 0.1, tuple(value.shape)).astype(ml_dtypes.bfloat16) + for name, value in inputs.items() + } + mask = np.array([[0, 0, 0, -np.inf], [0, 0, 0, 0]], dtype=np.float32) + feeds["bias"] = np.broadcast_to(mask, (2, 1, 2, 4)).copy() + if padded: + feeds["bias"][0, :, :, 0] = -np.inf + feeds["bias"] = feeds["bias"].astype(ml_dtypes.bfloat16) + tensors = { + name: torch.tensor(value.astype(np.float32), device="cuda").to(torch.bfloat16) + for name, value in feeds.items() + if name != "bias" + } + q = tensors["query"].reshape(2, 2, 4, 8).transpose(1, 2) + k = torch.cat( + (tensors["past_key"], tensors["key"].reshape(2, 2, 2, 8).transpose(1, 2)), dim=2 + ) + v = torch.cat( + (tensors["past_value"], tensors["value"].reshape(2, 2, 2, 8).transpose(1, 2)), dim=2 + ) + expected = ( + torch.nn.functional.scaled_dot_product_attention( + q, + k.repeat_interleave(2, dim=1), + v.repeat_interleave(2, dim=1), + attn_mask=torch.tensor(np.isfinite(feeds["bias"]), device="cuda"), + ) + .transpose(1, 2) + .reshape(2, 2, 32) + ) + ort.preload_dlls() + session = OnnxModelSession(ir.Model(graph, ir_version=12), device="cuda") + try: + result = session.run(feeds) + np.testing.assert_allclose( + result["output"].astype(np.float32), + expected.float().cpu().numpy(), + atol=1e-2, + rtol=1e-2, + ) + np.testing.assert_array_equal( + result["present_key"].astype(np.float32), k.float().cpu().numpy() + ) + np.testing.assert_array_equal( + result["present_value"].astype(np.float32), v.float().cpu().numpy() + ) + if padded: + feeds["past_value"][0, :, 0] = 1000 + changed = session.run(feeds) + np.testing.assert_array_equal(result["output"], changed["output"]) + finally: + session.close() + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="Native BF16 projection layout parity requires CUDA", +) +def test_bf16_down_projection_preserves_native_output_major_accumulation(): + class ProjectionTask(ModelTask): + def build(self, module, config): + graph, builder = _make_graph("projection") + x = builder.input("x", config.dtype, [1, 3, 6144]) + output = module(builder.op, x) + output.type = ir.TensorType(config.dtype) + output.shape = ir.Shape([1, 3, 2048]) + builder.add_output(output, "output") + return ModelPackage({"projection": _make_model(graph)}) + + config = dataclasses.replace(tiny_config(), dtype=ir.DataType.BFLOAT16) + graph = build_from_module( + _OutputMajorLinear(6144, 2048), + config, + task=ProjectionTask(), + execution_provider="onnx-standard", + )["projection"] + gemm = next(node for node in graph.graph if node.op_type == "Gemm") + assert gemm.attributes["transB"].as_int() == 1 + assert not any(node.op_type == "Transpose" for node in graph.graph) + rng = torch.Generator(device="cuda").manual_seed(74) + weight = torch.randn(2048, 6144, dtype=torch.bfloat16, device="cuda", generator=rng) + x = torch.randn(3, 6144, dtype=torch.bfloat16, device="cuda", generator=rng) + expected = torch.nn.functional.linear(x, weight).float().cpu().numpy()[None] + parameter = next( + value for value in graph.graph.initializers.values() if value.const_value is None + ) + parameter.const_value = ir.tensor(weight.float().cpu().numpy().astype(ml_dtypes.bfloat16)) + session = OnnxModelSession(graph, device="cuda") + try: + actual = session.run({"x": x.float().cpu().numpy().astype(ml_dtypes.bfloat16)[None]})[ + "output" + ].astype(np.float32) + finally: + session.close() + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="Mixed-output FP16 GEMM requires CUDA", +) +def test_fp16_projection_preserves_values_above_half_range(): + class ProjectionTask(ModelTask): + def build(self, module, config): + graph, builder = _make_graph("projection") + x = builder.input("x", ir.DataType.FLOAT16, [1, 2, 4]) + output = module(builder.op, x) + output.type = ir.TensorType(ir.DataType.FLOAT) + output.shape = ir.Shape([1, 2, 2]) + builder.add_output(output, "output") + return ModelPackage({"projection": _make_model(graph)}) + + config = dataclasses.replace(tiny_config(), dtype=ir.DataType.FLOAT16) + graph = build_from_module( + _FP32OutputLinear(4, 2), config, task=ProjectionTask(), execution_provider="cuda" + )["projection"] + weights = [ + value for value in graph.graph.initializers.values() if value.const_value is None + ] + assert len(weights) == 1 + assert weights[0].dtype == ir.DataType.FLOAT16 + weights[0].const_value = ir.tensor(np.full((2, 4), 1000, dtype=np.float16)) + session = OnnxModelSession(graph, device="cuda") + try: + actual = session.run({"x": np.full((1, 2, 4), 1000, dtype=np.float16)})["output"] + finally: + session.close() + assert actual.dtype == np.float32 + np.testing.assert_array_equal(actual, np.full((1, 2, 2), 4_000_000, dtype=np.float32)) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="FP16 residual stress test requires CUDA", +) +def test_fp16_backbone_residual_stream_does_not_overflow(): + class ConstantMLP(nn.Module): + def forward(self, op, hidden): + return op.Expand(op.CastLike(60000.0, hidden), op.Shape(hidden)) + + class ConstantLinearAttention(ConstantMLP): + def forward(self, op, hidden, conv_state, recurrent_state): + return super().forward(op, hidden), op.Abs(conv_state), op.Abs(recurrent_state) + + class ConstantAttention(ConstantMLP): + def forward(self, op, hidden, bias, positions, past): + return super().forward(op, hidden), (op.Abs(past[0]), op.Abs(past[1])) + + config = dataclasses.replace(tiny_config(), dtype=ir.DataType.FLOAT16) + module = FishOmniModel(config) + for layer in module.decoder.model.layers: + layer.mlp = ConstantMLP() + if layer.layer_type == "linear_attention": + layer.linear_attn = ConstantLinearAttention() + else: + layer.self_attn = ConstantAttention() + graph = build_from_module( + module, + config, + task=reduced_precision_prototype_task("fish-omni"), + execution_provider="cuda", + )["decoder"] + for value in graph.graph.initializers.values(): + if value.const_value is None: + value.const_value = ir.tensor(np.zeros(tuple(value.shape), dtype=np.float16)) + cache = {} + for value in graph.graph.inputs: + if value.name.startswith("past_key_values."): + shape = list(value.shape) + shape[0] = 1 + if value.name.endswith((".key", ".value")): + shape[2] = 0 + dtype = np.float32 if value.dtype == ir.DataType.FLOAT else np.float16 + cache[value.name] = np.zeros(shape, dtype=dtype) + session = OnnxModelSession(graph, device="cuda") + try: + result = session.run( + { + **cache, + "inputs_embeds": np.full((1, 1, config.hidden_size), 60000, dtype=np.float16), + "attention_mask": np.ones((1, 1), np.int64), + "position_ids": np.zeros((1, 1), np.int64), + } + ) + finally: + session.close() + # Two 60000-valued sublayer outputs per layer exceed FP16 residual range. + # Normalization of the finite FP32 residual is exactly one after FP16 rounding. + np.testing.assert_array_equal( + result["hidden_states"], np.ones((1, 1, config.hidden_size), dtype=np.float16) + ) + + +class _MoEOnlyTask(ModelTask): + def build(self, module, config): + graph, builder = _make_graph("experts") + x = builder.input("x", config.dtype, ["tokens", config.hidden_size]) + output = module(builder.op, x) + output.type = ir.TensorType(config.dtype) + output.shape = ir.Shape(["tokens", config.hidden_size]) + builder.add_output(output, "output") + return ModelPackage({"experts": _make_model(graph)}) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="FP16 sparse expert overflow test requires CUDA", +) +def test_unselected_fp16_expert_overflow_cannot_poison_output(): + config = dataclasses.replace( + tiny_config(), dtype=ir.DataType.FLOAT16, num_local_experts=2, num_experts_per_tok=1 + ) + model = build_from_module( + _FishMoEBlock(config), config, task=_MoEOnlyTask(), execution_provider="onnx-standard" + )["experts"] + for name, value in model.graph.initializers.items(): + if value.const_value is not None: + continue + data = np.zeros(tuple(value.shape), np.float16) + if name == "gate.weight": + data[0, 0], data[1, 0] = 1, -1 + elif "packed_experts" in name: + # Expert 1's finite projections overflow if executed; it is never selected. + data[1] = 100 + value.const_value = ir.tensor(data) + session = OnnxModelSession(model, device="cuda") + try: + output = session.run({"x": np.ones((1, config.hidden_size), np.float16)})["output"] + finally: + session.close() + np.testing.assert_array_equal(output, np.zeros_like(output)) + + +@pytest.mark.parametrize("tokens", [1, 16, 17, 31, 32, 33]) +def test_sparse_expert_chunk_boundaries_match_selected_only_math(tokens): + config = tiny_config() + model = build_from_module( + _FishMoEBlock(config), config, task=_MoEOnlyTask(), execution_provider="onnx-standard" + )["experts"] + rng = np.random.default_rng(42) + weights = {} + for name, value in model.graph.initializers.items(): + if value.const_value is None: + data = rng.normal(0, 0.1, tuple(value.shape)).astype(np.float32) + value.const_value = ir.tensor(data) + weights[name] = torch.tensor(data) + x = torch.tensor(rng.normal(0, 0.1, (tokens, config.hidden_size)).astype(np.float32)) + probabilities = torch.nn.functional.linear(x, weights["gate.weight"]).softmax(-1) + routing, ids = probabilities.topk(config.num_experts_per_tok, dim=-1) + routing /= routing.sum(-1, keepdim=True) + expected = torch.zeros_like(x) + for row in range(tokens): + for rank in range(config.num_experts_per_tok): + expert = ids[row, rank] + gate, up = torch.nn.functional.linear( + x[row], weights["packed_experts.gate_up_proj.weight"][expert] + ).chunk(2, dim=-1) + projected = torch.nn.functional.linear( + torch.nn.functional.silu(gate) * up, + weights["packed_experts.down_proj.weight"][expert], + ) + expected[row] += routing[row, rank] * projected + shared = torch.nn.functional.silu( + torch.nn.functional.linear(x, weights["shared_expert.gate_proj.weight"]) + ) * torch.nn.functional.linear(x, weights["shared_expert.up_proj.weight"]) + shared = torch.nn.functional.linear(shared, weights["shared_expert.down_proj.weight"]) + expected += ( + shared * torch.nn.functional.linear(x, weights["shared_expert_gate.weight"]).sigmoid() + ) + session = OnnxModelSession(model) + try: + actual = session.run({"x": x.numpy()})["output"] + finally: + session.close() + np.testing.assert_allclose(actual, expected.numpy(), atol=1e-4, rtol=1e-4) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 72ca3bfd6..504ac0aa5 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -19,6 +19,8 @@ from __future__ import annotations __all__ = [ + "FishOmniTask", + "FishOmniGenerationTask", "AdapterTask", "AudioCTCTask", "AudioFeatureExtractionTask", @@ -152,6 +154,7 @@ GGUFEmbeddingFeatureExtractionTask, GGUFEncoderFeatureExtractionTask, ) +from mobius.tasks._fish_omni import FishOmniGenerationTask, FishOmniTask from mobius.tasks._fun_asr_speech_language import FunASRSpeechLanguageTask from mobius.tasks._gemma3n import Gemma3nTask from mobius.tasks._gemma4 import ( @@ -239,6 +242,8 @@ # --------------------------------------------------------------------------- TASK_REGISTRY: dict[str, type[ModelTask]] = { + "fish-omni": FishOmniTask, + "fish-omni-generation": FishOmniGenerationTask, "adapter": AdapterTask, "audio-ctc": AudioCTCTask, "audio-feature-extraction": AudioFeatureExtractionTask, diff --git a/src/mobius/tasks/_fish_omni.py b/src/mobius/tasks/_fish_omni.py new file mode 100644 index 000000000..e0adbfbf2 --- /dev/null +++ b/src/mobius/tasks/_fish_omni.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Drama 3 neural graphs and ONNX sampling policies.""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir + +from mobius._build_context import ep_capabilities +from mobius._configs.fish_omni import FishOmniConfig +from mobius._model_package import ModelPackage +from mobius.generation._fish_omni import ( + FishAudioSamplingInputs, + build_fish_codebook_sampler, + build_fish_main_sampler, +) +from mobius.tasks._base import ComponentSpec, ModelTask, _make_graph, _make_model +from mobius.tasks._cache_utils import ( + _make_hybrid_cache_inputs, + _make_kv_cache_inputs, + _register_hybrid_cache_outputs, + _register_kv_cache_outputs, + _register_linear_attention_functions, +) + + +class FishOmniTask(ModelTask): + """Drama 3 frame embedding, hybrid decoder, and temporal codebook tower.""" + + model_roles: ClassVar[dict[str, str]] = { + "embedding": "embedding", + "decoder": "decoder", + "audio_decoder": "decoder", + } + components: ClassVar[ComponentSpec] = ComponentSpec( + embedding="embedding", + decoder="decoder", + audio_decoder="audio_decoder", + ) + audio_model_name = "audio_decoder" + + def _sampling_inputs(self, builder) -> FishAudioSamplingInputs | None: + return None + + def _validate_precision(self, config: FishOmniConfig) -> None: + if config.dtype == ir.DataType.FLOAT16 and ep_capabilities().name != "cuda": + raise ValueError( + "Drama 3 FP16 currently requires execution_provider='cuda' " + "for the native FP16-input/FP32-output projection" + ) + if config.native_fp8_experts and ep_capabilities().name != "cuda": + raise ValueError( + "Drama 3 native FP8 requires execution_provider='cuda'; " + "use --dequantize for portable FP32." + ) + if config.native_fp8_experts and config.dtype == ir.DataType.FLOAT: + raise ValueError( + "Drama 3 native FP8 export is disabled pending real-weight numerical " + "qualification; use --dequantize for the qualified dense FP32 route." + ) + if ( + config.dtype == ir.DataType.BFLOAT16 + and not config.native_fp8_experts + and ep_capabilities().name != "onnx-standard" + ): + raise ValueError( + "Drama 3 BF16 currently requires execution_provider='onnx-standard' " + "and a runtime with BF16 CUDA Attention support" + ) + if config.dtype in (ir.DataType.FLOAT16, ir.DataType.BFLOAT16): + raise ValueError( + "Drama 3 reduced-precision export is disabled pending real-weight numerical " + "qualification. Production-dimension FP16 and BF16 reference probes exceed " + "the unchanged 1e-2 logit gate; use FP32." + ) + + def build(self, module, config: FishOmniConfig) -> ModelPackage: + self._validate_precision(config) + self._validate_components(module) + assert config.tower is not None + batch, seq = "batch", "sequence_len" + past = ir.SymbolicDim("past_sequence_len") + models = {} + graph, builder = _make_graph("embedding") + frames = builder.input( + "frames", dtype=ir.DataType.INT64, shape=[batch, seq, config.num_codebooks + 1] + ) + embeds, unseen = module.embedding(builder.op, frames) + builder.add_output(embeds, "inputs_embeds") + builder.add_output(unseen, "shared_input") + models["embedding"] = _make_model(graph) + + graph, builder = _make_graph("decoder") + embeds = builder.input( + "inputs_embeds", dtype=config.dtype, shape=[batch, seq, config.hidden_size] + ) + mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "total_sequence_len"] + ) + positions = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq]) + decoder_caches = _make_hybrid_cache_inputs( + builder, config, config.dtype, ir.SymbolicDim(batch), past + ) + logits, hidden, present = module.decoder( + builder.op, embeds, mask, positions, decoder_caches + ) + builder.add_output(logits, "logits") + builder.add_output(hidden, "hidden_states") + _register_hybrid_cache_outputs(builder, present, config.layer_types or []) + models["decoder"] = _make_model(graph) + _register_linear_attention_functions(models["decoder"], config) + + graph, builder = _make_graph(self.audio_model_name) + hidden = builder.input( + "hidden_states", dtype=config.dtype, shape=[batch, seq, config.hidden_size] + ) + codes = builder.input( + "codes", dtype=ir.DataType.INT64, shape=[batch, seq, config.num_codebooks] + ) + semantic = builder.input("semantic_mask", dtype=ir.DataType.BOOL, shape=[batch, seq]) + unseen = builder.input( + "shared_input", dtype=config.dtype, shape=[batch, seq, config.hidden_size] + ) + mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "total_sequence_len"] + ) + positions = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq]) + tower = config.tower + tower_caches = _make_kv_cache_inputs( + builder, + tower.num_hidden_layers, + tower.num_key_value_heads, + tower.head_dim, + config.dtype, + ir.SymbolicDim(batch), + past, + ) + sampling = self._sampling_inputs(builder) + if sampling is None: + logits, present = module.audio_decoder( + builder.op, hidden, codes, semantic, unseen, mask, positions, tower_caches + ) + else: + # A second method must retain __call__'s module scope without + # cloning any tower weights or placing parameters in subgraphs. + decoder = module.audio_decoder + builder.push_module(decoder.name or "audio_decoder", type(decoder).__qualname__) + try: + logits, present, frame, counter, valid = decoder.forward_sampled( + builder.op, + hidden, + codes, + semantic, + unseen, + mask, + positions, + tower_caches, + sampling, + ) + finally: + builder.pop_module() + frame.type = ir.TensorType(ir.DataType.INT64) + frame.shape = ir.Shape([batch, 1, config.num_codebooks]) + counter.type, counter.shape = ir.TensorType(ir.DataType.INT64), ir.Shape([batch]) + valid.type, valid.shape = ir.TensorType(ir.DataType.BOOL), ir.Shape([batch]) + builder.add_output(frame, "frame_codes") + builder.add_output(counter, "next_counter") + builder.add_output(valid, "valid") + builder.add_output(logits, "codebook_logits") + _register_kv_cache_outputs(builder, present) + models[self.audio_model_name] = _make_model(graph) + for model in models.values(): + model.metadata_props["mobius.fish_omni.scope"] = ( + "Neural token graphs with ONNX sampling policies; caller orchestrates " + "temporal decoding and prompt formatting; waveform decoder exported separately" + ) + if config.native_fp8_experts: + models["decoder"].metadata_props["mobius.required_runtime_features"] = ( + "CUDA SM89+; fp8-qmoe=1; QMoE activation_quantization=dynamic_fp8" + ) + assert config.num_experts_per_tok is not None + models["decoder"].metadata_props["mobius.native_fp8.max_tokens_per_call"] = str( + 65535 // config.num_experts_per_tok + ) + return ModelPackage( + models, + config=config, + policy_components={ + "main_sampler": build_fish_main_sampler(config), + "codebook_sampler": build_fish_codebook_sampler(config), + }, + ) + + +class FishOmniGenerationTask(FishOmniTask): + """Single-pass temporal codebook generation with shared teacher-forced prefill.""" + + audio_model_name = "audio_generator" + model_roles: ClassVar[dict[str, str]] = { + "embedding": "embedding", + "decoder": "decoder", + "audio_generator": "decoder", + } + components: ClassVar[ComponentSpec] = ComponentSpec( + embedding="embedding", + decoder="decoder", + audio_generator="audio_decoder", + ) + + def _sampling_inputs(self, builder) -> FishAudioSamplingInputs: + return FishAudioSamplingInputs( + builder.input("generate", ir.DataType.BOOL, []), + builder.input("do_sample", ir.DataType.BOOL, []), + builder.input("temperature", ir.DataType.FLOAT, ["batch"]), + builder.input("top_k", ir.DataType.INT64, ["batch"]), + builder.input("top_p", ir.DataType.FLOAT, ["batch"]), + builder.input("seed", ir.DataType.INT64, ["batch"]), + builder.input("counter", ir.DataType.INT64, ["batch"]), + ) diff --git a/src/mobius/tasks/fish_codec.py b/src/mobius/tasks/fish_codec.py new file mode 100644 index 000000000..9e371281b --- /dev/null +++ b/src/mobius/tasks/fish_codec.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Independent FP32 Drama 3 waveform graph/export integration helper.""" + +from __future__ import annotations + +from pathlib import Path + +import onnx_ir as ir +import torch +from huggingface_hub import hf_hub_download + +from mobius._configs.fish_codec import ( + DRAMA3_CODEC_REPO, + DRAMA3_CODEC_REVISION, + FishCodecConfig, +) +from mobius._model_package import ModelPackage +from mobius.integrations._weight_loading import apply_weights +from mobius.models.fish_codec import FishCodecDecoder +from mobius.tasks._base import _make_graph, _make_model + + +def build_fish_codec_decoder( + config: FishCodecConfig | None = None, + *, + state_dict: dict[str, torch.Tensor] | None = None, +) -> ModelPackage: + """Build ``pkg["waveform_decoder"]`` without downloads or tracing. + + Contract: ``codes`` INT64 [batch,10,frames] -> ``waveform`` FP32 + [batch,1,frames*2048], 44100 Hz. Batch and frame count are dynamic; + 1 <= frames <= config.max_frames. Codes are nonnegative. Values above + each native vocabulary are max-clamped exactly as upstream does. + A graph built without ``state_dict`` has unbound learned initializers. + Only FP32 is currently validated; no reduced-precision claim is implied. + """ + config = config or FishCodecConfig() + module = FishCodecDecoder(config) + graph, builder = _make_graph(name="fish_codec_decoder") + codes = builder.input( + "codes", dtype=ir.DataType.INT64, shape=["batch", config.num_codebooks, "frames"] + ) + waveform = module(builder.op, codes) + waveform.type = ir.TensorType(ir.DataType.FLOAT) + waveform.shape = ir.Shape(["batch", 1, "samples"]) + builder.add_output(waveform, "waveform") + model = _make_model(graph) + model.metadata_props.update( + { + "sample_rate": str(config.sample_rate), + "samples_per_frame": str(config.samples_per_frame), + "max_frames": str(config.max_frames), + "codec_repo": DRAMA3_CODEC_REPO, + "codec_revision": DRAMA3_CODEC_REVISION, + } + ) + if state_dict is not None: + apply_weights(model, module.preprocess_weights(state_dict)) + return ModelPackage({"waveform_decoder": model}) + + +def load_drama3_codec_decoder( + checkpoint_path: str | Path | None = None, + *, + cache_dir: str | Path | None = None, +) -> ModelPackage: + """Load the pinned codec (or a caller-owned local copy) safely, then build. + + The 1.87 GB checkpoint includes unused encoding weights. Loading is + CPU memory-mapped and always ``weights_only=True``; arbitrary pickle + globals are never allowlisted. Local paths are explicitly caller-owned. + Save the resulting package with ``pkg.save(output_directory)`` or add + its ``waveform_decoder`` graph to the parent Drama 3 package. + """ + if checkpoint_path is None: + checkpoint_path = hf_hub_download( + DRAMA3_CODEC_REPO, + "codec.pth", + revision=DRAMA3_CODEC_REVISION, + cache_dir=cache_dir, + ) + state_dict = torch.load(checkpoint_path, map_location="cpu", mmap=True, weights_only=True) + if not isinstance(state_dict, dict) or not all( + isinstance(key, str) and isinstance(value, torch.Tensor) + for key, value in state_dict.items() + ): + raise ValueError("Expected codec.pth to contain a tensor state dictionary") + return build_fish_codec_decoder(state_dict=state_dict) diff --git a/testdata/cases/speech/drama-3.yaml b/testdata/cases/speech/drama-3.yaml new file mode 100644 index 000000000..8f8ff6782 --- /dev/null +++ b/testdata/cases/speech/drama-3.yaml @@ -0,0 +1,11 @@ +model_id: "fishaudio/drama-3" +model_type: "fish_omni" +revision: "b2b5ee998034b2975720906c815e7fb0e9f007bd" +task_type: "text-to-speech" +dtype: "float32" +level: "L4+L5" +inputs: + prompts: + - "Hello, this is a speech synthesis test." +skip_reason: "The generic TTS runner does not orchestrate Drama 3 temporal caches and sequential codebooks; dedicated token, ONNX policy and codec tests cover the exported graphs." +notes: "The original pinned FP32 checkpoint passed H200 CUDA prefill and 20-frame greedy parity (scripts/validate_drama3.py and testdata/golden/speech/drama-3_tokens_generation.json). The separate FP32 ONNX codec matches all 40960 real-reference samples on CPU (max_abs 1.848e-6). FP16 and BF16 production exports are disabled after real-weight reduced fixtures failed the unchanged1e-2 logit gate, despite synthetic prototype20-frame parity. Full-checkpoint reduced precision and native FP8 remain unqualified." diff --git a/testdata/golden/speech/drama-3_tokens_generation.json b/testdata/golden/speech/drama-3_tokens_generation.json new file mode 100644 index 000000000..d8dfc2f92 --- /dev/null +++ b/testdata/golden/speech/drama-3_tokens_generation.json @@ -0,0 +1,48 @@ +{ + "format": "mobius.drama3-token-golden.v1", + "scope": "Token generation only; no waveform codec or speech-quality claim.", + "model_id": "fishaudio/drama-3", + "revision": "b2b5ee998034b2975720906c815e7fb0e9f007bd", + "source_revision": "541566932abbaebcb0a40e562a1f846738383f38", + "dtype": "float32", + "expert_execution": "dense", + "reference_environment": { + "torch": "2.8.0+cu126", + "torch_cuda": "12.6", + "triton": "3.4.0" + }, + "prompt": "Hello, this is a speech synthesis test.", + "prompt_rows": 30, + "prefill_shape": [1, 30, 252288], + "prefill_last_top10": [58, 249087, 249489, 251042, 249484, 252004, 249296, 248300, 248538, 251637], + "prefill_last_top10_logits": [16.343189239501953, 13.476046562194824, 13.093083381652832, 12.944562911987305, 12.827812194824219, 12.753364562988281, 12.615554809570312, 12.454527854919434, 12.443135261535645, 12.27248477935791], + "prefill_logits_sha256": "371c5178385eff1ee0955ec71ed848a440f2dab26f280af5d15e7b9bbfab28a9", + "max_new_tokens": 20, + "do_sample": false, + "sampling": "argmax", + "tf32": false, + "finish_reason": "length", + "attempted_frames": 20, + "generated_frames": [ + [249087, 988, 379, 113, 519, 944, 87, 105, 269, 1016, 780], + [251908, 3809, 589, 90, 209, 970, 624, 130, 678, 123, 640], + [250710, 2611, 117, 437, 624, 933, 66, 880, 484, 863, 283], + [249077, 978, 227, 134, 492, 244, 318, 615, 583, 777, 782], + [252091, 3992, 935, 184, 178, 713, 120, 616, 332, 332, 696], + [249862, 1763, 151, 360, 745, 933, 794, 379, 1012, 936, 169], + [251047, 2948, 1012, 865, 856, 715, 615, 75, 835, 332, 108], + [251154, 3055, 746, 320, 684, 37, 615, 903, 308, 423, 737], + [251223, 3124, 193, 437, 924, 52, 802, 925, 475, 818, 31], + [249178, 1079, 163, 528, 210, 425, 732, 722, 475, 408, 162], + [248637, 538, 225, 509, 176, 653, 8, 483, 890, 33, 256], + [249519, 1420, 225, 970, 777, 718, 495, 826, 900, 263, 7], + [250032, 1933, 168, 677, 379, 584, 277, 961, 307, 367, 707], + [250032, 1933, 168, 677, 379, 584, 277, 769, 307, 367, 413], + [250032, 1933, 168, 613, 351, 252, 277, 344, 359, 328, 258], + [250032, 1933, 168, 613, 351, 187, 909, 344, 359, 328, 258], + [250032, 1933, 168, 121, 351, 59, 277, 74, 989, 431, 258], + [250032, 1933, 168, 121, 351, 499, 277, 74, 346, 152, 258], + [250032, 1933, 168, 121, 544, 651, 978, 173, 353, 109, 326], + [250032, 1933, 888, 121, 930, 692, 272, 173, 1008, 440, 362] + ] +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 25efe9fac..edcb7a6e5 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -72,9 +72,69 @@ WhisperConfig, YolosConfig, Zamba2Config, + _as_attribute_config, ) +from mobius._configs.fish_omni import FishOmniConfig from mobius.models import EsmConfig, ReUseConfig +# Tiny composite config for Drama 3's model-local parity and speech graph matrix. +FISH_OMNI_CONFIG = { + "model_type": "fish_omni", + "architectures": ["FishOmniModel"], + "speech_template_version": "reference-bracket-v1", + "semantic_start_token_id": 16, + "semantic_end_token_id": 31, + "eos_token_id": 2, + "text_config": { + "model_type": "fish_omni_text", + "hidden_size": 32, + "intermediate_size": 48, + "num_hidden_layers": 2, + "layer_types": ["linear_attention", "full_attention"], + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "vocab_size": 128, + "hidden_act": "silu", + "rms_norm_eps": 1e-6, + "max_position_embeddings": 256, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10000000.0, + "partial_rotary_factor": 0.25, + }, + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "num_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 16, + "shared_expert_intermediate_size": 16, + "tie_word_embeddings": False, + }, + "audio_decoder_config": { + "model_type": "fish_qwen3_audio_decoder", + "dim": 32, + "text_dim": 32, + "n_layer": 1, + "shared_enrichment_layers": 1, + "intermediate_size": 48, + "n_head": 4, + "n_local_heads": 2, + "head_dim": 8, + "norm_eps": 1e-6, + "rope_base": 10000000.0, + "rotary_percent": 0.25, + "vocab_size": 16, + "num_codebooks": 3, + "vq_backbone_codebooks": 1, + "vq_feature_placement": "backbone_partial", + "attention_qk_norm": False, + }, +} + # --------------------------------------------------------------------------- # Tiny model dimensions shared by all configs # --------------------------------------------------------------------------- @@ -3535,6 +3595,14 @@ def vl_overrides(model_type: str) -> dict: # Speech / TTS / Codec configs # --------------------------------------------------------------------------- SPEECH_CONFIGS: list[tuple[str, dict, bool]] = [ + ( + "fish_omni", + { + "_config_cls": FishOmniConfig, + **vars(FishOmniConfig.from_transformers(_as_attribute_config(FISH_OMNI_CONFIG))), + }, + True, + ), # --- Parakeet CTC (feature-input offline FastConformer) --- ( "parakeet_ctc", diff --git a/tests/build_graph/_support.py b/tests/build_graph/_support.py index c6ef23f35..a31b2458c 100644 --- a/tests/build_graph/_support.py +++ b/tests/build_graph/_support.py @@ -245,6 +245,7 @@ def _make_params(configs: list[tuple[str, dict, bool]]) -> list: "wav2vec2-conformer", "wavlm", # Audio/TTS dedicated tests + "fish_omni", "fun_asr", "mms", "qwen3_asr", diff --git a/tests/build_graph/speech_test.py b/tests/build_graph/speech_test.py index c8f600dff..cd16cad5f 100644 --- a/tests/build_graph/speech_test.py +++ b/tests/build_graph/speech_test.py @@ -12,6 +12,7 @@ import onnx_ir as ir import pytest from _test_configs import ( + FISH_OMNI_CONFIG, SPEECH_CONFIGS, TINY_HEAD_DIM, TINY_HEADS, @@ -45,6 +46,7 @@ _SPEECH_MODEL_PARAMS = _make_params(SPEECH_CONFIGS) _SPEECH_TASK_KEYS: dict[str, set[str]] = { + "fish-omni": {"embedding", "decoder", "audio_decoder"}, "speech-to-text": {"encoder", "decoder"}, "speech-language": {"audio_encoder", "embedding", "decoder"}, "codec": {"decoder", "encoder"}, @@ -80,6 +82,19 @@ } +def test_fish_omni_package_builds(): + """Drama 3's dedicated ABI is not the generic input_ids/logits TTS matrix.""" + from mobius._configs.fish_omni import FishOmniConfig + from mobius.integrations.transformers._config_resolver import _dict_to_pretrained_config + + config = FishOmniConfig.from_transformers(_dict_to_pretrained_config(FISH_OMNI_CONFIG)) + module = registry.get("fish_omni")(config) + package = build_from_module(module, config, task="fish-omni") + assert set(package) == {"embedding", "decoder", "audio_decoder"} + _assert_outputs_have_shapes_and_dtypes(package, "fish_omni") + _run_onnx_checker(package, "fish_omni") + + class TestBuildGraphWhisper: """Verify Whisper encoder-decoder builds with SpeechToTextTask.""" @@ -1229,6 +1244,10 @@ def test_component_quantization_builds(self, model_type: str, config_overrides: config = _base_config(**config_overrides) task = get_task(_default_task_for_model(model_type)) config = _with_component_quantization(config, task) + if model_type == "fish_omni": + with pytest.raises(ValueError, match="affine component quantization"): + registry.get(model_type)(config) + return package = build_from_module( registry.get(model_type)(config), diff --git a/tests/integration/_fish_codec_reference.py b/tests/integration/_fish_codec_reference.py new file mode 100644 index 000000000..e660c155e --- /dev/null +++ b/tests/integration/_fish_codec_reference.py @@ -0,0 +1,147 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Test-only, immutable upstream codec reference with no optional audio imports. + +Only the unused AudioSignal/BaseModel/CodecMixin imports are stubbed. All +decoder, RVQ, weight-normalization, Snake, transformer and ConvNeXt math runs +unmodified upstream source. No source or checkpoint is redistributed. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +import urllib.request +from pathlib import Path +from unittest.mock import patch + +import torch + +from mobius._configs.fish_codec import FISH_CODEC_SOURCE_REVISION, FishCodecConfig + +_DAC_REVISION = "c7cfc5d2647e26471dc394f95846a0830e7bec34" + + +def load_reference_sources(directory: Path): + sources = { + "dac.nn.layers": ( + "descriptinc/descript-audio-codec", + _DAC_REVISION, + "dac/nn/layers.py", + ), + "dac.nn.quantize": ( + "descriptinc/descript-audio-codec", + _DAC_REVISION, + "dac/nn/quantize.py", + ), + "_mobius_fish_codec_dac": ( + "fishaudio/fish-speech", + FISH_CODEC_SOURCE_REVISION, + "fish_speech/models/dac/modded_dac.py", + ), + "_mobius_fish_codec_rvq": ( + "fishaudio/fish-speech", + FISH_CODEC_SOURCE_REVISION, + "fish_speech/models/dac/rvq.py", + ), + } + paths = {} + for name, (repo, revision, path) in sources.items(): + target = directory / (name.replace(".", "_") + ".py") + urllib.request.urlretrieve( + f"https://raw.githubusercontent.com/{repo}/{revision}/{path}", target + ) + paths[name] = target + + stubs = { + name: types.ModuleType(name) + for name in ("audiotools", "audiotools.ml", "dac.model.base") + } + stubs["audiotools"].AudioSignal = object + stubs["audiotools.ml"].BaseModel = torch.nn.Module + stubs["dac.model.base"].CodecMixin = type("UnusedCodecMixin", (), {}) + modules = {} + with patch.dict(sys.modules, stubs): + for name, path in paths.items(): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + # dataclasses and TorchScript need the real source-backed module. + with patch.dict(sys.modules, {**modules, name: module}): + spec.loader.exec_module(module) + modules[name] = module + return modules["_mobius_fish_codec_dac"], modules["_mobius_fish_codec_rvq"] + + +def make_reference(config: FishCodecConfig, sources): + dac, rvq = sources + + class Reference(torch.nn.Module): + def __init__(self): + super().__init__() + args = dac.ModelArgs( + dim=config.latent_dim, + n_head=config.transformer_heads, + n_local_heads=config.transformer_heads, + head_dim=config.head_dim, + n_layer=config.transformer_layers, + intermediate_size=config.transformer_intermediate_size, + norm_eps=config.norm_eps, + rope_base=config.rope_base, + ) + post = dac.WindowLimitedTransformer( + args, + input_dim=config.latent_dim, + window_size=config.transformer_window, + causal=True, + ) + self.quantizer = rvq.DownsampleResidualVectorQuantize( + input_dim=config.latent_dim, + n_codebooks=config.residual_codebooks, + codebook_size=config.residual_codebook_size, + semantic_codebook_size=config.semantic_codebook_size, + codebook_dim=config.codebook_dim, + downsample_factor=config.quantizer_rates, + post_module=post, + ) + self.decoder = dac.Decoder( + config.latent_dim, + config.decoder_dim, + config.decoder_rates, + causal=True, + n_transformer_layers=[0] * len(config.decoder_rates), + ) + + def forward(self, codes): + # Upstream mutates/clamps indices in place. Do not mutate test feeds. + return dac.DAC.from_indices(self, codes.clone()) + + # Avoid allocating the upstream 32768-square unused causal_mask. The + # WindowLimitedTransformer.forward builds its own actual attention mask. + with torch.device("meta"): + model = Reference() + model.quantizer.post_module.causal_mask = torch.empty(0, device="meta") + model.quantizer.post_module.freqs_cis = torch.empty(0, device="meta") + model.to_empty(device="cpu") + model.quantizer.post_module.freqs_cis = dac.precompute_freqs_cis( + config.max_frames, config.head_dim, config.rope_base + ) + return model.eval() + + +def initialize_reference(model): + """Deterministic, well-conditioned weights for production-topology tiny tests.""" + generator = torch.Generator().manual_seed(17) + with torch.no_grad(): + for name, param in model.named_parameters(): + if name.endswith(".alpha"): + param.uniform_(0.5, 1.5, generator=generator) + elif name.endswith("norm.weight"): + param.fill_(1) + elif name.endswith((".weight_g", ".original0")): + param.fill_(0.2) + else: + param.normal_(0, 0.05, generator=generator) + return model diff --git a/tests/integration/fish_codec_test.py b/tests/integration/fish_codec_test.py new file mode 100644 index 000000000..3c687fdd7 --- /dev/null +++ b/tests/integration/fish_codec_test.py @@ -0,0 +1,129 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unmodified-upstream and real-checkpoint tests for the pinned waveform codec.""" + +from __future__ import annotations + +import dataclasses +import json +import os +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import pytest +import torch +from huggingface_hub import hf_hub_download + +from mobius._configs.fish_codec import ( + DRAMA3_CODEC_REPO, + DRAMA3_CODEC_REVISION, + FishCodecConfig, +) +from mobius._testing.ort_inference import OnnxModelSession +from mobius.tasks.fish_codec import build_fish_codec_decoder, load_drama3_codec_decoder +from tests.integration._fish_codec_reference import ( + initialize_reference, + load_reference_sources, + make_reference, +) + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module") +def sources(tmp_path_factory): + pytest.importorskip("einops") + return load_reference_sources(tmp_path_factory.mktemp("fish-codec-source")) + + +@pytest.mark.parametrize("batch,frames", [(1, 1), (2, 7), (1, 131)]) +def test_tiny_upstream_waveform_parity(sources, batch, frames): + config = FishCodecConfig( + latent_dim=16, + transformer_layers=2, + transformer_heads=2, + transformer_intermediate_size=32, + max_frames=160, + decoder_dim=32, + decoder_rates=(2, 2), + ) + reference = initialize_reference(make_reference(config, sources)) + state = reference.state_dict() + package = build_fish_codec_decoder(config, state_dict=state) + # Include above-native IDs to exercise the exact upstream upper clamping. + codes = np.random.default_rng(42).integers( + 0, 4096, size=(batch, config.num_codebooks, frames), dtype=np.int64 + ) + with torch.inference_mode(): + expected = reference(torch.from_numpy(codes)).numpy() + session = OnnxModelSession(package["waveform_decoder"]) + actual = session.run({"codes": codes})["waveform"] + assert actual.shape == (batch, 1, frames * config.samples_per_frame) + np.testing.assert_allclose(actual, expected, atol=1e-5, rtol=1e-4) + if frames > config.transformer_window: + # Reuse the same graph with a different time dimension and verify the + # whole decode stack is causal, not just its transformer attention. + prefix = session.run({"codes": codes[:, :, :3]})["waveform"] + np.testing.assert_allclose( + prefix, actual[:, :, : 3 * config.samples_per_frame], atol=1e-5, rtol=1e-4 + ) + + +@pytest.mark.skipif( + os.environ.get("MOBIUS_TEST_DRAMA3_CODEC") != "1", + reason="Opt in to the 1.87GB pinned real-codec download", +) +def test_real_checkpoint_waveform_parity(sources, tmp_path): + """Bounded 20-frame decode of the independently generated token golden. + + Set MOBIUS_DRAMA3_CODEC_PATH to reuse a local pinned checkpoint copy. + The default downloads only codec.pth, never the large language weights. + """ + path = os.environ.get("MOBIUS_DRAMA3_CODEC_PATH") + if path is None: + path = hf_hub_download(DRAMA3_CODEC_REPO, "codec.pth", revision=DRAMA3_CODEC_REVISION) + state = torch.load(path, map_location="cpu", mmap=True, weights_only=True) + config = dataclasses.replace(FishCodecConfig(), max_frames=256) + reference = make_reference(config, sources) + loaded = reference.load_state_dict(state, strict=False, assign=True) + assert not loaded.missing_keys + golden_path = ( + Path(__file__).parents[2] + / "testdata" + / "golden" + / "speech" + / "drama-3_tokens_generation.json" + ) + golden = json.loads(golden_path.read_text()) + frames = np.asarray(golden["generated_frames"], dtype=np.int64) + codes = np.ascontiguousarray(frames[:20, 1:].T[None]) + assert codes.shape == (1, 10, 20) + with torch.inference_mode(): + expected = reference(torch.from_numpy(codes)).numpy() + # Exercise the production loader with its full-size default configuration, + # then the public export API and a fresh CPU ORT session on the saved file. + package = load_drama3_codec_decoder(path) + output_dir = tmp_path / "export" + package.save(str(output_dir), progress_bar=False) + options = ort.SessionOptions() + options.intra_op_num_threads = 4 + session = ort.InferenceSession( + str(output_dir / "model.onnx"), + sess_options=options, + providers=["CPUExecutionProvider"], + ) + actual = session.run(["waveform"], {"codes": codes})[0] + assert actual.shape == (1, 1, 40960) + assert actual.dtype == expected.dtype == np.float32 + assert np.isfinite(actual).all() + assert np.sqrt(np.mean(expected**2)) > 1e-3 + np.testing.assert_allclose(actual, expected, atol=1e-4, rtol=1e-3) + error = actual.astype(np.float64) - expected + rmse = float(np.sqrt(np.mean(error**2))) + snr = float(10 * np.log10(np.mean(expected.astype(np.float64) ** 2) / np.mean(error**2))) + print( + f"Drama3 codec: samples={actual.size}, rate={config.sample_rate}, " + f"max_abs={np.max(np.abs(error)):.9g}, rmse={rmse:.9g}, snr_db={snr:.4f}" + ) diff --git a/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py new file mode 100644 index 000000000..5a90b1540 --- /dev/null +++ b/tests/integration/fish_omni_test.py @@ -0,0 +1,779 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Reduced Drama 3 parity against the pinned Fish Speech inference implementation.""" + +from __future__ import annotations + +import copy +import dataclasses +import importlib.util +import json +import os +import sys +import types +import urllib.request +from pathlib import Path + +import ml_dtypes +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch +from _test_configs import FISH_OMNI_CONFIG +from safetensors.torch import save_file + +from mobius import build_from_module +from mobius._configs.fish_omni import DRAMA3_SOURCE_REVISION, FishOmniConfig +from mobius._model_package import ModelPackage +from mobius._testing.ort_inference import OnnxModelSession +from mobius.generation._fish_omni import build_fish_codebook_sampler +from mobius.integrations.transformers._config_resolver import _dict_to_pretrained_config +from mobius.integrations.transformers._fish_omni_weights import ( + stream_fish_omni_safetensors_to_package, +) +from mobius.models import FishOmniModel +from mobius.models.fish_omni import _FishFP8Experts +from mobius.models.fish_omni_test import reduced_precision_prototype_task +from mobius.tasks import ModelTask +from mobius.tasks._base import _make_graph, _make_model + +pytestmark = pytest.mark.integration +_PREFIX = "fish_speech.models.text2semantic.fish_omni" + + +@pytest.mark.skipif( + os.environ.get("MOBIUS_TEST_DRAMA3_NATIVE_FP8") != "1", + reason="Opt in with an FP8-QMoE-enabled runtime implementing dynamic_fp8", +) +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) +@pytest.mark.parametrize("length", [1, 19, 33]) +def test_native_fp8_exported_experts_match_pinned_w8a8_source(upstream_fp8, dtype, length): + from onnx_ir import tensor_adapters + + from mobius.models.fish_omni_test import native_fp8_prototype_config + + assert torch.cuda.is_available() + assert torch.cuda.get_device_capability() >= (8, 9) + assert "fp8-qmoe=1" in ort.get_build_info() + config = dataclasses.replace( + native_fp8_prototype_config(dtype), + hidden_size=256, + moe_intermediate_size=256, + num_local_experts=4, + num_experts_per_tok=2, + ) + + class ExpertsTask(ModelTask): + def build(self, module, config): + graph, builder = _make_graph("experts") + x = builder.input("x", config.dtype, [1, length, config.hidden_size]) + logits = builder.input("router_logits", config.dtype, [1, length, 4]) + result = module(builder.op, x, None, None, logits) + result.type, result.shape = x.type, x.shape + builder.add_output(result, "output") + return ModelPackage({"experts": _make_model(graph)}) + + model = build_from_module( + _FishFP8Experts(config), config, task=ExpertsTask(), execution_provider="cuda" + )["experts"] + torch.manual_seed(94) + compute_dtype = { + ir.DataType.FLOAT: torch.float32, + ir.DataType.FLOAT16: torch.float16, + ir.DataType.BFLOAT16: torch.bfloat16, + }[dtype] + weights = {} + for projection in ("gate_proj", "up_proj", "down_proj"): + weights[projection + ".weight"] = torch.randn(4, 256, 256, device="cuda").to( + torch.float8_e4m3fn + ) + weights[projection + ".weight_scale_inv"] = ( + torch.rand(4, 2, 2, device="cuda") * 0.06 + 0.02 + ) + for name, value in model.graph.initializers.items(): + if value.const_value is None: + value.const_value = tensor_adapters.TorchTensor(weights[name].cpu(), name=name) + x = torch.randn(length, 256, device="cuda", dtype=compute_dtype) + logits = torch.randn(length, 4, device="cuda", dtype=compute_dtype) + probabilities = logits.float().softmax(-1) + coefficients, selected = probabilities.topk(2, dim=-1) + coefficients = (coefficients / coefficients.sum(-1, keepdim=True)).to(compute_dtype) + gate_up = torch.cat( + [ + weights["gate_proj.weight"].view(torch.uint8), + weights["up_proj.weight"].view(torch.uint8), + ], + dim=1, + ).view(torch.float8_e4m3fn) + scales = torch.cat( + [ + weights["gate_proj.weight_scale_inv"], + weights["up_proj.weight_scale_inv"], + ], + dim=1, + ) + expected = upstream_fp8.fp8_moe_forward( + x, + selected, + coefficients, + gate_up, + scales, + weights["down_proj.weight"], + weights["down_proj.weight_scale_inv"], + ) + array_dtype = { + ir.DataType.FLOAT: np.float32, + ir.DataType.FLOAT16: np.float16, + ir.DataType.BFLOAT16: ml_dtypes.bfloat16, + }[dtype] + session = OnnxModelSession(model, device="cuda") + try: + actual = session.run( + { + "x": x.float().cpu().numpy().astype(array_dtype)[None], + "router_logits": logits.float().cpu().numpy().astype(array_dtype)[None], + } + )["output"] + finally: + session.close() + assert actual.dtype == np.dtype(array_dtype) + np.testing.assert_allclose( + actual.astype(np.float32)[0], + expected.float().cpu().numpy(), + atol=2e-6 if dtype == ir.DataType.FLOAT else 1e-2, + rtol=2e-5 if dtype == ir.DataType.FLOAT else 1e-2, + ) + + +def _load_reference_module(name, tmp_path, monkeypatch): + url = ( + f"https://raw.githubusercontent.com/fishaudio/fish-speech/" + f"{DRAMA3_SOURCE_REVISION}/fish_speech/models/text2semantic/fish_omni/{name}.py" + ) + with urllib.request.urlopen(url, timeout=60) as response: + path = tmp_path / f"{name}.py" + path.write_bytes(response.read()) + spec = importlib.util.spec_from_file_location(f"{_PREFIX}.{name}", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def upstream_fp8(tmp_path, monkeypatch): + return _load_reference_module("fp8_moe", tmp_path, monkeypatch) + + +@pytest.fixture +def upstream(tmp_path, monkeypatch): + """Load unmodified reference files in an isolated temporary package namespace.""" + for name in ( + "fish_speech", + "fish_speech.models", + "fish_speech.models.text2semantic", + _PREFIX, + ): + package = types.ModuleType(name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, name, package) + modules = {} + for name in ("fp8_moe", "config", "layers", "moe", "backbone", "tower"): + modules[name] = _load_reference_module(name, tmp_path, monkeypatch) + return modules + + +def _checkpoint(backbone, tower): + weights = {} + for name, value in backbone.state_dict().items(): + if name.startswith("model."): + name = "model.language_model." + name.removeprefix("model.") + if name.endswith(".experts.w13_weight"): + prefix = name.removesuffix("w13_weight") + for expert, pair in enumerate(value): + gate, up = pair.chunk(2) + weights[f"{prefix}{expert}.gate_proj.weight"] = gate.contiguous() + weights[f"{prefix}{expert}.up_proj.weight"] = up.contiguous() + elif name.endswith(".experts.w2_weight"): + prefix = name.removesuffix("w2_weight") + for expert, down in enumerate(value): + weights[f"{prefix}{expert}.down_proj.weight"] = down.contiguous() + else: + weights[name] = value.contiguous() + for name, value in tower.state_dict().items(): + for local, source in ( + ("codebook_embeddings.", "audio_decoder.codebook_embeddings."), + ("shared_layers.", "audio_decoder_mtp_shared.layers."), + ("residual_layers.", "audio_decoder_mtp.layers."), + ("codebook_heads.", "audio_decoder_mtp.heads."), + ): + if name.startswith(local): + weights[source + name.removeprefix(local)] = value.contiguous() + break + else: + raise AssertionError(name) + return weights + + +def _empty_cache(model): + result = {} + for value in model.graph.inputs: + if not value.name.startswith("past_key_values."): + continue + shape = list(value.shape) + shape[0] = 1 + if value.name.endswith((".key", ".value")): + shape[2] = 0 + assert all(isinstance(dim, int) for dim in shape), (value.name, shape) + dtype = { + ir.DataType.FLOAT: np.float32, + ir.DataType.FLOAT16: np.float16, + ir.DataType.BFLOAT16: ml_dtypes.bfloat16, + }[value.dtype] + result[value.name] = np.zeros(shape, dtype=dtype) + return result + + +def _next_cache(outputs): + return { + key.replace("present.", "past_key_values.", 1): value + for key, value in outputs.items() + if key.startswith("present.") + } + + +@pytest.mark.parametrize("backbone_codebooks", [0, 1]) +@pytest.mark.parametrize("task", ["fish-omni", "fish-omni-generation"]) +@pytest.mark.parametrize( + "dtype,execution_provider", + [ + (ir.DataType.FLOAT, "default"), + pytest.param( + ir.DataType.BFLOAT16, + "onnx-standard", + marks=pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="BF16 numerical parity requires PyTorch and ORT CUDA", + ), + ), + pytest.param( + ir.DataType.FLOAT, + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="Fused sparse experts require CUDA", + ), + ), + pytest.param( + ir.DataType.FLOAT16, + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="FP16 mixed-output projection requires CUDA", + ), + ), + ], +) +def test_fish_omni_pipeline_prefill_and_cached_decode( + upstream, + tmp_path, + monkeypatch, + backbone_codebooks, + dtype, + execution_provider, + task, + prefill_length=4, +): + raw = copy.deepcopy(FISH_OMNI_CONFIG) + if prefill_length > 4: + raw["text_config"]["max_position_embeddings"] = 64 + raw["audio_decoder_config"]["vq_backbone_codebooks"] = backbone_codebooks + reference_config = upstream["config"].FishOmniConfig.from_dict(raw) + torch.manual_seed(42) + tolerance = 1e-4 if dtype == ir.DataType.FLOAT else 1e-2 + torch_dtype = { + ir.DataType.FLOAT: torch.float32, + ir.DataType.FLOAT16: torch.float16, + ir.DataType.BFLOAT16: torch.bfloat16, + }[dtype] + device = "cuda:0" if dtype != ir.DataType.FLOAT or execution_provider == "cuda" else "cpu" + backbone = upstream["backbone"].FishOmniBackbone( + reference_config.text, dtype=torch_dtype, device=device + ) + tower = upstream["tower"].FishOmniAudioDecoder( + reference_config, dtype=torch_dtype, device=device + ) + for module in (backbone, tower): + for parameter in module.parameters(): + torch.nn.init.normal_(parameter, mean=0.0, std=0.1) + module.eval() + config = dataclasses.replace( + FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)), dtype=dtype + ) + package = build_from_module( + FishOmniModel(config), + config, + task=reduced_precision_prototype_task(task) if dtype != ir.DataType.FLOAT else task, + execution_provider=execution_provider, + ) + save_file(_checkpoint(backbone, tower), tmp_path / "model.safetensors") + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + output = tmp_path / "onnx" + package.save(str(output), max_workers=1, progress_bar=False) + package = ModelPackage.load(str(output)) + assert set(package.policy_components) == {"main_sampler", "codebook_sampler"} + if device != "cpu": + ort.preload_dlls() + sessions = { + name: OnnxModelSession( + model, + device="cpu" if device == "cpu" else "cuda", + enable_profiling=name == "decoder" and execution_provider == "cuda", + profile_file_prefix=str(tmp_path / f"{name}_profile"), + ) + for name, model in package.items() + } + fast = task == "fish-omni-generation" + audio_name = "audio_generator" if fast else "audio_decoder" + sampling = ( + { + "generate": np.array(False), + "do_sample": np.array(False), + "temperature": np.ones(1, np.float32), + "top_k": np.array([config.tower.vocab_size], np.int64), + "top_p": np.ones(1, np.float32), + "seed": np.array([42], np.int64), + "counter": np.zeros(1, np.int64), + } + if fast + else {} + ) + capacity = 64 if prefill_length > 4 else 32 + backbone.setup_caches(1, capacity) + tower.setup_caches(capacity) + caches = {name: _empty_cache(model) for name, model in package.items()} + frames = torch.tensor( + [[1, 0, 0, 0], [16, 0, 4, 7], [30, 14, 3, 5], [3, 0, 0, 0]], device=device + ) + next_frames = torch.tensor( + [[16, 0, 2, 8], [20, 4, 3, 1], [2, 0, 0, 0], [18, 2, 6, 9]], device=device + ) + if prefill_length != 4: + indices = torch.arange(prefill_length, device=device) % 4 + frames, next_frames = frames[indices], next_frames[indices] + start = 0 + sampled_logits = [] + + def greedy(logits, *args, **kwargs): + sampled_logits.append(logits.detach().float().cpu().numpy()) + return logits.argmax(dim=-1) + + monkeypatch.setattr( + upstream["tower"], + "sample_top_k_top_p", + greedy, + ) + try: + # Multi-token prompt followed by 20 reused-state steps exercises both + # DeltaNet state and all shared/residual tower caches across time. + with torch.no_grad(): + for length in [prefill_length] + [1] * 20: + positions = torch.arange(start, start + length, device=device) + positions_np = positions.cpu().numpy()[None] + if start: + frames = next_frames[-1:].clone() + embedded = sessions["embedding"].run({"frames": frames.cpu().numpy()[None]}) + ref_embed = tower.embed_backbone(frames, backbone.embed(frames[:, 0])) + np.testing.assert_allclose( + embedded["inputs_embeds"][0].astype(np.float32), + ref_embed.float().cpu().numpy(), + atol=tolerance, + rtol=tolerance, + ) + mask = np.ones((1, start + length), dtype=np.int64) + outputs = sessions["decoder"].run( + { + **caches["decoder"], + "inputs_embeds": embedded["inputs_embeds"], + "position_ids": positions_np, + "attention_mask": mask, + } + ) + hidden = backbone(inputs_embeds=ref_embed[None], input_pos=positions) + np.testing.assert_allclose( + outputs["hidden_states"].astype(np.float32), + hidden.float().cpu().numpy(), + atol=tolerance, + rtol=tolerance, + err_msg=f"decoder hidden states at position {start}", + ) + np.testing.assert_allclose( + outputs["logits"].astype(np.float32), + backbone.compute_logits(hidden).float().cpu().numpy(), + atol=tolerance, + rtol=tolerance, + err_msg=f"decoder logits at position {start}", + ) + if start: + candidates = [ + config.eos_token_id, + *range( + config.semantic_start_token_id, config.semantic_end_token_id + 1 + ), + ] + main = candidates[ + int(outputs["logits"][0, -1, candidates].astype(np.float32).argmax()) + ] + next_codes = np.zeros((1, 1, config.num_codebooks), dtype=np.int64) + next_codes[0, 0, 0] = np.clip( + main - config.semantic_start_token_id, 0, config.tower.vocab_size - 1 + ) + semantic = np.array( + [ + [ + config.semantic_start_token_id + <= main + <= config.semantic_end_token_id + ] + ] + ) + for i in range(1 if fast else config.num_codebooks - 1): + tower_output = sessions[audio_name].run( + { + **sampling, + **({"generate": np.array(True)} if fast else {}), + **caches[audio_name], + "hidden_states": outputs["hidden_states"], + "shared_input": embedded["shared_input"], + "codes": next_codes, + "semantic_mask": semantic, + "position_ids": positions_np, + "attention_mask": mask, + } + ) + if fast: + assert tower_output["valid"].all() + next_codes = tower_output["frame_codes"] + else: + next_codes[0, 0, i + 1] = ( + tower_output["codebook_logits"][0, 0, i] + .astype(np.float32) + .argmax() + ) + sampled_logits.clear() + next_frames = tower.decode_step( + backbone.compute_logits(hidden)[0], + hidden[0], + positions, + frames, + torch.zeros(1, 16, dtype=torch.int64, device=device), + temperature=1.0, + top_k=16, + top_p=1.0, + ) + np.testing.assert_array_equal( + np.concatenate((np.array([[[main]]]), next_codes), axis=-1)[0], + next_frames.cpu().numpy(), + ) + np.testing.assert_allclose( + tower_output["codebook_logits"].astype(np.float32), + np.stack(sampled_logits[2:], axis=-2)[None], + atol=tolerance, + rtol=tolerance, + ) + caches["decoder"] = _next_cache(outputs) + caches[audio_name] = _next_cache(tower_output) + start += length + continue + semantic = tower.is_semantic(next_frames[:, 0]) + tower_output = sessions[audio_name].run( + { + **sampling, + **caches[audio_name], + "hidden_states": outputs["hidden_states"], + "shared_input": embedded["shared_input"], + "codes": next_frames[:, 1:].cpu().numpy()[None], + "semantic_mask": semantic.cpu().numpy()[None], + "position_ids": positions_np, + "attention_mask": mask, + } + ) + h = tower.run_shared(hidden[0], next_frames[:, 1], semantic, positions, frames) + expected = [] + for i, (layer, head) in enumerate( + zip(tower.residual_layers, tower.codebook_heads) + ): + h = layer(h, next_frames[:, i + 1], semantic, positions) + expected.append(head(h)) + expected = torch.stack(expected, dim=-2)[None].float().cpu().numpy() + np.testing.assert_allclose( + tower_output["codebook_logits"].astype(np.float32), + expected, + atol=tolerance, + rtol=tolerance, + ) + caches["decoder"] = _next_cache(outputs) + caches[audio_name] = _next_cache(tower_output) + start += length + if execution_provider == "cuda" and dtype == ir.DataType.FLOAT: + profile = json.loads( + Path(sessions["decoder"]._session.end_profiling()).read_text() + ) + actual_moe = sum( + item.get("args", {}).get("op_name") == "MoE" + and item.get("args", {}).get("provider") == "CUDAExecutionProvider" + for item in profile + ) + assert actual_moe == 21 * config.num_hidden_layers + finally: + for session in sessions.values(): + session.close() + + +def test_sampled_tower_matches_sequential_replays_and_rejects_multiframe_generate( + upstream, tmp_path +): + raw = copy.deepcopy(FISH_OMNI_CONFIG) + native_config = upstream["config"].FishOmniConfig.from_dict(raw) + torch.manual_seed(42) + backbone = upstream["backbone"].FishOmniBackbone(native_config.text, dtype=torch.float32) + tower = upstream["tower"].FishOmniAudioDecoder(native_config, dtype=torch.float32) + for module in (backbone, tower): + for parameter in module.parameters(): + torch.nn.init.normal_(parameter, mean=0, std=0.1) + config = FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)) + save_file(_checkpoint(backbone, tower), tmp_path / "model.safetensors") + models = {} + for task, name in ( + ("fish-omni", "audio_decoder"), + ("fish-omni-generation", "audio_generator"), + ): + package = build_from_module(FishOmniModel(config), config, task=task) + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + models[name] = package[name] + sessions = { + name: OnnxModelSession( + model, + enable_profiling=name == "audio_generator", + profile_file_prefix=str(tmp_path / f"{name}_profile"), + ) + for name, model in models.items() + } + sampler = OnnxModelSession(build_fish_codebook_sampler(config).model) + rng = np.random.default_rng(42) + inputs = { + "hidden_states": rng.normal(size=(1, 1, config.hidden_size)).astype(np.float32), + "shared_input": rng.normal(size=(1, 1, config.hidden_size)).astype(np.float32), + "semantic_mask": np.ones((1, 1), bool), + "attention_mask": np.ones((1, 1), np.int64), + "position_ids": np.zeros((1, 1), np.int64), + } + codes = np.zeros((1, 1, config.num_codebooks), np.int64) + codes[0, 0, 0] = 3 + controls = { + "temperature": np.array([0.8], np.float32), + "top_k": np.array([4], np.int64), + "top_p": np.array([0.9], np.float32), + "seed": np.array([42], np.int64), + } + initial_counter = np.array([17], np.int64) + try: + expected_codes = codes.copy() + counter = initial_counter + for index in range(config.num_codebooks - 1): + expected = sessions["audio_decoder"].run( + {**_empty_cache(models["audio_decoder"]), **inputs, "codes": expected_codes} + ) + sample = sampler.run( + { + **controls, + "counter": counter, + "logits": expected["codebook_logits"][:, 0, index], + } + ) + expected_codes[0, 0, index + 1] = sample["token"][0] + counter = sample["next_counter"] + actual = sessions["audio_generator"].run( + { + **_empty_cache(models["audio_generator"]), + **inputs, + **controls, + "codes": codes, + "counter": initial_counter, + "generate": np.array(True), + "do_sample": np.array(True), + } + ) + assert actual["valid"].all() + np.testing.assert_array_equal(actual["frame_codes"], expected_codes) + np.testing.assert_array_equal(actual["next_counter"], counter) + for name, value in expected.items(): + np.testing.assert_allclose(actual[name], value, atol=1e-4, rtol=1e-4) + doubled = {key: np.repeat(value, 2, axis=1) for key, value in inputs.items()} + doubled["position_ids"] = np.array([[0, 1]], np.int64) + invalid = sessions["audio_generator"].run( + { + **_empty_cache(models["audio_generator"]), + **doubled, + **controls, + "codes": np.repeat(codes, 2, axis=1), + "counter": initial_counter, + "generate": np.array(True), + "do_sample": np.array(False), + } + ) + assert not invalid["valid"].any() + assert (invalid["frame_codes"] == -1).all() + np.testing.assert_array_equal(invalid["next_counter"], initial_counter) + prefill = sessions["audio_generator"].run( + { + **_empty_cache(models["audio_generator"]), + **doubled, + **controls, + "temperature": np.array([np.nan], np.float32), + "top_p": np.array([np.nan], np.float32), + "codes": np.repeat(codes, 2, axis=1), + "counter": initial_counter, + "generate": np.array(False), + "do_sample": np.array(True), + } + ) + assert prefill["valid"].all() + np.testing.assert_array_equal(prefill["next_counter"], initial_counter) + profile = json.loads( + Path(sessions["audio_generator"]._session.end_profiling()).read_text() + ) + assert ( + sum(item.get("args", {}).get("op_name") == "TopK" for item in profile) + == config.num_codebooks - 1 + ) + finally: + sampler.close() + for session in sessions.values(): + session.close() + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="Native mixed-output tower requires CUDA", +) +def test_fp16_tower_matches_source_when_down_projections_exceed_half_range(upstream, tmp_path): + raw = copy.deepcopy(FISH_OMNI_CONFIG) + native_config = upstream["config"].FishOmniConfig.from_dict(raw) + backbone = upstream["backbone"].FishOmniBackbone( + native_config.text, dtype=torch.float16, device="cuda:0" + ) + tower = upstream["tower"].FishOmniAudioDecoder( + native_config, dtype=torch.float16, device="cuda:0" + ) + with torch.no_grad(): + for module in (backbone, tower): + for parameter in module.parameters(): + parameter.zero_() + for layer in [*tower.shared_layers, *tower.residual_layers]: + table = layer.semantic_embedding if layer.shared else layer.codebook_embedding + table.weight.fill_(1) + layer.eh_proj.weight[:, : native_config.audio_decoder.dim].fill_( + 1 / native_config.audio_decoder.dim + ) + mlp = layer.transformer_layer.feed_forward + mlp.w1.weight.fill_(1 / native_config.audio_decoder.dim) + mlp.w3.weight.fill_(1 / native_config.audio_decoder.dim) + mlp.w2.weight.fill_(10000) + for head in tower.codebook_heads: + head.weight.fill_(0.01) + config = dataclasses.replace( + FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)), + dtype=ir.DataType.FLOAT16, + ) + package = build_from_module( + FishOmniModel(config), + config, + task=reduced_precision_prototype_task("fish-omni"), + execution_provider="cuda", + ) + save_file(_checkpoint(backbone, tower), tmp_path / "model.safetensors") + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + graph = package["audio_decoder"] + tower.setup_caches(8) + hidden = torch.ones((1, config.hidden_size), dtype=torch.float16, device="cuda") + codes = torch.tensor([[3, 4, 5]], dtype=torch.int64, device="cuda") + frames = torch.cat( + (torch.tensor([[config.semantic_start_token_id + 3]], device="cuda"), codes), dim=-1 + ) + positions = torch.zeros(1, dtype=torch.int64, device="cuda") + semantic = torch.ones(1, dtype=torch.bool, device="cuda") + projected = [] + handles = [ + layer.transformer_layer.feed_forward.w2.register_forward_hook( + lambda _module, _inputs, output: projected.append(output.detach()) + ) + for layer in [*tower.shared_layers, *tower.residual_layers] + ] + with torch.no_grad(): + h = tower.run_shared(hidden, codes[:, 0], semantic, positions, frames) + expected = [] + for index, (layer, head) in enumerate( + zip(tower.residual_layers, tower.codebook_heads) + ): + h = layer(h, codes[:, index], semantic, positions) + expected.append(head(h).float().cpu().numpy()) + for handle in handles: + handle.remove() + assert projected + assert all( + value.dtype == torch.float32 and value.max().item() > 65504 for value in projected + ) + session = OnnxModelSession(graph, device="cuda") + try: + actual = session.run( + { + **_empty_cache(graph), + "hidden_states": np.ones((1, 1, config.hidden_size), np.float16), + "codes": codes.cpu().numpy()[None], + "semantic_mask": np.ones((1, 1), bool), + "shared_input": np.zeros((1, 1, config.hidden_size), np.float16), + "attention_mask": np.ones((1, 1), np.int64), + "position_ids": np.zeros((1, 1), np.int64), + } + )["codebook_logits"] + finally: + session.close() + assert np.isfinite(actual).all() + np.testing.assert_allclose( + actual.astype(np.float32), np.stack(expected, axis=-2)[None], atol=1e-2, rtol=1e-2 + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="Chunked reduced-precision source parity requires CUDA", +) +@pytest.mark.parametrize( + "dtype,ep", + [(ir.DataType.FLOAT16, "cuda"), (ir.DataType.BFLOAT16, "onnx-standard")], +) +@pytest.mark.parametrize("prefill_length", [17, 32]) +def test_chunked_sparse_prefill_and_reused_state_match_source( + upstream, tmp_path, monkeypatch, dtype, ep, prefill_length +): + test_fish_omni_pipeline_prefill_and_cached_decode( + upstream, + tmp_path, + monkeypatch, + 1, + dtype, + ep, + "fish-omni-generation", + prefill_length=prefill_length, + )