Release the GPU after every stage #11

2026-09-29T00:01:19-07:00

dejima-agent commented

2026-09-29 00:01:19 -07:00

Summary generated by LLM agent

songtrace serve is long-lived but only wants the GPU in bursts. Until now it kept everything it ever loaded, so a box shared with anything else had its VRAM held against a process that was mostly idle.

What was actually holding it

Three things, and all three had to be undone:

align cached the wav2vec2 model in a module-level dict keyed by device, and nothing ever cleared it. This is the big one — MMS_FA is the largest of the three models and it stayed resident from the first alignment until the process exited.
torchcrepe caches its model as attributes on its own infer() function (infer.model and infer.capacity), so pitch leaked a model it never declared.
PyTorch's caching allocator keeps freed blocks for reuse rather than handing them back to the driver, so even Demucs — whose Separator is a local and is collected normally — left its footprint unavailable to anything else. I checked demucs.pretrained: there is no in-memory model cache there, only the on-disk weight cache, so dropping the reference plus emptying the allocator is enough for that one.

The shape of the fix

A stage that caches a model registers how to drop it with @on_release, so models.py knows nothing about the pipeline and each stage owns the knowledge of what it cached. run_pipeline calls release() in a finally after every stage.

Per stage rather than per job, deliberately: each model is used by exactly one stage, so holding it past that buys nothing, and transcribe sits between separate and align waiting on your Whisper server — potentially minutes of a GPU nothing is using while Demucs' weights sit in it. The finally also covers a stage that raises, which previously left its model resident.

