Skip to content

fix: Nemotron VoiceChat TTS backbone runs bf16 on CUDA #2109

Closed

Description

@inureyes

Problem

The CUDA overlay src/lib/mlx-cpp/patches-cuda/dtype.cpp promotes bf16+f32 to bf16 on purpose (#636). #2093 made the code embeddings (tts/rvq.rs) and MoG head (tts/mog_head.rs) cast explicitly, but in src/models/nemotron_voicechat/tts/model.rs three bf16 tensors still meet the f32 backbone stream, so on CUDA they demote it while the reference runs f32:

  1. bos_emb and null_emb (loaded at :124-125, used at :298 multiply and :220 broadcast_to).
  2. audio_prompt_projection_W (loaded at :149, matmul at :295).
  3. The backbone's GemmaRMSNorm (src/lib/mlxcel-core/src/layers.rs:588-590, used via src/models/gemma3_backbone.rs:158 and per layer) builds 1 + w in the stored bf16 dtype. The CUDA rms_norm (patches-cuda/fast.cpp:98-140) takes out_type = result_type(x, weight), so it casts the f32 input down to bf16. The TTS-local OffsetRmsNorm (tts/norm_mlp.rs:93-95) already builds in f32 and is fine.

All three are reasoned from the op rules, not measured: there is no VoiceChat checkpoint on GB10.

Proposed fix

  • Add bos_emb, null_emb and audio_prompt_projection_W to promotes_to_f32 (model.rs:77); widening them is exact.
  • Do not add the norm keys: promoting w before 1 + w changes rounding on Metal and CPU. In the VoiceChat load path, widen the built bf16 1 + w to f32 instead. GemmaRMSNorm is shared by every Gemma model, so leave layers.rs unchanged.
  • Update f32_promotion_keeps_norms_and_the_subword_path_as_stored (tts/tts_tests.rs:443), which currently asserts that null_emb is not promoted.

Acceptance Criteria

Do it in one PR, run the checks once at the end, and push as few times as possible.

Verification

gpu-lock run --tag voicechat-dtype -- cargo test --release --features cuda --lib nemotron_voicechat -- --test-threads=1

Activity

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:modelsModel architectures, weights, loading, metadata
    platform:linuxLinux (CUDA / packaging) specific
    and removed on Oct 5, 2026
  2. added a commit that references this issue on Oct 5, 2026
    013f61e
  3. added a commit that references this issue on Oct 6, 2026
    6f3fa7b
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:modelsModel architectures, weights, loading, metadatamodeltype:audioAudio or speech modelplatform:linuxLinux (CUDA / packaging) specificpriority:lowLow prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions