A full fp16 cast on the audio_tokenizer saved ~3 GB VRAM but produced
silent audio: fp16 lacks the dynamic range LayerNorm/RMSNorm need to
keep their normalization statistics meaningful, and the decoder
collapsed to near-zero output.
Now: cast all weights to the resolved dtype, then explicitly cast
LayerNorm/GroupNorm/RMSNorm back to fp32. A forward pre-hook on every
leaf-ish module aligns incoming float tensors with each module's own
param dtype, so:
- fp32 wav input -> fp16 Linear: hook casts to fp16
- fp16 hidden state -> fp32 LayerNorm: hook casts to fp32
- fp32 norm output -> fp16 next layer: that layer's hook casts to fp16
Tested with a tiny mixed Linear/Conv1d/LayerNorm stand-in: fp32 inputs
flow through fp16 Linear, fp16 hidden states reach fp32 norm cleanly.
End-to-end audio quality has to be verified on the real model.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
