fix(demo): free GPU tensors between batches to prevent VRAM accumulation#374
fix(demo): free GPU tensors between batches to prevent VRAM accumulation#374SuperMarioYL wants to merge 1 commit into
Conversation
In transcribe_batch(), the large `inputs` and `output_ids` tensors were retained in scope until the next GC cycle, preventing the CUDA allocator from reclaiming their memory between batches. On a 24 GB GPU this causes OOM when processing 3+ files of 25-min audio even though each individual batch fits (issue microsoft#368). Changes: - Explicitly delete `inputs` and `output_ids` after the decode loop in transcribe_batch() and flush the allocator with empty_cache(). - Add _release_device_cache() helper (empty_cache + gc.collect) called after each batch in transcribe_with_batching() as well. - Add close() method to release the model and processor at end-of-use. - Wrap the transcribe_with_batching() call in main() with try/finally so close() runs even when an exception is raised mid-batch.
|
@microsoft-github-policy-service agree |
1 similar comment
|
@microsoft-github-policy-service agree |
rickthomasjr
left a comment
There was a problem hiding this comment.
Review: Approve — GPU VRAM memory leak fix
Problem addressed: VRAM accumulation across consecutive batch processing. This is a real issue (#368) where torch holds onto allocated device memory even after tensors are no longer referenced.
Changes are well-structured:
del inputs, output_idsafter each batch — explicit tensor release- New
_release_device_cache()— callstorch.cuda.empty_cache()andgc.collect() - New
close()method — ensures model/processor are freed when done try/finally: asr.close()in main — ensures cleanup even on error
Potential concern: torch.cuda.empty_cache() is an expensive operation. In tight batch loops it could add measurable overhead. However, since this is for processing "many files" (not intra-batch), the trade-off is appropriate. The alternative would be to only call empty_cache() every N batches.
Verdict: Approve. Addresses a genuine VRAM leak with a clean API (close pattern). The try/finally ensures no resource leak.
|
This has been approved and green for a few weeks now — is there anything blocking the merge, or anything else you'd like from me? Happy to rebase onto the latest base if that helps. |
Problem
When
--audio_filescontains multiple files,transcribe_with_batching()loops overbatches and calls
transcribe_batch()for each. Insidetranscribe_batch(), two largeGPU tensors are created:
inputs— the padded batch of speech features moved to deviceoutput_ids— the generated token ids frommodel.generate()Neither is explicitly freed before the function returns. Python's reference-count GC
does not guarantee immediate CUDA memory deallocation, so the allocator cannot reclaim
VRAM between iterations. After 2–3 batches of 25-min audio on a 24 GB GPU the process
OOMs even though each individual batch fits comfortably (issue #368).
Fix
Three targeted changes, no logic modifications:
transcribe_batch()—del inputs, output_idsimmediately after the decode loop,then flush the CUDA allocator via
torch.cuda.empty_cache()._release_device_cache()helper — 2-line private method (empty_cache+gc.collect)called from the three callsites that need it; avoids repeating the pair inline.
close()method — explicit teardown for callers who want to reclaim all VRAM afterprocessing (e.g. before loading a second model in the same process). The
main()entrypoint wraps inference in
try/finallyto call it on both normal exit and exception.Verification
transcribe_batchandtranscribe_with_batchingare unchanged.python demo/vibevoice_asr_inference_from_file.py ...) works as before.Fixes #368