⚡ DiT Layer Streaming (low VRAM)
Fit a 14B diffusion model into 3GB of VRAM without quantizing it
- model
- model
What this node actually does
Drop it right after your model loader - after the LoRAs - and it keeps the model's weights in system RAM instead of on the GPU, shipping each block across PCIe a split second before that block runs. Nothing about the math changes. That's the whole pitch, and it's the reason this node is interesting: ComfyUI already has on-demand offloading (Dynamic VRAM), and so do the block-swap packs, but WeightFerry claims bit-for-bit identical output to an unstreamed run. The author tested that claim block by block and generation by generation. I can't independently verify it, but the claim is unusually specific, which counts for something.
Numbers from the author's own bench on an RTX 4060 8GB: Cosmos-Predict2 14B in fp8 peaks at 2.9GB instead of 7GB at basically the same speed per step (6.56 vs 6.91 s/it), and Qwen-Image 20B in 3.3GB. Clamp the card to 2GB and Anima runs 1.65x faster than ComfyUI's own offloading. On a Colab L4 the wins mostly evaporate - Wan, Z-Image and Lumina land within ±7% of Dynamic VRAM, sometimes slower. So: a big win when you're genuinely starved, close to a wash when you're borderline.
Where it fits: it's the thing that makes Anima or a 14B video model possible on a 6–8GB card at all, and it buys headroom for the upscale pass later in the same graph. Anima is a 2B DiT that's slow for its size and wants 30–50 steps, so anything that gets it onto a small card without a GGUF detour is worth having.
How it works
Blocks are captured into pinned (page-locked) host memory once, on the first generation after the model loads - that's the "few extra seconds" you see. With the default double buffer, the next block's copy is queued while the current block computes, so the transfer hides behind arithmetic instead of stalling. resident_blocks decides how many blocks stay permanently in VRAM; the rest shuttle back and forth. If RAM is tight, tiering can spill streamed blocks to a disk cache and read ahead from there.
The arithmetic is ComfyUI's own code - this node moves weights, it doesn't reimplement blocks. UNet models (SDXL, SD 1.5) are explicitly not supported; it's a DiT-only engine.
The inputs that matter
model- the loaded model, post-LoRA.enabled- kill switch. Off means the model passes through untouched.auto_config(default on) - the one you should leave alone. It reads free VRAM, computes the block size, and picksresident_blocksandtieringfor you. It also decides not to stream when everything fits, which is the correct call - forced streaming of a model that fits made Anima slightly slower in the author's own measurement.vram_budget_gb- a soft cap on total ComfyUI VRAM so you can render while a game is open. The image doesn't change, only the speed. A VAE or upscale spike can still blow past it, once, with a warning.prefetch_slots- 3 is the fastest and costs one extra block of VRAM (about 132MB for Anima 2B); 2 is the frugal setting.buffer_mode- single buffer makes transfers synchronous, losing the overlap; only for misbehaving async transfers.sage_attention- injects SageAttention 2 into every block. It applies globally and is not undone when you switch it off, and it leaks to other clones of the same diffusion model (USDU, detailers). Skip it if you launch ComfyUI with--use-sage-attentionalready.
The deeper knobs (host_format, resident_mode, ram_budget_gb, sage_pv_fp16) are edge cases: an fp8 host format halves the transfer for huge, transfer-bound models without LoRA, and sage_pv_fp16 exists because SageAttention's default FP8 P·V kernel returns NaN when a channel of V is all zeros - seen for real with Cosmos 14B in fp8.
One output, model, going into your KSampler or the next enhancement node.
Install
Manager → search WeightFerry (registry name comfyui-weightferry), install, restart. Or:
cd ComfyUI/custom_nodes
git clone https://github.com/nyanmi-1828/ComfyUI-WeightFerry
requirements.txt needs nothing beyond ComfyUI itself. Optional extras you install yourself: sageattention, triton-windows on Windows for the fused Anima/Cosmos block kernels, and gguf only for the bundled converter scripts. You need ComfyUI 0.38.0+, an NVIDIA card (tested on Ada), and roughly the model's file size in free RAM - 16GB runs Anima fine. Be aware the pack changes global state on load: TF32 on, fused blocks, real/tors appended to KSampler's lists.
Where people get burned
"… MB of this process' GPU memory is in shared GPU memory." That's Windows spilling VRAM into system RAM instead of failing, and it makes generation crawl. auto_config reserves that much back on the next load, but closing the other GPU apps is the real fix.
"Anima fused block: Triton not available." Not fatal - the Triton kernel parts stay off and the images are the same, just slightly slower. Install triton-windows to get them.
Wan with --force-channels-last dies on "required rank 4 tensor to use channels_last format". That's a ComfyUI limitation with Wan's 5-D weights; launch without the flag.
Honest caveat: speed depends on your PCIe lanes and RAM bandwidth, so a riser card or a slow DDR4 kit will eat the gains. To get back to vanilla behaviour, Ctrl+B the node; the fused-block and GGUF-dequant changes have their own off-switches (WEIGHTFERRY_FUSED_BLOCK=0, WEIGHTFERRY_GGUF_TRITON=0).
Inputs (14)
| Name | Type | Default | Description |
|---|---|---|---|
| model | MODEL | — | |
| enabled | BOOLEAN | true | Enable layer streaming: DiT blocks are streamed between host RAM and the GPU while the model runs |
| pin_memory | BOOLEAN | true | True: asynchronous DMA from page-locked (pinned) RAM. False: pageable memory (transfers become effectively synchronous) |
| buffer_mode | COMBO | Double Buffer (Async Prefetch, Recommended) | Double: prefetch the next blocks asynchronously so transfers overlap the computation. Single: synchronous transfers |
| resident_blocks | INT | 81–128 | Blocks kept resident in VRAM (minimum 1: block 0 always stays). At least the block count: no streaming |
| sage_attention | BOOLEAN | false | Inject SageAttention 2 into every block. Note: it also applies to other clones sharing the diffusion model (USDU / detailers) and is not undone when switched off. Not needed with --use-sage-attention |
| host_formatopt | COMBO | Compute dtype, LoRA baked (fast with LoRA) | Format of the weights in host RAM. Compute dtype: LoRA baked in (more bytes to transfer, no per-forward work; faster with LoRA). Storage dtype (fp8): half the transfer, LoRA re-applied on the GPU every forward (for very large transfer-bound models without LoRA) |
| tieringopt | COMBO | RAM (all streamed blocks pinned in RAM) | Where the streamed blocks live. RAM: all in pinned RAM. Auto: measures free RAM, SSD speed and compute time and splits the blocks into a RAM tier and an SSD tier (a cache file read ahead by background threads); for machines with little RAM. With resident_mode = Self-managed the resident blocks' copy (read once per load) also goes to the SSD when RAM is short |
| ram_budget_gbopt | FLOAT | 0.00–512 | RAM tier limit (GB) with tiering Auto. 0: automatic (free RAM minus a 6 GB margin) |
| resident_modeopt | COMBO | ComfyUI (Dynamic VRAM) | How the resident blocks are managed. ComfyUI: left to ComfyUI's Dynamic VRAM. Self-managed: kept in dedicated GPU buffers with LoRA baked in; unaffected by Dynamic VRAM write-backs when VRAM is short, so runs give identical results, but their VRAM stays occupied for the whole generation |
| sage_pv_fp16opt | BOOLEAN | false | Use SageAttention's FP16 P*V kernel (global, not undone). The Ada default (FP8) returns NaN when a channel of V is all zeros (seen with Cosmos 14B in fp8). Also enabled by WEIGHTFERRY_SAGE_PV=fp16 |
| auto_configopt | BOOLEAN | true | Auto: chooses the resident block count from the free VRAM and the block size, and sets tiering to Auto. No streaming when every block fits in VRAM (faster). Overrides resident_blocks / tiering |
| prefetch_slotsopt | INT | 32–6 | GPU slots with Double Buffer (blocks prefetched + 1). 3: fastest (measured 0.94x of no streaming, one more slot of VRAM). 2: less VRAM (0.90x). One slot = one block (about 132 MB for Anima 2B) |
| vram_budget_gbopt | FLOAT | 0.00–512 | Most VRAM ComfyUI may use (GB); the automatic choices (resident blocks, activation streaming, caches, VAE tiles, ComfyUI's model loads) plan inside it, so only the speed changes, not the image. 0: WEIGHTFERRY_VRAM_BUDGET_GB, else no budget (all free VRAM, fastest) |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| model | MODEL | — |