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.
Defer transformer config access until LoRA conversion is actually requested so base inference keeps lazy DiTs unloaded through conditioning. Initialize transformer bookkeeping per pipeline and cover deferred no-LoRA construction, on-demand exclusion setup, and instance isolation.
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.
The existing tests build stages and pipelines by hand, so they would stay
green through both defects review found: a load_modules that never reaches the
deferral, and a stage constructor that reads a component's attributes and
materializes it during post_init.
Three tests now run the production paths. Two call the real
ComposedPipelineBase.load_modules with the component loader stubbed and a
counter on it, asserting that only opted-in names become proxies and that the
loader is never asked for them. The third builds the real MiniMax-H3 stages
over tracked proxies and asserts nothing materialized.
The third one was mutation-checked: adding a single transformer.patch_size
read to MiniMaxH3DenoisingStage.__init__ makes it fail, which is the habit
that defeats deferral in the shared DenoisingStage today.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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>
The schedule maps each deferred component to the last stage that holds it,
derived once after create_pipeline_stages. Every pipeline in the tree builds
its stages there, but nothing enforced it. A stage appended afterwards could
hold a component an earlier stage had already been told to free, and would
then be handed a released component mid-run with no error.
add_stage now rebuilds the schedule and says so, turning an invariant nothing
checked into a visible self-correction.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Same reason as the MiniMax H3 T2V example: this script builds its OffloadConfig
directly rather than going through FastVideoArgs.add_cli_args, so the flag needs
its own switch here to be reachable. FastH3 is the case the flag exists for,
since its components sum past what a 121 GiB unified-memory device can hold.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The example builds its OffloadConfig directly rather than going through
FastVideoArgs.add_cli_args, so the new flag needs its own switch to be
reachable from the motivating command line.
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