fix: wrap generate() in inference_mode() to stop a host-RAM leak from watermarking - #164
Open
malammar wants to merge 1 commit into
Open
fix: wrap generate() in inference_mode() to stop a host-RAM leak from watermarking#164malammar wants to merge 1 commit into
malammar wants to merge 1 commit into
Conversation
… watermarking perth's apply_watermark() (called internally by ChatterboxTTS.generate()/ ChatterboxTurboTTS.generate()) runs its own conv-based encoder forward pass with no no_grad()/inference_mode() guard of its own. The call site in engine.py only wrapped generate() in torch.autocast(), which doesn't disable gradient tracking -- so autograd built and retained a full backward graph for that forward pass on every single call, none of which is ever backpropped through. Under sustained load this leaks real host RAM that never gets freed. Fix: add torch.inference_mode() around the same call site. No vendored code needs to change -- inference_mode() as a context manager covers everything called inside it, including the internal watermark step.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
ChatterboxTTS.generate()/ChatterboxTurboTTS.generate()internally call perth'sapply_watermark(), which runs its own conv-based encoder forward pass with nono_grad()/inference_mode()guard of its own.engine.py's call site only wrapsgenerate()intorch.autocast(), which doesn't disable gradient tracking. Result: autograd builds and retains a full backward graph for the watermarking forward pass on every single call, none of which is ever backpropped through — a real host-RAM leak under sustained load.How it was found
Traced with
py-spy(caught worker processes mid-request insideapply_watermarkat peak memory) and confirmed withmemray attach(no code changes or restart needed — attaches to a live process) against a running deployment under sustained load.get_leaked_allocation_records()full stack traces showed the never-freed allocations rooted atgenerate()→apply_watermark()→ perth'sencoder.forward(), one 500MB-1GB+ allocation per call.Fix
Add
torch.inference_mode()around the existinggenerate()call site. Covers the internal watermark call too since it's a context manager, not a decorator — no vendored code needs to change.Verification
Measured against a 6-worker deployment (P102-100 GPUs), same load (concurrency 15, ~700 requests over 220s), same voice,
memray-tracked before and after an equivalent fix applied to a downstream fork with additional batching call sites that hit the same underlying bug:The small residual left isn't growing across the run and RSS stayed flat post-fix — looks like a one-time-per-shape cuDNN algorithm workspace cache (bounded, independent of grad mode), not the runaway per-request leak this fixes.
🤖 Generated with Claude Code