Repository navigation
fix: Nemotron VoiceChat TTS backbone runs bf16 on CUDA #2109
Copy link
Copy link
Closed
Labels
area:modelsModel architectures, weights, loading, metadataModel architectures, weights, loading, metadatamodeltype:audioAudio or speech modelAudio or speech modelplatform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:lowLow priorityLow prioritystatus:doneCompletedCompletedtype:bugBug fixes, error corrections, or issue resolutionsBug fixes, error corrections, or issue resolutions
Description
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:
bos_embandnull_emb(loaded at :124-125, used at :298multiplyand :220broadcast_to).audio_prompt_projection_W(loaded at :149,matmulat :295).- The backbone's
GemmaRMSNorm(src/lib/mlxcel-core/src/layers.rs:588-590, used viasrc/models/gemma3_backbone.rs:158and per layer) builds1 + win the stored bf16 dtype. The CUDArms_norm(patches-cuda/fast.cpp:98-140) takesout_type = result_type(x, weight), so it casts the f32 input down to bf16. The TTS-localOffsetRmsNorm(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_embandaudio_prompt_projection_Wtopromotes_to_f32(model.rs:77); widening them is exact. - Do not add the norm keys: promoting
wbefore1 + wchanges rounding on Metal and CPU. In the VoiceChat load path, widen the built bf161 + wto f32 instead.GemmaRMSNormis shared by every Gemma model, so leavelayers.rsunchanged. - Update
f32_promotion_keeps_norms_and_the_subword_path_as_stored(tts/tts_tests.rs:443), which currently asserts thatnull_embis not promoted.
Acceptance Criteria
- A new synthetic-weight test (bf16 weights) asserts that the backbone hidden stream is f32 after embedding assembly, after a backbone norm and after the projection. Under
--features cudait fails on 33c4505. - On non-CUDA builds, TTS output for a fixed seed is bit-identical before and after the fix. (argued from the code in PR fix: keep the VoiceChat TTS backbone stream f32 on CUDA #2115, not measured: the CPU-only build is blocked by fix: CPU-only Linux build fails to link kv_inplace_write #2108)
- The PR states that no real checkpoint was used for validation.
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=1Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
area:modelsModel architectures, weights, loading, metadataModel architectures, weights, loading, metadatamodeltype:audioAudio or speech modelAudio or speech modelplatform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:lowLow priorityLow prioritystatus:doneCompletedCompletedtype:bugBug fixes, error corrections, or issue resolutionsBug fixes, error corrections, or issue resolutions
Activity
fix: keep the VoiceChat TTS backbone stream f32 on CUDA
fix: keep the VoiceChat TTS backbone stream f32 on CUDA (lablup#2115)