fix(offload): per-rank shared memory for EP>1 to prevent expert weight corruption - #64
Open
cennn wants to merge 39 commits into
Open
fix(offload): per-rank shared memory for EP>1 to prevent expert weight corruption#64cennn wants to merge 39 commits into
cennn wants to merge 39 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
Replace ad-hoc ENGINE_CONFIG__EP_SIZE / EP_SIZE env var lookups with get_topology_dim(ep) which reads from the canonical topology key set by ParallelStateManager.
Replace ad-hoc ENGINE_CONFIG__EP_SIZE env var with the canonical MAGI_COMPILE_TOPOLOGY_KEY + get_topology_dim(), consistent with the production code refactor in c1bdb90.
gloo is always available in standard PyTorch installations.
- test_shm_memory_peak: use _create_empty_shm, _stream_copy_and_replace, _pack_params_flat, _split_flat_to_params from magi_compiler._api - test_ep_shared_memory: same, plus remove hand-written numpy tofile path - Remove unnecessary _skip_no_procfs / _skip_no_dist guards - Relax speed threshold to 1.5x (register_parameter overhead)
…ve_to_device - @staticmethod: no dependency on self, enables direct use in tests via PiecewiseCompileInterpreter._fix_graph_device_placement(module) - Recursive _move_to_device: handles nested list/tuple in one pass instead of separate if/elif branches, shorter and extensible
…methods Rename nested closures to descriptive class-level names and promote to @staticmethod on PiecewiseCompileInterpreter, enabling direct use in unit tests without constructing an interpreter instance: - _is_cpu_device → _device_is_cpu - _move_to_device → _recursive_to_device(val, target_device)
Delete local _is_cpu_device and _apply_full_fix from test file; import PiecewiseCompileInterpreter._device_is_cpu and _fix_graph_device_placement directly. Saves ~25 lines and tests actual production code paths.
_device_is_cpu, _recursive_to_device, fix_graph_device_placement are pure functions with no dependency on class state. Move them out of PiecewiseCompileInterpreter to module level for simpler import/test.
Replace hand-written _shm_write_read with production function, using mock.patch for MAGI_SHARED_BIN_PATH and pin_memory_in_place. Bug repro simply passes per_rank=False to the real function (simulating EP>1 with wrong shared-mmap path). Removes ~40 lines of reimplemented logic.
- 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
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.
Summary
Fix garbled video (花屏) caused by
_patch_cpu_offload_applyoverwriting EP-sharded expert weights.Root Cause
When
model_cpu_offload=TrueandEP_SIZE > 1, the original shared-memory logic had onlylocal_rank=0write model weights to a shared file, with all ranks reading that same file. This silently replaced each rank's unique expert shard with rank 0's data, destroying expert weight diversity and producing numerically corrupt output.Fix
_rank{N}.bin), preserving per-rank expert shards while keeping the speed of mmap +pin_memory_in_place._fix_graph_device_placementto fixexample_valuemetadata for all node types (not justget_attr/placeholder), including list/tuple values.ENGINE_CONFIG__EP_SIZEwithEP_SIZEfallback.Unit Test
tests/feature_tests/test_ep_shared_memory.py:test_shared_memory_overwrites_ep_shards— reproduces the bug (rank 1 loses its weights)test_ep_fix_preserves_per_rank_shards— verifies the fix (each rank keeps its own shard)Uses
torch.multiprocessing.spawn+ gloo backend, no GPU required.Verification
Base service with
compile+offloadon 8×H100 (EP=8, CP=8):