[WIP]: [Fix] Per-rank SHM for EP offload + streaming 1× host peak - #64
Open
cennn wants to merge 44 commits into
Open
[WIP]: [Fix] Per-rank SHM for EP offload + streaming 1× host peak#64cennn wants to merge 44 commits into
cennn wants to merge 44 commits into
Conversation
_patch_cpu_offload_apply created a single shared-memory file from local_rank=0 and had all ranks read it. With expert parallelism (EP>1), each rank holds a different expert shard; reading rank-0 data on every rank destroyed expert weight diversity and produced garbled video output. Fix: when EP_SIZE>1, fall back to per-rank pin_memory instead of cross-rank shared-memory dedup. Also: move model weights to CUDA before Dynamo tracing (_deep_cuda) so Dynamo captures the fused Triton kernel path instead of the decomposed Python fallback, and extend _fix_graph_device_placement to fix ALL FX nodes with CPU example_values (not just get_attr/placeholder).
…tion Two tests using torch.multiprocessing.spawn with gloo backend: 1. test_shared_memory_overwrites_ep_shards: Reproduces the bug — local_rank=0 writes expert weights to a shared file, all other ranks read it, silently overwriting their own expert shards with rank 0's data. 2. test_ep_fix_preserves_per_rank_shards: Verifies the fix — when EP_SIZE > 1, the shared-memory path is skipped and each rank retains its own expert weights.
ENGINE_CONFIG__EP_SIZE may not be set if the host framework (e.g. disagg_compute_runner) only sets EP_SIZE or configures ep_size programmatically. Fall back to EP_SIZE env var before defaulting to 1.
When EP_SIZE > 1, each rank holds a unique expert shard. The previous fix skipped shared memory entirely and used pin_memory, which was extremely slow for large models (~46GB per rank). Now each rank writes its own shared-memory file to /dev/shm and mmap-reads it back, preserving per-rank expert weights while keeping the speed benefit of shared memory + pin_memory_in_place on already-resident pages. For EP_SIZE <= 1, the original rank-0-writes-all-read scheme is retained (all ranks have identical weights). Also updates the regression test to verify the per-rank shm path.
cennn
changed the title
…_apply With EP>1, all 8 ranks simultaneously created ~43GB flat_buffer + wrote ~43GB to /dev/shm = ~87GB per rank x 8 = ~700GB, exceeding 512Gi container limit. Fix: serialize writes across ranks (one at a time) and write directly into mmap file (no flat_buffer). Peak memory drops from ~700GB to ~392GB.
1. _force_cpu: skip GPU roundtrip for CPU tensors when fn only changes device (not dtype). Reduces peak host memory during model.cuda() by avoiding temporary CUDA host allocations for every parameter. 2. MAGI_OFFLOAD_SKIP_SHM: when set to "1", skip shared memory creation and pin_memory_in_place entirely. Params remain as regular CPU tensors. This allows OffloadExecutor to work on memory-constrained nodes (e.g. 5090 with 512Gi container limit for 8x EP ranks) where the shm+pin overhead causes OOM.
When converting CPU example_value metadata to CUDA for Inductor, the .to(device) call strips torch.nn.Parameter wrapping. This caused OffloadExecutor to misidentify all model weights as regular input tensors, loading all 43.6GB onto GPU simultaneously instead of offloading per-submodule — OOM on 5090 (31GB VRAM). Re-wrap the converted FakeTensor in nn.Parameter to preserve type info.
- OffloadExecutor: log per-step H2D vs compute breakdown when MAGI_OFFLOAD_DEBUG=1 (cuda.synchronize between prefetch and compute for accurate wall-clock split) - _patch_cpu_offload_apply: support MAGI_OFFLOAD_PIN_BUDGET_GB env var Pin up to N GB of weights per rank via cudaHostRegister (no SHM copy) for faster async H2D while staying within host memory budget
Log start/end time for each rank during staggered pin, plus OS memlock limit. Helps debug slow NFS page faults during cudaHostRegister.
Replace sequential 1-rank-at-a-time pinning with parallel waves. Auto-detects max concurrent ranks: total_ram/2 / per_rank_param_size. Override via MAGI_OFFLOAD_PIN_CONCURRENCY env var. 512GB node, 43.36GB/rank → concurrency=5, 2 waves instead of 8. Expected pin time: ~6min vs ~22min sequential.
The per-submodule cuda.synchronize() barriers prevented H2D/compute pipeline overlap, reducing production throughput. Profiling data has been collected; this debug scaffolding is no longer needed.
- Revert offload_warpper.py (all changes were unused imports after debug removal) - Remove dead offload() __dict__ branch (call sites only pass tuple/dict) - Extract 80-line inline pin logic into _staggered_pin_memory() helper - magi_backend.py: use module-level os/magi_logger instead of inline imports - Use %-formatting instead of f-string for logger calls
…fload_apply - _shm_path(): centralize /dev/shm path construction - _pack_params_flat(): copy named tensors into contiguous buffer - _split_flat_to_params(): split flat buffer back to named param views - _create_shm_tensor(): create mmap file + pack in one call - Unify EP>1 and EP<=1 serialization (both use mmap now, remove numpy bf16/fp8 workaround) - _patch_cpu_offload_apply SHM logic: ~135 lines -> ~45 lines
- Merge per_rank_shm and shared branches into single helper - Collapse 3-way dispatch to 2-way (skip_shm vs materialize) - Single cleanup point for del/gc.collect() - Eliminate duplicated pin/append/split/remove/load_state_dict
1 task
…tracing After _deep_cuda removal (e0c7277), Dynamo traces with CPU tensors and specialises .to(x.device) as .to(device('cpu')) — a hardcoded literal in the FX graph. _fix_graph_device_placement already moves example_values to CUDA, but these baked .to(cpu) nodes remained, causing index_select(CUDA, CPU) → BackendCompilerFailed during PiecewiseCompileInterpreter.run(). Extend _fix_graph_device_placement to also rewrite: - call_method('to', device('cpu')) → call_method('to', device('cuda')) - call_function(..., device='cpu') → call_function(..., device='cuda') Add regression test (test_fix_to_cpu_in_graph.py) that: 1. Confirms metadata-only fix still produces the device mismatch 2. Verifies the full rewrite resolves the error 3. Ensures .to(dtype) calls are not affected
- Add gc.collect() to _batch_materialize for symmetry with streaming - Move _split_flat_to_params to top-level import, remove inline import - Remove unused _assign_param import
512MB model caused subprocess timeout in CI Docker containers. 256MB still clearly distinguishes batch (0.6x growth) from streaming (0.08x growth). Threshold 0.4x accommodates smaller model sizes where VmHWM overhead is proportionally lower.
…ss isolation Major changes: - _batch_materialize faithfully replicates the ORIGINAL production code (flat_buffer + tofile + from_file) instead of writing directly to mmap - Switch metric from VmHWM to peak RssAnon via background polling thread. VmHWM could not distinguish batch from streaming (both create mmap). RssAnon cleanly isolates the flat_buffer overhead: batch=0.59x vs streaming=0.00x - Replace mp.fork+Manager with subprocess.run for CI Docker compatibility (fork in multi-threaded environments caused 120s deadlocks in CI) - Fix copyright year in test_ep_shared_memory.py (2025 -> 2026)
Replace RssAnon (per-CPU batched counter, ~40% under-reporting) with smaps_rollup Anonymous (page-table walk, exact). Remove polling thread in favor of deterministic reads at known peak points. Tighten thresholds to 0.8x batch / 0.1x streaming.
cennn
changed the title
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.
Motivation
model_cpu_offload=Truematerializes weights into/dev/shmso ranks on one node can share pages. Onmain, rank 0 writes one file and every rank maps it. That is correct when all ranks hold the same weights. With EP>1 each rank holds a different expert shard, so the shared file silently replaces every shard with rank 0's — MoE routing still runs, but the experts are numerically wrong (garbled video).The same
mainpath also builds a fullflat_buffer = torch.zeros(total_numel)beforetofile(), so host anonymous memory peaks at ~2× model size. That is the reason concurrent per-rank materialize used to look unsafe.Separately, offload tracing leaves
.to(device("cpu"))literals and CPUexample_values in the FX graph.mainonly rewrites factorydevice=kwargs, so Inductor still hits mixed-device FakeTensor errors (aten.index_selectoncuda+cpu).What this PR does
Keep EP shards.
ep_sizecomes fromMAGI_COMPILE_TOPOLOGY_KEY(get_topology_dim("ep")). EP>1 writes_rank{N}.binper rank; EP≤1 still shares one file (rank 0 writes, others map). Same mmap +pin_memory_in_placeH2D path asmain.Stream instead of batch-copy.
_stream_copy_and_replacecopies one parameter into the mmap and immediately swaps it into the module. Peak extra Anonymous memory is ~1/N of the model, not a second full copy. Dense / EP=1 is not slower.Rewrite offload-traced CPU devices.
fix_graph_device_placementremaps.to(cpu)/device=cpuand recursively moves CPUexample_values to the current CUDA device. Metadata-only is not enough — the.to(cpu)node itself must move.Evidence
Minimal EP repro (
test_ep_shared_memory.py, 2 ranks, gloo, no GPU):per_rank=False→ both ranks end with identical weights (bug);per_rank=True→ each rank keeps its own shard.Host peak (
test_shm_memory_peak.py, 256 MB dummy,smaps_rollupAnonymousat the known peak site):main-styleflat_bufferDevice mismatch (
test_fix_to_cpu_in_graph.py): graphx.index_select(0, mapping.to(cpu))— metadata-only rewrite raises FakeTensor device error; full rewrite runs on CUDA and leaves no residual.to(cpu).Cluster: 8×5090 base compile+offload warmup completes on the SHM path; all-5090 E2E (
textencCPU + base/sr/vae 5090) produced valid videos.Other
_force_cpuskips the GPU round-trip when the tensor is already on CPU andfndoes not change dtype (commonmodel.cuda()offload apply).fmt_host_mem()one-liner after_force_cpu/ SHM /gc.collect.