From 2703b89bf49aac4f6f74c01dd83f3c6d9a8cfcdc Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 19:56:46 -0700 Subject: [PATCH 01/22] Add FP32 Drama 3 token graph export Export the Fish Omni frame embedding, hybrid MoE backbone, and temporal codebook tower as three ONNX graphs. Pin checkpoint/source revisions and stream block-FP8 expert reconstruction with FP32 scale arithmetic and transactional source coverage validation. Cover graph/config compatibility, streaming refusals, persisted CLI output, intrinsic runtime metadata, and reduced upstream prefill plus 20-frame cached greedy parity. Reject unqualified FP16/BF16 and native FP8 rather than advertising unsupported inference; keep the waveform codec and host sampler external. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/getting-started.md | 69 ++++ pyproject.toml | 1 + src/mobius/_configs/fish_omni.py | 200 +++++++++++ src/mobius/_registry.py | 9 + src/mobius/integrations/_weight_loading.py | 26 +- .../integrations/onnx_genai/auto_export.py | 11 + .../integrations/ort_genai/auto_export.py | 6 + .../integrations/transformers/_builder.py | 30 +- .../transformers/_fish_omni_weights.py | 122 +++++++ .../transformers/_fish_omni_weights_test.py | 193 +++++++++++ src/mobius/models/__init__.py | 2 + src/mobius/models/fish_omni.py | 327 ++++++++++++++++++ src/mobius/models/fish_omni_test.py | 118 +++++++ src/mobius/tasks/__init__.py | 3 + src/mobius/tasks/_fish_omni.py | 108 ++++++ testdata/cases/speech/drama-3.yaml | 11 + tests/_test_configs.py | 59 ++++ tests/build_graph/_support.py | 1 + tests/build_graph/speech_test.py | 14 + tests/integration/fish_omni_test.py | 289 ++++++++++++++++ 20 files changed, 1593 insertions(+), 6 deletions(-) create mode 100644 src/mobius/_configs/fish_omni.py create mode 100644 src/mobius/integrations/transformers/_fish_omni_weights.py create mode 100644 src/mobius/integrations/transformers/_fish_omni_weights_test.py create mode 100644 src/mobius/models/fish_omni.py create mode 100644 src/mobius/models/fish_omni_test.py create mode 100644 src/mobius/tasks/_fish_omni.py create mode 100644 testdata/cases/speech/drama-3.yaml create mode 100644 tests/integration/fish_omni_test.py diff --git a/docs/getting-started.md b/docs/getting-started.md index 2c6fdf539..7859e7d2b 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -94,6 +94,75 @@ 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. Native +`bsgemm-moe` execution is **not** claimed. Plan for roughly 144 GB of dense FP32 +weights, plus the source checkpoint and runtime allocations. Only FP32 is +currently accepted. FP16 is rejected because released residuals/down projections +exceed its range; BF16 is rejected because reduced upstream numerical parity +does not pass the `1e-2` gate. Tolerances are not relaxed to qualify it. + +| 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. 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. + +Sampling (including repetition-aware resampling), reference audio encoding, +the separately distributed `codec.pth`, waveform decoding, and ORT GenAI +generation integration are external. `--runtime onnx-genai` writes graph-derived +advisory component contracts, not an executable speech pipeline; +`--runtime ort-genai` is explicitly rejected. No full-checkpoint waveform or +generation-golden claim is made by this token-graph export. Fish Audio's +checkpoint/source license is separate from Mobius's MIT license. + ### 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/src/mobius/_configs/fish_omni.py b/src/mobius/_configs/fish_omni.py new file mode 100644 index 000000000..5eea3e739 --- /dev/null +++ b/src/mobius/_configs/fish_omni.py @@ -0,0 +1,200 @@ +# 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 separately distributed codec is external.""" + + 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 + + @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.dtype != ir.DataType.FLOAT: + raise ValueError( + "Drama 3 currently requires dtype='f32'. FP16 residuals/down projections " + "overflow; BF16 has not passed the upstream numerical parity gate." + ) + 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/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index 47aebecca..f966c5bf2 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -102,6 +102,10 @@ class StreamingWeightSource: mode: Literal["direct", "fp8_scalar", "fp8_block_128"] = "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 @dataclasses.dataclass(frozen=True) @@ -715,9 +719,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 tensor if tensor.dtype == target_dtype else tensor.to(target_dtype) if source.mode == "fp8_scalar": if not source_dtype.startswith("F8"): raise ValueError( @@ -1049,12 +1053,24 @@ def validate_source( 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 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 + ) 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, ) else: if len(expected_shape) != 3 or len(source.experts) != expected_shape[0]: 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..feaa3c66d 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -2034,6 +2034,12 @@ 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 host-owned next-frame codebook sampling, temporal " + "audio-decoder caches and an external waveform codec; an ONNX Runtime " + "GenAI pipeline configuration is not implemented." + ) 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..b00f8f5dd 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 != "fish-omni" and not isinstance(task, FishOmniTask): + raise ValueError( + "Drama 3 requires task='fish-omni'; 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: + raise ValueError( + "Drama 3 native bsgemm-moe FP8 execution is not supported. " + "Use keep_quantized=False (--dequantize) for streaming dense conversion." + ) + 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..79014c30a --- /dev/null +++ b/src/mobius/integrations/transformers/_fish_omni_weights.py @@ -0,0 +1,122 @@ +# 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 ( + StreamingWeightPlan, + StreamingWeightSource, + stream_preprocessed_safetensors_to_model, +) + + +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.") + 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 = {} + for name, value in initializers.items(): + if value.const_value is not None: + continue + source = _source_name(name) + if source not in index: + raise ValueError(f"Missing Drama 3 checkpoint tensor {source!r}") + source_dtype = index[source][2] + expert = ".mlp.experts." in 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") + targets[name] = StreamingWeightSource( + source, + mode="fp8_block_128", + scale_name=source.removesuffix(".weight") + ".weight_scale_inv", + reconstruction_dtype=torch.float32, + ) + else: + if source_dtype not in {"F32", "BF16", "F16"}: + raise ValueError(f"Unexpected Drama 3 weight dtype {source_dtype}: {source}") + targets[name] = StreamingWeightSource(source) + return StreamingWeightPlan( + targets=targets, + report={"architecture": "fish_omni", "codec_included": False}, + ) + + +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..96d9f7ad4 --- /dev/null +++ b/src/mobius/integrations/transformers/_fish_omni_weights_test.py @@ -0,0 +1,193 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import copy +import dataclasses +import json + +import onnx_ir as ir +import pytest +import torch +from safetensors.torch import 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 tiny_config +from tests._test_configs import FISH_OMNI_CONFIG + + +def _checkpoint(package): + torch.manual_seed(19) + return { + _source_name(name): torch.randn(tuple(value.shape), dtype=torch.float32) * 0.02 + for model in package.values() + for name, value in model.graph.initializers.items() + if value.const_value is None + } + + +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 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) + + +@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") + + +def test_fish_omni_cli_persists_three_weighted_models(tmp_path): + 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", + "--dequantize", + "--output", + str(output), + ] + ) + for component in ("embedding", "decoder", "audio_decoder"): + 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 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_omni.py b/src/mobius/models/fish_omni.py new file mode 100644 index 000000000..752dd3f03 --- /dev/null +++ b/src/mobius/models/fish_omni.py @@ -0,0 +1,327 @@ +# 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; prompt formatting, sampling, +and the separately distributed audio codec remain caller responsibilities. +""" + +from __future__ import annotations + +import math +from typing import ClassVar + +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig +from mobius._configs.fish_omni import FishOmniConfig +from mobius.components import ( + Embedding, + Linear, + OffsetRMSNorm, + SoftmaxTopKGate, + apply_rotary_pos_emb, + create_attention_bias, + initialize_rope, +) +from mobius.components._attention import _apply_attention +from mobius.models.qwen35 import Qwen35MoEDecoderLayer, Qwen35MoETextModel + + +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 + + +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 = Qwen35MoETextModel(config) + assert config.num_local_experts is not None and config.num_experts_per_tok is not None + for layer in self.model.layers: + assert isinstance(layer, Qwen35MoEDecoderLayer) + layer.mlp.gate = _FishRouter( + config.hidden_size, + config.num_local_experts, + config.num_experts_per_tok, + norm_topk_prob=config.norm_topk_prob, + ) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward(self, op, inputs_embeds, attention_mask, position_ids, past_key_values): + 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, + ) + 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.Div( + op.Add(text, code_sum(1, count + 1)), op.CastLike(math.sqrt(1 + count), 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 = apply_rotary_pos_emb( + op, + q, + positions, + num_heads=cfg.num_attention_heads, + rotary_embedding_dim=self.rotary_dim, + interleaved=True, + ) + k = apply_rotary_pos_emb( + op, + k, + positions, + num_heads=cfg.num_key_value_heads, + rotary_embedding_dim=self.rotary_dim, + interleaved=True, + ) + attended, pk, pv = _apply_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, + ) + return self.wo(op, op.Mul(attended, op.Sigmoid(gate))), (pk, pv) + + +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) + self.w2 = Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, op, hidden): + gate = self.w1(op, hidden) + return self.w2(op, op.Mul(op.Mul(gate, op.Sigmoid(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) + + def forward(self, op, hidden, bias, positions, past): + 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 self.final_layernorm(op, output), 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 + self.shared_layers = nn.ModuleList( + [ + _TowerLayer(config.tower, shared=True) + for _ in range(config.shared_enrichment_layers) + ] + ) + self.residual_layers = nn.ModuleList( + [_TowerLayer(config.tower, 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) + ] + ) + rotary_emb = initialize_rope(config.tower) + assert rotary_emb is not None + self.rotary_emb = rotary_emb + + def 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) + 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 + + +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. Sampling, reference-bracket prompt formatting, and waveform codec + execution are external. 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"), + } + + 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..1fafc4ff1 --- /dev/null +++ b/src/mobius/models/fish_omni_test.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import copy +import dataclasses +import json + +import onnx_ir as ir +import pytest + +from mobius import build_from_module +from mobius._configs.fish_omni import FishOmniConfig +from mobius._registry import registry +from mobius.integrations.transformers._config_resolver import ( + _config_from_hf, + _dict_to_pretrained_config, +) +from mobius.models import FishOmniModel +from tests._test_configs import FISH_OMNI_CONFIG + + +def tiny_config(): + return FishOmniConfig.from_transformers( + _dict_to_pretrained_config(copy.deepcopy(FISH_OMNI_CONFIG)) + ) + + +def test_fish_omni_build(): + dtype = ir.DataType.FLOAT + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task="fish-omni") + assert set(package) == {"decoder", "embedding", "audio_decoder"} + assert {v.name for v in package["embedding"].graph.outputs} == { + "inputs_embeds", + "shared_input", + } + decoder = package["decoder"] + 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"] + 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() + ) + + +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 + + +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_fish_omni_rejects_unqualified_precision(dtype): + with pytest.raises(ValueError, match="requires dtype='f32'"): + FishOmniModel(dataclasses.replace(tiny_config(), dtype=dtype)) + + +@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)) + + +def test_fish_omni_ort_genai_metadata_rejects_incomplete_runtime(tmp_path): + from mobius.integrations.ort_genai import write_ort_genai_config + + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task="fish-omni") + 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() + + +def test_fish_omni_onnx_genai_metadata_is_graph_derived_and_advisory(tmp_path): + from mobius.integrations.onnx_genai import write_onnx_genai_config + + config = tiny_config() + package = build_from_module(FishOmniModel(config), config, task="fish-omni") + 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] diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 72ca3bfd6..c8601ab5e 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -19,6 +19,7 @@ from __future__ import annotations __all__ = [ + "FishOmniTask", "AdapterTask", "AudioCTCTask", "AudioFeatureExtractionTask", @@ -152,6 +153,7 @@ GGUFEmbeddingFeatureExtractionTask, GGUFEncoderFeatureExtractionTask, ) +from mobius.tasks._fish_omni import FishOmniTask from mobius.tasks._fun_asr_speech_language import FunASRSpeechLanguageTask from mobius.tasks._gemma3n import Gemma3nTask from mobius.tasks._gemma4 import ( @@ -239,6 +241,7 @@ # --------------------------------------------------------------------------- TASK_REGISTRY: dict[str, type[ModelTask]] = { + "fish-omni": FishOmniTask, "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..2c8cf0c11 --- /dev/null +++ b/src/mobius/tasks/_fish_omni.py @@ -0,0 +1,108 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Three-graph Drama 3 token-generation ABI.""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir + +from mobius._configs.fish_omni import FishOmniConfig +from mobius._model_package import ModelPackage +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", + ) + + def build(self, module, config: FishOmniConfig) -> ModelPackage: + 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("audio_decoder") + 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, + ) + logits, present = module.audio_decoder( + builder.op, hidden, codes, semantic, unseen, mask, positions, tower_caches + ) + builder.add_output(logits, "codebook_logits") + _register_kv_cache_outputs(builder, present) + models["audio_decoder"] = _make_model(graph) + for model in models.values(): + model.metadata_props["mobius.fish_omni.scope"] = ( + "Token generation only; external sampling, prompt formatting and waveform codec" + ) + return ModelPackage(models, config=config) diff --git a/testdata/cases/speech/drama-3.yaml b/testdata/cases/speech/drama-3.yaml new file mode 100644 index 000000000..b27a6c5a5 --- /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: "Token-graph export only: full-checkpoint waveform and generation goldens require the external Fish Speech codec and sampler." +notes: "Pinned L2 model is registered; reduced upstream prefill and temporal-cache parity is covered in tests/integration/fish_omni_test.py. No fabricated L4/L5 goldens." diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 25efe9fac..6b300516f 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -75,6 +75,65 @@ ) from mobius.models import EsmConfig, ReUseConfig +# Drama 3 has a temporal codebook ABI and is covered by model-local tests, +# rather than the generic input_ids/logits causal-LM 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 # --------------------------------------------------------------------------- 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..281d1e15a 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, @@ -80,6 +81,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.""" diff --git a/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py new file mode 100644 index 000000000..976f45967 --- /dev/null +++ b/tests/integration/fish_omni_test.py @@ -0,0 +1,289 @@ +# 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 importlib.util +import sys +import types +import urllib.request + +import numpy as np +import pytest +import torch +from safetensors.torch import save_file + +from mobius import build_from_module +from mobius._configs.fish_omni import DRAMA3_SOURCE_REVISION, FishOmniConfig +from mobius._testing.ort_inference import OnnxModelSession +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 tests._test_configs import FISH_OMNI_CONFIG + +pytestmark = pytest.mark.integration +_PREFIX = "fish_speech.models.text2semantic.fish_omni" + + +@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"): + 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) + modules[name] = module + 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 = [ + dim if isinstance(dim, int) else (0 if str(dim) == "past_sequence_len" else 1) + for dim in value.shape + ] + result[value.name] = np.zeros(shape, dtype=np.float32) + 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]) +def test_fish_omni_pipeline_prefill_and_cached_decode( + upstream, tmp_path, monkeypatch, backbone_codebooks +): + raw = copy.deepcopy(FISH_OMNI_CONFIG) + raw["audio_decoder_config"]["vq_backbone_codebooks"] = backbone_codebooks + reference_config = upstream["config"].FishOmniConfig.from_dict(raw) + torch.manual_seed(42) + tolerance = 1e-4 + backbone = upstream["backbone"].FishOmniBackbone( + reference_config.text, dtype=torch.float32 + ) + tower = upstream["tower"].FishOmniAudioDecoder(reference_config, dtype=torch.float32) + for module in (backbone, tower): + for parameter in module.parameters(): + torch.nn.init.normal_(parameter, mean=0.0, std=0.1) + module.eval() + config = FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)) + package = build_from_module(FishOmniModel(config), config, task="fish-omni") + save_file(_checkpoint(backbone, tower), tmp_path / "model.safetensors") + stream_fish_omni_safetensors_to_package(package, str(tmp_path), config, revision=None) + sessions = {name: OnnxModelSession(model) for name, model in package.items()} + backbone.setup_caches(1, 32) + tower.setup_caches(32) + 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]]) + next_frames = torch.tensor([[16, 0, 2, 8], [20, 4, 3, 1], [2, 0, 0, 0], [18, 2, 6, 9]]) + start = 0 + sampled_logits = [] + + def greedy(logits, *args, **kwargs): + sampled_logits.append(logits.detach().float().numpy()) + return logits.argmax(dim=-1) + + monkeypatch.setattr( + upstream["tower"], + "sample_top_k_top_p", + greedy, + ) + try: + # Multi-token prompt followed by two reused-state steps exercises both + # DeltaNet state and all shared/residual tower caches across time. + with torch.no_grad(): + for length in [4] + [1] * 20: + positions = torch.arange(start, start + length) + if start: + frames = next_frames[-1:].clone() + embedded = sessions["embedding"].run({"frames": frames.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().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.numpy()[None], + "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().numpy(), + atol=tolerance, + rtol=tolerance, + ) + np.testing.assert_allclose( + outputs["logits"].astype(np.float32), + backbone.compute_logits(hidden).float().numpy(), + atol=tolerance, + rtol=tolerance, + ) + 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(config.num_codebooks - 1): + tower_output = sessions["audio_decoder"].run( + { + **caches["audio_decoder"], + "hidden_states": outputs["hidden_states"], + "shared_input": embedded["shared_input"], + "codes": next_codes, + "semantic_mask": semantic, + "position_ids": positions.numpy()[None], + "attention_mask": mask, + } + ) + 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), + 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.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_decoder"] = _next_cache(tower_output) + start += length + continue + semantic = tower.is_semantic(next_frames[:, 0]) + tower_output = sessions["audio_decoder"].run( + { + **caches["audio_decoder"], + "hidden_states": outputs["hidden_states"], + "shared_input": embedded["shared_input"], + "codes": next_frames[:, 1:].numpy()[None], + "semantic_mask": semantic.numpy()[None], + "position_ids": positions.numpy()[None], + "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().numpy() + np.testing.assert_allclose( + tower_output["codebook_logits"].astype(np.float32), + expected, + atol=tolerance, + rtol=tolerance, + ) + caches["decoder"] = _next_cache(outputs) + caches["audio_decoder"] = _next_cache(tower_output) + start += length + finally: + for session in sessions.values(): + session.close() From dfec0b319afd894d97bd71d4e8ea33cb4bfb4828 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 20:15:00 -0700 Subject: [PATCH 02/22] Resolve Drama 3 test configs from the worktree The H200 validation host has an unrelated installed tests package that shadows the repository's namespace-package imports. Put the shared test config directory on pytest's local path and use the repository-standard _test_configs import, so the model suites work in shared Python environments. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- conftest.py | 5 ++++- .../integrations/transformers/_fish_omni_weights_test.py | 2 +- src/mobius/models/fish_omni_test.py | 2 +- tests/integration/fish_omni_test.py | 4 ++-- 4 files changed, 8 insertions(+), 5 deletions(-) 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/src/mobius/integrations/transformers/_fish_omni_weights_test.py b/src/mobius/integrations/transformers/_fish_omni_weights_test.py index 96d9f7ad4..943872231 100644 --- a/src/mobius/integrations/transformers/_fish_omni_weights_test.py +++ b/src/mobius/integrations/transformers/_fish_omni_weights_test.py @@ -10,6 +10,7 @@ import onnx_ir as ir import pytest import torch +from _test_configs import FISH_OMNI_CONFIG from safetensors.torch import save_file from mobius import build_from_module @@ -25,7 +26,6 @@ ) from mobius.models import FishOmniModel from mobius.models.fish_omni_test import tiny_config -from tests._test_configs import FISH_OMNI_CONFIG def _checkpoint(package): diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py index 1fafc4ff1..85016a708 100644 --- a/src/mobius/models/fish_omni_test.py +++ b/src/mobius/models/fish_omni_test.py @@ -9,6 +9,7 @@ import onnx_ir as ir import pytest +from _test_configs import FISH_OMNI_CONFIG from mobius import build_from_module from mobius._configs.fish_omni import FishOmniConfig @@ -18,7 +19,6 @@ _dict_to_pretrained_config, ) from mobius.models import FishOmniModel -from tests._test_configs import FISH_OMNI_CONFIG def tiny_config(): diff --git a/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py index 976f45967..1b1a308c5 100644 --- a/tests/integration/fish_omni_test.py +++ b/tests/integration/fish_omni_test.py @@ -14,6 +14,7 @@ import numpy as np import pytest import torch +from _test_configs import FISH_OMNI_CONFIG from safetensors.torch import save_file from mobius import build_from_module @@ -24,7 +25,6 @@ stream_fish_omni_safetensors_to_package, ) from mobius.models import FishOmniModel -from tests._test_configs import FISH_OMNI_CONFIG pytestmark = pytest.mark.integration _PREFIX = "fish_speech.models.text2semantic.fish_omni" @@ -153,7 +153,7 @@ def greedy(logits, *args, **kwargs): greedy, ) try: - # Multi-token prompt followed by two reused-state steps exercises both + # 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 [4] + [1] * 20: From 1d854aef53a66434e89e8351ce47a09bfc9a923e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 20:52:37 -0700 Subject: [PATCH 03/22] Include Drama 3 in speech coverage discovery Wire the existing tiny composite config into the shared speech matrix so the repository-wide L1/L3 audit discovers its real coverage. Exercise the dedicated task through the normal graph checks, and assert that unsupported affine component quantization is rejected instead of claiming an unimplemented path. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/mobius/_configs/fish_omni.py | 5 +++++ tests/_test_configs.py | 13 +++++++++++-- tests/build_graph/speech_test.py | 5 +++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/src/mobius/_configs/fish_omni.py b/src/mobius/_configs/fish_omni.py index 5eea3e739..e56013b55 100644 --- a/src/mobius/_configs/fish_omni.py +++ b/src/mobius/_configs/fish_omni.py @@ -190,6 +190,11 @@ def from_transformers(cls, config, parent_config=None) -> FishOmniConfig: 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 the FP32 graph with explicit FP8 checkpoint dequantization." + ) if self.dtype != ir.DataType.FLOAT: raise ValueError( "Drama 3 currently requires dtype='f32'. FP16 residuals/down projections " diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 6b300516f..edcb7a6e5 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -72,11 +72,12 @@ WhisperConfig, YolosConfig, Zamba2Config, + _as_attribute_config, ) +from mobius._configs.fish_omni import FishOmniConfig from mobius.models import EsmConfig, ReUseConfig -# Drama 3 has a temporal codebook ABI and is covered by model-local tests, -# rather than the generic input_ids/logits causal-LM matrix. +# Tiny composite config for Drama 3's model-local parity and speech graph matrix. FISH_OMNI_CONFIG = { "model_type": "fish_omni", "architectures": ["FishOmniModel"], @@ -3594,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/speech_test.py b/tests/build_graph/speech_test.py index 281d1e15a..cd16cad5f 100644 --- a/tests/build_graph/speech_test.py +++ b/tests/build_graph/speech_test.py @@ -46,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"}, @@ -1243,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), From 53b0a3b9c9f26570702af110c0c842a732f1cb72 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 21:15:43 -0700 Subject: [PATCH 04/22] Exercise Drama 3 parity through saved packages Round-trip the weighted package before numerical parity so component-qualified symbolic dimensions are covered. Initialize empty KV caches by their documented sequence axis rather than assuming a bare symbolic dimension name; recurrent and convolution states retain the actual batch size. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/integration/fish_omni_test.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py index 1b1a308c5..38450c471 100644 --- a/tests/integration/fish_omni_test.py +++ b/tests/integration/fish_omni_test.py @@ -19,6 +19,7 @@ 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.integrations.transformers._config_resolver import _dict_to_pretrained_config from mobius.integrations.transformers._fish_omni_weights import ( @@ -97,10 +98,11 @@ def _empty_cache(model): for value in model.graph.inputs: if not value.name.startswith("past_key_values."): continue - shape = [ - dim if isinstance(dim, int) else (0 if str(dim) == "past_sequence_len" else 1) - for dim in value.shape - ] + 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) result[value.name] = np.zeros(shape, dtype=np.float32) return result @@ -134,6 +136,9 @@ def test_fish_omni_pipeline_prefill_and_cached_decode( package = build_from_module(FishOmniModel(config), config, task="fish-omni") 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)) sessions = {name: OnnxModelSession(model) for name, model in package.items()} backbone.setup_caches(1, 32) tower.setup_caches(32) From a5818d880d5e11d31b63a13edaa62ac3080d2e14 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 21:42:47 -0700 Subject: [PATCH 05/22] Record full-checkpoint Drama 3 H200 parity Validate the complete pinned FP32 export against an independently invoked Fish Speech pipeline on H200 CUDA. All prefill/cached vectors agree within 1e-4 absolute error and all 20 greedy frames match exactly; profiling confirms CUDA execution. Commit the native token golden and a staged reproducible validator, with smoke tests for cache namespaces, logging and revision identity. Update documentation and distinguish verified token generation from the still-external waveform pipeline. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/getting-started.md | 30 +- scripts/validate_drama3.py | 561 ++++++++++++++++++ scripts/validate_drama3_test.py | 103 ++++ testdata/cases/speech/drama-3.yaml | 4 +- .../speech/drama-3_tokens_generation.json | 42 ++ 5 files changed, 734 insertions(+), 6 deletions(-) create mode 100644 scripts/validate_drama3.py create mode 100644 scripts/validate_drama3_test.py create mode 100644 testdata/golden/speech/drama-3_tokens_generation.json diff --git a/docs/getting-started.md b/docs/getting-started.md index 7859e7d2b..64067f931 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -138,7 +138,9 @@ 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. Padding must be removed before running the recurrent backbone; +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 @@ -159,9 +161,29 @@ Sampling (including repetition-aware resampling), reference audio encoding, the separately distributed `codec.pth`, waveform decoding, and ORT GenAI generation integration are external. `--runtime onnx-genai` writes graph-derived advisory component contracts, not an executable speech pipeline; -`--runtime ort-genai` is explicitly rejected. No full-checkpoint waveform or -generation-golden claim is made by this token-graph export. Fish Audio's -checkpoint/source license is separate from Mobius's MIT license. +`--runtime ort-genai` is explicitly rejected. Fish Audio's checkpoint/source +license is separate from Mobius's MIT license. + +The complete pinned checkpoint 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. + +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. ### Build from a GGUF file diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py new file mode 100644 index 000000000..90e68b8d2 --- /dev/null +++ b/scripts/validate_drama3.py @@ -0,0 +1,561 @@ +# 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 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 +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 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) + caches[port.name] = np.zeros(shape, dtype=np.float32) + 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): + from mobius import build + + start = time.monotonic() + event("export_start") + package = build(str(checkpoint), revision=REVISION, dtype="f32", keep_quantized=False) + 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, + "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 reference(checkpoint): + os.environ["NVIDIA_TF32_OVERRIDE"] = "0" + os.environ["FISH_DRAMA3_FP8_EXPERTS"] = "0" + os.environ["FISH_DRAMA3_CUDA_GRAPH"] = "0" + import numpy as np + import torch + + 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=torch.float32, + ) + event("reference_loaded", seconds=time.monotonic() - start) + backbone, tower = model.backbone, model.audio_decoder + 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 array(tensor): + return tensor.detach().float().cpu().numpy() + + 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": array(output)[None], + "shared_input": array(unseen)[None], + } + ) + return output + + def forward(*args, **kwargs): + hidden = original_forward(*args, **kwargs) + states[-1]["hidden_states"] = array(hidden) + states[-1]["logits"] = array(backbone.compute_logits(hidden)) + 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(array(head(output))) + + 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": "float32", + "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()) + 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-4, + rtol=1e-4, + ) + + +def compare(): + os.environ["NVIDIA_TF32_OVERRIDE"] = "0" + import numpy as np + import onnxruntime as ort + + 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 = {} + 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 = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC + 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) + sessions[name] = ort.InferenceSession( + str(ROOT / "onnx" / name / "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): + session = sessions[name] + output = session.run(None, feeds) + names = [port.name for port in session.get_outputs()] + assert len(names) == len(output) + return dict(zip(names, output)) + + 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) + 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]) + np.testing.assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, 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 + ) + (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].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]]) + for index in range(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, + }, + ) + codes[0, 0, index + 1] = tower_output["codebook_logits"][0, 0, index].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": "float32", + "onnxruntime_version": ort.__version__, + "provider": "CUDAExecutionProvider", + "ort_graph_optimization": "ORT_ENABLE_BASIC", + "atol": 1e-3, + "rtol": 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("--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 + if args.stage == "download": + download() + elif args.stage == "export": + export(checkpoint_path()) + elif args.stage == "reference": + reference(checkpoint_path()) + else: + compare() diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py new file mode 100644 index 000000000..86a9f0ab6 --- /dev/null +++ b/scripts/validate_drama3_test.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Network-free checks for the full-checkpoint validation driver.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +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", + } + + +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_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["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_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/testdata/cases/speech/drama-3.yaml b/testdata/cases/speech/drama-3.yaml index b27a6c5a5..92b8f5dd5 100644 --- a/testdata/cases/speech/drama-3.yaml +++ b/testdata/cases/speech/drama-3.yaml @@ -7,5 +7,5 @@ level: "L4+L5" inputs: prompts: - "Hello, this is a speech synthesis test." -skip_reason: "Token-graph export only: full-checkpoint waveform and generation goldens require the external Fish Speech codec and sampler." -notes: "Pinned L2 model is registered; reduced upstream prefill and temporal-cache parity is covered in tests/integration/fish_omni_test.py. No fabricated L4/L5 goldens." +skip_reason: "The generic TTS waveform golden runner requires the external codec and sampler; token-level goldens are validated separately." +notes: "The full pinned checkpoint passed H200 CUDA prefill and 20-frame greedy parity. Reproduce with scripts/validate_drama3.py; the token reference fixture is testdata/golden/speech/drama-3_tokens_generation.json. No waveform or reduced-precision claim." 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..916cf0c4e --- /dev/null +++ b/testdata/golden/speech/drama-3_tokens_generation.json @@ -0,0 +1,42 @@ +{ + "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", + "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.343191146850586, 13.476040840148926, 13.0930814743042, 12.944562911987305, 12.827808380126953, 12.753363609313965, 12.615551948547363, 12.4545259475708, 12.443134307861328, 12.272481918334961], + "prefill_logits_sha256": "e7cd85ba35c64b51f4d895b85bd1f567d4494f413ddfcd6cfb1381c583fa2c90", + "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] + ] +} From e939a768ac824e4efcbf23d2c9a452f31e8868fd Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 22:15:50 -0700 Subject: [PATCH 06/22] Add native Drama 3 waveform codec decoder Implement the pinned Fish Speech decode path declaratively: ten RVQ code channels, BF16-rounded RoPE in the FP32 causal transformer, ConvNeXt upsampling, and causal DAC waveform synthesis. Reuse shared convolution and normalization components while binding both checkpoint weight-normalization formats strictly. Expose build_fish_codec_decoder and safe, pinned load_drama3_codec_decoder helpers independently of the token graph registry. No shared Drama 3 model/config, exports, or documentation files are changed. Verify tiny single-frame, batched and attention-window-crossing parity against unmodified upstream, dynamic causal prefix reuse, safe checkpoint loading, and a saved/reloaded real checkpoint export. Twenty golden token frames decode to 40960 samples at 44100 Hz: max absolute error 1.84774399e-6, RMSE 1.18384925e-7, SNR 117.7911 dB. FP16/BF16 codec execution is not claimed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/_configs/fish_codec.py | 79 ++++++ src/mobius/models/fish_codec.py | 289 +++++++++++++++++++++ src/mobius/models/fish_codec_test.py | 121 +++++++++ src/mobius/tasks/fish_codec.py | 90 +++++++ tests/integration/_fish_codec_reference.py | 147 +++++++++++ tests/integration/fish_codec_test.py | 129 +++++++++ 6 files changed, 855 insertions(+) create mode 100644 src/mobius/_configs/fish_codec.py create mode 100644 src/mobius/models/fish_codec.py create mode 100644 src/mobius/models/fish_codec_test.py create mode 100644 src/mobius/tasks/fish_codec.py create mode 100644 tests/integration/_fish_codec_reference.py create mode 100644 tests/integration/fish_codec_test.py 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/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/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/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}" + ) From f07dd960a344a8a32bd23f5867c3e05587c331b2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 23:35:51 -0700 Subject: [PATCH 07/22] Add Drama 3 ONNX sampling and experimental BF16 execution Export native drop-crossing top-p and semantic repetition sampling as explicit-state ONNX policies. Preserve source FP32 nonlinear, normalization, rotary and routed-sum boundaries, and select unmasked fused attention only after proving padding is absent. Gate BF16 to the numerically verified standard-ONNX route, leaving FP16 and dynamic FP8 unsupported. Document the separately exported waveform codec and extend the full-checkpoint validator without claiming unmeasured BF16 fidelity. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 58 +++- scripts/validate_drama3.py | 86 ++++-- scripts/validate_drama3_test.py | 61 ++++ src/mobius/_configs/fish_omni.py | 7 +- src/mobius/components/__init__.py | 3 +- src/mobius/components/_activations.py | 15 +- src/mobius/components/_activations_test.py | 33 +++ src/mobius/components/_gated_deltanet.py | 7 +- src/mobius/components/_gated_deltanet_test.py | 30 ++ src/mobius/components/_rotary_embedding.py | 2 + src/mobius/functions/causal_conv.py | 5 +- src/mobius/generation/_fish_omni.py | 207 ++++++++++++++ src/mobius/generation/_fish_omni_test.py | 197 +++++++++++++ src/mobius/models/fish_omni.py | 263 ++++++++++++++++-- src/mobius/models/fish_omni_test.py | 130 ++++++++- src/mobius/tasks/_fish_omni.py | 24 +- tests/integration/fish_omni_test.py | 91 ++++-- 17 files changed, 1119 insertions(+), 100 deletions(-) create mode 100644 src/mobius/generation/_fish_omni.py create mode 100644 src/mobius/generation/_fish_omni_test.py diff --git a/docs/getting-started.md b/docs/getting-started.md index 64067f931..254392f23 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -122,10 +122,16 @@ 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. Native `bsgemm-moe` execution is **not** claimed. Plan for roughly 144 GB of dense FP32 -weights, plus the source checkpoint and runtime allocations. Only FP32 is -currently accepted. FP16 is rejected because released residuals/down projections -exceed its range; BF16 is rejected because reduced upstream numerical parity -does not pass the `1e-2` gate. Tolerances are not relaxed to qualify it. +weights, plus the source checkpoint and runtime allocations. FP16 is rejected +until its FP32 residual/mixed-output projection path is implemented. An opt-in +BF16 dense-weight path uses `dtype="bf16", execution_provider="onnx-standard"`; +its reduced upstream pipeline passes full-logit parity at `1e-2` and 20 exact +cached greedy frames on CUDA. Full-checkpoint BF16 validation is pending. +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 does not initialize on +CUDA, so use the standard-ONNX export path for this dtype. | Graph | Inputs | Outputs | |---|---|---| @@ -157,9 +163,47 @@ 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. -Sampling (including repetition-aware resampling), reference audio encoding, -the separately distributed `codec.pth`, waveform decoding, and ORT GenAI -generation integration are external. `--runtime onnx-genai` writes graph-derived +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. diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index 90e68b8d2..b9ff513c2 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -34,6 +34,7 @@ EMBEDDING_DEVICE = 0 AUDIO_DEVICE = 1 DECODER_DEVICE = 2 +DTYPE = "f32" GOLDEN = ( pathlib.Path(__file__).resolve().parents[1] / "testdata" @@ -48,6 +49,7 @@ def event(kind, **details): def empty_cache_ports(ports): + import ml_dtypes import numpy as np caches = {} @@ -60,7 +62,11 @@ def empty_cache_ports(ports): assert len(shape) == 4 shape[2] = 0 assert all(isinstance(dim, int) for dim in shape), (port.name, shape) - caches[port.name] = np.zeros(shape, dtype=np.float32) + type_name = getattr(port, "type", "tensor(float)") + dtypes = {"tensor(float)": np.float32, "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 @@ -139,7 +145,13 @@ def export(checkpoint): start = time.monotonic() event("export_start") - package = build(str(checkpoint), revision=REVISION, dtype="f32", keep_quantized=False) + package = build( + str(checkpoint), + revision=REVISION, + dtype=DTYPE, + keep_quantized=False, + execution_provider="onnx-standard" if DTYPE == "bf16" else "default", + ) event("graph_and_bindings_ready", seconds=time.monotonic() - start) package.save(str(ROOT / "onnx"), max_workers=1) event("export_complete", seconds=time.monotonic() - start) @@ -147,6 +159,7 @@ def export(checkpoint): json.dumps( { "revision": REVISION, + "dtype": DTYPE, "seconds": time.monotonic() - start, "weight_loading_report": package.weight_loading_report, "files": { @@ -181,7 +194,7 @@ def reference(checkpoint): checkpoint, devices=[f"cuda:{device}" for device in REFERENCE_DEVICES], tower_device=f"cuda:{REFERENCE_TOWER_DEVICE}", - dtype=torch.float32, + dtype=torch.bfloat16 if DTYPE == "bf16" else torch.float32, ) event("reference_loaded", seconds=time.monotonic() - start) backbone, tower = model.backbone, model.audio_decoder @@ -271,7 +284,7 @@ def capture(_module, _inputs, output): "model_id": "fishaudio/drama-3", "revision": REVISION, "source_revision": SOURCE_REVISION, - "dtype": "float32", + "dtype": "bfloat16" if DTYPE == "bf16" else "float32", "prompt": PROMPT, "prompt_rows": len(prompt), "prefill_shape": list(logits.shape), @@ -299,6 +312,7 @@ def validate_reference_golden(golden): import numpy as np committed = json.loads(GOLDEN.read_text()) + assert golden["dtype"] == ("bfloat16" if DTYPE == "bf16" else "float32"), "dtype" for field in ( "model_id", "revision", @@ -319,15 +333,32 @@ def validate_reference_golden(golden): np.testing.assert_allclose( golden["prefill_last_top10_logits"], committed["prefill_last_top10_logits"], - atol=1e-4, - rtol=1e-4, + atol=1e-2 if DTYPE == "bf16" else 1e-4, + rtol=1e-2 if DTYPE == "bf16" 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() + } + 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 compare(): os.environ["NVIDIA_TF32_OVERRIDE"] = "0" import numpy as np import onnxruntime as ort + import onnxruntime_easy as ort_easy golden = json.loads((ROOT / "reference" / "golden.json").read_text()) validate_reference_golden(golden) @@ -350,7 +381,7 @@ def compare(): options.enable_profiling = True options.profile_file_prefix = str(ROOT / "decoder-profile") event("ort_session_start", component=name, device=device) - sessions[name] = ort.InferenceSession( + sessions[name] = ort_easy.EasySession( str(ROOT / "onnx" / name / "model.onnx"), sess_options=options, providers=[ @@ -368,11 +399,10 @@ def compare(): event("ort_session_ready", component=name, seconds=time.monotonic() - start) def run(name, feeds): - session = sessions[name] - output = session.run(None, feeds) - names = [port.name for port in session.get_outputs()] - assert len(names) == len(output) - return dict(zip(names, output)) + return run_session(sessions[name], feeds) + + expected_type = "tensor(bfloat16)" if DTYPE == "bf16" else "tensor(float)" + assert sessions["decoder"].get_outputs()[0].type == expected_type def empty_cache(name): return empty_cache_ports(sessions[name].get_inputs()) @@ -388,6 +418,7 @@ def next_cache(output): 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 @@ -397,7 +428,10 @@ def check(name, actual, expected): "mean_abs_diff": float(difference.mean()) if difference.size else 0.0, } event("comparison", name=name, **metrics[name]) - np.testing.assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, err_msg=name) + tolerance = 1e-2 if DTYPE == "bf16" 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"] @@ -480,7 +514,9 @@ def check(name, actual, expected): state["prefill_tower_logits"], ) tower_cache = next_cache(prefilled) - main = int(candidates[decoded["logits"][0, -1, candidates].argmax()]) + 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]]) @@ -497,7 +533,9 @@ def check(name, actual, expected): "attention_mask": mask, }, ) - codes[0, 0, index + 1] = tower_output["codebook_logits"][0, 0, index].argmax() + 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"] ) @@ -514,12 +552,12 @@ def check(name, actual, expected): "status": "passed", "revision": REVISION, "source_revision": SOURCE_REVISION, - "dtype": "float32", + "dtype": "bfloat16" if DTYPE == "bf16" else "float32", "onnxruntime_version": ort.__version__, "provider": "CUDAExecutionProvider", "ort_graph_optimization": "ORT_ENABLE_BASIC", - "atol": 1e-3, - "rtol": 1e-3, + "atol": 1e-2 if DTYPE == "bf16" else 1e-3, + "rtol": 1e-2 if DTYPE == "bf16" else 1e-3, "prompt_rows": golden["prompt_rows"], "generated_frames": len(generated), "attempted_frames": golden["attempted_frames"], @@ -538,6 +576,8 @@ def check(name, actual, expected): 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", "bf16"), default="f32") + 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) @@ -551,6 +591,16 @@ def check(name, actual, expected): EMBEDDING_DEVICE = args.embedding_device AUDIO_DEVICE = args.audio_device DECODER_DEVICE = args.decoder_device + DTYPE = args.dtype + if args.golden is not None: + GOLDEN = args.golden.resolve() + elif DTYPE == "bf16": + GOLDEN = GOLDEN.with_name("drama-3_bf16_tokens_generation.json") + if args.stage == "compare" and not GOLDEN.is_file(): + parser.error( + "BF16 comparison requires --golden pointing to the independent " + "reference stage's golden.json until a BF16 golden is committed" + ) if args.stage == "download": download() elif args.stage == "export": diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py index 86a9f0ab6..0be14c9c3 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -8,6 +8,8 @@ import json from types import SimpleNamespace +import ml_dtypes +import numpy as np import pytest import validate_drama3 from validate_drama3 import ( @@ -60,6 +62,65 @@ def test_unknown_nonbatch_cache_dimension_is_an_error(): ) +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_committed_token_golden_is_pinned_and_complete(): golden = json.loads(GOLDEN.read_text()) validate_reference_golden(golden) diff --git a/src/mobius/_configs/fish_omni.py b/src/mobius/_configs/fish_omni.py index e56013b55..2fb29c626 100644 --- a/src/mobius/_configs/fish_omni.py +++ b/src/mobius/_configs/fish_omni.py @@ -20,7 +20,7 @@ @dataclasses.dataclass class FishOmniConfig(ArchitectureConfig): - """Drama 3 token generation; the separately distributed codec is external.""" + """Drama 3 token generation; the waveform codec is exported separately.""" tower: ArchitectureConfig | None = None num_codebooks: int = 10 @@ -195,10 +195,9 @@ def validate(self) -> None: "Drama 3 affine component quantization is unsupported; " "use the FP32 graph with explicit FP8 checkpoint dequantization." ) - if self.dtype != ir.DataType.FLOAT: + if self.dtype not in (ir.DataType.FLOAT, ir.DataType.BFLOAT16): raise ValueError( - "Drama 3 currently requires dtype='f32'. FP16 residuals/down projections " - "overflow; BF16 has not passed the upstream numerical parity gate." + "Drama 3 FP16 requires an FP32 residual/mixed-output projection path." ) if self.tower is None: raise ValueError("Drama 3 requires an audio tower configuration") diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index d9cee1d06..4792905f9 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -121,6 +121,7 @@ "SequenceMambaBlock", "SequenceSelectiveScan", "SiLU", + "Float32SiLU", "Siglip2NaFlexVisionEmbeddings", "Siglip2NaFlexVisionModel", "SigmoidTopKGate", @@ -165,7 +166,7 @@ "siglip2_naflex_attention_mask", ] -from mobius.components._activations import SiLU, get_activation +from mobius.components._activations import Float32SiLU, 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..ba33ca6c2 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,10 @@ 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(op.Swish(op.Cast(x, to=ir.DataType.FLOAT)), x) diff --git a/src/mobius/components/_activations_test.py b/src/mobius/components/_activations_test.py index cfe2e39be..c755850b0 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") + swish = next(node for node in graph if node.op_type == "Swish") + assert swish.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..12f8863f3 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -130,6 +130,8 @@ def __init__( name="sin_cache", data=ir.tensor(sin_cache_data), ) + self.cos_cache._keep_float32 = True # type: ignore[attr-defined] + self.sin_cache._keep_float32 = True # type: ignore[attr-defined] def _cast_embeddings( self, diff --git a/src/mobius/functions/causal_conv.py b/src/mobius/functions/causal_conv.py index 80502d99c..cf0fd073d 100644 --- a/src/mobius/functions/causal_conv.py +++ b/src/mobius/functions/causal_conv.py @@ -128,7 +128,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( + op.Swish(op.Cast(conv_out, to=ir.DataType.FLOAT)), + conv_out, + ) elif activation == "none": output = conv_out else: diff --git a/src/mobius/generation/_fish_omni.py b/src/mobius/generation/_fish_omni.py new file mode 100644 index 000000000..f99e97e34 --- /dev/null +++ b/src/mobius/generation/_fish_omni.py @@ -0,0 +1,207 @@ +# 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 + +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, +) + + +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/models/fish_omni.py b/src/mobius/models/fish_omni.py index 752dd3f03..703d5527f 100644 --- a/src/mobius/models/fish_omni.py +++ b/src/mobius/models/fish_omni.py @@ -4,12 +4,14 @@ """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; prompt formatting, sampling, -and the separately distributed audio codec remain caller responsibilities. +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 @@ -19,16 +21,192 @@ from mobius._configs import ArchitectureConfig from mobius._configs.fish_omni import FishOmniConfig from mobius.components import ( + MLP, Embedding, + Float32SiLU, Linear, OffsetRMSNorm, + PostGatedRMSNorm, + Qwen35Attention, SoftmaxTopKGate, apply_rotary_pos_emb, create_attention_bias, initialize_rope, ) from mobius.components._attention import _apply_attention -from mobius.models.qwen35 import Qwen35MoEDecoderLayer, Qwen35MoETextModel +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 + return rope + + +def _sigmoid_float32(op, value): + return op.CastLike(op.Sigmoid(op.Cast(value, to=ir.DataType.FLOAT)), 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), + op.Swish(op.Cast(gate, to=ir.DataType.FLOAT)), + ) + 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): + super().__init__(config, gate=gate) + if self.experts is not None: + for expert in self.experts: + assert isinstance(expert, MLP) + expert.act_fn = Float32SiLU() + self.shared_expert.act_fn = Float32SiLU() + + def forward(self, op, hidden_states): + if self.experts is None: + return super().forward(op, hidden_states) + routing, selected = self.gate(op, hidden_states) + combined = None + for index, expert in enumerate(self.experts): + output = expert(op, hidden_states) + match = op.CastLike(op.Equal(selected, index), routing) + weight = op.ReduceSum(op.Mul(routing, match), [-1], keepdims=True) + contribution = op.Mul( + op.Cast(output, to=ir.DataType.FLOAT), + op.Cast(weight, to=ir.DataType.FLOAT), + ) + combined = contribution if combined is None else op.Add(combined, contribution) + assert combined is not None + routed = op.CastLike(combined, hidden_states) + 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 _FishRouter(SoftmaxTopKGate): @@ -43,21 +221,42 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return op.CastLike(weights, hidden_states), indices -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 = Qwen35MoETextModel(config) - assert config.num_local_experts is not None and config.num_experts_per_tok is not None - for layer in self.model.layers: - assert isinstance(layer, Qwen35MoEDecoderLayer) - layer.mlp.gate = _FishRouter( +class _FishDecoderLayer(Qwen35MoEDecoderLayer): + def __init__(self, config, layer_idx): + super().__init__(config, layer_idx) + 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) + + +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) self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) def forward(self, op, inputs_embeds, attention_mask, position_ids, past_key_values): @@ -106,8 +305,10 @@ def code_sum(first: int, end: int): count = cfg.vq_backbone_codebooks mixed = op.Div( - op.Add(text, code_sum(1, count + 1)), op.CastLike(math.sqrt(1 + count), text) + op.Cast(op.Add(text, code_sum(1, count + 1)), to=ir.DataType.FLOAT), + 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) @@ -139,23 +340,23 @@ def forward(self, op, hidden, bias, positions, past): 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 = apply_rotary_pos_emb( + q = _rotate_float32( op, q, positions, - num_heads=cfg.num_attention_heads, - rotary_embedding_dim=self.rotary_dim, + heads=cfg.num_attention_heads, + rotary_dim=self.rotary_dim, interleaved=True, ) - k = apply_rotary_pos_emb( + k = _rotate_float32( op, k, positions, - num_heads=cfg.num_key_value_heads, - rotary_embedding_dim=self.rotary_dim, + heads=cfg.num_key_value_heads, + rotary_dim=self.rotary_dim, interleaved=True, ) - attended, pk, pv = _apply_attention( + attended, pk, pv = _apply_fish_attention( op, q, k, @@ -166,8 +367,9 @@ def forward(self, op, hidden, bias, positions, past): 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, op.Sigmoid(gate))), (pk, pv) + return self.wo(op, op.Mul(attended, _sigmoid_float32(op, gate))), (pk, pv) class _TowerMLP(nn.Module): @@ -176,10 +378,11 @@ def __init__(self, config): self.w1 = Linear(config.hidden_size, config.intermediate_size, bias=False) self.w3 = Linear(config.hidden_size, config.intermediate_size, bias=False) 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(op.Mul(gate, op.Sigmoid(gate)), self.w3(op, hidden))) + return self.w2(op, op.Mul(self.act_fn(op, gate), self.w3(op, hidden))) class _TowerTransformer(nn.Module): @@ -233,14 +436,15 @@ 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(config.tower, shared=True) + _TowerLayer(tower_config, shared=True) for _ in range(config.shared_enrichment_layers) ] ) self.residual_layers = nn.ModuleList( - [_TowerLayer(config.tower, shared=False) for _ in range(config.num_codebooks - 1)] + [_TowerLayer(tower_config, shared=False) for _ in range(config.num_codebooks - 1)] ) self.codebook_heads = nn.ModuleList( [ @@ -248,9 +452,7 @@ def __init__(self, config: FishOmniConfig): for _ in range(config.num_codebooks - 1) ] ) - rotary_emb = initialize_rope(config.tower) - assert rotary_emb is not None - self.rotary_emb = rotary_emb + self.rotary_emb = _float32_rope(config.tower) def forward( self, @@ -302,8 +504,9 @@ 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. Sampling, reference-bracket prompt formatting, and waveform codec - execution are external. FP8 experts require explicit streaming dequantization. + 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" diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py index 85016a708..0ff977341 100644 --- a/src/mobius/models/fish_omni_test.py +++ b/src/mobius/models/fish_omni_test.py @@ -7,18 +7,25 @@ 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 from mobius import build_from_module from mobius._configs.fish_omni import FishOmniConfig 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 def tiny_config(): @@ -27,11 +34,17 @@ def tiny_config(): ) -def test_fish_omni_build(): - dtype = ir.DataType.FLOAT - config = tiny_config() - package = build_from_module(FishOmniModel(config), config, task="fish-omni") +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.BFLOAT16]) +def test_fish_omni_build(dtype): + config = dataclasses.replace(tiny_config(), dtype=dtype) + package = build_from_module( + FishOmniModel(config), + config, + task="fish-omni", + execution_provider="onnx-standard" if dtype == ir.DataType.BFLOAT16 else "default", + ) 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", @@ -42,6 +55,8 @@ def test_fish_omni_build(): 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 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) @@ -65,12 +80,18 @@ def test_fish_omni_registry_config_resolution(): assert resolved.quantization is None -@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16]) def test_fish_omni_rejects_unqualified_precision(dtype): - with pytest.raises(ValueError, match="requires dtype='f32'"): + with pytest.raises(ValueError, match="FP16 requires"): FishOmniModel(dataclasses.replace(tiny_config(), dtype=dtype)) +def test_fish_omni_bfloat16_rejects_unqualified_ep(): + config = dataclasses.replace(tiny_config(), dtype=ir.DataType.BFLOAT16) + with pytest.raises(ValueError, match="requires execution_provider='onnx-standard'"): + build_from_module(FishOmniModel(config), config, task="fish-omni") + + @pytest.mark.parametrize( ("section", "key", "value"), [ @@ -116,3 +137,100 @@ def test_fish_omni_onnx_genai_metadata_is_graph_derived_and_advisory(tmp_path): 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() diff --git a/src/mobius/tasks/_fish_omni.py b/src/mobius/tasks/_fish_omni.py index 2c8cf0c11..9dbe15129 100644 --- a/src/mobius/tasks/_fish_omni.py +++ b/src/mobius/tasks/_fish_omni.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Three-graph Drama 3 token-generation ABI.""" +"""Drama 3 neural graphs and ONNX sampling policies.""" from __future__ import annotations @@ -9,8 +9,13 @@ 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 ( + 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, @@ -36,6 +41,11 @@ class FishOmniTask(ModelTask): ) def build(self, module, config: FishOmniConfig) -> ModelPackage: + if config.dtype == ir.DataType.BFLOAT16 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" + ) self._validate_components(module) assert config.tower is not None batch, seq = "batch", "sequence_len" @@ -103,6 +113,14 @@ def build(self, module, config: FishOmniConfig) -> ModelPackage: models["audio_decoder"] = _make_model(graph) for model in models.values(): model.metadata_props["mobius.fish_omni.scope"] = ( - "Token generation only; external sampling, prompt formatting and waveform codec" + "Neural token graphs with ONNX sampling policies; caller orchestrates " + "temporal decoding and prompt formatting; waveform decoder exported separately" ) - return ModelPackage(models, config=config) + return ModelPackage( + models, + config=config, + policy_components={ + "main_sampler": build_fish_main_sampler(config), + "codebook_sampler": build_fish_codebook_sampler(config), + }, + ) diff --git a/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py index 38450c471..4bb4f13fe 100644 --- a/tests/integration/fish_omni_test.py +++ b/tests/integration/fish_omni_test.py @@ -6,12 +6,16 @@ from __future__ import annotations import copy +import dataclasses import importlib.util import sys import types import urllib.request +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 @@ -103,7 +107,8 @@ def _empty_cache(model): if value.name.endswith((".key", ".value")): shape[2] = 0 assert all(isinstance(dim, int) for dim in shape), (value.name, shape) - result[value.name] = np.zeros(shape, dtype=np.float32) + dtype = ml_dtypes.bfloat16 if value.dtype == ir.DataType.BFLOAT16 else np.float32 + result[value.name] = np.zeros(shape, dtype=dtype) return result @@ -116,40 +121,75 @@ def _next_cache(outputs): @pytest.mark.parametrize("backbone_codebooks", [0, 1]) +@pytest.mark.parametrize( + "dtype", + [ + ir.DataType.FLOAT, + pytest.param( + ir.DataType.BFLOAT16, + 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", + ), + ), + ], +) def test_fish_omni_pipeline_prefill_and_cached_decode( - upstream, tmp_path, monkeypatch, backbone_codebooks + upstream, tmp_path, monkeypatch, backbone_codebooks, dtype ): raw = copy.deepcopy(FISH_OMNI_CONFIG) raw["audio_decoder_config"]["vq_backbone_codebooks"] = backbone_codebooks reference_config = upstream["config"].FishOmniConfig.from_dict(raw) torch.manual_seed(42) - tolerance = 1e-4 + tolerance = 1e-4 if dtype == ir.DataType.FLOAT else 1e-2 + torch_dtype = torch.float32 if dtype == ir.DataType.FLOAT else torch.bfloat16 + device = "cpu" if dtype == ir.DataType.FLOAT else "cuda:0" backbone = upstream["backbone"].FishOmniBackbone( - reference_config.text, dtype=torch.float32 + reference_config.text, dtype=torch_dtype, device=device + ) + tower = upstream["tower"].FishOmniAudioDecoder( + reference_config, dtype=torch_dtype, device=device ) - tower = upstream["tower"].FishOmniAudioDecoder(reference_config, dtype=torch.float32) for module in (backbone, tower): for parameter in module.parameters(): torch.nn.init.normal_(parameter, mean=0.0, std=0.1) module.eval() - config = FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)) - package = build_from_module(FishOmniModel(config), config, task="fish-omni") + config = dataclasses.replace( + FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)), dtype=dtype + ) + package = build_from_module( + FishOmniModel(config), + config, + task="fish-omni", + execution_provider="default" if dtype == ir.DataType.FLOAT else "onnx-standard", + ) 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)) - sessions = {name: OnnxModelSession(model) for name, model in package.items()} + assert set(package.policy_components) == {"main_sampler", "codebook_sampler"} + if dtype != ir.DataType.FLOAT: + ort.preload_dlls() + sessions = { + name: OnnxModelSession(model, device="cpu" if dtype == ir.DataType.FLOAT else "cuda") + for name, model in package.items() + } backbone.setup_caches(1, 32) tower.setup_caches(32) 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]]) - next_frames = torch.tensor([[16, 0, 2, 8], [20, 4, 3, 1], [2, 0, 0, 0], [18, 2, 6, 9]]) + 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 + ) start = 0 sampled_logits = [] def greedy(logits, *args, **kwargs): - sampled_logits.append(logits.detach().float().numpy()) + sampled_logits.append(logits.detach().float().cpu().numpy()) return logits.argmax(dim=-1) monkeypatch.setattr( @@ -162,14 +202,15 @@ def greedy(logits, *args, **kwargs): # DeltaNet state and all shared/residual tower caches across time. with torch.no_grad(): for length in [4] + [1] * 20: - positions = torch.arange(start, start + length) + 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.numpy()[None]}) + 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().numpy(), + ref_embed.float().cpu().numpy(), atol=tolerance, rtol=tolerance, ) @@ -178,22 +219,24 @@ def greedy(logits, *args, **kwargs): { **caches["decoder"], "inputs_embeds": embedded["inputs_embeds"], - "position_ids": positions.numpy()[None], + "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().numpy(), + 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().numpy(), + backbone.compute_logits(hidden).float().cpu().numpy(), atol=tolerance, rtol=tolerance, + err_msg=f"decoder logits at position {start}", ) if start: candidates = [ @@ -226,7 +269,7 @@ def greedy(logits, *args, **kwargs): "shared_input": embedded["shared_input"], "codes": next_codes, "semantic_mask": semantic, - "position_ids": positions.numpy()[None], + "position_ids": positions_np, "attention_mask": mask, } ) @@ -241,14 +284,14 @@ def greedy(logits, *args, **kwargs): hidden[0], positions, frames, - torch.zeros(1, 16, dtype=torch.int64), + 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.numpy(), + next_frames.cpu().numpy(), ) np.testing.assert_allclose( tower_output["codebook_logits"].astype(np.float32), @@ -266,9 +309,9 @@ def greedy(logits, *args, **kwargs): **caches["audio_decoder"], "hidden_states": outputs["hidden_states"], "shared_input": embedded["shared_input"], - "codes": next_frames[:, 1:].numpy()[None], - "semantic_mask": semantic.numpy()[None], - "position_ids": positions.numpy()[None], + "codes": next_frames[:, 1:].cpu().numpy()[None], + "semantic_mask": semantic.cpu().numpy()[None], + "position_ids": positions_np, "attention_mask": mask, } ) @@ -279,7 +322,7 @@ def greedy(logits, *args, **kwargs): ): h = layer(h, next_frames[:, i + 1], semantic, positions) expected.append(head(h)) - expected = torch.stack(expected, dim=-2)[None].float().numpy() + expected = torch.stack(expected, dim=-2)[None].float().cpu().numpy() np.testing.assert_allclose( tower_output["codebook_logits"].astype(np.float32), expected, From d1218b605773d9933e446d2dc0060c1d46ffdd15 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 26 Oct 2026 23:38:46 -0700 Subject: [PATCH 08/22] Update Drama 3 golden runner qualification scope Keep the generic waveform-runner waiver tied to missing temporal orchestration rather than already-exported codec and sampling graphs. Distinguish real FP32 evidence from experimental BF16 and unfinished FP16/FP8 routes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- testdata/cases/speech/drama-3.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/testdata/cases/speech/drama-3.yaml b/testdata/cases/speech/drama-3.yaml index 92b8f5dd5..7cbbda7a6 100644 --- a/testdata/cases/speech/drama-3.yaml +++ b/testdata/cases/speech/drama-3.yaml @@ -7,5 +7,5 @@ level: "L4+L5" inputs: prompts: - "Hello, this is a speech synthesis test." -skip_reason: "The generic TTS waveform golden runner requires the external codec and sampler; token-level goldens are validated separately." -notes: "The full pinned checkpoint passed H200 CUDA prefill and 20-frame greedy parity. Reproduce with scripts/validate_drama3.py; the token reference fixture is testdata/golden/speech/drama-3_tokens_generation.json. No waveform or reduced-precision claim." +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 full 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). Experimental BF16 has reduced upstream 20-frame parity only; full-checkpoint BF16, FP16 and native FP8 remain unqualified." From 62a0eb482f0bb1ef3af8760f9885126e085a25f2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 00:47:46 -0700 Subject: [PATCH 09/22] Optimize Drama 3 sparse experts and single-pass audio generation Use packed, lazily reconstructed FP32 expert banks with the existing CUDA MoE ABI instead of evaluating every expert. Add an opt-in generation task that prefills the same temporal caches and samples conditioned codebooks in one tower traversal, with greedy/prefill RNG work skipped. Preserve portable exports, source rounding, cache semantics and fail-stop validity flags. Add a request benchmark with zero-copy tensor views and device-resident state, separate initialization/prefill/steady-frame/waveform metrics, and contention and real-time-factor gates. Source parity, stochastic replay equivalence and actual CUDA MoE placement are covered. No uncontended full-model latency claim is made while the shared H200 devices are occupied. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 69 ++ scripts/benchmark_drama3.py | 608 ++++++++++++++++++ scripts/benchmark_drama3_test.py | 187 ++++++ src/mobius/_testing/ort_inference.py | 4 +- src/mobius/generation/_fish_omni.py | 63 ++ .../integrations/ort_genai/auto_export.py | 7 +- .../integrations/transformers/_builder.py | 7 +- .../transformers/_fish_omni_weights.py | 58 +- .../transformers/_fish_omni_weights_test.py | 44 +- src/mobius/models/fish_omni.py | 206 +++++- src/mobius/models/fish_omni_test.py | 59 +- src/mobius/tasks/__init__.py | 4 +- src/mobius/tasks/_fish_omni.py | 70 +- tests/integration/fish_omni_test.py | 214 +++++- 14 files changed, 1541 insertions(+), 59 deletions(-) create mode 100644 scripts/benchmark_drama3.py create mode 100644 scripts/benchmark_drama3_test.py diff --git a/docs/getting-started.md b/docs/getting-started.md index 254392f23..5cec8ca1e 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -163,6 +163,47 @@ 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. It still +stores dense FP32 weights: this is **not native FP8**, and it does not remove the +large model-memory requirement. The portable/default export remains available. + +`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 BF16 standard-ONNX route also supports this one-pass task, but remains +experimentally qualified as described above. + 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`, @@ -229,6 +270,34 @@ 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. +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. +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 --sample --frames 100 --output drama3-perf.json +``` + +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/scripts/benchmark_drama3.py b/scripts/benchmark_drama3.py new file mode 100644 index 000000000..45effa342 --- /dev/null +++ b/scripts/benchmark_drama3.py @@ -0,0 +1,608 @@ +# 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): + 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", + }, + ), + "CPUExecutionProvider", + ] + if device == "cuda" + else ["CPUExecutionProvider"] + ) + self.session = ort.InferenceSession(str(path), options, providers=providers) + self.device, self.device_id = device, device_id + 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 + binding.bind_input( + name, + device_type=value.owner.device_name(), + device_id=self.device_id, + 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) + 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 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) + 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 + element_type = { + "tensor(float)": ir.DataType.FLOAT, + "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) + 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, + "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, + "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("--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..8515cd0df --- /dev/null +++ b/scripts/benchmark_drama3_test.py @@ -0,0 +1,187 @@ +# 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 + +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.tasks.fish_codec import build_fish_codec_decoder + + +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_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.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.BFLOAT16 and device == "cpu": + pytest.skip("Stock ORT CPU has no BF16 Attention kernel") + 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=task, + execution_provider="onnx-standard" if dtype == ir.DataType.BFLOAT16 else "default", + ) + torch.manual_seed(42) + weights = { + _source_name(name): torch.randn(tuple(value.shape)) * 0.1 + for model in package.values() + for name, value in model.graph.initializers.items() + if value.const_value is None + } + # 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)) + with_codec = task == "fish-omni-generation" and device == "cpu" + 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 []), + ] + ) + result = benchmark(args) + assert result["generated_frames"] == 4 + 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") + 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/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/generation/_fish_omni.py b/src/mobius/generation/_fish_omni.py index f99e97e34..e6e12a62a 100644 --- a/src/mobius/generation/_fish_omni.py +++ b/src/mobius/generation/_fish_omni.py @@ -9,6 +9,8 @@ from __future__ import annotations +from typing import NamedTuple + import onnx_ir as ir from onnxscript import GraphBuilder, OpBuilder @@ -23,6 +25,67 @@ ) +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))) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index feaa3c66d..33c1b2a39 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -2036,9 +2036,10 @@ def write_ort_genai_config( ) if getattr(config, "model_type", None) == "fish_omni": raise ValueError( - "Drama 3 requires host-owned next-frame codebook sampling, temporal " - "audio-decoder caches and an external waveform codec; an ONNX Runtime " - "GenAI pipeline configuration is not implemented." + "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( diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index b00f8f5dd..8a2c27429 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -651,9 +651,12 @@ def build_transformers_model( if model_type == "fish_omni": from mobius.tasks import FishOmniTask - if task != "fish-omni" and not isinstance(task, FishOmniTask): + if task not in ("fish-omni", "fish-omni-generation") and not isinstance( + task, FishOmniTask + ): raise ValueError( - "Drama 3 requires task='fish-omni'; it is not a generic causal LM" + "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") diff --git a/src/mobius/integrations/transformers/_fish_omni_weights.py b/src/mobius/integrations/transformers/_fish_omni_weights.py index 79014c30a..ae0123627 100644 --- a/src/mobius/integrations/transformers/_fish_omni_weights.py +++ b/src/mobius/integrations/transformers/_fish_omni_weights.py @@ -13,12 +13,31 @@ 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): + if source not in index: + raise ValueError(f"Missing Drama 3 checkpoint tensor {source!r}") + source_dtype = index[source][2] + 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") + return StreamingWeightSource( + source, + mode="fp8_block_128", + scale_name=source.removesuffix(".weight") + ".weight_scale_inv", + reconstruction_dtype=torch.float32, + ) + if source_dtype not in {"F32", "BF16", "F16"}: + raise ValueError(f"Unexpected Drama 3 weight dtype {source_dtype}: {source}") + return StreamingWeightSource(source) + + def _source_name(target: str) -> str: if target == "embedding.embed_tokens.weight": return "model.language_model.embed_tokens.weight" @@ -26,6 +45,7 @@ def _source_name(target: str) -> str: 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" @@ -45,28 +65,34 @@ def build_fish_omni_streaming_plan( config: FishOmniConfig, ) -> StreamingWeightPlan: """Map every learned parameter and classify each independently scaled expert.""" - targets = {} + targets: dict[str, StreamingWeightSource | StreamingExpertBankSource] = {} for name, value in initializers.items(): if value.const_value is not None: continue source = _source_name(name) - if source not in index: - raise ValueError(f"Missing Drama 3 checkpoint tensor {source!r}") - source_dtype = index[source][2] - expert = ".mlp.experts." in 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") - targets[name] = StreamingWeightSource( - source, - mode="fp8_block_128", - scale_name=source.removesuffix(".weight") + ".weight_scale_inv", - reconstruction_dtype=torch.float32, + if any( + source.endswith(f".experts.{projection}.weight") + for projection in ("gate_proj", "up_proj", "down_proj") + ): + assert config.num_local_experts is not None + prefix, projection = source.rsplit(".experts.", 1) + targets[name] = StreamingExpertBankSource( + tuple( + ( + _classify_source( + f"{prefix}.experts.{expert}.{projection}", + index, + config, + expert=True, + ), + ) + for expert in range(config.num_local_experts) + ) ) else: - if source_dtype not in {"F32", "BF16", "F16"}: - raise ValueError(f"Unexpected Drama 3 weight dtype {source_dtype}: {source}") - targets[name] = StreamingWeightSource(source) + targets[name] = _classify_source( + source, index, config, expert=".mlp.experts." in source + ) return StreamingWeightPlan( targets=targets, report={"architecture": "fish_omni", "codec_included": False}, diff --git a/src/mobius/integrations/transformers/_fish_omni_weights_test.py b/src/mobius/integrations/transformers/_fish_omni_weights_test.py index 943872231..c7201f8c3 100644 --- a/src/mobius/integrations/transformers/_fish_omni_weights_test.py +++ b/src/mobius/integrations/transformers/_fish_omni_weights_test.py @@ -131,6 +131,33 @@ def test_fish_omni_fp8_scales_are_reconstructed_in_fp32(tmp_path, dtype): 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 = [] @@ -166,7 +193,15 @@ def test_fish_omni_requires_explicit_dequantization_before_weight_io(monkeypatch _builder.build_transformers_model("fishaudio/drama-3") -def test_fish_omni_cli_persists_three_weighted_models(tmp_path): +@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") @@ -179,12 +214,17 @@ def test_fish_omni_cli_persists_three_weighted_models(tmp_path): str(tmp_path), "--dtype", "f32", + "--task", + task, + "--execution-provider", + ep, "--dequantize", "--output", str(output), ] ) - for component in ("embedding", "decoder", "audio_decoder"): + 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() diff --git a/src/mobius/models/fish_omni.py b/src/mobius/models/fish_omni.py index 703d5527f..39346d2bb 100644 --- a/src/mobius/models/fish_omni.py +++ b/src/mobius/models/fish_omni.py @@ -18,6 +18,7 @@ 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 ( @@ -34,6 +35,12 @@ initialize_rope, ) from mobius.components._attention import _apply_attention +from mobius.generation._fish_omni import ( + FishAudioSamplingInputs, + _all_rows, + _finite, + sample_codebook_or_greedy, +) from mobius.models.qwen35 import ( Qwen35MoEBlock, Qwen35MoEDecoderLayer, @@ -181,7 +188,19 @@ 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 = ( + _FishCudaExperts(config) if config.dtype == ir.DataType.FLOAT else None + ) if self.experts is not None: for expert in self.experts: assert isinstance(expert, MLP) @@ -191,7 +210,12 @@ def __init__(self, config, gate=None): def forward(self, op, hidden_states): if self.experts is None: return super().forward(op, hidden_states) - routing, selected = self.gate(op, hidden_states) + routing, selected, logits = self.gate(op, hidden_states) + caps = ep_capabilities() + if self.packed_experts is not None and caps.name == "cuda" and caps.supports_fused_moe: + # Floating CUDA MoE performs its own softmax/top-k from RAW logits. + routed = self.packed_experts(op, hidden_states, logits) + return self._add_shared_expert(op, hidden_states, routed) combined = None for index, expert in enumerate(self.experts): output = expert(op, hidden_states) @@ -204,11 +228,77 @@ def forward(self, op, hidden_states): combined = contribution if combined is None else op.Add(combined, contribution) assert combined is not None routed = op.CastLike(combined, hidden_states) + 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 _FishCudaExperts(nn.Module): + """Source-layout expert banks consumed sparsely by ORT's floating CUDA MoE.""" + + def __init__(self, config): + super().__init__() + assert config.num_local_experts is not None + assert config.moe_intermediate_size is not None + assert config.dtype == ir.DataType.FLOAT + 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.gate_proj = _FishExpertBank( + [self.num_experts, self.intermediate_size, self.hidden_size] + ) + self.up_proj = _FishExpertBank( + [self.num_experts, 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, router_logits): + # CUDA MoE declares [E,H,I]/[E,I,H], but consumes the original + # output-major HF bytes. Reshape changes the ABI shape, NOT the layout. + # This matches upstream test_moe_cuda.py's separate-FC3 Mixtral path. + gate = op.Reshape( + self.gate_proj(op), [self.num_experts, self.hidden_size, self.intermediate_size] + ) + up = op.Reshape( + self.up_proj(op), [self.num_experts, self.hidden_size, self.intermediate_size] + ) + down = op.Reshape( + self.down_proj(op), [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, + None, + down, + None, + up, + activation_type="silu", + k=self.top_k, + normalize_routing_weights=self.normalize, + swiglu_fusion=0, + _domain="com.microsoft", + ) + output.type = ir.TensorType(ir.DataType.FLOAT) + 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])) @@ -218,7 +308,7 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): ) if self.norm_topk_prob: weights = op.Div(weights, op.ReduceSum(weights, [-1], keepdims=True)) - return op.CastLike(weights, hidden_states), indices + return op.CastLike(weights, hidden_states), indices, logits class _FishDecoderLayer(Qwen35MoEDecoderLayer): @@ -464,6 +554,43 @@ def forward( 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( @@ -483,21 +610,89 @@ def forward( 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, - op.Gather(codes, i, axis=-1), + incoming, 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 + 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): @@ -519,6 +714,7 @@ class FishOmniModel(nn.Module): "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): diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py index 0ff977341..514d742c8 100644 --- a/src/mobius/models/fish_omni_test.py +++ b/src/mobius/models/fish_omni_test.py @@ -80,6 +80,55 @@ def test_fish_omni_registry_config_resolution(): 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 * 3 + assert all(len(value.shape) == 3 for value in banks) + assert not any(".experts.0." in name for name in graph.initializers) + + +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.BFLOAT16]) +def test_fish_single_pass_task_reuses_tower_weights_and_layer_count(dtype): + config = dataclasses.replace(tiny_config(), dtype=dtype) + ep = "onnx-standard" if dtype == ir.DataType.BFLOAT16 else "default" + teacher = build_from_module( + FishOmniModel(config), config, task="fish-omni", execution_provider=ep + )["audio_decoder"] + generated = build_from_module( + FishOmniModel(config), config, task="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.FLOAT16]) def test_fish_omni_rejects_unqualified_precision(dtype): with pytest.raises(ValueError, match="FP16 requires"): @@ -111,22 +160,24 @@ def test_fish_omni_rejects_incompatible_configs(section, key, value): FishOmniConfig.from_transformers(_dict_to_pretrained_config(raw)) -def test_fish_omni_ort_genai_metadata_rejects_incomplete_runtime(tmp_path): +@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="fish-omni") + 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() -def test_fish_omni_onnx_genai_metadata_is_graph_derived_and_advisory(tmp_path): +@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="fish-omni") + 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"} diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index c8601ab5e..504ac0aa5 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -20,6 +20,7 @@ __all__ = [ "FishOmniTask", + "FishOmniGenerationTask", "AdapterTask", "AudioCTCTask", "AudioFeatureExtractionTask", @@ -153,7 +154,7 @@ GGUFEmbeddingFeatureExtractionTask, GGUFEncoderFeatureExtractionTask, ) -from mobius.tasks._fish_omni import FishOmniTask +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 ( @@ -242,6 +243,7 @@ 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 index 9dbe15129..1de51720c 100644 --- a/src/mobius/tasks/_fish_omni.py +++ b/src/mobius/tasks/_fish_omni.py @@ -13,6 +13,7 @@ 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, ) @@ -39,6 +40,10 @@ class FishOmniTask(ModelTask): decoder="decoder", audio_decoder="audio_decoder", ) + audio_model_name = "audio_decoder" + + def _sampling_inputs(self, builder) -> FishAudioSamplingInputs | None: + return None def build(self, module, config: FishOmniConfig) -> ModelPackage: if config.dtype == ir.DataType.BFLOAT16 and ep_capabilities().name != "onnx-standard": @@ -80,7 +85,7 @@ def build(self, module, config: FishOmniConfig) -> ModelPackage: models["decoder"] = _make_model(graph) _register_linear_attention_functions(models["decoder"], config) - graph, builder = _make_graph("audio_decoder") + graph, builder = _make_graph(self.audio_model_name) hidden = builder.input( "hidden_states", dtype=config.dtype, shape=[batch, seq, config.hidden_size] ) @@ -105,12 +110,40 @@ def build(self, module, config: FishOmniConfig) -> ModelPackage: ir.SymbolicDim(batch), past, ) - logits, present = module.audio_decoder( - builder.op, hidden, codes, semantic, unseen, mask, positions, tower_caches - ) + 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["audio_decoder"] = _make_model(graph) + 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 " @@ -124,3 +157,30 @@ def build(self, module, config: FishOmniConfig) -> ModelPackage: "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/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py index 4bb4f13fe..6ae46d4b6 100644 --- a/tests/integration/fish_omni_test.py +++ b/tests/integration/fish_omni_test.py @@ -8,9 +8,11 @@ import copy import dataclasses import importlib.util +import json import sys import types import urllib.request +from pathlib import Path import ml_dtypes import numpy as np @@ -25,6 +27,7 @@ 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, @@ -121,22 +124,33 @@ def _next_cache(outputs): @pytest.mark.parametrize("backbone_codebooks", [0, 1]) +@pytest.mark.parametrize("task", ["fish-omni", "fish-omni-generation"]) @pytest.mark.parametrize( - "dtype", + "dtype,execution_provider", [ - ir.DataType.FLOAT, + (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", + ), + ), ], ) def test_fish_omni_pipeline_prefill_and_cached_decode( - upstream, tmp_path, monkeypatch, backbone_codebooks, dtype + upstream, tmp_path, monkeypatch, backbone_codebooks, dtype, execution_provider, task ): raw = copy.deepcopy(FISH_OMNI_CONFIG) raw["audio_decoder_config"]["vq_backbone_codebooks"] = backbone_codebooks @@ -144,7 +158,7 @@ def test_fish_omni_pipeline_prefill_and_cached_decode( torch.manual_seed(42) tolerance = 1e-4 if dtype == ir.DataType.FLOAT else 1e-2 torch_dtype = torch.float32 if dtype == ir.DataType.FLOAT else torch.bfloat16 - device = "cpu" if dtype == ir.DataType.FLOAT else "cuda:0" + 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 ) @@ -161,8 +175,8 @@ def test_fish_omni_pipeline_prefill_and_cached_decode( package = build_from_module( FishOmniModel(config), config, - task="fish-omni", - execution_provider="default" if dtype == ir.DataType.FLOAT else "onnx-standard", + task=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) @@ -170,12 +184,32 @@ def test_fish_omni_pipeline_prefill_and_cached_decode( 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 dtype != ir.DataType.FLOAT: + if device != "cpu": ort.preload_dlls() sessions = { - name: OnnxModelSession(model, device="cpu" if dtype == ir.DataType.FLOAT else "cuda") + 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 {} + ) backbone.setup_caches(1, 32) tower.setup_caches(32) caches = {name: _empty_cache(model) for name, model in package.items()} @@ -261,10 +295,12 @@ def greedy(logits, *args, **kwargs): ] ] ) - for i in range(config.num_codebooks - 1): - tower_output = sessions["audio_decoder"].run( + for i in range(1 if fast else config.num_codebooks - 1): + tower_output = sessions[audio_name].run( { - **caches["audio_decoder"], + **sampling, + **({"generate": np.array(True)} if fast else {}), + **caches[audio_name], "hidden_states": outputs["hidden_states"], "shared_input": embedded["shared_input"], "codes": next_codes, @@ -273,11 +309,15 @@ def greedy(logits, *args, **kwargs): "attention_mask": mask, } ) - next_codes[0, 0, i + 1] = ( - tower_output["codebook_logits"][0, 0, i] - .astype(np.float32) - .argmax() - ) + 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], @@ -300,13 +340,14 @@ def greedy(logits, *args, **kwargs): rtol=tolerance, ) caches["decoder"] = _next_cache(outputs) - caches["audio_decoder"] = _next_cache(tower_output) + caches[audio_name] = _next_cache(tower_output) start += length continue semantic = tower.is_semantic(next_frames[:, 0]) - tower_output = sessions["audio_decoder"].run( + tower_output = sessions[audio_name].run( { - **caches["audio_decoder"], + **sampling, + **caches[audio_name], "hidden_states": outputs["hidden_states"], "shared_input": embedded["shared_input"], "codes": next_frames[:, 1:].cpu().numpy()[None], @@ -330,8 +371,141 @@ def greedy(logits, *args, **kwargs): rtol=tolerance, ) caches["decoder"] = _next_cache(outputs) - caches["audio_decoder"] = _next_cache(tower_output) + caches[audio_name] = _next_cache(tower_output) start += length + if execution_provider == "cuda": + 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() From 2a25cd90a1ee45c4e1652490c525c3e3a0ab10ca Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 02:46:46 -0700 Subject: [PATCH 10/22] Isolate selected experts and enforce real-weight precision gates Replace portable all-expert execution with selected-pair gate/up and down BMMs, using bounded 16-token prefill chunks and an unpadded decode path. Preserve root-scope streamed banks and FP32 weighted reduction. This avoids inactive expert overflow contaminating routed results and reduces wasted computation. Prepare and stress-test FP16 residuals, native mixed-output projections and opt-in logged weight saturation, but keep both reduced-precision production exports closed: independent pinned production-dimension real-weight probes failed the unchanged 1e-2 logit gate despite synthetic successes. Qualification-only task subclasses carry explicit unqualified metadata. Correct scalar reciprocal and SiLU opmath without relaxing tolerances. Account for reconstruction dtype and saturation scratch in streaming reports. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 37 ++- scripts/benchmark_drama3.py | 1 + scripts/benchmark_drama3_test.py | 45 ++- scripts/validate_drama3.py | 43 +-- src/mobius/_configs/fish_omni.py | 8 +- src/mobius/components/__init__.py | 3 +- src/mobius/components/_activations.py | 8 +- src/mobius/components/_activations_test.py | 4 +- src/mobius/functions/causal_conv.py | 3 +- src/mobius/functions/causal_conv_test.py | 6 +- src/mobius/integrations/_weight_loading.py | 67 +++- .../transformers/_fish_omni_weights.py | 29 +- .../transformers/_fish_omni_weights_test.py | 114 ++++++- src/mobius/models/fish_omni.py | 236 ++++++++++++--- src/mobius/models/fish_omni_test.py | 285 +++++++++++++++++- src/mobius/tasks/_fish_omni.py | 16 +- testdata/cases/speech/drama-3.yaml | 2 +- tests/integration/fish_omni_test.py | 165 +++++++++- 18 files changed, 934 insertions(+), 138 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 5cec8ca1e..69c059b4f 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -122,16 +122,26 @@ 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. Native `bsgemm-moe` execution is **not** claimed. Plan for roughly 144 GB of dense FP32 -weights, plus the source checkpoint and runtime allocations. FP16 is rejected -until its FP32 residual/mixed-output projection path is implemented. An opt-in -BF16 dense-weight path uses `dtype="bf16", execution_provider="onnx-standard"`; -its reduced upstream pipeline passes full-logit parity at `1e-2` and 20 exact -cached greedy frames on CUDA. Full-checkpoint BF16 validation is pending. +weights, plus the source checkpoint and runtime allocations. FP16 and BF16 +production exports remain **explicitly disabled**. Test-only qualification +prototypes pass synthetic full-logit parity at `1e-2` and 20 exact cached greedy +frames, but stricter pinned real-weight probes exposed remaining production-size +rounding differences. Full-checkpoint reduced-precision validation is pending. +The FP16 qualification +prototype implements FP32 residual streams, native FP16-input/FP32-output +`GemmFloat8` tower down projections, and logged FP32 saturation before backbone +weight casts. Overflow and synthetic-source tests pass, but a pinned real-weight, +production-dimension reduced fixture still exceeds the unchanged `1e-2` logit +gate (FP16: 3 of 12288 prefill-tower values, maximum violating error `0.01171875`; +BF16: 373 values, maximum violating error `0.0625`). +That is a failed qualification, not permission to widen tolerances or advertise +reduced-precision support. Checkpoint storage is never modified; other adapters retain their +conversion behavior unless they explicitly opt into FP16 saturation. 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 does not initialize on -CUDA, so use the standard-ONNX export path for this dtype. +ORT 1.28's native BF16 `CausalConvWithState` fallback also does not initialize +on CUDA; qualification prototypes use standard-ONNX expansion for this dtype. | Graph | Inputs | Outputs | |---|---|---| @@ -182,7 +192,11 @@ 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. It still stores dense FP32 weights: this is **not native FP8**, and it does not remove the -large model-memory requirement. The portable/default export remains available. +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` @@ -201,8 +215,11 @@ Outputs include `frame_codes` INT64 `[B,1,10]`, `valid` BOOL `[B]`, 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 BF16 standard-ONNX route also supports this one-pass task, but remains -experimentally qualified as described above. +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. diff --git a/scripts/benchmark_drama3.py b/scripts/benchmark_drama3.py index 45effa342..544b6d037 100644 --- a/scripts/benchmark_drama3.py +++ b/scripts/benchmark_drama3.py @@ -273,6 +273,7 @@ def benchmark(args): dtype = sessions["decoder"].session.get_outputs()[0].type element_type = { "tensor(float)": ir.DataType.FLOAT, + "tensor(float16)": ir.DataType.FLOAT16, "tensor(bfloat16)": ir.DataType.BFLOAT16, }[dtype] diff --git a/scripts/benchmark_drama3_test.py b/scripts/benchmark_drama3_test.py index 8515cd0df..eed8c8175 100644 --- a/scripts/benchmark_drama3_test.py +++ b/scripts/benchmark_drama3_test.py @@ -36,6 +36,7 @@ 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 @@ -78,7 +79,9 @@ def test_zero_copy_contiguous_views_bind_the_correct_bytes(tmp_path): @pytest.mark.parametrize("task", ["fish-omni", "fish-omni-generation"]) @pytest.mark.parametrize("sample", [False, True]) -@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.BFLOAT16]) +@pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] +) @pytest.mark.parametrize( "device", [ @@ -96,8 +99,8 @@ def test_zero_copy_contiguous_views_bind_the_correct_bytes(tmp_path): def test_benchmark_runs_actual_cached_graphs_without_claiming_full_model_perf( tmp_path, task, sample, device, dtype ): - if dtype == ir.DataType.BFLOAT16 and device == "cpu": - pytest.skip("Stock ORT CPU has no BF16 Attention kernel") + 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)) ) @@ -105,16 +108,36 @@ def test_benchmark_runs_actual_cached_graphs_without_claiming_full_model_perf( package = build_from_module( FishOmniModel(config), config, - task=task, - execution_provider="onnx-standard" if dtype == ir.DataType.BFLOAT16 else "default", + 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 = { - _source_name(name): torch.randn(tuple(value.shape)) * 0.1 - for model in package.values() - for name, value in model.graph.initializers.items() - if value.const_value is None - } + 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 diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index b9ff513c2..941bbf47f 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -35,6 +35,7 @@ AUDIO_DEVICE = 1 DECODER_DEVICE = 2 DTYPE = "f32" +DTYPE_NAMES = {"f32": "float32", "f16": "float16", "bf16": "bfloat16"} GOLDEN = ( pathlib.Path(__file__).resolve().parents[1] / "testdata" @@ -63,7 +64,11 @@ def empty_cache_ports(ports): 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(bfloat16)": ml_dtypes.bfloat16} + 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]) @@ -150,7 +155,7 @@ def export(checkpoint): revision=REVISION, dtype=DTYPE, keep_quantized=False, - execution_provider="onnx-standard" if DTYPE == "bf16" else "default", + execution_provider={"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) @@ -194,7 +199,7 @@ def reference(checkpoint): checkpoint, devices=[f"cuda:{device}" for device in REFERENCE_DEVICES], tower_device=f"cuda:{REFERENCE_TOWER_DEVICE}", - dtype=torch.bfloat16 if DTYPE == "bf16" else torch.float32, + 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 @@ -284,7 +289,7 @@ def capture(_module, _inputs, output): "model_id": "fishaudio/drama-3", "revision": REVISION, "source_revision": SOURCE_REVISION, - "dtype": "bfloat16" if DTYPE == "bf16" else "float32", + "dtype": DTYPE_NAMES[DTYPE], "prompt": PROMPT, "prompt_rows": len(prompt), "prefill_shape": list(logits.shape), @@ -312,7 +317,7 @@ def validate_reference_golden(golden): import numpy as np committed = json.loads(GOLDEN.read_text()) - assert golden["dtype"] == ("bfloat16" if DTYPE == "bf16" else "float32"), "dtype" + assert golden["dtype"] == DTYPE_NAMES[DTYPE], "dtype" for field in ( "model_id", "revision", @@ -333,8 +338,8 @@ def validate_reference_golden(golden): np.testing.assert_allclose( golden["prefill_last_top10_logits"], committed["prefill_last_top10_logits"], - atol=1e-2 if DTYPE == "bf16" else 1e-4, - rtol=1e-2 if DTYPE == "bf16" else 1e-4, + atol=1e-2 if DTYPE != "f32" else 1e-4, + rtol=1e-2 if DTYPE != "f32" else 1e-4, ) @@ -401,7 +406,11 @@ def compare(): def run(name, feeds): return run_session(sessions[name], feeds) - expected_type = "tensor(bfloat16)" if DTYPE == "bf16" else "tensor(float)" + expected_type = { + "f32": "tensor(float)", + "f16": "tensor(float16)", + "bf16": "tensor(bfloat16)", + }[DTYPE] assert sessions["decoder"].get_outputs()[0].type == expected_type def empty_cache(name): @@ -428,7 +437,7 @@ def check(name, actual, expected): "mean_abs_diff": float(difference.mean()) if difference.size else 0.0, } event("comparison", name=name, **metrics[name]) - tolerance = 1e-2 if DTYPE == "bf16" else 1e-3 + tolerance = 1e-2 if DTYPE != "f32" else 1e-3 np.testing.assert_allclose( actual, expected, atol=tolerance, rtol=tolerance, err_msg=name ) @@ -552,12 +561,12 @@ def check(name, actual, expected): "status": "passed", "revision": REVISION, "source_revision": SOURCE_REVISION, - "dtype": "bfloat16" if DTYPE == "bf16" else "float32", + "dtype": DTYPE_NAMES[DTYPE], "onnxruntime_version": ort.__version__, "provider": "CUDAExecutionProvider", "ort_graph_optimization": "ORT_ENABLE_BASIC", - "atol": 1e-2 if DTYPE == "bf16" else 1e-3, - "rtol": 1e-2 if DTYPE == "bf16" else 1e-3, + "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"], @@ -576,7 +585,7 @@ def check(name, actual, expected): 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", "bf16"), default="f32") + parser.add_argument("--dtype", choices=("f32", "f16", "bf16"), default="f32") 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) @@ -594,12 +603,12 @@ def check(name, actual, expected): DTYPE = args.dtype if args.golden is not None: GOLDEN = args.golden.resolve() - elif DTYPE == "bf16": - GOLDEN = GOLDEN.with_name("drama-3_bf16_tokens_generation.json") + elif DTYPE != "f32": + GOLDEN = GOLDEN.with_name(f"drama-3_{DTYPE}_tokens_generation.json") if args.stage == "compare" and not GOLDEN.is_file(): parser.error( - "BF16 comparison requires --golden pointing to the independent " - "reference stage's golden.json until a BF16 golden is committed" + "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() diff --git a/src/mobius/_configs/fish_omni.py b/src/mobius/_configs/fish_omni.py index 2fb29c626..b58b0bbb9 100644 --- a/src/mobius/_configs/fish_omni.py +++ b/src/mobius/_configs/fish_omni.py @@ -193,12 +193,10 @@ def validate(self) -> None: if self.quantization is not None or self.component_quantization: raise ValueError( "Drama 3 affine component quantization is unsupported; " - "use the FP32 graph with explicit FP8 checkpoint dequantization." - ) - if self.dtype not in (ir.DataType.FLOAT, ir.DataType.BFLOAT16): - raise ValueError( - "Drama 3 FP16 requires an FP32 residual/mixed-output projection path." + "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.tower is None: raise ValueError("Drama 3 requires an audio tower configuration") self.tower.validate() diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 4792905f9..d3c43a191 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -122,6 +122,7 @@ "SequenceSelectiveScan", "SiLU", "Float32SiLU", + "float32_silu", "Siglip2NaFlexVisionEmbeddings", "Siglip2NaFlexVisionModel", "SigmoidTopKGate", @@ -166,7 +167,7 @@ "siglip2_naflex_attention_mask", ] -from mobius.components._activations import Float32SiLU, 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 ba33ca6c2..4a2db66fc 100644 --- a/src/mobius/components/_activations.py +++ b/src/mobius/components/_activations.py @@ -124,4 +124,10 @@ 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(op.Swish(op.Cast(x, to=ir.DataType.FLOAT)), x) + 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 c755850b0..8b164a967 100644 --- a/src/mobius/components/_activations_test.py +++ b/src/mobius/components/_activations_test.py @@ -42,8 +42,8 @@ def test_float32_silu_rounds_only_at_output(self, dtype, array_dtype, torch_dtyp output = Float32SiLU()(builder.op, x) output.type, output.shape = x.type, x.shape builder.add_output(output, "output") - swish = next(node for node in graph if node.op_type == "Swish") - assert swish.inputs[0].producer().attributes["to"].as_int() == ir.DataType.FLOAT + 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)) diff --git a/src/mobius/functions/causal_conv.py b/src/mobius/functions/causal_conv.py index cf0fd073d..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" @@ -129,7 +130,7 @@ def body(op, input_val, weight_val, bias_val, conv_state_val): # Step 5: Apply activation. if activation in ("silu", "swish"): output = op.CastLike( - op.Swish(op.Cast(conv_out, to=ir.DataType.FLOAT)), + float32_silu(op, conv_out), conv_out, ) elif activation == "none": 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/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index f966c5bf2..8a77ba488 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -106,6 +106,8 @@ class StreamingWeightSource: # 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) @@ -704,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]], @@ -721,7 +746,7 @@ def _materialize_preprocessed_source( name=source.source_name, target_dtype=source.reconstruction_dtype, ) - return tensor if tensor.dtype == target_dtype 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( @@ -738,14 +763,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}") @@ -1067,10 +1092,20 @@ def validate_source( 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 + reconstruction_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]: @@ -1092,10 +1127,25 @@ 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 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( @@ -1105,9 +1155,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/transformers/_fish_omni_weights.py b/src/mobius/integrations/transformers/_fish_omni_weights.py index ae0123627..cecd523eb 100644 --- a/src/mobius/integrations/transformers/_fish_omni_weights.py +++ b/src/mobius/integrations/transformers/_fish_omni_weights.py @@ -20,7 +20,7 @@ ) -def _classify_source(source, index, config, *, expert): +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] @@ -32,10 +32,11 @@ def _classify_source(source, index, config, *, expert): 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) + return StreamingWeightSource(source, saturate_float16=saturate_float16) def _source_name(target: str) -> str: @@ -70,28 +71,42 @@ def build_fish_omni_streaming_plan( 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}.weight") - for projection in ("gate_proj", "up_proj", "down_proj") + for projection in ("gate_proj", "up_proj", "down_proj", "gate_up_proj") ): 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}.{projection}", + 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 + source, + index, + config, + expert=".mlp.experts." in source, + saturate_float16=saturate, ) return StreamingWeightPlan( targets=targets, diff --git a/src/mobius/integrations/transformers/_fish_omni_weights_test.py b/src/mobius/integrations/transformers/_fish_omni_weights_test.py index c7201f8c3..bb80b02bc 100644 --- a/src/mobius/integrations/transformers/_fish_omni_weights_test.py +++ b/src/mobius/integrations/transformers/_fish_omni_weights_test.py @@ -11,7 +11,7 @@ import pytest import torch from _test_configs import FISH_OMNI_CONFIG -from safetensors.torch import save_file +from safetensors.torch import load_file, save_file from mobius import build_from_module from mobius.__main__ import main @@ -30,12 +30,29 @@ def _checkpoint(package): torch.manual_seed(19) - return { - _source_name(name): torch.randn(tuple(value.shape), dtype=torch.float32) * 0.02 - for model in package.values() - for name, value in model.graph.initializers.items() - if value.const_value is None - } + 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 def test_fish_omni_streams_all_components(tmp_path): @@ -47,7 +64,25 @@ def test_fish_omni_streams_all_components(tmp_path): for model in package.values(): for name, value in model.graph.initializers.items(): assert value.const_value is not None, name - if name.startswith( + 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.", @@ -231,3 +266,66 @@ def test_fish_omni_cli_persists_three_weighted_models(tmp_path, task, ep): ) 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/fish_omni.py b/src/mobius/models/fish_omni.py index 39346d2bb..f6c4ebb74 100644 --- a/src/mobius/models/fish_omni.py +++ b/src/mobius/models/fish_omni.py @@ -22,7 +22,6 @@ from mobius._configs import ArchitectureConfig from mobius._configs.fish_omni import FishOmniConfig from mobius.components import ( - MLP, Embedding, Float32SiLU, Linear, @@ -32,9 +31,11 @@ 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, @@ -128,7 +129,7 @@ def forward(self, op, hidden_states, gate): scaled = op.Mul(rounded, self.weight) gated = op.Mul( op.Cast(scaled, to=ir.DataType.FLOAT), - op.Swish(op.Cast(gate, to=ir.DataType.FLOAT)), + float32_silu(op, gate), ) return op.CastLike(gated, hidden_states) @@ -198,36 +199,12 @@ def __init__(self, config, gate=None): norm_topk_prob=config.norm_topk_prob, ) super().__init__(config, gate=gate) - self.packed_experts = ( - _FishCudaExperts(config) if config.dtype == ir.DataType.FLOAT else None - ) - if self.experts is not None: - for expert in self.experts: - assert isinstance(expert, MLP) - expert.act_fn = Float32SiLU() + self.packed_experts = _FishExperts(config) self.shared_expert.act_fn = Float32SiLU() def forward(self, op, hidden_states): - if self.experts is None: - return super().forward(op, hidden_states) routing, selected, logits = self.gate(op, hidden_states) - caps = ep_capabilities() - if self.packed_experts is not None and caps.name == "cuda" and caps.supports_fused_moe: - # Floating CUDA MoE performs its own softmax/top-k from RAW logits. - routed = self.packed_experts(op, hidden_states, logits) - return self._add_shared_expert(op, hidden_states, routed) - combined = None - for index, expert in enumerate(self.experts): - output = expert(op, hidden_states) - match = op.CastLike(op.Equal(selected, index), routing) - weight = op.ReduceSum(op.Mul(routing, match), [-1], keepdims=True) - contribution = op.Mul( - op.Cast(output, to=ir.DataType.FLOAT), - op.Cast(weight, to=ir.DataType.FLOAT), - ) - combined = contribution if combined is None else op.Add(combined, contribution) - assert combined is not None - routed = op.CastLike(combined, 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): @@ -236,41 +213,152 @@ def _add_shared_expert(self, op, hidden_states, routed): return op.Add(routed, op.Mul(shared, gate)) -class _FishCudaExperts(nn.Module): - """Source-layout expert banks consumed sparsely by ORT's floating CUDA MoE.""" +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.dtype == ir.DataType.FLOAT + 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_proj = _FishExpertBank( [self.num_experts, self.intermediate_size, self.hidden_size] ) self.up_proj = _FishExpertBank( [self.num_experts, self.intermediate_size, self.hidden_size] ) + 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, router_logits): + 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_proj(op), + self.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_bank, up_bank, down_bank): # CUDA MoE declares [E,H,I]/[E,I,H], but consumes the original # output-major HF bytes. Reshape changes the ABI shape, NOT the layout. # This matches upstream test_moe_cuda.py's separate-FC3 Mixtral path. gate = op.Reshape( - self.gate_proj(op), [self.num_experts, self.hidden_size, self.intermediate_size] - ) - up = op.Reshape( - self.up_proj(op), [self.num_experts, self.hidden_size, self.intermediate_size] + gate_bank, [self.num_experts, self.hidden_size, self.intermediate_size] ) + up = op.Reshape(up_bank, [self.num_experts, self.hidden_size, self.intermediate_size]) down = op.Reshape( - self.down_proj(op), [self.num_experts, self.intermediate_size, self.hidden_size] + down_bank, [self.num_experts, self.intermediate_size, self.hidden_size] ) output = op.MoE( op.Reshape(hidden_states, [-1, self.hidden_size]), @@ -314,6 +402,7 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): 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( @@ -331,6 +420,28 @@ def __init__(self, config, layer_idx): 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): @@ -348,8 +459,11 @@ def __init__(self, config: FishOmniConfig): super().__init__() self.model = _FishTextModel(config) self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + 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, @@ -358,6 +472,8 @@ def forward(self, op, inputs_embeds, attention_mask, position_ids, past_key_valu 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 @@ -394,9 +510,9 @@ def code_sum(first: int, end: int): return op.ReduceSum(embeds, [-2], keepdims=False) count = cfg.vq_backbone_codebooks - mixed = op.Div( + mixed = op.Mul( op.Cast(op.Add(text, code_sum(1, count + 1)), to=ir.DataType.FLOAT), - math.sqrt(1 + count), + 1.0 / math.sqrt(1 + count), ) mixed = op.CastLike(mixed, text) unseen = code_sum(count + 1, cfg.num_codebooks) @@ -462,12 +578,42 @@ def forward(self, op, hidden, bias, positions, past): return self.wo(op, op.Mul(attended, _sigmoid_float32(op, gate))), (pk, pv) +class _FP32OutputLinear(Linear): + """Native FP16-input GEMM with FP32 output, without promoting the weight bank.""" + + 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 = op.GemmFloat8( + op.Reshape(hidden, [-1, self.in_features]), + 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) + 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) + + 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) - self.w2 = Linear(config.intermediate_size, config.hidden_size, bias=False) + self.w2 = ( + _FP32OutputLinear(config.intermediate_size, config.hidden_size) + if config.dtype == ir.DataType.FLOAT16 + else Linear(config.intermediate_size, config.hidden_size, bias=False) + ) self.act_fn = Float32SiLU() def forward(self, op, hidden): @@ -482,8 +628,16 @@ def __init__(self, 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 ) @@ -516,7 +670,7 @@ def forward(self, op, hidden, codes, semantic, bias, positions, past, shared_inp op, op.Concat(self.enorm(op, embedded), self.hnorm(op, hidden), axis=-1) ) output, present = self.transformer_layer(op, mixed, bias, positions, past) - return self.final_layernorm(op, output), present + return op.CastLike(self.final_layernorm(op, output), mixed), present class FishOmniAudioDecoder(nn.Module): diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py index 514d742c8..df68bb275 100644 --- a/src/mobius/models/fish_omni_test.py +++ b/src/mobius/models/fish_omni_test.py @@ -14,10 +14,12 @@ import pytest import torch from _test_configs import FISH_OMNI_CONFIG -from onnxscript import GraphBuilder +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 ( @@ -25,7 +27,9 @@ _dict_to_pretrained_config, ) from mobius.models import FishOmniModel -from mobius.models.fish_omni import _apply_fish_attention +from mobius.models.fish_omni import _apply_fish_attention, _FishMoEBlock, _FP32OutputLinear +from mobius.tasks import FishOmniGenerationTask, FishOmniTask, ModelTask +from mobius.tasks._base import _make_graph, _make_model def tiny_config(): @@ -34,14 +38,52 @@ def tiny_config(): ) -@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.BFLOAT16]) +class _ReducedPrecisionQualificationOnly: + def _validate_precision(self, config): + if config.dtype not in (ir.DataType.FLOAT16, ir.DataType.BFLOAT16): + return super()._validate_precision(config) + expected = "cuda" if config.dtype == ir.DataType.FLOAT16 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-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() + + +@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="fish-omni", - execution_provider="onnx-standard" if dtype == ir.DataType.BFLOAT16 else "default", + 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"} @@ -94,15 +136,31 @@ def test_fish_omni_cuda_build_uses_sparse_fused_expert_banks(): assert not any(".experts.0." in name for name in graph.initializers) -@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.BFLOAT16]) +@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 = "onnx-standard" if dtype == ir.DataType.BFLOAT16 else "default" + ep = { + ir.DataType.FLOAT: "default", + ir.DataType.FLOAT16: "cuda", + ir.DataType.BFLOAT16: "onnx-standard", + }[dtype] teacher = build_from_module( - FishOmniModel(config), config, task="fish-omni", execution_provider=ep + 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="fish-omni-generation", execution_provider=ep + 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 @@ -129,18 +187,34 @@ def test_fish_single_pass_task_reuses_tower_weights_and_layer_count(dtype): ) -@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16]) +@pytest.mark.parametrize("dtype", [ir.DataType.DOUBLE]) def test_fish_omni_rejects_unqualified_precision(dtype): - with pytest.raises(ValueError, match="FP16 requires"): + with pytest.raises(ValueError, match="requires floating"): FishOmniModel(dataclasses.replace(tiny_config(), dtype=dtype)) -def test_fish_omni_bfloat16_rejects_unqualified_ep(): - config = dataclasses.replace(tiny_config(), dtype=ir.DataType.BFLOAT16) - with pytest.raises(ValueError, match="requires execution_provider='onnx-standard'"): +@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"), [ @@ -285,3 +359,186 @@ def test_fish_attention_preserves_padding_and_cache(padded): 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="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/_fish_omni.py b/src/mobius/tasks/_fish_omni.py index 1de51720c..b43cf519b 100644 --- a/src/mobius/tasks/_fish_omni.py +++ b/src/mobius/tasks/_fish_omni.py @@ -45,12 +45,26 @@ class FishOmniTask(ModelTask): def _sampling_inputs(self, builder) -> FishAudioSamplingInputs | None: return None - def build(self, module, config: FishOmniConfig) -> ModelPackage: + 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.dtype == ir.DataType.BFLOAT16 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" diff --git a/testdata/cases/speech/drama-3.yaml b/testdata/cases/speech/drama-3.yaml index 7cbbda7a6..8f8ff6782 100644 --- a/testdata/cases/speech/drama-3.yaml +++ b/testdata/cases/speech/drama-3.yaml @@ -8,4 +8,4 @@ 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 full 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). Experimental BF16 has reduced upstream 20-frame parity only; full-checkpoint BF16, FP16 and native FP8 remain unqualified." +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/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py index 6ae46d4b6..e4ec5d909 100644 --- a/tests/integration/fish_omni_test.py +++ b/tests/integration/fish_omni_test.py @@ -33,6 +33,7 @@ stream_fish_omni_safetensors_to_package, ) from mobius.models import FishOmniModel +from mobius.models.fish_omni_test import reduced_precision_prototype_task pytestmark = pytest.mark.integration _PREFIX = "fish_speech.models.text2semantic.fish_omni" @@ -110,7 +111,11 @@ def _empty_cache(model): if value.name.endswith((".key", ".value")): shape[2] = 0 assert all(isinstance(dim, int) for dim in shape), (value.name, shape) - dtype = ml_dtypes.bfloat16 if value.dtype == ir.DataType.BFLOAT16 else np.float32 + 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 @@ -147,17 +152,39 @@ def _next_cache(outputs): 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 + 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 = torch.float32 if dtype == ir.DataType.FLOAT else torch.bfloat16 + 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 @@ -175,7 +202,7 @@ def test_fish_omni_pipeline_prefill_and_cached_decode( package = build_from_module( FishOmniModel(config), config, - task=task, + 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") @@ -210,8 +237,9 @@ def test_fish_omni_pipeline_prefill_and_cached_decode( if fast else {} ) - backbone.setup_caches(1, 32) - tower.setup_caches(32) + 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 @@ -219,6 +247,9 @@ def test_fish_omni_pipeline_prefill_and_cached_decode( 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 = [] @@ -235,7 +266,7 @@ def greedy(logits, *args, **kwargs): # 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 [4] + [1] * 20: + for length in [prefill_length] + [1] * 20: positions = torch.arange(start, start + length, device=device) positions_np = positions.cpu().numpy()[None] if start: @@ -373,7 +404,7 @@ def greedy(logits, *args, **kwargs): caches["decoder"] = _next_cache(outputs) caches[audio_name] = _next_cache(tower_output) start += length - if execution_provider == "cuda": + if execution_provider == "cuda" and dtype == ir.DataType.FLOAT: profile = json.loads( Path(sessions["decoder"]._session.end_profiling()).read_text() ) @@ -509,3 +540,121 @@ def test_sampled_tower_matches_sequential_replays_and_rejects_multiframe_generat 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, + ) From dc69e5daa83e220ac8282f3468d743c33e235aac Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 08:53:28 -0700 Subject: [PATCH 11/22] Scope Fish RoPE precision and preserve source projection rounding Limit FP32 rotary-cache preservation to the Fish model instead of changing every architecture. Restore generic dtype casting and GQA rewrites that caused eight CI failures; all 52 affected regression tests now pass. Use explicit FP32 exponential/division sigmoid at native rounding boundaries and preserve BF16 tower down-projection output-major layout with standard Gemm. Controlled real-weight probes now match BF16 prefill and initial cached calls exactly, but later RMS reduction drift still fails the unchanged gate. Keep both reduced-precision production admissions closed and document failed stages. Add exact negative-tail/projection regressions and fix the multi-device validation bridge's current CUDA context. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 15 ++- scripts/validate_drama3.py | 12 ++- scripts/validate_drama3_test.py | 35 +++++++ src/mobius/components/_rotary_embedding.py | 2 - src/mobius/models/fish_omni.py | 48 ++++++--- src/mobius/models/fish_omni_test.py | 109 ++++++++++++++++++++- 6 files changed, 200 insertions(+), 21 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 69c059b4f..22e8f6b9d 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -130,10 +130,19 @@ rounding differences. Full-checkpoint reduced-precision validation is pending. The FP16 qualification prototype implements FP32 residual streams, native FP16-input/FP32-output `GemmFloat8` tower down projections, and logged FP32 saturation before backbone -weight casts. Overflow and synthetic-source tests pass, but a pinned real-weight, -production-dimension reduced fixture still exceeds the unchanged `1e-2` logit -gate (FP16: 3 of 12288 prefill-tower values, maximum violating error `0.01171875`; +weight casts. Overflow and synthetic-source tests pass, but the initial pinned +real-weight, production-dimension reduced fixture exceeded the unchanged `1e-2` +logit gate (FP16: 3 of 12288 prefill-tower values, maximum violating error `0.01171875`; BF16: 373 values, maximum violating error `0.0625`). +Subsequent source-fed probes identified negative-tail sigmoid rounding and +down-projection accumulation order. Explicit FP32 exponential/division sigmoid +and output-major BF16 `Gemm` restore bit-exact reduced-fixture BF16 prefill +and the first two cached calls, but the third cached call still fails +(62 of 4096 values, maximum violating error `0.03125`). Its first divergence +is a one-ULP FP32 RMS variance reduction difference before BF16 rounding. +These are diagnostic improvements, not completed reduced-precision qualification. +The updated FP16 fixture still fails prefill (7 of 12288 values, maximum +violating error `0.0166015625`). That is a failed qualification, not permission to widen tolerances or advertise reduced-precision support. Checkpoint storage is never modified; other adapters retain their conversion behavior unless they explicitly opt into FP16 saturation. diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index 941bbf47f..449af4064 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -13,6 +13,7 @@ import argparse import collections +import contextlib import hashlib import importlib.util import json @@ -353,7 +354,16 @@ def run_session(session, feeds): name: _numpy_to_ort_value(np.ascontiguousarray(value) if value.ndim else value) for name, value in feeds.items() } - values = session(**converted) + 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)} diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py index 0be14c9c3..4644364c2 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -6,6 +6,7 @@ from __future__ import annotations import json +from contextlib import contextmanager from types import SimpleNamespace import ml_dtypes @@ -121,6 +122,40 @@ def test_runtime_bridge_preserves_types_and_output_names(tmp_path, 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)] + + def test_committed_token_golden_is_pinned_and_complete(): golden = json.loads(GOLDEN.read_text()) validate_reference_golden(golden) diff --git a/src/mobius/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 12f8863f3..4e286d40b 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -130,8 +130,6 @@ def __init__( name="sin_cache", data=ir.tensor(sin_cache_data), ) - self.cos_cache._keep_float32 = True # type: ignore[attr-defined] - self.sin_cache._keep_float32 = True # type: ignore[attr-defined] def _cast_embeddings( self, diff --git a/src/mobius/models/fish_omni.py b/src/mobius/models/fish_omni.py index f6c4ebb74..1aa642be6 100644 --- a/src/mobius/models/fish_omni.py +++ b/src/mobius/models/fish_omni.py @@ -52,11 +52,17 @@ 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): - return op.CastLike(op.Sigmoid(op.Cast(value, to=ir.DataType.FLOAT)), 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( @@ -578,16 +584,33 @@ def forward(self, op, hidden, bias, positions, past): return self.wo(op, op.Mul(attended, _sigmoid_float32(op, gate))), (pk, pv) -class _FP32OutputLinear(Linear): - """Native FP16-input GEMM with FP32 output, without promoting the weight bank.""" +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( - op.Reshape(hidden, [-1, self.in_features]), + hidden, self.weight, transA=0, transB=1, @@ -597,11 +620,7 @@ def forward(self, op, hidden): _domain="com.microsoft", ) projected.type = ir.TensorType(ir.DataType.FLOAT) - 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) + return projected class _TowerMLP(nn.Module): @@ -609,11 +628,12 @@ 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) - self.w2 = ( - _FP32OutputLinear(config.intermediate_size, config.hidden_size) - if config.dtype == ir.DataType.FLOAT16 - else Linear(config.intermediate_size, config.hidden_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): diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py index df68bb275..54dd6ad2e 100644 --- a/src/mobius/models/fish_omni_test.py +++ b/src/mobius/models/fish_omni_test.py @@ -27,7 +27,13 @@ _dict_to_pretrained_config, ) from mobius.models import FishOmniModel -from mobius.models.fish_omni import _apply_fish_attention, _FishMoEBlock, _FP32OutputLinear +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 @@ -68,6 +74,62 @@ def reduced_precision_prototype_task(task): return _GenerationProbe() if task == "fish-omni-generation" else _TeacherProbe() +@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] ) @@ -99,6 +161,7 @@ def test_fish_omni_build(dtype): 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) @@ -361,6 +424,50 @@ def test_fish_attention_preserves_padding_and_cache(padded): 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(), From 8735c1ed3060cec05ac9f1fcf659a79f69530571 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 09:24:20 -0700 Subject: [PATCH 12/22] Eliminate whole expert-bank copies in Drama 3 CUDA execution Serialize gate/up banks in half-split SwiGLU order and omit FC3. The separate-FC3 floating MoE path otherwise allocates and copies the full 2 GiB expert bank on every layer invocation. Verify the swiglu/fusion=2 ABI numerically on ORT 1.28; preserve routing, gate/up ordering and shared-expert semantics. Use output-major Gemm for the vocabulary projection, avoiding a 2 GiB runtime transpose above ORT's constant-fold limit. Add no-FC3/layout and Fish-only rotary-cache assertions. Extend the full-checkpoint validator to export/select the one-pass task and verify greedy RNG/validity while preserving the original replay path. The complete pinned pre-fused CUDA and single-pass pipeline now passes independent 30-row prefill/all-vector and 20 exact complete-frame comparison. Maximum decoder/tower logit errors are 5.150e-5/4.673e-5; the profile confirms 40 actual CUDA MoE executions. Scoped regressions pass; no wall-clock performance claim is inferred. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 24 +++++- scripts/validate_drama3.py | 71 ++++++++++++++--- scripts/validate_drama3_test.py | 88 ++++++++++++++++++++++ src/mobius/components/_rotary_embedding.py | 9 ++- src/mobius/models/fish_omni.py | 36 ++++----- src/mobius/models/fish_omni_test.py | 19 ++++- 6 files changed, 209 insertions(+), 38 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 22e8f6b9d..19e658a03 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -199,7 +199,11 @@ 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. It still +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 @@ -275,12 +279,20 @@ 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 complete pinned checkpoint was validated on H200 with ORT 1.28 CUDA and +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 also passes that +independent complete-checkpoint case: all vectors for the 30-row prefill and +20 exact complete frames, with maximum decoder/tower logit differences +`5.150e-5` / `4.673e-5` and hidden difference `2.289e-5`. +The prefill profile asserts 40 actual CUDA MoE executions. This fresh check +qualifies the changed fast token graph for this case; it is not a latency or +speech-quality measurement. + To reproduce on a sufficiently large CUDA host, select free devices and run: ```bash @@ -295,6 +307,14 @@ 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. For performance measurement, `scripts/benchmark_drama3.py` runs actual cached requests with I/O binding: recurrent/KV state and neural intermediates remain on diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index 449af4064..3754e454b 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -146,7 +146,7 @@ def load_upstream(): return result -def export(checkpoint): +def export(checkpoint, *, task="fish-omni", execution_provider=None): from mobius import build start = time.monotonic() @@ -156,7 +156,9 @@ def export(checkpoint): revision=REVISION, dtype=DTYPE, keep_quantized=False, - execution_provider={"f32": "default", "f16": "cuda", "bf16": "onnx-standard"}[DTYPE], + task=task, + execution_provider=execution_provider + or {"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) @@ -369,7 +371,32 @@ def run_session(session, feeds): return {port.name: _ort_value_to_numpy(value) for port, value in zip(ports, values)} -def compare(): +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): os.environ["NVIDIA_TF32_OVERRIDE"] = "0" import numpy as np import onnxruntime as ort @@ -380,6 +407,12 @@ def compare(): 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), @@ -396,8 +429,9 @@ def compare(): options.enable_profiling = True options.profile_file_prefix = str(ROOT / "decoder-profile") event("ort_session_start", component=name, device=device) + component = audio_model_name if name == "audio_decoder" else name sessions[name] = ort_easy.EasySession( - str(ROOT / "onnx" / name / "model.onnx"), + str(ROOT / "onnx" / component / "model.onnx"), sess_options=options, providers=[ ( @@ -413,7 +447,9 @@ def compare(): ) event("ort_session_ready", component=name, seconds=time.monotonic() - start) - def run(name, feeds): + 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 = { @@ -539,7 +575,8 @@ def check(name, actual, expected): 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]]) - for index in range(count - 1): + single_pass = audio_model_name == "audio_generator" + for index in range(1 if single_pass else count - 1): tower_output = run( "audio_decoder", { @@ -551,10 +588,16 @@ def check(name, actual, expected): "position_ids": positions[:, -1:], "attention_mask": mask, }, + generate=True, ) - codes[0, 0, index + 1] = ( - tower_output["codebook_logits"][0, 0, index].astype(np.float32).argmax() - ) + 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"] ) @@ -572,6 +615,7 @@ def check(name, actual, expected): "revision": REVISION, "source_revision": SOURCE_REVISION, "dtype": DTYPE_NAMES[DTYPE], + "audio_component": audio_model_name, "onnxruntime_version": ort.__version__, "provider": "CUDAExecutionProvider", "ort_graph_optimization": "ORT_ENABLE_BASIC", @@ -596,6 +640,11 @@ def check(name, actual, expected): 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( + "--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("--golden", type=pathlib.Path) parser.add_argument("--reference-devices", default="0,1,2,3") parser.add_argument("--reference-tower-device", default=4, type=int) @@ -623,8 +672,8 @@ def check(name, actual, expected): if args.stage == "download": download() elif args.stage == "export": - export(checkpoint_path()) + export(checkpoint_path(), task=args.task, execution_provider=args.execution_provider) elif args.stage == "reference": reference(checkpoint_path()) else: - compare() + compare(audio_model_name=args.audio_component) diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py index 4644364c2..7f98cb9f3 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -156,6 +156,94 @@ def get_outputs(self): 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) 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/models/fish_omni.py b/src/mobius/models/fish_omni.py index 1aa642be6..1cabc513e 100644 --- a/src/mobius/models/fish_omni.py +++ b/src/mobius/models/fish_omni.py @@ -234,12 +234,6 @@ def __init__(self, config): self.normalize = int(config.norm_topk_prob) self.dtype = config.dtype self.act_fn = Float32SiLU() - self.gate_proj = _FishExpertBank( - [self.num_experts, self.intermediate_size, self.hidden_size] - ) - self.up_proj = _FishExpertBank( - [self.num_experts, self.intermediate_size, self.hidden_size] - ) self.gate_up_proj = _FishExpertBank( [self.num_experts, 2 * self.intermediate_size, self.hidden_size] ) @@ -255,8 +249,7 @@ def forward(self, op, hidden_states, weights, selected, router_logits): op, hidden_states, router_logits, - self.gate_proj(op), - self.up_proj(op), + self.gate_up_proj(op), self.down_proj(op), ) gate_up, down = self.gate_up_proj(op), self.down_proj(op) @@ -355,29 +348,30 @@ def _selected(self, op, hidden, weights, selected, gate_up_bank, down_bank): ) return op.CastLike(summed, hidden) - def _fused(self, op, hidden_states, router_logits, gate_bank, up_bank, down_bank): - # CUDA MoE declares [E,H,I]/[E,I,H], but consumes the original - # output-major HF bytes. Reshape changes the ABI shape, NOT the layout. - # This matches upstream test_moe_cuda.py's separate-FC3 Mixtral path. - gate = op.Reshape( - gate_bank, [self.num_experts, self.hidden_size, self.intermediate_size] + 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] ) - up = op.Reshape(up_bank, [self.num_experts, self.hidden_size, 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, + gate_up, None, down, None, - up, - activation_type="silu", + 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=0, + swiglu_fusion=2, _domain="com.microsoft", ) output.type = ir.TensorType(ir.DataType.FLOAT) @@ -464,7 +458,9 @@ class FishOmniBackbone(nn.Module): def __init__(self, config: FishOmniConfig): super().__init__() self.model = _FishTextModel(config) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + # 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): diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py index 54dd6ad2e..e4e2c7a2c 100644 --- a/src/mobius/models/fish_omni_test.py +++ b/src/mobius/models/fish_omni_test.py @@ -154,6 +154,10 @@ def test_fish_omni_build(dtype): "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 @@ -170,6 +174,14 @@ def test_fish_omni_build(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(): @@ -194,8 +206,13 @@ def test_fish_omni_cuda_build_uses_sparse_fused_expert_banks(): 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 * 3 + 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) From c6b34d076defe3cc7d5276b6960d7e763745527c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 11:14:17 -0700 Subject: [PATCH 13/22] Wire guarded native FP8 and qualify real-time Drama 3 speech Prepare typed E4M3FN expert banks, unchanged FP32 grids and explicit dynamic_fp8 QMoE emission. Add a byte-preserving streaming classification that rejects dtype conversion, malformed scales and unsafe reinterpretation. Validate every finite FP8 byte including negative zero through save/load, and compare actual exporter-emitted FLOAT/FP16/BF16 decode/prefill experts against pinned Triton on the new runtime. Keep every production native-FP8/low-dtype admission closed. Complete native-FP8/FP32 prefill fails the unchanged hidden-state gate despite all 40 real expert layers passing with identical source inputs. A 1.192e-7 upstream input difference changes an FP8 code and amplifies across layers. Separate reference execution modes and preserve their identities instead of widening tolerances or substituting weight-only FP8. Add explicit CPU embedding/CUDA codec placement, nonzero-device binding guards and measured cuDNN NC1D codec tuning. Dense FP32 real speech reaches RTF0.770 for100greedy frames and RTF0.901 for53sampled frames on the separately qualified combined runtime, with identical sampled-frame hashes. Document initialization/memory costs, runtime provenance, both ORT PRs and failed precision qualifications. Final shared/source suite295pass/17skip,9native operator-source cases pass, and scoped lint/typing pass. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 75 ++++++++- scripts/benchmark_drama3.py | 76 ++++++++- scripts/benchmark_drama3_test.py | 65 +++++++- scripts/validate_drama3.py | 54 +++++-- scripts/validate_drama3_test.py | 31 ++++ src/mobius/_configs/fish_omni.py | 21 +++ src/mobius/integrations/_weight_loading.py | 45 +++++- .../integrations/transformers/_builder.py | 5 +- .../transformers/_fish_omni_weights.py | 19 ++- .../transformers/_fish_omni_weights_test.py | 76 ++++++++- src/mobius/models/fish_omni.py | 70 ++++++++- src/mobius/models/fish_omni_test.py | 75 ++++++++- src/mobius/tasks/_fish_omni.py | 24 ++- tests/integration/fish_omni_test.py | 145 ++++++++++++++++-- 14 files changed, 729 insertions(+), 52 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 19e658a03..e33b9ba33 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -152,6 +152,26 @@ 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`). This is **not full-model low-dtype +qualification**: production native-FP8 admission also remains closed. A complete +native-FP8/FP32 probe fails the unchanged 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. The independent source +reference repeats bit-exactly. This is a pipeline qualification failure, not a +reason to substitute weight-only FP8 or widen tolerances. The custom +CUDA build needs `onnxruntime_USE_FP8_QMOE=ON`, SM89+, and the new mode; old +weight-only FP8 is not an equivalent fallback. 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]` | @@ -293,6 +313,43 @@ The prefill profile asserts 40 actual CUDA MoE executions. This fresh check qualifies the changed fast token graph for this case; it is not a latency or speech-quality measurement. +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 @@ -315,19 +372,35 @@ 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 --sample --frames 100 --output drama3-perf.json +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. diff --git a/scripts/benchmark_drama3.py b/scripts/benchmark_drama3.py index 544b6d037..544178bee 100644 --- a/scripts/benchmark_drama3.py +++ b/scripts/benchmark_drama3.py @@ -63,7 +63,16 @@ def codebook_view(value, index): 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): + 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 @@ -85,6 +94,7 @@ def __init__(self, path, *, device, device_id=0, profile=False, profile_director "device_id": device_id, "use_tf32": "0", "arena_extend_strategy": "kSameAsRequested", + "cudnn_conv1d_pad_to_nc1d": str(int(cudnn_conv1d_pad_to_nc1d)), }, ), "CPUExecutionProvider", @@ -93,7 +103,7 @@ def __init__(self, path, *, device, device_id=0, profile=False, profile_director else ["CPUExecutionProvider"] ) self.session = ort.InferenceSession(str(path), options, providers=providers) - self.device, self.device_id = device, device_id + 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=()): @@ -106,10 +116,11 @@ def run(self, feeds, *, fetches=None, cpu_outputs=()): 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=value.owner.device_name(), - device_id=self.device_id, + 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, @@ -125,7 +136,9 @@ def run(self, feeds, *, fetches=None, cpu_outputs=()): 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) + 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() @@ -225,6 +238,12 @@ def greedy_model(config, dtype): 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 @@ -251,7 +270,15 @@ def benchmark(args): "profile_directory": args.output.parent / "drama3-profiles", } sessions = { - name: BoundSession(args.package / name / "model.onnx", **kwargs) + 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 = { @@ -271,6 +298,11 @@ def benchmark(args): 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, @@ -457,7 +489,14 @@ def request(limit): path = args.codec / "waveform_decoder" / "model.onnx" if not path.is_file(): path = args.codec / "model.onnx" - sessions["waveform_decoder"] = BoundSession(path, **kwargs) + 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" @@ -498,6 +537,15 @@ def decode_waveform(frames): "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, @@ -512,6 +560,9 @@ def decode_waveform(frames): "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, @@ -589,6 +640,17 @@ def parser(): 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) diff --git a/scripts/benchmark_drama3_test.py b/scripts/benchmark_drama3_test.py index eed8c8175..74fda11f8 100644 --- a/scripts/benchmark_drama3_test.py +++ b/scripts/benchmark_drama3_test.py @@ -8,6 +8,7 @@ import copy import dataclasses import json +from types import SimpleNamespace import numpy as np import onnx_ir as ir @@ -40,6 +41,27 @@ 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"): @@ -56,6 +78,29 @@ def test_scalar_boolean_keeps_its_rank(): 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) @@ -150,7 +195,13 @@ def test_benchmark_runs_actual_cached_graphs_without_claiming_full_model_perf( 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)) - with_codec = task == "fish-omni-generation" and device == "cpu" + 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( @@ -192,14 +243,26 @@ def test_benchmark_runs_actual_cached_graphs_without_claiming_full_model_perf( *(["--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 diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index 3754e454b..19a1a3c9d 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -36,6 +36,7 @@ AUDIO_DEVICE = 1 DECODER_DEVICE = 2 DTYPE = "f32" +EXPERT_EXECUTION = "dense" DTYPE_NAMES = {"f32": "float32", "f16": "float16", "bf16": "bfloat16"} GOLDEN = ( pathlib.Path(__file__).resolve().parents[1] @@ -146,7 +147,7 @@ def load_upstream(): return result -def export(checkpoint, *, task="fish-omni", execution_provider=None): +def export(checkpoint, *, task="fish-omni", execution_provider=None, native_fp8=False): from mobius import build start = time.monotonic() @@ -155,10 +156,14 @@ def export(checkpoint, *, task="fish-omni", execution_provider=None): str(checkpoint), revision=REVISION, dtype=DTYPE, - keep_quantized=False, + keep_quantized=native_fp8, task=task, execution_provider=execution_provider - or {"f32": "default", "f16": "cuda", "bf16": "onnx-standard"}[DTYPE], + 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) @@ -168,6 +173,7 @@ def export(checkpoint, *, task="fish-omni", execution_provider=None): { "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": { @@ -181,9 +187,9 @@ def export(checkpoint, *, task="fish-omni", execution_provider=None): ) -def reference(checkpoint): +def reference(checkpoint, *, native_fp8=False): os.environ["NVIDIA_TF32_OVERRIDE"] = "0" - os.environ["FISH_DRAMA3_FP8_EXPERTS"] = "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 @@ -206,6 +212,8 @@ def reference(checkpoint): ) 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] = [] @@ -293,6 +301,7 @@ def capture(_module, _inputs, output): "revision": REVISION, "source_revision": SOURCE_REVISION, "dtype": DTYPE_NAMES[DTYPE], + "expert_execution": "dynamic_fp8" if native_fp8 else "dense", "prompt": PROMPT, "prompt_rows": len(prompt), "prefill_shape": list(logits.shape), @@ -320,6 +329,8 @@ 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" assert golden["dtype"] == DTYPE_NAMES[DTYPE], "dtype" for field in ( "model_id", @@ -458,6 +469,15 @@ def run(name, feeds, *, generate=False): "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()) @@ -539,6 +559,11 @@ def check(name, actual, expected): ) > 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) ) @@ -615,6 +640,7 @@ def check(name, actual, expected): "revision": REVISION, "source_revision": SOURCE_REVISION, "dtype": DTYPE_NAMES[DTYPE], + "expert_execution": EXPERT_EXECUTION, "audio_component": audio_model_name, "onnxruntime_version": ort.__version__, "provider": "CUDAExecutionProvider", @@ -640,6 +666,7 @@ def check(name, actual, expected): 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" ) @@ -660,20 +687,27 @@ def check(name, actual, expected): 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": - GOLDEN = GOLDEN.with_name(f"drama-3_{DTYPE}_tokens_generation.json") + 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( - "Reduced-precision comparison requires --golden pointing to the independent " + "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) + export( + checkpoint_path(), + task=args.task, + execution_provider=args.execution_provider, + native_fp8=args.native_fp8, + ) elif args.stage == "reference": - reference(checkpoint_path()) + reference(checkpoint_path(), native_fp8=args.native_fp8) else: compare(audio_model_name=args.audio_component) diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py index 7f98cb9f3..93109036c 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -270,6 +270,37 @@ def test_reference_mismatch_fails_before_model_loading(change): 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("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() diff --git a/src/mobius/_configs/fish_omni.py b/src/mobius/_configs/fish_omni.py index b58b0bbb9..348a79d3a 100644 --- a/src/mobius/_configs/fish_omni.py +++ b/src/mobius/_configs/fish_omni.py @@ -29,6 +29,7 @@ class FishOmniConfig(ArchitectureConfig): 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: @@ -197,6 +198,26 @@ def validate(self) -> None: ) 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/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index 8a77ba488..67fe00073 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -96,10 +96,10 @@ @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 @@ -112,7 +112,7 @@ class StreamingWeightSource: @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, ...], ...] @@ -737,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) @@ -876,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 @@ -922,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 = ( @@ -955,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( @@ -1071,7 +1092,9 @@ 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 " @@ -1081,6 +1104,8 @@ def validate_source( 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 = ( @@ -1117,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 " @@ -1130,6 +1157,8 @@ def validate_source( 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 = ( diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 8a2c27429..2b9f3c378 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -661,10 +661,7 @@ def build_transformers_model( 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: - raise ValueError( - "Drama 3 native bsgemm-moe FP8 execution is not supported. " - "Use keep_quantized=False (--dequantize) for streaming dense conversion." - ) + config = dataclasses.replace(config, native_fp8_experts=True) from mobius.tasks import get_task diff --git a/src/mobius/integrations/transformers/_fish_omni_weights.py b/src/mobius/integrations/transformers/_fish_omni_weights.py index cecd523eb..0a81cf36f 100644 --- a/src/mobius/integrations/transformers/_fish_omni_weights.py +++ b/src/mobius/integrations/transformers/_fish_omni_weights.py @@ -24,9 +24,15 @@ 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", @@ -75,8 +81,9 @@ def build_fish_omni_streaming_plan( source.startswith("model.language_model.") or source == "lm_head.weight" ) if any( - source.endswith(f".experts.{projection}.weight") + 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) @@ -110,7 +117,15 @@ def build_fish_omni_streaming_plan( ) return StreamingWeightPlan( targets=targets, - report={"architecture": "fish_omni", "codec_included": False}, + 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", + }, ) diff --git a/src/mobius/integrations/transformers/_fish_omni_weights_test.py b/src/mobius/integrations/transformers/_fish_omni_weights_test.py index bb80b02bc..e82531d3c 100644 --- a/src/mobius/integrations/transformers/_fish_omni_weights_test.py +++ b/src/mobius/integrations/transformers/_fish_omni_weights_test.py @@ -7,6 +7,7 @@ import dataclasses import json +import numpy as np import onnx_ir as ir import pytest import torch @@ -25,7 +26,11 @@ stream_fish_omni_safetensors_to_package, ) from mobius.models import FishOmniModel -from mobius.models.fish_omni_test import tiny_config +from mobius.models.fish_omni_test import ( + native_fp8_prototype_config, + reduced_precision_prototype_task, + tiny_config, +) def _checkpoint(package): @@ -55,6 +60,75 @@ def _checkpoint(package): 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") diff --git a/src/mobius/models/fish_omni.py b/src/mobius/models/fish_omni.py index 1cabc513e..6f9929125 100644 --- a/src/mobius/models/fish_omni.py +++ b/src/mobius/models/fish_omni.py @@ -205,7 +205,9 @@ def __init__(self, config, gate=None): norm_topk_prob=config.norm_topk_prob, ) super().__init__(config, gate=gate) - self.packed_experts = _FishExperts(config) + 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): @@ -378,6 +380,72 @@ def _fused(self, op, hidden_states, router_logits, gate_up_bank, down_bank): 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__() diff --git a/src/mobius/models/fish_omni_test.py b/src/mobius/models/fish_omni_test.py index e4e2c7a2c..d1b613a19 100644 --- a/src/mobius/models/fish_omni_test.py +++ b/src/mobius/models/fish_omni_test.py @@ -46,9 +46,16 @@ def tiny_config(): class _ReducedPrecisionQualificationOnly: def _validate_precision(self, config): - if config.dtype not in (ir.DataType.FLOAT16, ir.DataType.BFLOAT16): + 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 else "onnx-standard" + 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") @@ -56,7 +63,9 @@ def build(self, module, config): package = super().build(module, config) for model in package.values(): model.metadata_props["mobius.precision_validation"] = ( - "unqualified-reduced-precision-prototype" + "unqualified-native-fp8-prototype" + if config.native_fp8_experts + else "unqualified-reduced-precision-prototype" ) return package @@ -74,6 +83,66 @@ def reduced_precision_prototype_task(task): 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", diff --git a/src/mobius/tasks/_fish_omni.py b/src/mobius/tasks/_fish_omni.py index b43cf519b..e0adbfbf2 100644 --- a/src/mobius/tasks/_fish_omni.py +++ b/src/mobius/tasks/_fish_omni.py @@ -51,7 +51,21 @@ def _validate_precision(self, config: FishOmniConfig) -> None: "Drama 3 FP16 currently requires execution_provider='cuda' " "for the native FP16-input/FP32-output projection" ) - if config.dtype == ir.DataType.BFLOAT16 and ep_capabilities().name != "onnx-standard": + 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" @@ -163,6 +177,14 @@ def build(self, module, config: FishOmniConfig) -> ModelPackage: "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, diff --git a/tests/integration/fish_omni_test.py b/tests/integration/fish_omni_test.py index e4ec5d909..5a90b1540 100644 --- a/tests/integration/fish_omni_test.py +++ b/tests/integration/fish_omni_test.py @@ -9,6 +9,7 @@ import dataclasses import importlib.util import json +import os import sys import types import urllib.request @@ -33,12 +34,142 @@ 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.""" @@ -53,19 +184,7 @@ def upstream(tmp_path, monkeypatch): monkeypatch.setitem(sys.modules, name, package) modules = {} for name in ("fp8_moe", "config", "layers", "moe", "backbone", "tower"): - 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) - modules[name] = module + modules[name] = _load_reference_module(name, tmp_path, monkeypatch) return modules From f6ea60ace828652c7f7e4c7d128c4f3f8b693ac2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 18:11:39 -0700 Subject: [PATCH 14/22] Pin Drama 3 reference validation to the declared Torch release Require upstream's torch==2.8.0 before reference model loading and record Torch, CUDA and Triton provenance. Reject environment mismatches between versioned goldens while preserving legacy golden compatibility.\n\nDocument fresh real-weight precision results and source-boundary investigations separately from earlier Torch 2.10 diagnostics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 64 +++++++++++++++++---------------- scripts/validate_drama3.py | 19 ++++++++++ scripts/validate_drama3_test.py | 43 ++++++++++++++++++++++ 3 files changed, 95 insertions(+), 31 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index e33b9ba33..f20ddbc3f 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -120,32 +120,31 @@ 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. Native -`bsgemm-moe` execution is **not** claimed. Plan for roughly 144 GB of dense FP32 +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 **explicitly disabled**. Test-only qualification -prototypes pass synthetic full-logit parity at `1e-2` and 20 exact cached greedy -frames, but stricter pinned real-weight probes exposed remaining production-size -rounding differences. Full-checkpoint reduced-precision validation is pending. -The FP16 qualification +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. Overflow and synthetic-source tests pass, but the initial pinned -real-weight, production-dimension reduced fixture exceeded the unchanged `1e-2` -logit gate (FP16: 3 of 12288 prefill-tower values, maximum violating error `0.01171875`; -BF16: 373 values, maximum violating error `0.0625`). -Subsequent source-fed probes identified negative-tail sigmoid rounding and -down-projection accumulation order. Explicit FP32 exponential/division sigmoid -and output-major BF16 `Gemm` restore bit-exact reduced-fixture BF16 prefill -and the first two cached calls, but the third cached call still fails -(62 of 4096 values, maximum violating error `0.03125`). Its first divergence -is a one-ULP FP32 RMS variance reduction difference before BF16 rounding. -These are diagnostic improvements, not completed reduced-precision qualification. -The updated FP16 fixture still fails prefill (7 of 12288 values, maximum -violating error `0.0166015625`). -That is a failed qualification, not permission to widen tolerances or advertise -reduced-precision support. Checkpoint storage is never modified; other adapters retain their -conversion behavior unless they explicitly opt into FP16 saturation. +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. Earlier Torch 2.10 diagnostics are retained separately: +CUDA reduction order differs between these releases. + +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. + 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. @@ -158,17 +157,20 @@ 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`). This is **not full-model low-dtype -qualification**: production native-FP8 admission also remains closed. A complete -native-FP8/FP32 probe fails the unchanged prefill hidden-state gate +(`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. The independent source -reference repeats bit-exactly. This is a pipeline qualification failure, not a -reason to substitute weight-only FP8 or widen tolerances. The custom -CUDA build needs `onnxruntime_USE_FP8_QMOE=ON`, SM89+, and the new mode; old -weight-only FP8 is not an equivalent fallback. The runtime rejects more than +those discontinuities amplify across subsequent layers. Qualification against +the declared Torch 2.8 environment is in progress. 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. The next differences +are inside DeltaNet gating, normalization and recurrent arithmetic. +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. diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index 19a1a3c9d..8595d3fa3 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -187,6 +187,14 @@ def export(checkpoint, *, task="fish-omni", execution_provider=None, native_fp8= ) +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(checkpoint, *, native_fp8=False): os.environ["NVIDIA_TF32_OVERRIDE"] = "0" os.environ["FISH_DRAMA3_FP8_EXPERTS"] = "1" if native_fp8 else "0" @@ -194,6 +202,7 @@ def reference(checkpoint, *, native_fp8=False): 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 @@ -302,6 +311,11 @@ def capture(_module, _inputs, output): "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), @@ -331,6 +345,10 @@ def validate_reference_golden(golden): 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", @@ -641,6 +659,7 @@ def check(name, actual, expected): "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", diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py index 93109036c..ead50bd42 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -5,6 +5,7 @@ from __future__ import annotations +import copy import json from contextlib import contextmanager from types import SimpleNamespace @@ -277,6 +278,48 @@ def test_native_fp8_reference_cannot_be_compared_as_dense(): 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) + + +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()) + 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 From e94ddcd69167e39082a6af17c474afb11ba19df3 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 18:38:08 -0700 Subject: [PATCH 15/22] Requalify Drama 3 golden against the declared Torch 2.8 environment Regenerate the independent full-checkpoint FP32 oracle with Torch 2.8.0+cu126 and Triton 3.4.0. Preserve all 20 exact frame IDs while refreshing full-logit provenance and recording the reference environment. Verify all prefill/cached vectors and frames on stock ORT 1.28 and the separately preserved QMoE/XOR/Pow ORT 1.32 runtime. Record actual numerical bounds and retain the original performance configuration separately. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 25 +++++++++++++------ scripts/validate_drama3_test.py | 7 ++++++ .../speech/drama-3_tokens_generation.json | 10 ++++++-- 3 files changed, 32 insertions(+), 10 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index f20ddbc3f..293fe3a65 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -135,7 +135,9 @@ The pinned Fish Speech source declares `torch==2.8.0`. Run 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. Earlier Torch 2.10 diagnostics are retained separately: -CUDA reduction order differs between these releases. +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 @@ -307,13 +309,20 @@ 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 also passes that -independent complete-checkpoint case: all vectors for the 30-row prefill and -20 exact complete frames, with maximum decoder/tower logit differences -`5.150e-5` / `4.673e-5` and hidden difference `2.289e-5`. -The prefill profile asserts 40 actual CUDA MoE executions. This fresh check -qualifies the changed fast token graph for this case; it is not a latency or -speech-quality measurement. +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, diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py index ead50bd42..242274f9c 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -250,6 +250,12 @@ def test_committed_token_golden_is_pinned_and_complete(): 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 @@ -308,6 +314,7 @@ 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: diff --git a/testdata/golden/speech/drama-3_tokens_generation.json b/testdata/golden/speech/drama-3_tokens_generation.json index 916cf0c4e..d8dfc2f92 100644 --- a/testdata/golden/speech/drama-3_tokens_generation.json +++ b/testdata/golden/speech/drama-3_tokens_generation.json @@ -5,12 +5,18 @@ "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.343191146850586, 13.476040840148926, 13.0930814743042, 12.944562911987305, 12.827808380126953, 12.753363609313965, 12.615551948547363, 12.4545259475708, 12.443134307861328, 12.272481918334961], - "prefill_logits_sha256": "e7cd85ba35c64b51f4d895b85bd1f567d4494f413ddfcd6cfb1381c583fa2c90", + "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", From 74a13a811febe95ac448c78506d1f5fc9ee2bedd Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 19:37:21 -0700 Subject: [PATCH 16/22] Reject nonfinite Drama 3 source references before persisting goldens Check source embeddings, shared features, decoder hidden states/logits and tower heads before writing reference tensors. Report the failing stage and execution dtype instead of recording invalid JSON or an unusable numerical oracle. Document complete native-FP8 BF16 and FP16 source captures in the declared Torch 2.8 environment. Both produce 20 finite frames; matching ONNX qualification remains separate. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 14 ++++++++++---- scripts/validate_drama3.py | 28 ++++++++++++++++++++-------- scripts/validate_drama3_test.py | 22 ++++++++++++++++++++++ 3 files changed, 52 insertions(+), 12 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 293fe3a65..57064cbfb 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -134,7 +134,9 @@ 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. Earlier Torch 2.10 diagnostics are retained separately: +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. @@ -166,11 +168,15 @@ hidden-state gate 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. Source-ordered reductions +the declared Torch 2.8 environment is in progress. Complete native-FP8 source runs +in BF16 and FP16 each produce 20 finite frames; matching whole-model ONNX +comparisons are 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. The next differences -are inside DeltaNet gating, normalization and recurrent arithmetic. +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. 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 diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index 8595d3fa3..9ec3f7902 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -195,6 +195,17 @@ def validate_reference_torch_version(version: str) -> None: ) +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" @@ -231,9 +242,6 @@ def reference(checkpoint, *, native_fp8=False): original_prefill = tower.prefill original_decode = tower.decode_step - def array(tensor): - return tensor.detach().float().cpu().numpy() - def embed(frames, text): output = original_embed(frames, text) unseen = tower.embedding_sum( @@ -247,16 +255,18 @@ def embed(frames, text): states.append( { "input_frames": frames.detach().cpu().numpy(), - "inputs_embeds": array(output)[None], - "shared_input": array(unseen)[None], + "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"] = array(hidden) - states[-1]["logits"] = array(backbone.compute_logits(hidden)) + 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 @@ -275,7 +285,9 @@ def decode(*args, **kwargs): def capture_head(head): def capture(_module, _inputs, output): - active_heads.append(array(head(output))) + active_heads.append( + reference_array(head(output), name=f"tower.head_{len(active_heads)}.logits") + ) return capture diff --git a/scripts/validate_drama3_test.py b/scripts/validate_drama3_test.py index 242274f9c..72556d396 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -9,6 +9,7 @@ import json from contextlib import contextmanager from types import SimpleNamespace +from unittest.mock import Mock import ml_dtypes import numpy as np @@ -295,6 +296,27 @@ def test_reference_rejects_a_different_torch_release(version): 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"} From b8bba3fa329a1d38d33271f1021a4b3143afbf72 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 19:56:43 -0700 Subject: [PATCH 17/22] Document validated CUDA and WebGPU Log1p fusion Link the independently committed provider implementation, required optimization levels and exact source-softplus integration result. Record native WebGPU FP32/FP16 software-Vulkan coverage and browser/hardware validation scope. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/getting-started.md b/docs/getting-started.md index 57064cbfb..80d2164b2 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -177,6 +177,17 @@ 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. +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 From 1f2f54ba0481702cc5773445dfd81b46b33937a0 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 21:33:43 -0700 Subject: [PATCH 18/22] docs: record native BF16 precision and CUDA library findings Record the complete Torch 2.8 native BF16 diagnostic failure and its first expert-boundary divergence. Distinguish the cuBLAS 12.6.4 reference from the 13.6 runtime before matched-library replay; keep production precision admission closed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 80d2164b2..25a2a3cee 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -149,6 +149,13 @@ hidden-state inputs, the BF16 tower matches bit-for-bit through prefill and thre 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. @@ -169,8 +176,16 @@ 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; matching whole-model ONNX -comparisons are pending. Source-ordered reductions +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. +CUDA-library-aligned BF16 replay and the FP16 whole-model comparison 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 From 6e8afc7ed895e4c4e5cb4067b8ac9f58c0c38621 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 22:10:55 -0700 Subject: [PATCH 19/22] feat: record explicit ORT optimization in speech validation Allow disabled/basic/extended/all comparison levels while preserving the BASIC default. Apply the level to every component and record it in events and reports so provider-fusion qualification cannot be mislabeled. Cover session configuration and report metadata without network access. The full FP32 checkpoint also passes EXTENDED against the pinned Torch 2.8 oracle with all 20 exact frames and 40 CUDA MoE executions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 14 +++++++ scripts/validate_drama3.py | 34 +++++++++++++--- scripts/validate_drama3_test.py | 69 +++++++++++++++++++++++++++++++++ 3 files changed, 112 insertions(+), 5 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 25a2a3cee..0d446f299 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -184,6 +184,12 @@ in nine BF16 values (maximum `0.000244140625`); later layers amplify the differe 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 +[justinchuby/onnx-shape-inference#135](https://github.com/justinchuby/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. CUDA-library-aligned BF16 replay and the FP16 whole-model comparison remain pending. Source-ordered reductions and [a CUDA reciprocal-square-root fast path](https://github.com/microsoft/onnxruntime/pull/33662) @@ -197,6 +203,14 @@ 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 diff --git a/scripts/validate_drama3.py b/scripts/validate_drama3.py index 9ec3f7902..181f321d1 100644 --- a/scripts/validate_drama3.py +++ b/scripts/validate_drama3.py @@ -38,6 +38,12 @@ 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" @@ -437,12 +443,16 @@ def run_tower_session(session, feeds, *, generate=False): return result -def compare(*, audio_model_name=None): +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()) @@ -463,13 +473,18 @@ def compare(*, audio_model_name=None): options = ort.SessionOptions() options.intra_op_num_threads = 8 options.inter_op_num_threads = 1 - options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC + 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) + 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"), @@ -675,7 +690,7 @@ def check(name, actual, expected): "audio_component": audio_model_name, "onnxruntime_version": ort.__version__, "provider": "CUDAExecutionProvider", - "ort_graph_optimization": "ORT_ENABLE_BASIC", + "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"], @@ -703,6 +718,12 @@ def check(name, actual, expected): ) 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) @@ -741,4 +762,7 @@ def check(name, actual, expected): elif args.stage == "reference": reference(checkpoint_path(), native_fp8=args.native_fp8) else: - compare(audio_model_name=args.audio_component) + 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 index 72556d396..a3220c169 100644 --- a/scripts/validate_drama3_test.py +++ b/scripts/validate_drama3_test.py @@ -33,6 +33,75 @@ def test_logging_accepts_a_metric_name(capsys): } +@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( From e8f374f05323c49cfb11ab0fe341cb01c0309e47 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 Oct 2026 22:57:48 -0700 Subject: [PATCH 20/22] docs: record normal precision builds and FP16 model gate Link the upstream symbolic-constraint and GemmFloat8 fixes, record unchanged expert-boundary tensors after normal builds, and distinguish the CUDA13 full-FP16 failure from pending matched-library qualification. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 0d446f299..706097268 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -186,12 +186,18 @@ 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 -[justinchuby/onnx-shape-inference#135](https://github.com/justinchuby/onnx-shape-inference/pull/135). +[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. -CUDA-library-aligned BF16 replay and the FP16 whole-model comparison remain -pending. Source-ordered reductions +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 From b7064b2cbb3144078869bd587c425ed986b3f473 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 28 Oct 2026 02:45:56 -0700 Subject: [PATCH 21/22] docs: record matched runtime and source backend isolation Record qualified CUDA12.6 operator evidence without admitting failed native models. Distinguish the actual source MEA backend from forced Flash probes, document source-exact non-expert diagnostic controls, and identify the remaining QMoE scale/compute differences truthfully. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 34 ++++++++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 706097268..bbe42df21 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -152,13 +152,21 @@ 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. +controls reproduce the same distinction. A separately built CUDA 12.6 runtime +now reproduces all three unchanged ORT matmul forms and four direct cuBLAS +controls exactly. Its source is `dceae081a48ac45b32af785b057c0a30b8aa6856`, wheel +SHA256 `9baea5c2f5e300c39b903fb99364917a861aef291c90583360003d04c82589b5`. +QMoE, XOR, Pow, full-head Flash and Log1p operator gates pass on that runtime; +whole-model reduced-precision gates still fail. Matching the ORT version alone +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. +when padding is present. The unpadded reduced-precision prototype omits redundant +causal masks to enable Flash. Profiling the actual pinned backbone instead shows +CUTLASS memory-efficient attention with its explicit BOOL causal mask. A separate +native MEA prototype retains that mask and matches the actual source attention +fixture bit-for-bit; a forced-Flash operator comparison is not source-backend +qualification. ORT 1.28's native BF16 `CausalConvWithState` fallback also does not initialize on CUDA; qualification prototypes use standard-ONNX expansion for this dtype. @@ -197,13 +205,23 @@ the normal FP16 build completes without skipped shape passes and preserves all 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 +original/normal/Log1p variants also fail their unchanged prefill gates; matching +the numerical libraries does not by itself qualify native models. +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. +cache inputs. Native unit-lower triangular solves now remove the first prefix +discrepancy: the first expert inputs and logits are bit-exact for both dtypes. +With source-output controls for the experts, runtime CUDA RoPE, native triangular +solves and source MEA reproduce all 80 prefill expert-input/logit vectors and the +final hidden tensor exactly. This is diagnostic isolation, never whole-model +qualification. The remaining earliest difference is inside QMoE: matching the +source sigmoid sequence corrects 301 activation-scale discrepancies, while +down-projection accumulation is still under investigation. The captured Triton +down kernel uses register-local FP8-to-FP16 conversion and FP16 MMA; that source +compute mode is distinct from native FP8 MMA and must be reported as such. 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). From 77d4f43df5df2e04e08eb9fd9fe8ba21027857da Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 28 Oct 2026 04:07:10 -0700 Subject: [PATCH 22/22] Document source-compatible BF16 cached precision boundary Record exact FP8-storage/FP16-tile weighted-down and reduction evidence, full prefill and seventeen-frame agreement, and the isolated state-17 layer-12 QMoE discrepancy. Keep production admission and native hardware-FP8 claims unchanged. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- docs/getting-started.md | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index bbe42df21..106dcd5ad 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -218,10 +218,22 @@ With source-output controls for the experts, runtime CUDA RoPE, native triangula solves and source MEA reproduce all 80 prefill expert-input/logit vectors and the final hidden tensor exactly. This is diagnostic isolation, never whole-model qualification. The remaining earliest difference is inside QMoE: matching the -source sigmoid sequence corrects 301 activation-scale discrepancies, while -down-projection accumulation is still under investigation. The captured Triton -down kernel uses register-local FP8-to-FP16 conversion and FP16 MMA; that source -compute mode is distinct from native FP8 MMA and must be reported as such. +source sigmoid sequence corrects 301 activation-scale discrepancies. The captured +Triton down kernel uses register-local FP8-to-FP16 conversion and FP16 MMA; +matching its fragment order reproduces all 491520 first-case weighted-down +FP32 values exactly. Matching the source's four-accumulator reduction also makes +that case's BF16 output exact. This source-compatible mode retains FP8 storage +and activation quantization but uses FP16 MMA, not native hardware FP8 MMA. + +With that mode, native triangular solves, source MEA and runtime source-style +RoPE in both decoder and tower, a private complete BF16 pipeline now matches +decoder and teacher-forced tower prefill exactly, followed by 17 exact generated +frames. It still fails the unchanged hidden-state gate at cached state 17. +Actual-cache replay isolates the first difference to one BF16 QMoE output in +layer 12 (`0.000244140625` maximum error), despite exact expert inputs/logits +and all 30 recurrent caches matching bit-for-bit after prefill and state 16. +This is not a qualified reduced-precision export: the public admission guards +remain closed, and native hardware FP8 MMA is separately unqualified. 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).