release() drops the registered caches, runs a gc.collect() (a module graph is cyclic enough that refcounting alone will not get the weights to the allocator's free list), then empties the allocator. It will not initialise CUDA just to clean up after a run that never touched it — there is an is_initialized() guard, so a CPU-only run does not acquire a context on its way out.

The cost is reloading each model once per job instead of once per process. That is seconds against stages that run for minutes, and it is the trade the shared box wants.

How you can check it

release() answers the bytes it handed back, and the stage logs a line when that is not zero:

[separate] released 3704 MiB of GPU memory

Only when something was actually freed, so it is at most three lines per job and it is evidence rather than noise. Watching nvidia-smi across a run should show the same steps down.

What this does not do

The CUDA context itself — a few hundred MB that belongs to the process until it exits — stays. Releasing that would mean running each stage in a subprocess, which is a much larger change than this and not obviously worth it against the model weights and activations, which are where the gigabytes are. The README says so, and says to stop the container if you need every last MB.

Testing

13 tests for models.py and 6 for the pipeline wiring: callbacks registered and run in order, align's dict cleared, torchcrepe's function attributes deleted (and both the never-imported and nothing-cached cases), the allocator emptied only when CUDA is available and initialised, the freed-byte count, release after every stage, release when a stage raises, no release for a cached stage, and the log line appearing only when something was freed.

Full suite: 302 passed, 1 skipped. tests/test_audio_io.py and tests/test_separate.py were not run — they import torch, which is not installable in the environment I had.

The honest limit: none of this was verified against a real GPU. There is no CUDA and no torch here, so the tests drive a stand-in torch module and prove the plumbing — that the caches are cleared, that empty_cache() is called exactly when it should be, and that the guards hold — not that VRAM actually drops on your box. The log line above is there so that part is checkable in one run rather than requiring you to race nvidia-smi. Worth confirming on a real job before this is relied on.

**Summary generated by LLM agent** `songtrace serve` is long-lived but only wants the GPU in bursts. Until now it kept everything it ever loaded, so a box shared with anything else had its VRAM held against a process that was mostly idle. ## What was actually holding it Three things, and all three had to be undone: - `align` cached the wav2vec2 model in a module-level dict keyed by device, and nothing ever cleared it. This is the big one — MMS_FA is the largest of the three models and it stayed resident from the first alignment until the process exited. - `torchcrepe` caches its model as attributes on its own `infer()` function (`infer.model` and `infer.capacity`), so `pitch` leaked a model it never declared. - PyTorch's caching allocator keeps freed blocks for reuse rather than handing them back to the driver, so even Demucs — whose `Separator` is a local and is collected normally — left its footprint unavailable to anything else. I checked `demucs.pretrained`: there is no in-memory model cache there, only the on-disk weight cache, so dropping the reference plus emptying the allocator is enough for that one. ## The shape of the fix A stage that caches a model registers how to drop it with `@on_release`, so `models.py` knows nothing about the pipeline and each stage owns the knowledge of what it cached. `run_pipeline` calls `release()` in a `finally` after every stage. Per stage rather than per job, deliberately: each model is used by exactly one stage, so holding it past that buys nothing, and `transcribe` sits between `separate` and `align` waiting on your Whisper server — potentially minutes of a GPU nothing is using while Demucs' weights sit in it. The `finally` also covers a stage that raises, which previously left its model resident. `release()` drops the registered caches, runs a `gc.collect()` (a module graph is cyclic enough that refcounting alone will not get the weights to the allocator's free list), then empties the allocator. It will not initialise CUDA just to clean up after a run that never touched it — there is an `is_initialized()` guard, so a CPU-only run does not acquire a context on its way out. The cost is reloading each model once per job instead of once per process. That is seconds against stages that run for minutes, and it is the trade the shared box wants. ## How you can check it `release()` answers the bytes it handed back, and the stage logs a line when that is not zero: ``` [separate] released 3704 MiB of GPU memory ``` Only when something was actually freed, so it is at most three lines per job and it is evidence rather than noise. Watching `nvidia-smi` across a run should show the same steps down. ## What this does not do The CUDA context itself — a few hundred MB that belongs to the process until it exits — stays. Releasing that would mean running each stage in a subprocess, which is a much larger change than this and not obviously worth it against the model weights and activations, which are where the gigabytes are. The README says so, and says to stop the container if you need every last MB. ## Testing 13 tests for `models.py` and 6 for the pipeline wiring: callbacks registered and run in order, `align`'s dict cleared, `torchcrepe`'s function attributes deleted (and both the never-imported and nothing-cached cases), the allocator emptied only when CUDA is available *and* initialised, the freed-byte count, release after every stage, release when a stage raises, no release for a cached stage, and the log line appearing only when something was freed. Full suite: 302 passed, 1 skipped. `tests/test_audio_io.py` and `tests/test_separate.py` were not run — they import torch, which is not installable in the environment I had. The honest limit: **none of this was verified against a real GPU.** There is no CUDA and no torch here, so the tests drive a stand-in torch module and prove the plumbing — that the caches are cleared, that `empty_cache()` is called exactly when it should be, and that the guards hold — not that VRAM actually drops on your box. The log line above is there so that part is checkable in one run rather than requiring you to race `nvidia-smi`. Worth confirming on a real job before this is relied on.

dejima-agent added 1 commit 2026-09-29 00:01:20 -07:00

Release the GPU after every stage eac236893f

songtrace serve is long-lived but only wants the GPU in bursts, and until now
it kept everything it ever loaded. align cached wav2vec2 in a module dict that
was never cleared, torchcrepe caches its model on its own infer() function,
and even once the last reference to a model goes, PyTorch keeps the freed
blocks in its allocator rather than returning them to the driver. On a shared
box that is the whole GPU held against one idle process.

Stages now hand it back as soon as they are done with it. A stage that caches
a model registers how to drop it with @on_release, so models.py needs to know
nothing about the pipeline, and run_pipeline calls release() in a finally
after each stage — not once per job, because transcribe sits between separate
and align waiting on the Whisper server, which can be minutes of a GPU nothing
is using.

release() answers the bytes it handed back and the stage logs a line when that
is not zero, which is the only part of this visible from outside. It will not
initialise CUDA just to clean up after a run that never touched it.

The cost is reloading each model once per job rather than once per process,
seconds against stages that run for minutes. The CUDA context itself stays
until the process exits; releasing that would mean a subprocess per stage.

branderson approved these changes 2026-10-06 13:44:40 -07:00

branderson merged commit 4e4aa07516 into main

2026-10-06 13:44:45 -07:00

branderson referenced this issue from a commit

2026-10-15 14:17:07 -07:00

Merge pull request 'Release the GPU after every stage' (#11) from gpu-unload into main

dejima-agent referenced this issue from a commit

2026-10-15 17:44:36 -07:00

Merge pull request 'Release the GPU after every stage' (#11) from gpu-unload into main

Sign in to join this conversation.