Single-GPU CPU offloading via reasoner/generator split#11
Single-GPU CPU offloading via reasoner/generator split#11rahul-steiger-nv wants to merge 6 commits into
Conversation
|
|
||
| self.language_model.init_weights(buffer_device=buffer_device) | ||
|
|
||
| def offload_module_groups(self) -> dict[str, list[torch.nn.Module]]: |
There was a problem hiding this comment.
Are these changes for model/vfm already in internal version?
There was a problem hiding this comment.
@gekurian Hi George, this MR has added lots of contents in model/vfm/mot and omni_mot_model.py, do you think these contents could be synced into internal codebase?
|
Some conflicts might be solved. THX! |
I fixed the conflicts |
7a9253c
|
Hi @rahul-steiger-nv Let's resolve the conflicts because we have re-synced from internal. We could merge this ASAP. |
Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
|
@foreverlms I resolved the conflicts |
|
https://huggingface.co/SimpleTuner/cosmos3-component-reasoning-layers-bf16-super for training implementations (eg. if anyone wishes to implement this way in musubi, diffusion-pipe or elsewhere) i've extracted the layers for reasoning into their own component repo which allows simpletuner to retrieve these and generate kv tensors that get persisted to disk unless the user decides to disable the text encoder cache. calling it a text encoder is merely because it's reusing the same caching infrastructure inside simpletuner, they are not text encoder components. but at least this way the main model weights do not need to have their reasoning layers included during training time. you can only do this though if you don't use text input augmentation (or if you pre-permute the text inputs and cache them AOT) |
4 participants
Summary
Opt-in, single-GPU CPU offloading for in-tree inference so Cosmos3-Nano fits on
smaller GPUs. The reasoner tower runs once as a prefill that caches per-layer K/V;
the denoise loop then runs generator-only with the reasoner offloaded to CPU.
Default-off — the joint path is unchanged.
Enable with
--offload-stages reasoner generator vae.How it works
MoTDecoderLayer.forwardgainsund_only(prefill) /gen_only(denoise) branches;
ReasonerMemoryStatecaches understanding K/V in prefill andreplays it during generator-only denoise. Empty sub-sequences true-skip their
projections, so offloaded weights are never invoked.
OffloadPipelinepacks each group (reasoner/generator/vae) intopinned CPU memory and stages one at a time into a single reusable GPU arena —
reasoner around prefill, generator around denoise, VAE around decode.
metathrough init and materializedirectly on CPU (never touch the GPU); freed device memory is returned before the
arena is allocated.
Constraints
Single-GPU only, incompatible with CUDA graphs (both guarded);
two_wayattentionwith
video_temporal_causal=False.Testing
Validated on Cosmos3-Nano, RTX 5090 (32 GB): 720p up to 33 frames and
480p up to 389 frames.
Future work
Layerwise offloading could push resolution/frames higher, at the cost of more
staging overhead and runtime — left as a follow-up.