[Perf][Qwen3-TTS] Free unused decoder in Talker SpeechTokenizer to VRAM (#2429)

Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
Sy03andYueqian Lin authored and GitHub committed 2026-04-05 19:04:37 +00:00
1 parent 025408f693
commit f6cfacdd16
1 file changed
+8 -3
@@ -1124,14 +1124,19 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
speech_tokenizer_dir,
torch_dtype=torch.bfloat16,
)
# Prefer GPU for encoder if available; otherwise keep CPU.
# Only move encoder to GPU; the decoder is unused by Talker (which
# only calls tok.encode()) and would otherwise waste bf16 VRAM.
# NOTE: after this point the tokenizer instance is encode-only;
# calling tok.decode() will fail because tok.model.decoder is None.
dev = next(self.parameters()).device
if dev.type != "cpu":
try:
tok.model.to(dev)
del tok.model.decoder
tok.model.decoder = None
tok.model.encoder.to(dev)
tok.device = dev
except Exception as e:
raise RuntimeError(f"Failed to move speech tokenizer to {dev}: {e}") from e
raise RuntimeError(f"Failed to move speech tokenizer encoder to {dev}: {e}") from e
else:
tok.device = dev
self._speech_tokenizer = tok