diff --git a/lib/model_state.py b/lib/model_state.py index ebda2a8..87f1cbd 100644 --- a/lib/model_state.py +++ b/lib/model_state.py @@ -241,6 +241,42 @@ def _fix_moss_model_config_token_ids(processor) -> None: cls._moss_tts_lang_attr_fallback_patched = True +def _cast_inputs_to_param_dtype_hook(module, args): + """Forward pre-hook: cast floating tensor inputs to module's param dtype. + + AutoProcessor loads the audio_tokenizer in fp32. After we cast its weights + to fp16/bf16 to save VRAM, raw fp32 wav inputs flow through Linear/Conv + layers whose weights are now half-precision -> RuntimeError "expected + scalar type Float but found Half". The hook fires before each leaf + module's forward and aligns input dtype to the module's own parameters. + """ + if not args: + return args + params = list(module.parameters(recurse=False)) + if not params: + return args + target_dtype = params[0].dtype + new_args = tuple( + a.to(target_dtype) + if isinstance(a, torch.Tensor) + and a.is_floating_point() + and a.dtype != target_dtype + else a + for a in args + ) + return new_args + + +def _patch_audio_tokenizer_dtype(audio_tokenizer) -> None: + """Register the dtype-cast pre-hook on every leaf-ish module of the audio + tokenizer. Covers both encode (voice ref) and decode paths without + needing to know the model's submodule names. + """ + for module in audio_tokenizer.modules(): + if list(module.parameters(recurse=False)): + module.register_forward_pre_hook(_cast_inputs_to_param_dtype_hook) + + def resolve_attn_impl(requested: str, device: str) -> str: if requested == "auto": return "sdpa" if device == "cuda" else "eager" @@ -314,6 +350,7 @@ def get_or_load( processor.audio_tokenizer = processor.audio_tokenizer.to( resolved_at_device, dtype=resolved_dtype ) + _patch_audio_tokenizer_dtype(processor.audio_tokenizer) else: processor.audio_tokenizer = processor.audio_tokenizer.to(resolved_at_device) diff --git a/tests/test_model_state.py b/tests/test_model_state.py index ea0e0fb..944877b 100644 --- a/tests/test_model_state.py +++ b/tests/test_model_state.py @@ -133,6 +133,30 @@ def test_audio_tokenizer_cast_to_dtype_on_cuda(): fake_at.to.assert_called_with("cuda", dtype=torch.float16) +def test_dtype_cast_hook_aligns_fp32_input_with_fp16_layer(): + """End-to-end: a Linear with fp16 weights normally errors on fp32 input, + but with the hook installed it casts the input and produces fp16 output. + """ + import torch + layer = torch.nn.Linear(4, 4).to(torch.float16) + inp_fp32 = torch.randn(2, 4, dtype=torch.float32) + with pytest.raises(RuntimeError, match=r"expected scalar type"): + layer(inp_fp32) + model_state._patch_audio_tokenizer_dtype(layer) + out = layer(inp_fp32) + assert out.dtype == torch.float16 + + +def test_dtype_cast_hook_passes_through_matching_dtype(): + """Hook is a no-op when input dtype already matches param dtype.""" + import torch + layer = torch.nn.Linear(4, 4).to(torch.float16) + model_state._patch_audio_tokenizer_dtype(layer) + inp_fp16 = torch.randn(2, 4, dtype=torch.float16) + out = layer(inp_fp16) + assert out.dtype == torch.float16 + + def test_audio_tokenizer_no_dtype_cast_on_cpu(): """When the audio_tokenizer is on CPU, do not force fp16 — fp16 ops are slow on most CPUs and the VRAM concern doesn't apply off-GPU.