Run Lance inference on low-RAM / multi-GPU consumer hardware (e.g. 8 GB RAM + 5×3060)#43
Open
johbau wants to merge 9 commits into
Open
Run Lance inference on low-RAM / multi-GPU consumer hardware (e.g. 8 GB RAM + 5×3060)#43johbau wants to merge 9 commits into
johbau wants to merge 9 commits into
Conversation
Lets inference_lance.py run on hosts with less system RAM than the model needs to materialize in fp32 on CPU (~12 GB for Lance_3B). On main, the first call `Qwen2ForCausalLM(llm_config)` allocates a freshly-init'd fp32 3B model on CPU and OOM-kills an 8 GB host before any GPU code runs. Changes: - Build LLM / ViT / Lance wrapper under `accelerate.init_empty_weights()` so every nn.Parameter is shape-only on the meta device — near-zero CPU RAM during construction. - Replace `safetensors.safe_open()` with a hand-rolled reader that does plain seek+read of one tensor at a time. safe_open mmaps the whole 12 GB file, which Linux refuses on a host with strict overcommit / no swap (ENOMEM). Peak CPU RAM during load is one tensor at a time. - Pass `dtype=torch.bfloat16` to `set_module_tensor_to_device` so loaded values aren't silently upcast back to the meta tensor's fp32 default. Without this the model lives at fp32 on the GPU, doubling VRAM and breaking the bf16 autocast path (fp32 weights * bf16 activations → fp32 output, then index-put into bf16 destination crashes). - Replace numpy fp64 sin-cos position embeddings with a torch fp32 port that computes on the param's device. PositionEmbedding3D._init_weights used to peak around 4 GB of CPU RAM building 3 intermediate arrays of shape (t*h*w, ~D/3); the GPU version contributes ~zero CPU. - Materialize any params left on `meta` after the load (the popped latent_pos_embed sin-cos buffer) on the target device and re-init. - Set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` so the per-tensor streaming pattern doesn't fragment the CUDA caching allocator into a state where large allocations fail despite plenty of free VRAM. Peak CPU RSS during load stays under ~2 GB. The model loads in bf16 across whatever device the runner picks (cuda:LOCAL_RANK). Multi-GPU sharding for hosts that can't fit the model on one card is a separate follow-up — see SHARDED_LOAD.md. See LOW_RAM_LOAD.md for the full memory profile and per-file rationale. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Builds on the low-RAM streaming load (previous commit) to enable single-
process, model-parallel inference across N GPUs. Lets Lance run on hosts
where no single card has enough VRAM but the aggregate does — e.g.
5 × RTX 3060 (60 GB) for Lance_3B + ViT + VAE.
Changes:
- `_build_lance_device_map(model, num_gpus)` spreads LLM transformer
layers across cuda:0..N-1 with cuda:0 getting a reduced share (it also
hosts embed/lm_head/ViT/VAE/connectors). A safety net pins any
uncovered parameter to cuda:0 so future top-level MoT siblings don't
break dispatch.
- `accelerate.dispatch_model` installs pre/post forward hooks that move
activations between cards as needed. The streaming loader from the
previous commit already routes each tensor onto its target shard at
load time; this commit just attaches the runtime hooks.
- Replace flex_attention with eager-SDPA in all three call sites in
lance.py. flex_attention's BlockMask captures device-specific tensors
that dynamo refuses to combine with Q/K/V from a different shard.
qwen2_navit.py already has an isinstance(attention_mask, List) → SDPA
branch that crosses devices cleanly via accelerate hooks. A new helper
`_flex_mask_to_dense_list` evaluates the flex mask function on a
(q_idx, kv_idx) meshgrid to produce that List.
- Align parent-class Python combine sites that dispatch_model's hooks
can't reach:
- Pin ViT to cuda:0 so its output matches embed_tokens for the inline
`masked_scatter` in validation_video_to_text.
- Move `packed_sequence` back to the index tensor's device after the
Qwen2Model layer loop, so the final norm/lm_head indexing combine
works.
- Handle both Tensor and List in the per-layer
`attention_mask.to(device=…)` call.
- Launcher and config:
- inference_lance.sh: NUM_GPUS=5 default (now means shard count, not
data-parallel rank count), `--num_processes 1`, forwards
`--shard_num_gpus $NUM_GPUS`.
- InferenceArguments: new `shard_num_gpus: int = 0`. 0 = use all
visible GPUs; >0 caps to that many.
Behavior on a 1-GPU host is unchanged — the device map collapses to
"everything on cuda:0" and dispatch_model is skipped. The dense-mask
SDPA replacement runs unconditionally; flex_attention can be gated
behind a flag if a single-card user wants the compiled kernel back.
Memory profile on a 5 × 3060 / 8 GB RAM host:
cuda:0 ~6 GB (3 LLM layers + ViT + VAE + embed + lm_head + extras)
cuda:1 ~3 GB (8 LLM layers)
cuda:2 ~3 GB (8 LLM layers)
cuda:3 ~3 GB (8 LLM layers)
cuda:4 ~3 GB (8 LLM layers)
Smoke test (x2t_image, 768 res, 6 cases) completes successfully. About
67 s per understanding batch — slow because activations shuttle across
PCIe between cards and SDPA is eager. The point is fitting the model on
this hardware, not throughput.
See SHARDED_LOAD.md for the full rationale and per-file summary.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The diffusion + VAE-decode generation path uses forward_inference (KVcache) and the WanVideoVAE, neither exercised by the understanding path. Three model-parallel fixes were needed to run t2i/t2v across multiple GPUs: - forward_inference gen-mode norm (qwen2_navit.py): inference-mode twin of the forward_train post-loop fix. After the decoder-layer loop the sequence is on the last shard, but the index tensors + norm/norm_moe_gen are on cuda:0. Added a .to(packed_text_indexes.device) guard before the gen-mode index-put. Runs every diffusion timestep. (The KVcache path attends via flash_attn_varlen_func, not flex_attention, so the dense-mask change is not exercised here.) - VAE device override (modeling/vae/wan/model.py): WanVideoVAE hard-coded get_device()=cuda:0. It now accepts a device= (default get_device(), so single-GPU is unchanged) used by configure_vae_model/vae_encode/vae_decode, so the VAE can live on a less-crowded card. - Dedicated VAE card (inference_lance.py): the video VAE decode's conv activations (~9 GB at 480^2 / 17 frames) won't fit on a card that also holds LLM layers. For generation tasks, _build_lance_device_map now takes reserve_last_for_vae and shards the LLM across the first N-1 cards, leaving the last card empty of LLM weights; the VAE is built there. Confirmed on the 8 GB-RAM / 5x3060 host: t2i produces a clean PNG and t2v produces a valid 480x480 17-frame h264 mp4 (1.42 s). Note: 768^2 video decode exceeds a single 12 GB card even when dedicated — that needs VAE decode tiling (not implemented). The launcher's VIDEO_HEIGHT/WIDTH default to 768, so pass --VIDEO_HEIGHT 480 --VIDEO_WIDTH 480 for t2v on 12 GB cards. See SHARDED_LOAD.md. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Convenience for bounded generation smoke tests (one output instead of all prompts in the example JSON, which apply_inference_defaults expands to validation_max_samples=100000): - inference_lance.sh: new --DATASET_CONFIG passthrough that forwards --val_dataset_config_file to the Python side. - config/examples/t2i_single.json, t2v_single.json: first prompt of the corresponding example file, so a smoke test generates a single image/video. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Lifts the VAE-decode resolution ceiling (previously ~480-512^2 on a 12 GB card; 768^2 OOMs even on a card dedicated to the VAE). The decode is already streamed temporally, so the memory peak is a single frame's full-resolution conv activations — a spatial problem, not a weight one (so LLM-style layer sharding wouldn't help). Tiling the latent spatially and feather-blending the per-tile decodes bounds per-tile memory and lifts the ceiling. - WanVideoVAE._tiled_decode / _should_tile (modeling/vae/wan/model.py): slice the latent [1,48,t,h,w] into overlapping spatial tiles, decode each via the existing self.vae.decode (which resets its own temporal feat_cache, so each tile is a correct independent temporal stream), and blend into the output with a linear edge ramp + weight-sum normalization. Reuses the validated decode per tile — no Decoder3d rewrite. - Config: vae_tile_size (0=auto above ~512^2 latent, >0=force, <0=disable) and vae_tile_overlap in InferenceArguments; plumbed through inference_lance.py and inference_lance.sh (--VAE_TILE / --VAE_TILE_OVERLAP). - Blend arithmetic unit-tested off-GPU: reconstructing exact-tile decodes matches the source to ~1e-16 across divisible/non-divisible/768^2-latent cases with full coverage (wsum >= 1). TILED_VAE.md documents the design (Approach A: single-GPU tiling, implemented; Approach B: multi-GPU tile distribution, proposal) and the pending in-container validation (480^2 parity, 768^2 memory, seams). SHARDED_LOAD.md cross-links it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
video_edit / image_edit crashed with "Expected all tensors on the same
device, cuda:4 and cuda:0" at the source-latent index-put in the generation
loop (lance.py validation_gen / validation_gen_KVcache):
curr_padded_latent[idx] = x_t[idx]
curr_padded_latent is the VAE-encoded *source* video — built from
val_padded_latent via vae_encode on the reserved VAE card (cuda:4) — while
x_t (the diffusion latent), the index tensors, and the whole denoising loop
run on `device` (cuda:0). Only editing/i2v carry a source latent, so t2v/t2i
never hit this.
Fix: move curr_padded_latent to `device` at the cat/cast point (both gen
methods), so all subsequent combines stay on one device.
Validated: video_edit runs end-to-end on the 8 GB-RAM / 5x3060 box at
--RESOLUTION video_192p (3 example clips). Higher resolution still needs
encode-side VAE tiling (vae_encode is not tiled) — separate follow-up.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Mirror of the decode-side tiling, for the encode path that editing / i2v tasks use (make_padded_latent -> vae_encode). Without it, encoding a source video at higher RESOLUTION overruns the dedicated VAE card (the conv activations at full spatial resolution), so editing only worked at video_192p. - WanVideoVAE._tiled_encode / _should_tile_encode (modeling/vae/wan/model.py): slice the pixel input into latent-cell-aligned spatial tiles (pixel span = latent tile x downsample_spatial, so tiles abut exactly), encode each via the existing self.vae.encode (returns mu/log_var; resets its own temporal feat_cache per call), feather-blend mu and log_var into full-resolution latent canvases with weight-sum normalization, then reparameterize once on the stitched tensors in vae_encode (so per-tile sampling noise never seams). - Auto-enables when the latent grid (H // downsample_spatial) exceeds the same ~512^2 threshold as decode; reuses the --VAE_TILE / --VAE_TILE_OVERLAP knobs. - torch.cuda.synchronize(self.device) at the end of _tiled_encode: the VAE is on a non-default card (cuda:N-1) while the current device is cuda:0, so the many per-tile kernels run async and a downstream empty_cache() in make_padded_latent raced them into an illegal memory access (masked only by CUDA_LAUNCH_BLOCKING=1). The sync flushes the VAE card before the result is consumed — encode tiling then works in normal async mode. Validated on the 8 GB-RAM / 5x3060 box: - Off-GPU: stitch arithmetic exact to ~1e-16, full coverage. - Encode parity (768^2, deterministic mu): production default (tile 32 / overlap 8) = 0.57% rel mean err, converging to the bf16 floor (~0.43%) by overlap 16; tiling is exact-to-bf16, no halo-crop needed. - End-to-end: video_edit at video_480p (752x560, latent 47 -> tiles) produces a correct, seam-free edit in normal async mode. Scope: lifts the spatial/resolution ceiling for editing. Does NOT address the low-res high-frame-count OOM (temporal axis) nor the multi-clip earlyoom (system-RAM accumulation across clips) — both separate. Docs: TILED_VAE_ENCODE.md (design + validation results); TILED_VAE.md renamed to TILED_VAE_DECODE.md so decode/encode docs are a matching pair. Reference updates to the new decode filename in SHARDED_LOAD.md, inference_lance.sh, config/config_factory.py, and model.py comments. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move the meta-init streaming-load helpers out of inference_lance.py into a shared module so the Gradio app can reuse the exact same load path (one source of truth): resolve_lance_checkpoint, build_lance_device_map, device_for_param, the safetensors header/tensor readers, stream_load_into, materialize_remaining_meta, plus POPPED_FROM_CHECKPOINT. inference_lance.py imports them (aliased to its prior underscore names) — no behavior change to the CLI path. Pending: re-run a CLI task on-box to confirm the refactor is byte-identical in behavior. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Make lance_gradio.py use the same meta-init + streaming + multi-GPU sharding + VAE-tiling path as the CLI (via common/model/sharded_load.py), so the web UI runs on the 8 GB-RAM / 5x3060 box. Previously it used the stock single-GPU loader (fp32 CPU materialization + model.to(device)) and also imported a function (init_from_model_path_if_needed) our branch had already removed — so it could not load on this hardware (or even import). - LanceT2VV2TPipeline takes a *shard GPU set* (not a single device_id): meta-init LLM/ViT/Lance under init_empty_weights(), stream weights onto the shards at bf16 via stream_load_into, build_lance_device_map (+ reserve last card for the VAE on generation), dispatch_model, materialize_remaining_meta, skip init_moe. No model.to() (it would collapse the shards). VAE built on the reserved card with tiling (vae_tile_size/overlap). - PipelinePool now holds ONE model sharded across all GPUs, not one full replica per GPU (concurrency stays 1). unload() no longer .cpu()s the dispatched model (accelerate forbids whole-model moves) and frees every shard's cache. - PYTORCH_CUDA_ALLOC_CONF default gains expandable_segments:True (needed by the streaming load). New lance_gradio.sh launcher (--gpus 0,1,2,3,4, binds 0.0.0.0:7860); note podman_run must publish the UI port. - Guards a contiguous-GPU requirement (device_map uses logical cuda:0..N-1). Validated: all 7 task types (x2t_image, x2t_video, t2i, t2v, i2v, image_edit, video_edit) run end-to-end through the Gradio UI on the 8 GB-RAM / 5x3060 box. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Runs Lance inference on an 8 GB-RAM host with 5× RTX 3060 (12 GB each) — far
below the README's ≥40 GB-VRAM requirement. The headline constraint here is system
RAM, not VRAM: the stock path materializes the 3B model in fp32 on CPU (~12 GB) before
moving it to GPU, which OOM-kills an 8 GB host during weight init, before any GPU code
runs. This PR removes that CPU spike and shards the model across the small cards.
Validated end-to-end on this box across all seven task types: x2t_image, x2t_video, t2i,
t2v (incl. 768² via tiled decode), i2v, image_edit, and video_edit — via both the CLI
and the Gradio web UI.
Headline: low system-RAM load
Build the model under
accelerate.init_empty_weights()(meta device) and stream thesafetensors checkpoint with plain
seek/read— nommap(which a low-RAM/no-swap hostrefuses for a 12 GB file) and no full fp32 copy — casting each tensor to bf16 and placing
it directly on its target GPU. Peak CPU RSS stays under ~2 GB instead of ~12 GB. This is
the piece that makes Lance load at all on a small-RAM box, independent of how many GPUs
you have.
Supporting pieces
device_map+accelerate.dispatch_modelspread theLLM layers across N GPUs (
--shard_num_gpus), with a dedicated card for the VAE ongeneration. flex-attention is swapped for an eager-SDPA dense-mask path that's safe
across shards.
forward_inference(KV-cache) gen path, and for the encoded source latent in source-conditioned tasks
(image_edit / video_edit / i2v) so they don't crash on a cross-device index-put.
resolution ceiling (768²+ on a 12 GB card); auto-enables above ~512², off below.
resolution), plus a
synchronizethat makes the sharded encode safe in async mode.common/model/sharded_load.py(one source of truth), andlance_gradio.pynow usesthe identical meta-init + streaming + sharding + tiling load. Its pipeline pool holds
one model sharded across all GPUs (not one replica per GPU); a
lance_gradio.shlauncher binds it for remote access. All seven task types were exercised through the
UI on the 8 GB-RAM / 5×3060 box.
No behavior change for capable hosts
--shard_num_gpusdefaults to all visible GPUs and collapses to single-GPU(
dispatch_modelskipped) when there's one card; VAE tiling auto-disables below ~512².Hosts with enough RAM/VRAM behave as before — and load faster, at half the at-rest VRAM,
since weights stay bf16. Design notes in
LOW_RAM_LOAD.md,SHARDED_LOAD.md,TILED_VAE_DECODE.md,TILED_VAE_ENCODE.md.Related work
Several memory-reduction efforts exist; this one is distinguished by targeting host
RAM and high-res tiling across all tasks:
UND→GEN→VAE load/offload, for T2I. Complementary: it could run within a shard here.
VAE-on-a-separate-GPU + eager-SDPA approach (for Turing 2080 Ti / FP16, T2V +
understanding), and adds MLP token-chunking. The sharding overlaps; what's new here is
the low-system-RAM streaming load, VAE encode/decode tiling for 768², and full
task coverage (editing + i2v). The two are composable.