GB10 unified memory made the CPU-then-copy path take minutes for
the 35B FastH3 transformer. When host offload is off, read
safetensors onto the accelerator.
Register pipeline-level compilation as a lazy materialization transform so every freshly loaded component receives the same compile setup. Keep whole-module compilation behind the proxy, preserve release scheduling, and add regression coverage for conditional and whole-module compile paths.
Review of the first version found that the flag was unsafe as a global
default. Releasing a component and loading it again is only correct when
nothing outside the loader has changed it, and two habits in this tree break
that without raising:
* mutating a component after load. LongCatPipeline.initialize_pipeline turns
on block-sparse attention and writes parameters into every transformer
block. That runs once, so a re-materialized component silently comes back
with the feature off and the second generation quietly degrades.
* reading a component's attributes while stages are built. The shared
DenoisingStage.__init__ derives the attention backend from
transformer.hidden_size, which materializes the DiT during post_init and
defeats the deferral it was meant to gain.
_lazy_module_names is now empty in the base class, so an unchecked pipeline
gets no deferral and says so. MiniMax-H3 opts in to the four components this
PR measured, and nothing else changes behaviour.
Also from review:
Releasing on the failure path. A stage's hook frees only what that stage is
the last user of, and it ran only after a successful forward. Both the stage
and the whole run now release on the way out, so the retry a memory
constrained caller attempts does not start from a worse position than the
request that just failed. A failing release cannot replace the exception
being propagated.
Walking into nested stages. Cosmos25AutoDenoisingStage keeps the transformer
inside child stages, so a one-level scan called it unreferenced and never
freed it. The scan now recurses through stages and containers with cycle
protection.
Keeping the proxy when torch.compile is skipped. The FSDP check ran after the
proxy had already been replaced by the real module, so an FSDP-wrapped
component lost both the compile and its release hook.
Not attaching the activation trace to a deferred component, since the hook
manager pins every module it wraps.
test_parser.py asserted on a whole serialized config dict, so adding a field
to OffloadConfig broke it. Grepping the field name could not find that; only
running the suite could.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A pipeline materializes every component before the first stage runs, so peak
memory is the sum of all components even though no two are needed at the same
moment. The CPU offload flags cannot help with this: they act after loading,
and on a unified-memory device moving weights to the host frees nothing because
it is the same pool.
Add lazy_module_load, off by default. Heavy components become a LazyModule
proxy that loads on first use, and the pipeline installs a release hook on the
last stage that holds each one. Peak becomes the largest overlapping set
instead of the sum.
Measured on MiniMax-H3 r16, 121 GiB GB10, 1 GPU, 192x320, 4 steps. Peak CUDA
allocated is 57.7 GiB, reached during conditioning where the text encoder and
video VAE overlap. The four deferred components account for 96.3 GiB together,
which is what stays resident without the flag: text encoder 48.0, DiT 37.8,
video VAE 9.7, audio VAE 0.8. Generation completed and wrote a video.
Details worth flagging:
The release hook lives on PipelineStage.__call__, not in the pipeline stage
loop, so pipelines that override forward still free.
The proxy forwards __class__, so isinstance stays honest. A proxy answering
False to isinstance(module, FSDPModule) would take the wrong branch silently.
Self-returning methods hand back the proxy rather than the component. Stages
write self.vae = self.vae.to(device) in a dozen places; returning the component
there would replace the proxy with a reference the pipeline cannot release, and
the run would look normal while freeing nothing.
Releasing is a latency cost, never a correctness one: a released component
reloads on next access. Training keeps everything resident and warns if the
flag is set. If no stage holds a deferred component the pipeline warns rather
than silently doing nothing.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Aryan Kumar