From a7d488a898b6e342bb002d0b0c8dfb88809b19fa Mon Sep 17 00:00:00 2001 From: zack Date: Sat, 19 Sep 2026 19:43:52 +0800 Subject: [PATCH 1/3] [Perf][CosyVoice3] Fold frozen HiFT generator weight norms after checkpoint load CosyVoice3's HiFT vocoder keeps 77 weight-norm parametrizations on frozen inference weights, so every decode recomputes g*v/||v|| per layer: profiled as 77 aten::_weight_norm + 77 CUDA weight_norm_fwd_first_dim_kernel events per chunk (RFC 6870, workstream C4). The shipped remove_weight_norm() was unusable: it calls the legacy torch.nn.utils spelling on modules built with the parametrizations API (ValueError on any current torch), invokes a non-existent SourceModuleHnNSF.remove_weight_norm(), and walks source_downs layers that were never weight-normed. Nothing called it. Fix and wire it into the load path: - per-layer _fold_weight_norm() supports both APIs (parametrizations API; legacy WeightNorm forward-pre-hook via isinstance, same idea as the merged indextts2 _strip_weight_norm), returns 0/1, is idempotent, and normalizes the materialized weight to a frozen nn.Parameter regardless of the caller's grad mode (no_grad/inference_mode callers would otherwise get a buffer or an inference tensor); - ResBlock/HiFTGenerator.remove_weight_norm() return the folded count so a silent no-op shows up in the log (same convention as minimax_music3); - CosyVoice3Code2Wav.load_weights() folds after the strict checkpoint load and device move; - f0_predictor's 5 weight norms are deliberately left parametrized (pinned-CPU precision path; RFC C5, to be done with PR 7518). step_audio2 and glm_tts import HiFTGenerator but never fold, so they are unaffected. One-time post-load transform: folding rewrites the state_dict schema (parametrizations.weight.original0/1 -> weight). vllm-omni has no reload_weights path and sleep/wake restores physical pages via CuMemAllocator without re-invoking the loader, so nothing in-tree loads twice; a second load of the original checkpoint would fail the strict load_state_dict loudly. Validation (Runpod, vllm 0.29.0, torch 2.13.0+cu130, base d4ffde1a6ef): - tests/model_executor/models/cosyvoice3/: 114 passed (RTX 2000 Ada host); shared-consumer checks: minicpmo cuda-graph test 13 passed, step_audio2 and glm_tts imports OK - pre-commit run --files <3 changed files>: all gates passed (ruff, typos, mypy-3.10, CI marks, SPDX, forbidden imports, torch.cuda guard) - new L1 regressions (16 cases): fold exactness atol=rtol=0 across parametrized and legacy APIs, 22050 and 24000 Hz, finalize and streaming; fold counts, idempotency, Parameter-ness (including folding under inference_mode); shipped-config generator folds exactly 77 with only the 5 f0 layers left; load_weights folds loaded (not initial) weights for modern and legacy checkpoint key schemas. All 16 fail on unpatched d4ffde1a6ef. - real checkpoint (FunAudioLLM/Fun-CosyVoice3-0.5B-2512, GPU): production load_weights logs 'Folded 77 weight-norm layers'; inference() on the real hift weights is bit-identical before vs after fold; generator 0 / f0 5 parametrizations left; folded weight is a Parameter on CUDA - CUDA profiler (real-sized generator, decode path): weight-norm events 231 -> 0; inference output bit-identical (torch.equal=True) - wall-clock A/B, paired interleaved rounds, alternating arm order, bootstrap 95% CI: RTX A4500 (200 rounds, initial fold revision): decode() 3.83 ms [3.74, 3.89] chunk 41 (200/200 faster) and 4.51 ms [4.33, 4.62] chunk 191 (197/200); RTX 2000 Ada (30 rounds, final head): decode() 2.82 and 2.35 ms (30/30 each). Full inference() path is dominated by unchanged f0_predictor jitter (C2/7518 scope); no e2e RTF claim made (7521's harness cannot resolve this effect on shared GPUs). Refs: RFC 6870 (C4). Related: PR 6927 (generic fold helper for MiniCPM-o), PR 7518 (C2/C5 f0 predictor). Signed-off-by: zack --- .../test_cosyvoice3_hift_weight_norm.py | 204 ++++++++++++++++++ .../cosyvoice3/code2wav_core/hifigan.py | 67 ++++-- .../models/cosyvoice3/cosyvoice3_code2wav.py | 15 +- 3 files changed, 272 insertions(+), 14 deletions(-) create mode 100644 tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py new file mode 100644 index 00000000000..1c4d1fa4396 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU regressions for folding frozen HiFT generator weight normalization.""" + +import pytest +import torch +from torch import nn +from torch.nn.utils import parametrize + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cpu, + # The legacy weight_norm API warns on construction; the "legacy" params + # below exercise it on purpose. + pytest.mark.filterwarnings("ignore:.*weight_norm.*:FutureWarning"), +] + + +@pytest.fixture(autouse=True) +def fixed_seed(): + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + yield + + +@pytest.fixture(params=["parametrized", "legacy"]) +def weight_norm_api(request, monkeypatch): + if request.param == "legacy": + monkeypatch.setattr(hifigan, "weight_norm", nn.utils.weight_norm) + return request.param + + +def _has_weight_norm(module): + return parametrize.is_parametrized(module, "weight") or ( + hasattr(module, "weight_g") and hasattr(module, "weight_v") + ) + + +def _make_hift(causal=True, sampling_rate=24000): + # Note: pairing the non-causal generator with CausalConvRNNF0Predictor is + # not a shipped configuration; it exists only to exercise the shared + # ResBlock/conv fold code on both generator classes. + cls = hifigan.CausalHiFTGenerator if causal else hifigan.HiFTGenerator + return cls( + base_channels=32, + sampling_rate=sampling_rate, + # Keep the real CosyVoice3 upsample ratio and SineGen2 (24 kHz), + # while shrinking channels and residual blocks for CPU tests. + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + resblock_kernel_sizes=[3], + resblock_dilation_sizes=[[1, 3, 5]], + source_resblock_kernel_sizes=[3, 3, 3], + source_resblock_dilation_sizes=[[1, 3, 5]] * 3, + f0_predictor=hifigan.CausalConvRNNF0Predictor(cond_channels=16), + ).eval() + + +@pytest.mark.parametrize("causal", [False, True]) +def test_resblock_fold_is_exact_and_idempotent(weight_norm_api, causal): + block = hifigan.ResBlock(channels=4, causal=causal).eval() + x = torch.randn(1, 4, 12) + conv = block.convs1[0] + # Guard: the fixture really selected the intended API, so the legacy + # cases cannot silently re-test the parametrizations path. + if weight_norm_api == "legacy": + assert hasattr(conv, "weight_g") and not parametrize.is_parametrized(conv, "weight") + else: + assert parametrize.is_parametrized(conv, "weight") + + with torch.no_grad(): + before = block(x) + # Fold in normal grad mode, like the production load path (vLLM's loader + # does not wrap load_weights in no_grad/inference_mode): the materialized + # weight must be a real Parameter. + first = block.remove_weight_norm() + assert isinstance(conv.weight, nn.Parameter) + second = block.remove_weight_norm() + with torch.no_grad(): + once = block(x) + twice = block(x) + + assert (first, second) == (6, 0) # 3 dilations x {convs1, convs2}; no-op after + assert not any(_has_weight_norm(module) for module in block.modules()) + torch.testing.assert_close(once, before, rtol=0, atol=0) + torch.testing.assert_close(twice, before, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "causal,sampling_rate,finalize", + [(False, 22050, True), (True, 22050, True), (True, 24000, True), (True, 24000, False)], +) +def test_hift_fold_preserves_audio_and_leaves_f0_unchanged(weight_norm_api, causal, sampling_rate, finalize): + hift = _make_hift(causal, sampling_rate) + mel = torch.randn(1, 80, 20) + if weight_norm_api == "legacy": + assert hasattr(hift.conv_pre, "weight_g") and not parametrize.is_parametrized(hift.conv_pre, "weight") + normalized = {name: module for name, module in hift.named_modules() if _has_weight_norm(module)} + f0_state = {name: value.clone() for name, value in hift.f0_predictor.state_dict().items()} + assert any(name.startswith("source_resblocks.") for name in normalized) + assert any(name.startswith("resblocks.") for name in normalized) + # The C5-scope layers must start out weight-normed, or the "f0 untouched" + # asserts below would pass vacuously. + assert any(name.startswith("f0_predictor.") for name in normalized) + + # Non-causal HiFT draws fresh source noise on each invocation. + rng = torch.get_rng_state() + kwargs = {"finalize": finalize} if causal else {} + before = hift.inference(mel, **kwargs)[0] + # 3 ups + conv_pre + conv_post + (3 resblocks + 3 source resblocks) x 6 convs + assert hift.remove_weight_norm() == 41 + assert hift.remove_weight_norm() == 0 + torch.set_rng_state(rng) + after = hift.inference(mel, **kwargs)[0] + + assert before.numel() > 0 + assert torch.isfinite(before).all() + torch.testing.assert_close(after, before, rtol=0, atol=0) + for name, module in normalized.items(): + # F0 is a separate RFC workstream; do not silently fold it here. + assert _has_weight_norm(module) == name.startswith("f0_predictor.") + for name, value in hift.f0_predictor.state_dict().items(): + torch.testing.assert_close(value, f0_state[name], rtol=0, atol=0) + + +@pytest.mark.parametrize("legacy_checkpoint", [False, True]) +def test_load_weights_folds_loaded_not_initial_weights(tmp_path, legacy_checkpoint): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + # Exercise the real loader with a small vocoder and synthetic checkpoints. + # __new__ skips CosyVoice3Code2Wav.__init__ (which would build the DiT and + # flow decoder); load_weights only touches .flow_model and .hift. + model = CosyVoice3Code2Wav.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.flow_model = nn.Linear(2, 2) + model.hift = _make_hift() + state = {name: value.clone() for name, value in model.hift.state_dict().items()} + state["conv_pre.parametrizations.weight.original0"].mul_(2) + expected_weight = model.hift.conv_pre.weight.detach().clone() * 2 + flow_state = {name: torch.ones_like(value) for name, value in model.flow_model.state_dict().items()} + if legacy_checkpoint: + state = { + name.replace("parametrizations.weight.original0", "weight_g").replace( + "parametrizations.weight.original1", "weight_v" + ): value + for name, value in state.items() + } + torch.save(flow_state, tmp_path / "flow.pt") + torch.save({"generator." + name: value for name, value in state.items()}, tmp_path / "hift.pt") + + model.load_weights(str(tmp_path), torch.device("cpu")) + + torch.testing.assert_close(model.hift.conv_pre.weight, expected_weight, rtol=0, atol=0) + # Production loads in normal grad mode, so the folded weight is a Parameter. + assert isinstance(model.hift.conv_pre.weight, nn.Parameter) + assert not any( + _has_weight_norm(module) for name, module in model.hift.named_modules() if not name.startswith("f0_predictor.") + ) + assert _has_weight_norm(model.hift.f0_predictor.condnet[0]) + assert not model.hift.training + assert not model.flow_model.training + for name, value in model.flow_model.state_dict().items(): + torch.testing.assert_close(value, flow_state[name], rtol=0, atol=0) + + +def test_fold_under_inference_mode_stays_usable_outside_it(): + """Folding inside inference_mode must not leak inference tensors: the + materialized weight stays a frozen Parameter and forward keeps working + in a normal/no_grad context (pre-hardening this raised RuntimeError).""" + block = hifigan.ResBlock(channels=4).eval() + x = torch.randn(1, 4, 12) + with torch.inference_mode(): + assert block.remove_weight_norm() == 6 + assert isinstance(block.convs1[0].weight, nn.Parameter) + assert not block.convs1[0].weight.requires_grad + with torch.no_grad(): + out = block(x) + assert torch.isfinite(out).all() + + +def test_shipped_config_generator_folds_exactly_77_layers(): + """Pin the RFC #6870 C4 count against the shipped config, so silent + config drift (or a fold that quietly stops covering layers) shows up + here: conv_pre 1 + ups 3 + 9 resblocks x 6 + 3 source resblocks x 6 + + conv_post 1 = 77. Construction only (no forward), so the full + 512-channel model stays cheap on CPU.""" + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + + hift_cfg = dict(CosyVoice3Config().hift) + f0_cfg = hift_cfg.pop("f0_predictor") + f0 = hifigan.CausalConvRNNF0Predictor( + num_class=f0_cfg["num_class"], + in_channels=f0_cfg["in_channels"], + cond_channels=f0_cfg["cond_channels"], + ) + generator = hifigan.CausalHiFTGenerator(f0_predictor=f0, **hift_cfg) + + assert generator.remove_weight_norm() == 77 + assert generator.remove_weight_norm() == 0 + remaining = [name for name, m in generator.named_modules() if _has_weight_norm(m)] + assert len(remaining) == 5 + assert all(name.startswith("f0_predictor.") for name in remaining) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index 2b8f3e61e24..48daae8ed3c 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -12,7 +12,9 @@ import torch.nn.functional as F from scipy.signal import get_window from torch.nn import Conv1d, ConvTranspose1d -from torch.nn.utils import remove_weight_norm +from torch.nn.utils import parametrize +from torch.nn.utils import remove_weight_norm as remove_legacy_weight_norm +from torch.nn.utils.weight_norm import WeightNorm try: from torch.nn.utils.parametrizations import weight_norm @@ -41,6 +43,36 @@ def get_padding(kernel_size, dilation=1): return int((kernel_size * dilation - dilation) / 2) +def _fold_weight_norm(module: nn.Module) -> int: + """Fold one weight-normalized layer into a plain ``weight``, in place. + + Supports both APIs: the modern ``torch.nn.utils.parametrizations`` API + (what this module is built with when torch provides it) and the legacy + ``WeightNorm`` forward-pre-hook. Returns 1 when a weight norm was + folded and 0 when there was nothing to fold (e.g. on a second call). + """ + if parametrize.is_parametrized(module, "weight"): + # Removes every parametrization registered on "weight"; only + # weight_norm is applied to the layers in this module. + parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) + weight = module.weight + if not isinstance(weight, nn.Parameter): + # Under no_grad/inference_mode, leave_parametrized registers a plain + # (possibly inference) tensor instead of a Parameter. Normalize so + # callers in any grad mode get the same semantics and the weight + # stays usable outside the fold's grad context. + with torch.inference_mode(False): + frozen = nn.Parameter(weight.detach().clone(), requires_grad=False) + del module.weight + module.register_parameter("weight", frozen) + return 1 + for hook in list(module._forward_pre_hooks.values()): + if isinstance(hook, WeightNorm): + remove_legacy_weight_norm(module, name=hook.name) + return 1 + return 0 + + class ResBlock(torch.nn.Module): """Residual block module in HiFiGAN/BigVGAN.""" @@ -92,10 +124,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = xt + x return x - def remove_weight_norm(self): + def remove_weight_norm(self) -> int: + """Fold this block's conv weight norms; returns the folded count.""" + folded = 0 for idx in range(len(self.convs1)): - remove_weight_norm(self.convs1[idx]) - remove_weight_norm(self.convs2[idx]) + folded += _fold_weight_norm(self.convs1[idx]) + folded += _fold_weight_norm(self.convs2[idx]) + return folded def _carry_phase_at_boundary( @@ -570,18 +605,24 @@ def __init__( ) self.f0_predictor = f0_predictor - def remove_weight_norm(self): + def remove_weight_norm(self) -> int: + """Fold the generator's frozen weight norms into plain weights. + + Returns how many convolutions were folded. ``source_downs`` and + ``m_source`` carry no weight norm; ``f0_predictor`` is excluded: + it is pinned to CPU for precision, so it is left parametrized and + folding it is tracked separately in RFC #6870 (C5). + """ + folded = 0 for layer in self.ups: - remove_weight_norm(layer) + folded += _fold_weight_norm(layer) for block in self.resblocks: - block.remove_weight_norm() - remove_weight_norm(self.conv_pre) - remove_weight_norm(self.conv_post) - self.m_source.remove_weight_norm() - for layer in self.source_downs: - remove_weight_norm(layer) + folded += block.remove_weight_norm() + folded += _fold_weight_norm(self.conv_pre) + folded += _fold_weight_norm(self.conv_post) for block in self.source_resblocks: - block.remove_weight_norm() + folded += block.remove_weight_norm() + return folded def _stft(self, x): if x.device.type == "npu": diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 1713467c8b3..c5b17b07a53 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -518,5 +518,18 @@ def load_weights(self, model_dir: str, device: torch.device) -> None: k.replace("generator.", ""): v for k, v in torch.load(hift_path, map_location=device).items() } self.hift.load_state_dict(hift_state_dict, strict=True) - self.hift.to(device).eval() + self.hift.to(device) + # Inference weights are frozen: fold g*v/||v|| once here instead of + # recomputing it in every convolution on every streamed chunk. + # Must stay after load_state_dict: folding rewrites the state_dict + # keys (parametrizations.weight.original0/1 -> weight), so this is a + # one-time post-load transform. A second load_weights() with the + # original checkpoint fails the strict load_state_dict loudly; no + # in-tree path reloads (sleep/wake restores physical pages, and + # level-2 wake is NotImplementedError in async_omni.py). + folded = self.hift.remove_weight_norm() + logger.info("Folded %d weight-norm layers in HiFT generator", folded) + if folded == 0: + logger.warning("HiFT generator had no weight-norm layers to fold; check config drift") + self.hift.eval() logger.info(f"Loaded hift weights from {hift_path}") From 2e021d07370d8318147e83540fdecace148c0bc9 Mon Sep 17 00:00:00 2001 From: zack Date: Wed, 23 Sep 2026 14:41:20 +0800 Subject: [PATCH 2/3] [Bugfix] Keep folded CosyVoice3 HiFT weights usable by autograd Apply the shared normalizer suggested by @0z5a to both weight-norm APIs so folding under inference_mode cannot leave inference Parameters. Preserve existing ordinary Parameter flags. Cover both APIs across grad contexts, cold loading into a legacy-created vocoder, the loader's transition to eval, and multichunk finalization. Keep the generator-only 77-layer scope and five F0 norms unchanged. Related: #7871 Signed-off-by: zack --- .../test_cosyvoice3_hift_weight_norm.py | 86 +++++++++++++++---- .../cosyvoice3/code2wav_core/hifigan.py | 34 +++++--- .../models/cosyvoice3/cosyvoice3_code2wav.py | 11 +-- 3 files changed, 94 insertions(+), 37 deletions(-) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py index 1c4d1fa4396..449131e2438 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py @@ -125,8 +125,12 @@ def test_hift_fold_preserves_audio_and_leaves_f0_unchanged(weight_norm_api, caus torch.testing.assert_close(value, f0_state[name], rtol=0, atol=0) -@pytest.mark.parametrize("legacy_checkpoint", [False, True]) -def test_load_weights_folds_loaded_not_initial_weights(tmp_path, legacy_checkpoint): +@pytest.mark.parametrize( + "weight_norm_api,legacy_checkpoint", + [("parametrized", False), ("parametrized", True), ("legacy", True)], + indirect=["weight_norm_api"], +) +def test_load_weights_folds_loaded_not_initial_weights(tmp_path, weight_norm_api, legacy_checkpoint): from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav # Exercise the real loader with a small vocoder and synthetic checkpoints. @@ -135,10 +139,18 @@ def test_load_weights_folds_loaded_not_initial_weights(tmp_path, legacy_checkpoi model = CosyVoice3Code2Wav.__new__(CosyVoice3Code2Wav) nn.Module.__init__(model) model.flow_model = nn.Linear(2, 2) - model.hift = _make_hift() + # Start in training mode to check the loader's transition to eval. + model.hift = _make_hift().train() + assert model.hift.training + conv = model.hift.conv_pre + assert hasattr(conv, "weight_g") == (weight_norm_api == "legacy") state = {name: value.clone() for name, value in model.hift.state_dict().items()} - state["conv_pre.parametrizations.weight.original0"].mul_(2) - expected_weight = model.hift.conv_pre.weight.detach().clone() * 2 + gain = "weight_g" if weight_norm_api == "legacy" else "parametrizations.weight.original0" + state[f"conv_pre.{gain}"].mul_(2) + expected_weight = conv.weight.detach().clone() * 2 + assert not torch.equal(conv.weight, expected_weight) + # No forward before loading: a legacy fold must recompute from the loaded + # g/v parameters, not reuse the cached effective weight from construction. flow_state = {name: torch.ones_like(value) for name, value in model.flow_model.state_dict().items()} if legacy_checkpoint: state = { @@ -155,6 +167,7 @@ def test_load_weights_folds_loaded_not_initial_weights(tmp_path, legacy_checkpoi torch.testing.assert_close(model.hift.conv_pre.weight, expected_weight, rtol=0, atol=0) # Production loads in normal grad mode, so the folded weight is a Parameter. assert isinstance(model.hift.conv_pre.weight, nn.Parameter) + assert not model.hift.conv_pre.weight.is_inference() assert not any( _has_weight_norm(module) for name, module in model.hift.named_modules() if not name.startswith("f0_predictor.") ) @@ -165,19 +178,60 @@ def test_load_weights_folds_loaded_not_initial_weights(tmp_path, legacy_checkpoi torch.testing.assert_close(value, flow_state[name], rtol=0, atol=0) -def test_fold_under_inference_mode_stays_usable_outside_it(): - """Folding inside inference_mode must not leak inference tensors: the - materialized weight stays a frozen Parameter and forward keeps working - in a normal/no_grad context (pre-hardening this raised RuntimeError).""" +@pytest.mark.parametrize( + "fold_context", + [torch.enable_grad, torch.no_grad, torch.inference_mode], + ids=["grad", "no_grad", "inference_mode"], +) +def test_fold_stays_usable_outside_grad_context(weight_norm_api, fold_context): + """Both APIs must leave non-inference Parameters usable by autograd.""" block = hifigan.ResBlock(channels=4).eval() - x = torch.randn(1, 4, 12) - with torch.inference_mode(): - assert block.remove_weight_norm() == 6 - assert isinstance(block.convs1[0].weight, nn.Parameter) - assert not block.convs1[0].weight.requires_grad + x = torch.randn(1, 4, 12, requires_grad=True) + assert hasattr(block.convs1[0], "weight_g") == (weight_norm_api == "legacy") with torch.no_grad(): - out = block(x) - assert torch.isfinite(out).all() + before = block(x) + with fold_context(): + assert block.remove_weight_norm() == 6 + assert block.remove_weight_norm() == 0 + for conv in (*block.convs1, *block.convs2): + assert isinstance(conv.weight, nn.Parameter) + assert not conv.weight.is_inference() + # no_grad would hide the failure when autograd tries to save a weight. + out = block(x) + torch.testing.assert_close(out, before, rtol=0, atol=0) + out.sum().backward() + assert x.grad is not None + assert torch.isfinite(x.grad).all() + + +@pytest.mark.parametrize("sampling_rate", [22050, 24000]) +def test_fold_preserves_multichunk_audio_and_finalize(weight_norm_api, sampling_rate): + from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_model + + hift = _make_hift(sampling_rate=sampling_rate) + model = _make_model(hift, window_len=64) + mel = torch.randn(1, 80, 192) + chunks = mel.split(24, dim=-1) + + def stream(): + cache = None + audio = [] + window_shifted = False + for i, chunk in enumerate(chunks): + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=i == len(chunks) - 1) + audio.append(speech) + if cache is not None: + window_shifted |= cache["mel_offset"] > 0 + assert window_shifted + assert cache is None + return torch.cat(audio, dim=-1) + + before = stream() + assert hift.remove_weight_norm() == 41 + after = stream() + assert before.numel() > 0 + assert torch.isfinite(before).all() + torch.testing.assert_close(after, before, rtol=0, atol=0) def test_shipped_config_generator_folds_exactly_77_layers(): diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index 48daae8ed3c..8ad1bfb92f3 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -43,6 +43,22 @@ def get_padding(kernel_size, dilation=1): return int((kernel_size * dilation - dilation) / 2) +def _normalize_folded_weight(module: nn.Module, name: str = "weight") -> None: + """Keep ordinary Parameters; replace plain or inference tensors. + + In inference_mode, the legacy API can leave an inference Parameter, + while the parametrization API can leave a plain inference tensor. + Rebuild those outside inference_mode so they remain usable afterward. + """ + weight = getattr(module, name) + if isinstance(weight, nn.Parameter) and not weight.is_inference(): + return + with torch.inference_mode(False): + frozen = nn.Parameter(weight.detach().clone(), requires_grad=False) + delattr(module, name) + module.register_parameter(name, frozen) + + def _fold_weight_norm(module: nn.Module) -> int: """Fold one weight-normalized layer into a plain ``weight``, in place. @@ -55,20 +71,12 @@ def _fold_weight_norm(module: nn.Module) -> int: # Removes every parametrization registered on "weight"; only # weight_norm is applied to the layers in this module. parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) - weight = module.weight - if not isinstance(weight, nn.Parameter): - # Under no_grad/inference_mode, leave_parametrized registers a plain - # (possibly inference) tensor instead of a Parameter. Normalize so - # callers in any grad mode get the same semantics and the weight - # stays usable outside the fold's grad context. - with torch.inference_mode(False): - frozen = nn.Parameter(weight.detach().clone(), requires_grad=False) - del module.weight - module.register_parameter("weight", frozen) + _normalize_folded_weight(module) return 1 for hook in list(module._forward_pre_hooks.values()): if isinstance(hook, WeightNorm): remove_legacy_weight_norm(module, name=hook.name) + _normalize_folded_weight(module, hook.name) return 1 return 0 @@ -609,9 +617,9 @@ def remove_weight_norm(self) -> int: """Fold the generator's frozen weight norms into plain weights. Returns how many convolutions were folded. ``source_downs`` and - ``m_source`` carry no weight norm; ``f0_predictor`` is excluded: - it is pinned to CPU for precision, so it is left parametrized and - folding it is tracked separately in RFC #6870 (C5). + ``m_source`` carry no weight norm. ``f0_predictor`` is excluded: + causal inference runs it on CPU for precision, and folding it is + tracked separately in RFC #6870 (C5). """ folded = 0 for layer in self.ups: diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index c5b17b07a53..e01e8284234 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -519,14 +519,9 @@ def load_weights(self, model_dir: str, device: torch.device) -> None: } self.hift.load_state_dict(hift_state_dict, strict=True) self.hift.to(device) - # Inference weights are frozen: fold g*v/||v|| once here instead of - # recomputing it in every convolution on every streamed chunk. - # Must stay after load_state_dict: folding rewrites the state_dict - # keys (parametrizations.weight.original0/1 -> weight), so this is a - # one-time post-load transform. A second load_weights() with the - # original checkpoint fails the strict load_state_dict loudly; no - # in-tree path reloads (sleep/wake restores physical pages, and - # level-2 wake is NotImplementedError in async_omni.py). + # Fold after loading and device placement to avoid recomputing weights + # for every chunk. Folding changes state_dict keys, so reloading the + # original checkpoint into this instance is unsupported. folded = self.hift.remove_weight_norm() logger.info("Folded %d weight-norm layers in HiFT generator", folded) if folded == 0: From 5c7a151f7b05949fd73c0a96893272e71031fe4d Mon Sep 17 00:00:00 2001 From: zack Date: Fri, 25 Sep 2026 11:21:51 +0800 Subject: [PATCH 3/3] [Perf] Fold CosyVoice3 F0 weight norms after checkpoint loading Materialize F0 weights on CPU in FP32 using the shared folding helpers, while keeping generator and F0 folding as separate entry points. Add joint C4/C5 loading and streaming regressions for modern and legacy weight normalization, including strict loading and idempotence. Signed-off-by: zack --- .../test_cosyvoice3_f0_weight_norm.py | 283 ++++++++++++++++++ .../test_cosyvoice3_hift_weight_norm.py | 22 +- .../cosyvoice3/code2wav_core/hifigan.py | 8 + .../models/cosyvoice3/cosyvoice3_code2wav.py | 7 + 4 files changed, 312 insertions(+), 8 deletions(-) create mode 100644 tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py new file mode 100644 index 00000000000..97525d4f4ec --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""F0 folding and joint C4/C5 loading regressions with the existing CPU F0 policy.""" + +import pytest +import torch +from torch import nn +from torch.nn.utils import parametrize + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.filterwarnings("ignore:.*weight_norm.*:FutureWarning"), +] + + +@pytest.fixture(autouse=True) +def fixed_seed(): + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(5) + yield + + +@pytest.fixture(params=["parametrized", "legacy"]) +def weight_norm_api(request, monkeypatch): + if request.param == "legacy": + monkeypatch.setattr(hifigan, "weight_norm", nn.utils.weight_norm) + return request.param + + +@pytest.fixture( + params=[torch.enable_grad, torch.no_grad, torch.inference_mode], ids=["grad", "no_grad", "inference_mode"] +) +def fold_context(request): + return request.param + + +@pytest.fixture( + params=[ + pytest.param("cpu", marks=pytest.mark.cpu), + pytest.param("cuda", marks=pytest.mark.cuda), + ] +) +def runtime(request): + """Only the generator may run on CUDA; production F0 still runs on CPU.""" + on_cuda = request.param == "cuda" + if on_cuda and not torch.cuda.is_available(): + pytest.skip("CUDA is required for the generator/device-placement regressions") + with torch.random.fork_rng(devices=[0] if on_cuda else []): + torch.manual_seed(5) + with torch.backends.cudnn.flags(benchmark=False, deterministic=True, allow_tf32=False): + yield torch.device("cuda:0" if on_cuda else "cpu") + + +def _has_weight_norm(module): + return parametrize.is_parametrized(module, "weight") or ( + hasattr(module, "weight_g") and hasattr(module, "weight_v") + ) + + +def _f0_layers(f0): + return [layer for layer in f0.condnet if isinstance(layer, hifigan.CausalConv1d)] + + +def _make_hift(): + # Source noise/phase tensors are not in state_dict. Give independent + # instances identical construction RNG without changing the caller's RNG. + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(5) + return hifigan.CausalHiFTGenerator( + base_channels=32, + sampling_rate=24000, + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + resblock_kernel_sizes=[3], + resblock_dilation_sizes=[[1, 3, 5]], + source_resblock_kernel_sizes=[3, 3, 3], + source_resblock_dilation_sizes=[[1, 3, 5]] * 3, + f0_predictor=hifigan.CausalConvRNNF0Predictor(cond_channels=16), + ) + + +def _make_loader(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + # Real load_weights and vocoder; only the unrelated Flow is reduced. + model = CosyVoice3Code2Wav.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.flow_model = nn.Linear(2, 2) + model.hift = _make_hift().train() + model._hift_window_len = 64 + return model + + +def _write_checkpoint(model, directory, legacy_checkpoint=False): + state = {name: value.detach().clone() for name, value in model.hift.state_dict().items()} + # Change conv_pre and F0 gains to detect stale constructor weights without + # amplifying every generator layer and overflowing its magnitude exp(). + for name, value in state.items(): + if name.startswith(("conv_pre.", "f0_predictor.")) and name.endswith( + ("weight_g", "parametrizations.weight.original0") + ): + value.mul_(2) + if legacy_checkpoint: + state = { + name.replace("parametrizations.weight.original0", "weight_g").replace( + "parametrizations.weight.original1", "weight_v" + ): value + for name, value in state.items() + } + torch.save(model.flow_model.state_dict(), directory / "flow.pt") + torch.save({"generator." + name: value for name, value in state.items()}, directory / "hift.pt") + return state + + +@pytest.mark.cpu +@pytest.mark.parametrize("finalize", [False, True]) +def test_f0_fold_is_exact_and_idempotent(weight_norm_api, finalize): + f0 = hifigan.CausalConvRNNF0Predictor(cond_channels=16).eval() + layers = _f0_layers(f0) + assert len(layers) == 5 and all(_has_weight_norm(layer) for layer in layers) + assert hasattr(layers[0], "weight_g") == (weight_norm_api == "legacy") + classifier = {name: value.detach().clone() for name, value in f0.classifier.state_dict().items()} + context = f0.left_context_frames + mel = torch.randn(1, 80, 24) + with torch.no_grad(): + before = f0(mel, finalize=finalize) + assert f0.remove_weight_norm() == 5 + weights = [layer.weight for layer in layers] + assert f0.remove_weight_norm() == 0 + assert all(layer.weight is weight for layer, weight in zip(layers, weights)) + assert not any(_has_weight_norm(module) for module in f0.modules()) + assert f0.left_context_frames == context + with torch.no_grad(): + after = f0(mel, finalize=finalize) + assert before.numel() > 0 and torch.isfinite(before).all() + torch.testing.assert_close(after, before, rtol=0, atol=0) + for name, value in f0.classifier.state_dict().items(): + torch.testing.assert_close(value, classifier[name], rtol=0, atol=0) + + +@pytest.mark.cpu +def test_f0_fold_preserves_autograd_usability(weight_norm_api, fold_context): + f0 = hifigan.CausalConvRNNF0Predictor(cond_channels=16).eval() + mel = torch.randn(1, 80, 24, requires_grad=True) + with torch.no_grad(): + before = f0(mel) + with fold_context(): + assert f0.remove_weight_norm() == 5 + assert f0.remove_weight_norm() == 0 + for layer in _f0_layers(f0): + assert isinstance(layer.weight, nn.Parameter) + assert not layer.weight.is_inference() + out = f0(mel) + torch.testing.assert_close(out, before, rtol=0, atol=0) + out.sum().backward() + assert mel.grad is not None and torch.isfinite(mel.grad).all() + + +@pytest.mark.cpu +def test_f0_fold_leaves_generator_norms_unchanged(weight_norm_api): + hift = _make_hift() + generator_norms = { + name + for name, module in hift.named_modules() + if not name.startswith("f0_predictor.") and _has_weight_norm(module) + } + assert len(generator_norms) == 41 + assert sum(_has_weight_norm(module) for module in hift.f0_predictor.modules()) == 5 + assert hift.f0_predictor.remove_weight_norm() == 5 + assert {name for name, module in hift.named_modules() if _has_weight_norm(module)} == generator_norms + + +@pytest.mark.parametrize( + "weight_norm_api,legacy_checkpoint", + [("parametrized", False), ("parametrized", True), ("legacy", True)], + indirect=["weight_norm_api"], +) +def test_load_folds_loaded_generator_and_f0_weights( + tmp_path, weight_norm_api, legacy_checkpoint, runtime, fold_context +): + device = runtime + model = _make_loader() + assert model.hift.training and model.hift.f0_predictor.training and model.flow_model.training + assert hasattr(model.hift.f0_predictor.condnet[0], "weight_g") == (weight_norm_api == "legacy") + state = _write_checkpoint(model, tmp_path, legacy_checkpoint) + # Do not forward the loader instance before loading. A separate unfolded + # reference refreshes legacy cached weights from the changed checkpoint. + reference = _make_hift() + reference.load_state_dict(state, strict=True) + reference.to(device).eval() + reference.f0_predictor.to(device="cpu", dtype=torch.float32) + normalized = {name: module for name, module in reference.named_modules() if _has_weight_norm(module)} + assert len(normalized) == 46 # 41 reduced-fixture generator layers + 5 F0 + mel = torch.randn(1, 80, 24, device=device) + expected = reference.inference(mel, finalize=True)[0] + + with fold_context(): + model.load_weights(str(tmp_path), device) + + assert not model.hift.training and not model.hift.f0_predictor.training and not model.flow_model.training + assert not any(_has_weight_norm(module) for module in model.hift.modules()) + for name, ref_layer in normalized.items(): + layer = model.hift.get_submodule(name) + assert isinstance(layer.weight, nn.Parameter) and not layer.weight.is_inference() + assert layer.weight.device == (torch.device("cpu") if name.startswith("f0_predictor.") else device) + assert layer.weight.dtype == torch.float32 + torch.testing.assert_close(layer.weight, ref_layer.weight, rtol=0, atol=0) + assert {p.device for p in model.hift.f0_predictor.parameters()} == {torch.device("cpu")} + assert {p.dtype for p in model.hift.f0_predictor.parameters()} == {torch.float32} + assert model.hift.remove_weight_norm() == 0 + assert model.hift.f0_predictor.remove_weight_norm() == 0 + actual = model.hift.inference(mel, finalize=True)[0] + assert expected.numel() > 0 and torch.isfinite(expected).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.cpu +def test_invalid_checkpoint_does_not_fold_either_component(tmp_path, weight_norm_api, monkeypatch): + model = _make_loader() + state = _write_checkpoint(model, tmp_path) + state.pop("f0_predictor.classifier.bias") + torch.save({"generator." + name: value for name, value in state.items()}, tmp_path / "hift.pt") + + def unexpected_fold(*args, **kwargs): + pytest.fail("Folding must not run before strict checkpoint validation") + + monkeypatch.setattr(model.hift, "remove_weight_norm", unexpected_fold) + monkeypatch.setattr(model.hift.f0_predictor, "remove_weight_norm", unexpected_fold) + with pytest.raises(RuntimeError, match="Missing key"): + model.load_weights(str(tmp_path), torch.device("cpu")) + assert sum(_has_weight_norm(module) for module in model.hift.modules()) == 46 + + +@pytest.mark.parametrize("reference_fold_generator", [False, True], ids=["unfolded", "c4_only"]) +def test_joint_fold_preserves_streaming_audio_and_caches(tmp_path, weight_norm_api, runtime, reference_fold_generator): + device = runtime + model = _make_loader() + state = _write_checkpoint(model, tmp_path) + reference = _make_loader() + reference.hift.load_state_dict(state, strict=True) + reference.hift.to(device).eval() + reference.hift.f0_predictor.to(device="cpu", dtype=torch.float32) + # Compare both against no folding and against C4 alone, isolating C5's + # incremental effect while exercising the real loader/streaming path. + if reference_fold_generator: + assert reference.hift.remove_weight_norm() == 41 + model.load_weights(str(tmp_path), device) + assert not any(_has_weight_norm(module) for module in model.hift.modules()) + expected_norms = 5 if reference_fold_generator else 46 + assert sum(_has_weight_norm(module) for module in reference.hift.modules()) == expected_norms + chunks = torch.randn(1, 80, 144, device=device).split(24, dim=-1) + cache = ref_cache = None + shifted = False + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + expected, ref_cache = reference._stream_hift_from_feat(chunk, cache_state=ref_cache, finalize=finalize) + actual, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + assert expected.numel() > 0 and torch.isfinite(expected).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + if finalize: + assert cache is None and ref_cache is None + else: + assert cache.keys() == ref_cache.keys() + for key, value in ref_cache.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(cache[key], value, rtol=0, atol=0) + else: + assert cache[key] == value + shifted |= cache["mel_offset"] > 0 + assert shifted + + +@pytest.mark.cpu +def test_shipped_f0_config_folds_exactly_five_layers(weight_norm_api): + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + + f0 = hifigan.CausalConvRNNF0Predictor(**CosyVoice3Config().hift["f0_predictor"]) + assert sum(_has_weight_norm(module) for module in f0.modules()) == 5 + assert f0.remove_weight_norm() == 5 + assert f0.remove_weight_norm() == 0 + assert not any(_has_weight_norm(module) for module in f0.modules()) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py index 449131e2438..713d437ca0a 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py @@ -119,7 +119,7 @@ def test_hift_fold_preserves_audio_and_leaves_f0_unchanged(weight_norm_api, caus assert torch.isfinite(before).all() torch.testing.assert_close(after, before, rtol=0, atol=0) for name, module in normalized.items(): - # F0 is a separate RFC workstream; do not silently fold it here. + # The generator entry point must leave F0 to its separate fold call. assert _has_weight_norm(module) == name.startswith("f0_predictor.") for name, value in hift.f0_predictor.state_dict().items(): torch.testing.assert_close(value, f0_state[name], rtol=0, atol=0) @@ -171,7 +171,7 @@ def test_load_weights_folds_loaded_not_initial_weights(tmp_path, weight_norm_api assert not any( _has_weight_norm(module) for name, module in model.hift.named_modules() if not name.startswith("f0_predictor.") ) - assert _has_weight_norm(model.hift.f0_predictor.condnet[0]) + assert not any(_has_weight_norm(module) for module in model.hift.f0_predictor.modules()) assert not model.hift.training assert not model.flow_model.training for name, value in model.flow_model.state_dict().items(): @@ -234,12 +234,13 @@ def stream(): torch.testing.assert_close(after, before, rtol=0, atol=0) -def test_shipped_config_generator_folds_exactly_77_layers(): - """Pin the RFC #6870 C4 count against the shipped config, so silent - config drift (or a fold that quietly stops covering layers) shows up - here: conv_pre 1 + ups 3 + 9 resblocks x 6 + 3 source resblocks x 6 - + conv_post 1 = 77. Construction only (no forward), so the full - 512-channel model stays cheap on CPU.""" +def test_shipped_config_folds_77_generator_and_five_f0_layers(weight_norm_api): + """Pin separate C4/C5 counts for both APIs on the shipped configuration. + + Generator: conv_pre 1 + ups 3 + 9 resblocks x 6 + 3 source resblocks x 6 + + conv_post 1 = 77. F0 adds five condnet convolutions. Construction only + (no forward), so the full 512-channel model stays cheap on CPU. + """ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config hift_cfg = dict(CosyVoice3Config().hift) @@ -251,8 +252,13 @@ def test_shipped_config_generator_folds_exactly_77_layers(): ) generator = hifigan.CausalHiFTGenerator(f0_predictor=f0, **hift_cfg) + assert sum(_has_weight_norm(module) for module in generator.modules()) == 82 assert generator.remove_weight_norm() == 77 assert generator.remove_weight_norm() == 0 remaining = [name for name, m in generator.named_modules() if _has_weight_norm(m)] assert len(remaining) == 5 assert all(name.startswith("f0_predictor.") for name in remaining) + assert f0.remove_weight_norm() == 5 + assert f0.remove_weight_norm() == 0 + assert generator.remove_weight_norm() == 0 + assert not any(_has_weight_norm(module) for module in generator.modules()) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index 8ad1bfb92f3..acd8427cd75 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -1118,6 +1118,14 @@ def __init__(self, num_class: int = 1, in_channels: int = 80, cond_channels: int if isinstance(layer, CausalConv1d) and layer.causal_type == "left" ) + def remove_weight_norm(self) -> int: + """Fold the five F0 convolutions after loading on the execution device. + + Repeated calls are no-ops. Load original checkpoints into a fresh + predictor, before folding changes the state-dict keys. + """ + return sum(_fold_weight_norm(layer) for layer in self.condnet) + def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor: if finalize is True: x = self.condnet[0](x) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index e01e8284234..86511f8fe64 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -526,5 +526,12 @@ def load_weights(self, model_dir: str, device: torch.device) -> None: logger.info("Folded %d weight-norm layers in HiFT generator", folded) if folded == 0: logger.warning("HiFT generator had no weight-norm layers to fold; check config drift") + # F0 inference runs on CPU for causal precision. Materialize its + # normalized weights there in FP32, not on the generator's device. + self.hift.f0_predictor.to(device="cpu", dtype=torch.float32) + f0_folded = self.hift.f0_predictor.remove_weight_norm() + logger.info("Folded %d weight-norm layers in HiFT F0 predictor", f0_folded) + if f0_folded == 0: + logger.warning("HiFT F0 predictor had no weight-norm layers to fold; check config drift") self.hift.eval() logger.info(f"Loaded hift weights from {hift_path}")