From 3ab37665be05ccdc6086c470c9c2b8a4f0af807d Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 27 Aug 2026 18:11:47 +0800 Subject: [PATCH 01/54] fix(offload): skip shared-memory weight dedup when EP>1 _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). --- magi_compiler/_api.py | 28 +++++++++++++++++--- magi_compiler/magi_backend/magi_backend.py | 30 ++++++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 7bf5f21..d6ddef4 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -253,8 +253,28 @@ def new_call(*args, **kwargs): state = getattr(instance, state_attr) if state.compile_config.offload_config.model_cpu_offload and state.jit_compiled_code is None: - args = offload(args) - kwargs = offload(kwargs) + # Move ALL tensors in the model to CUDA before first compile call + # so Dynamo traces with CUDA tensors and captures the correct + # (Triton) code path. Without this, CPU weights cause + # x.is_cuda=False during tracing, leading to a decomposed Python + # path with 10x more subgraph boundaries and catastrophic + # numerical divergence in MoE routing. + def _deep_cuda(mod): + for name, child in mod.named_children(): + _deep_cuda(child) + for name, buf in list(mod._buffers.items()): + if buf is not None and buf.device.type == "cpu": + mod._buffers[name] = buf.cuda() + for name, param in list(mod._parameters.items()): + if param is not None and param.device.type == "cpu": + mod._parameters[name] = torch.nn.Parameter(param.data.cuda(), requires_grad=param.requires_grad) + for name in list(vars(mod)): + v = getattr(mod, name) + if isinstance(v, torch.Tensor) and not isinstance(v, torch.nn.Parameter) and v.device.type == "cpu": + setattr(mod, name, v.cuda()) + _deep_cuda(instance) + import gc; gc.collect(); torch.cuda.empty_cache() + # Do NOT offload args — keep them on CUDA for correct tracing if torch.compiler.is_compiling(): return old_method(*args, **kwargs) @@ -538,7 +558,8 @@ def _force_cpu(t): _orig_apply(self, _force_cpu) # create shared memory tensors for all parameters/buffers on CPU - if dist.is_initialized(): + ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", "1")) + if dist.is_initialized() and ep_size <= 1: local_rank = int(os.environ.get("LOCAL_RANK", 0)) full_state_dict = self.state_dict() @@ -572,7 +593,6 @@ def _force_cpu(t): if dtype == torch.bfloat16: flat_buffer.view(torch.int16).numpy().tofile(shared_bin_path) elif dtype.itemsize == 1 and dtype.is_floating_point: - # fp8 flat_buffer.view(torch.uint8).numpy().tofile(shared_bin_path) else: flat_buffer.numpy().tofile(shared_bin_path) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 239b224..b8b1dad 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -306,7 +306,37 @@ def _fix_graph_device_placement(self, module: torch.nn.Module): node.update_kwarg('device', target_device) needs_recompile = True + # Fix ALL nodes with CPU example_values — not just get_attr/placeholder. + # model_cpu_offload traces with CPU FakeTensors; Inductor autotuning + # creates benchmark tensors on the example_value device, and the + # codegen may pick CPU-specific paths if it sees CPU metadata. + cpu_fix_count = 0 + for node in module.graph.nodes: + ev = node.meta.get('example_value') + if ev is not None: + if hasattr(ev, 'device') and str(ev.device) == 'cpu': + node.meta['example_value'] = ev.to(target_device) + needs_recompile = True + cpu_fix_count += 1 + elif isinstance(ev, (list, tuple)): + fixed_list = [] + any_fixed = False + for item in ev: + if hasattr(item, 'device') and str(item.device) == 'cpu': + fixed_list.append(item.to(target_device)) + any_fixed = True + cpu_fix_count += 1 + else: + fixed_list.append(item) + if any_fixed: + node.meta['example_value'] = type(ev)(fixed_list) + needs_recompile = True + if needs_recompile: + import os as _os + if _os.environ.get('RANK', '0') == '0': + from magi_compiler.utils import magi_logger + magi_logger.info(f'[fix_device] fixed {cpu_fix_count} CPU example_values to cuda:{target_device}') module.recompile() @observe_lifecycle("piecewise_compile") From 18ab8d59d7628b718a2dc52f546adbb9d0e19b79 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 27 Aug 2026 18:59:39 +0800 Subject: [PATCH 02/54] test(offload): add regression test for EP shared-memory weight corruption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- tests/feature_tests/test_ep_shared_memory.py | 207 +++++++++++++++++++ 1 file changed, 207 insertions(+) create mode 100644 tests/feature_tests/test_ep_shared_memory.py diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py new file mode 100644 index 0000000..36c5255 --- /dev/null +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -0,0 +1,207 @@ +# Copyright (c) 2025 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Regression tests for the EP shared-memory weight corruption bug. + +Root cause: _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's data +on every rank destroyed expert weight diversity and produced garbled output. + +Fix: when EP_SIZE > 1, fall back to per-rank pin_memory instead of +cross-rank shared-memory dedup. + +These tests use torch.multiprocessing.spawn with 2 GPUs to reproduce the +exact multi-rank scenario without a real model. +""" + +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn + + +_skip_no_dist = pytest.mark.skipif( + not hasattr(dist, "is_gloo_available") or not dist.is_gloo_available(), + reason="requires gloo backend", +) + + +class FakeExpertBlock(nn.Module): + """Tiny module simulating an EP-sharded expert block.""" + + def __init__(self, num_experts: int, dim: int): + super().__init__() + self.expert_weight = nn.Parameter(torch.randn(num_experts, dim, dtype=torch.bfloat16)) + + def forward(self, x): + return x @ self.expert_weight.T + + +def _extract_shared_memory_logic(module, local_rank, shared_dir): + """ + Reproduces the ORIGINAL (buggy) shared-memory logic from + _patch_cpu_offload_apply: only local_rank==0 writes the file, + all ranks read the same file and load_state_dict(assign=True). + """ + full_state_dict = module.state_dict() + grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} + for name, tensor in full_state_dict.items(): + dt = tensor.dtype + grouped.setdefault(dt, []).append((name, tensor)) + + shared_state = {} + for dtype, param_list in grouped.items(): + dtype_str = str(dtype).split(".")[-1] + shared_path = os.path.join(shared_dir, f"shared_{dtype_str}.bin") + total_numel = sum(t.numel() for _, t in param_list) + + if local_rank == 0: + flat = torch.zeros(total_numel, dtype=dtype) + off = 0 + for _, t in param_list: + n = t.numel() + flat[off : off + n].copy_(t.view(-1)) + off += n + if dtype == torch.bfloat16: + flat.view(torch.int16).numpy().tofile(shared_path) + else: + flat.numpy().tofile(shared_path) + del flat + + dist.barrier() + + giant = torch.from_file(shared_path, shared=True, size=total_numel, dtype=dtype, device="cpu") + off = 0 + for name, orig in param_list: + n = orig.numel() + shared_state[name] = giant[off : off + n].view(orig.shape) + off += n + + dist.barrier() + + module.load_state_dict(shared_state, assign=True) + + +# ─────────────────────────────────────────────────────────────── +# Worker functions for spawn +# ─────────────────────────────────────────────────────────────── + +def _worker_bug_repro(rank, world_size, shared_dir, seed_per_rank, result_file): + """Reproduces the bug: all ranks get rank-0's weights after shared-memory dedup.""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "29501" + os.environ["LOCAL_RANK"] = str(rank) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + + torch.manual_seed(seed_per_rank[rank]) + model = FakeExpertBlock(num_experts=4, dim=8) + original_weight = model.expert_weight.data.clone() + + _extract_shared_memory_logic(model, local_rank=rank, shared_dir=shared_dir) + + weight_after = model.state_dict()["expert_weight"] + matches_own = torch.equal(weight_after, original_weight) + + torch.save({"rank": rank, "matches_own": matches_own, "weight": weight_after.clone()}, f"{result_file}_{rank}.pt") + dist.destroy_process_group() + + +def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_file): + """Verifies the fix: with EP>1, skip shared-memory → each rank keeps its own weights.""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "29502" + os.environ["LOCAL_RANK"] = str(rank) + os.environ["ENGINE_CONFIG__EP_SIZE"] = str(world_size) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + + torch.manual_seed(seed_per_rank[rank]) + model = FakeExpertBlock(num_experts=4, dim=8) + original_weight = model.expert_weight.data.clone() + + ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", "1")) + if ep_size <= 1: + _extract_shared_memory_logic(model, local_rank=rank, shared_dir=shared_dir) + else: + pass # pin_memory path: weights stay as-is + + weight_after = model.state_dict()["expert_weight"] + matches_own = torch.equal(weight_after, original_weight) + + torch.save({"rank": rank, "matches_own": matches_own, "weight": weight_after.clone()}, f"{result_file}_{rank}.pt") + dist.destroy_process_group() + + +# ─────────────────────────────────────────────────────────────── +# Tests +# ─────────────────────────────────────────────────────────────── + +@_skip_no_dist +def test_shared_memory_overwrites_ep_shards(): + """ + BUG REPRO: with original shared-memory logic, rank 1's expert weights + are silently overwritten by rank 0's data. + """ + world_size = 2 + seeds = {0: 42, 1: 123} + + with tempfile.TemporaryDirectory() as tmpdir: + result_file = os.path.join(tmpdir, "result") + mp.spawn(_worker_bug_repro, args=(world_size, tmpdir, seeds, result_file), nprocs=world_size, join=True) + + r0 = torch.load(f"{result_file}_0.pt", weights_only=True) + r1 = torch.load(f"{result_file}_1.pt", weights_only=True) + + assert r0["matches_own"], "rank 0 should keep its own weights (it wrote the file)" + assert not r1["matches_own"], ( + "BUG REPRO FAILED: rank 1 should have LOST its weights " + "(overwritten by rank 0's shared-memory file), but it still matches. " + "The bug may have been fixed upstream — update this test." + ) + assert torch.equal(r0["weight"], r1["weight"]), ( + "After shared-memory dedup, both ranks should have identical weights " + "(rank 0's data). This is the core of the EP corruption bug." + ) + + +@_skip_no_dist +def test_ep_fix_preserves_per_rank_shards(): + """ + FIX VERIFIED: with EP_SIZE > 1, shared-memory is skipped and each rank + keeps its own expert shard. + """ + world_size = 2 + seeds = {0: 42, 1: 123} + + with tempfile.TemporaryDirectory() as tmpdir: + result_file = os.path.join(tmpdir, "result") + mp.spawn(_worker_fix_verified, args=(world_size, tmpdir, seeds, result_file), nprocs=world_size, join=True) + + r0 = torch.load(f"{result_file}_0.pt", weights_only=True) + r1 = torch.load(f"{result_file}_1.pt", weights_only=True) + + assert r0["matches_own"], "rank 0 should keep its own weights" + assert r1["matches_own"], ( + "FIX FAILED: rank 1 should keep its own weights when EP > 1, " + "but they were overwritten." + ) + assert not torch.equal(r0["weight"], r1["weight"]), ( + "With EP > 1, each rank should have DIFFERENT expert weights. " + "If they're equal, the shared-memory path ran despite EP > 1." + ) From 9e7505856e25d5b40c0abfd7df5bed2e88e503a6 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 27 Aug 2026 19:22:14 +0800 Subject: [PATCH 03/54] fix(offload): read EP_SIZE env var as fallback for ep_size detection 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. --- magi_compiler/_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index d6ddef4..adf36fd 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -558,7 +558,7 @@ def _force_cpu(t): _orig_apply(self, _force_cpu) # create shared memory tensors for all parameters/buffers on CPU - ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", "1")) + ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", os.environ.get("EP_SIZE", "1"))) if dist.is_initialized() and ep_size <= 1: local_rank = int(os.environ.get("LOCAL_RANK", 0)) full_state_dict = self.state_dict() From e0c72779e1ed0fe1cce07e89f3e84e953840ac43 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 27 Aug 2026 20:45:14 +0800 Subject: [PATCH 04/54] fix(offload): use per-rank shared memory for EP>1 instead of pin_memory 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. --- magi_compiler/_api.py | 55 ++++++++++---------- tests/feature_tests/test_ep_shared_memory.py | 50 +++++++++++------- 2 files changed, 59 insertions(+), 46 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index adf36fd..dc8b238 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -253,28 +253,8 @@ def new_call(*args, **kwargs): state = getattr(instance, state_attr) if state.compile_config.offload_config.model_cpu_offload and state.jit_compiled_code is None: - # Move ALL tensors in the model to CUDA before first compile call - # so Dynamo traces with CUDA tensors and captures the correct - # (Triton) code path. Without this, CPU weights cause - # x.is_cuda=False during tracing, leading to a decomposed Python - # path with 10x more subgraph boundaries and catastrophic - # numerical divergence in MoE routing. - def _deep_cuda(mod): - for name, child in mod.named_children(): - _deep_cuda(child) - for name, buf in list(mod._buffers.items()): - if buf is not None and buf.device.type == "cpu": - mod._buffers[name] = buf.cuda() - for name, param in list(mod._parameters.items()): - if param is not None and param.device.type == "cpu": - mod._parameters[name] = torch.nn.Parameter(param.data.cuda(), requires_grad=param.requires_grad) - for name in list(vars(mod)): - v = getattr(mod, name) - if isinstance(v, torch.Tensor) and not isinstance(v, torch.nn.Parameter) and v.device.type == "cpu": - setattr(mod, name, v.cuda()) - _deep_cuda(instance) - import gc; gc.collect(); torch.cuda.empty_cache() - # Do NOT offload args — keep them on CUDA for correct tracing + args = offload(args) + kwargs = offload(kwargs) if torch.compiler.is_compiling(): return old_method(*args, **kwargs) @@ -559,7 +539,7 @@ def _force_cpu(t): # create shared memory tensors for all parameters/buffers on CPU ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", os.environ.get("EP_SIZE", "1"))) - if dist.is_initialized() and ep_size <= 1: + if dist.is_initialized(): local_rank = int(os.environ.get("LOCAL_RANK", 0)) full_state_dict = self.state_dict() @@ -574,15 +554,24 @@ def _force_cpu(t): shared_state_dict = {} self._magi_giant_buffers = [] + # EP <= 1: all ranks have identical weights → rank 0 writes, all read (saves RAM). + # EP > 1: each rank holds a unique expert shard → every rank writes its own file. + per_rank_shm = ep_size > 1 + writer_rank = local_rank if per_rank_shm else 0 + dist.barrier() for dtype, param_list in grouped_params.items(): dtype_str = str(dtype).split(".")[-1] - shared_bin_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{self.__class__.__name__}.bin" + cls_name = self.__class__.__name__ + if per_rank_shm: + shared_bin_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}_rank{local_rank}.bin" + else: + shared_bin_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}.bin" total_numel = sum(t.numel() for _, t in param_list) - if local_rank == 0: + if local_rank == writer_rank: flat_buffer = torch.zeros(total_numel, dtype=dtype) offset = 0 for _, tensor in param_list: @@ -621,8 +610,12 @@ def _force_cpu(t): offset += numel dist.barrier() - if local_rank == 0 and os.path.exists(shared_bin_path): - os.remove(shared_bin_path) + if per_rank_shm: + if os.path.exists(shared_bin_path): + os.remove(shared_bin_path) + else: + if local_rank == 0 and os.path.exists(shared_bin_path): + os.remove(shared_bin_path) self.load_state_dict(shared_state_dict, assign=True) @@ -646,4 +639,12 @@ def offload(obj): return {k: offload(v) for k, v in obj.items()} if isinstance(obj, (list, tuple)): return type(obj)(offload(i) for i in obj) + if hasattr(obj, "__dict__"): + import copy + obj = copy.copy(obj) + for attr_name in list(vars(obj)): + attr_val = getattr(obj, attr_name) + if isinstance(attr_val, torch.Tensor) and attr_val.is_cuda: + setattr(obj, attr_name, attr_val.cpu()) + return obj return obj diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 36c5255..670cdd2 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -20,13 +20,15 @@ (EP > 1), each rank holds a different expert shard; reading rank-0's data on every rank destroyed expert weight diversity and produced garbled output. -Fix: when EP_SIZE > 1, fall back to per-rank pin_memory instead of -cross-rank shared-memory dedup. +Fix: when EP_SIZE > 1, each rank writes its OWN shared-memory file so that +expert shards are preserved. When EP_SIZE <= 1, the original rank-0-writes +all-read scheme is safe (weights are identical across ranks). -These tests use torch.multiprocessing.spawn with 2 GPUs to reproduce the -exact multi-rank scenario without a real model. +These tests use torch.multiprocessing.spawn with 2 workers and the gloo +backend to reproduce the exact multi-rank scenario without a real model. """ +import gc import os import tempfile @@ -54,11 +56,12 @@ def forward(self, x): return x @ self.expert_weight.T -def _extract_shared_memory_logic(module, local_rank, shared_dir): +def _shm_write_read(module, local_rank, shared_dir, per_rank): """ - Reproduces the ORIGINAL (buggy) shared-memory logic from - _patch_cpu_offload_apply: only local_rank==0 writes the file, - all ranks read the same file and load_state_dict(assign=True). + Core shared-memory logic extracted from _patch_cpu_offload_apply. + + per_rank=False → original buggy path (rank 0 writes, all read same file). + per_rank=True → fixed path (each rank writes its own file). """ full_state_dict = module.state_dict() grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} @@ -66,13 +69,16 @@ def _extract_shared_memory_logic(module, local_rank, shared_dir): dt = tensor.dtype grouped.setdefault(dt, []).append((name, tensor)) + writer_rank = local_rank if per_rank else 0 shared_state = {} + for dtype, param_list in grouped.items(): dtype_str = str(dtype).split(".")[-1] - shared_path = os.path.join(shared_dir, f"shared_{dtype_str}.bin") + suffix = f"_rank{local_rank}" if per_rank else "" + shared_path = os.path.join(shared_dir, f"shared_{dtype_str}{suffix}.bin") total_numel = sum(t.numel() for _, t in param_list) - if local_rank == 0: + if local_rank == writer_rank: flat = torch.zeros(total_numel, dtype=dtype) off = 0 for _, t in param_list: @@ -84,6 +90,7 @@ def _extract_shared_memory_logic(module, local_rank, shared_dir): else: flat.numpy().tofile(shared_path) del flat + gc.collect() dist.barrier() @@ -95,6 +102,12 @@ def _extract_shared_memory_logic(module, local_rank, shared_dir): off += n dist.barrier() + if per_rank: + if os.path.exists(shared_path): + os.remove(shared_path) + else: + if local_rank == 0 and os.path.exists(shared_path): + os.remove(shared_path) module.load_state_dict(shared_state, assign=True) @@ -103,6 +116,7 @@ def _extract_shared_memory_logic(module, local_rank, shared_dir): # Worker functions for spawn # ─────────────────────────────────────────────────────────────── + def _worker_bug_repro(rank, world_size, shared_dir, seed_per_rank, result_file): """Reproduces the bug: all ranks get rank-0's weights after shared-memory dedup.""" os.environ["MASTER_ADDR"] = "127.0.0.1" @@ -114,7 +128,7 @@ def _worker_bug_repro(rank, world_size, shared_dir, seed_per_rank, result_file): model = FakeExpertBlock(num_experts=4, dim=8) original_weight = model.expert_weight.data.clone() - _extract_shared_memory_logic(model, local_rank=rank, shared_dir=shared_dir) + _shm_write_read(model, local_rank=rank, shared_dir=shared_dir, per_rank=False) weight_after = model.state_dict()["expert_weight"] matches_own = torch.equal(weight_after, original_weight) @@ -124,7 +138,7 @@ def _worker_bug_repro(rank, world_size, shared_dir, seed_per_rank, result_file): def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_file): - """Verifies the fix: with EP>1, skip shared-memory → each rank keeps its own weights.""" + """Verifies the fix: EP>1 uses per-rank shm files, each rank keeps its own weights.""" os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = "29502" os.environ["LOCAL_RANK"] = str(rank) @@ -136,10 +150,7 @@ def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_fil original_weight = model.expert_weight.data.clone() ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", "1")) - if ep_size <= 1: - _extract_shared_memory_logic(model, local_rank=rank, shared_dir=shared_dir) - else: - pass # pin_memory path: weights stay as-is + _shm_write_read(model, local_rank=rank, shared_dir=shared_dir, per_rank=(ep_size > 1)) weight_after = model.state_dict()["expert_weight"] matches_own = torch.equal(weight_after, original_weight) @@ -152,6 +163,7 @@ def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_fil # Tests # ─────────────────────────────────────────────────────────────── + @_skip_no_dist def test_shared_memory_overwrites_ep_shards(): """ @@ -183,8 +195,8 @@ def test_shared_memory_overwrites_ep_shards(): @_skip_no_dist def test_ep_fix_preserves_per_rank_shards(): """ - FIX VERIFIED: with EP_SIZE > 1, shared-memory is skipped and each rank - keeps its own expert shard. + FIX VERIFIED: with EP_SIZE > 1, per-rank shared-memory files ensure each + rank keeps its own expert shard intact. """ world_size = 2 seeds = {0: 42, 1: 123} @@ -203,5 +215,5 @@ def test_ep_fix_preserves_per_rank_shards(): ) assert not torch.equal(r0["weight"], r1["weight"]), ( "With EP > 1, each rank should have DIFFERENT expert weights. " - "If they're equal, the shared-memory path ran despite EP > 1." + "If they're equal, the per-rank shm path did not work correctly." ) From ff596c7fbc748c6987765cbe23195ca5c1c50411 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Fri, 28 Aug 2026 00:59:06 +0800 Subject: [PATCH 05/54] fix: stagger per-rank shm writes to prevent OOM in _patch_cpu_offload_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. --- magi_compiler/_api.py | 131 +++++++++++++++++++++++++++--------------- 1 file changed, 84 insertions(+), 47 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index dc8b238..1d60285 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -554,70 +554,107 @@ def _force_cpu(t): shared_state_dict = {} self._magi_giant_buffers = [] - # EP <= 1: all ranks have identical weights → rank 0 writes, all read (saves RAM). - # EP > 1: each rank holds a unique expert shard → every rank writes its own file. per_rank_shm = ep_size > 1 - writer_rank = local_rank if per_rank_shm else 0 - dist.barrier() + if per_rank_shm: + # EP > 1: each rank holds a unique expert shard. + # Stagger writes so only one rank at a time allocates the mmap + # file on /dev/shm (~43 GB). Without staggering, all 8 ranks + # would simultaneously need old_params + shm_file = ~87 GB each, + # totalling ~700 GB and exceeding the 512 Gi container limit. + world_size = dist.get_world_size() + for turn in range(world_size): + if local_rank == turn: + for dtype, param_list in grouped_params.items(): + dtype_str = str(dtype).split(".")[-1] + cls_name = self.__class__.__name__ + shm_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}_rank{local_rank}.bin" + total_numel = sum(t.numel() for _, t in param_list) + elem_size = torch.empty(0, dtype=dtype).element_size() + + with open(shm_path, "wb") as f: + f.truncate(total_numel * elem_size) + + giant = torch.from_file(shm_path, shared=True, size=total_numel, dtype=dtype, device="cpu") + offset = 0 + for _, tensor in param_list: + numel = tensor.numel() + giant[offset : offset + numel].copy_(tensor.view(-1)) + offset += numel + + pin_memory_in_place(giant) + self._magi_giant_buffers.append(giant) + + offset = 0 + for name, original_tensor in param_list: + numel = original_tensor.numel() + shared_param = giant[offset : offset + numel].view(original_tensor.shape) + if original_tensor.requires_grad: + shared_param.requires_grad_(True) + shared_state_dict[name] = shared_param + offset += numel + + if os.path.exists(shm_path): + os.remove(shm_path) + + self.load_state_dict(shared_state_dict, assign=True) + del full_state_dict, grouped_params + gc.collect() + + dist.barrier() - for dtype, param_list in grouped_params.items(): - dtype_str = str(dtype).split(".")[-1] - cls_name = self.__class__.__name__ - if per_rank_shm: - shared_bin_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}_rank{local_rank}.bin" - else: + else: + # EP <= 1: all ranks have identical weights; rank 0 writes, all read. + dist.barrier() + for dtype, param_list in grouped_params.items(): + dtype_str = str(dtype).split(".")[-1] + cls_name = self.__class__.__name__ shared_bin_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}.bin" + total_numel = sum(t.numel() for _, t in param_list) - total_numel = sum(t.numel() for _, t in param_list) + if local_rank == 0: + flat_buffer = torch.zeros(total_numel, dtype=dtype) + offset = 0 + for _, tensor in param_list: + numel = tensor.numel() + flat_buffer[offset : offset + numel].copy_(tensor.view(-1)) + offset += numel - if local_rank == writer_rank: - flat_buffer = torch.zeros(total_numel, dtype=dtype) - offset = 0 - for _, tensor in param_list: - numel = tensor.numel() - flat_buffer[offset : offset + numel].copy_(tensor.view(-1)) - offset += numel + if dtype == torch.bfloat16: + flat_buffer.view(torch.int16).numpy().tofile(shared_bin_path) + elif dtype.itemsize == 1 and dtype.is_floating_point: + flat_buffer.view(torch.uint8).numpy().tofile(shared_bin_path) + else: + flat_buffer.numpy().tofile(shared_bin_path) - if dtype == torch.bfloat16: - flat_buffer.view(torch.int16).numpy().tofile(shared_bin_path) - elif dtype.itemsize == 1 and dtype.is_floating_point: - flat_buffer.view(torch.uint8).numpy().tofile(shared_bin_path) - else: - flat_buffer.numpy().tofile(shared_bin_path) + del flat_buffer + gc.collect() - del flat_buffer - gc.collect() + dist.barrier() - dist.barrier() + giant_shared_tensor = torch.from_file( + shared_bin_path, shared=True, size=total_numel, dtype=dtype, device="cpu" + ) + self._magi_giant_buffers.append(giant_shared_tensor) - giant_shared_tensor = torch.from_file( - shared_bin_path, shared=True, size=total_numel, dtype=dtype, device="cpu" - ) - self._magi_giant_buffers.append(giant_shared_tensor) + pin_memory_in_place(giant_shared_tensor) - pin_memory_in_place(giant_shared_tensor) - - offset = 0 - for name, original_tensor in param_list: - numel = original_tensor.numel() - shared_param = giant_shared_tensor[offset : offset + numel].view(original_tensor.shape) + offset = 0 + for name, original_tensor in param_list: + numel = original_tensor.numel() + shared_param = giant_shared_tensor[offset : offset + numel].view(original_tensor.shape) - if original_tensor.requires_grad: - shared_param.requires_grad_(True) + if original_tensor.requires_grad: + shared_param.requires_grad_(True) - shared_state_dict[name] = shared_param - offset += numel + shared_state_dict[name] = shared_param + offset += numel - dist.barrier() - if per_rank_shm: - if os.path.exists(shared_bin_path): - os.remove(shared_bin_path) - else: + dist.barrier() if local_rank == 0 and os.path.exists(shared_bin_path): os.remove(shared_bin_path) - self.load_state_dict(shared_state_dict, assign=True) + self.load_state_dict(shared_state_dict, assign=True) else: From 4ec50acfc660065b246390bb1c4af1f0c0db1048 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Fri, 28 Aug 2026 20:13:44 +0800 Subject: [PATCH 06/54] feat(offload): optimize _force_cpu + add MAGI_OFFLOAD_SKIP_SHM bypass 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. --- magi_compiler/_api.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 1d60285..1a989c6 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -532,7 +532,22 @@ def _cpu_apply(self, fn): return _orig_apply(self, fn) # move all parameters/buffers to CPU + # Optimized: skip GPU roundtrip when tensor is already on CPU and fn + # only changes device (not dtype). The roundtrip was originally needed + # for cases where fn includes dtype conversion (e.g. model.to(dtype=fp16)), + # but the common offload path is just model.cuda() with no dtype change. + _dtype_target_cache: dict = {} + def _force_cpu(t): + if t.device.type == "cpu": + dt = t.dtype + if dt not in _dtype_target_cache: + probe = torch.empty(0, dtype=dt, device="cpu") + _dtype_target_cache[dt] = fn(probe).dtype + target_dt = _dtype_target_cache[dt] + if target_dt == dt: + return t + return t.to(dtype=target_dt) return fn(t).cpu() _orig_apply(self, _force_cpu) @@ -556,7 +571,18 @@ def _force_cpu(t): per_rank_shm = ep_size > 1 - if per_rank_shm: + skip_shm = os.environ.get("MAGI_OFFLOAD_SKIP_SHM", "0") == "1" + if skip_shm: + magi_logger.info( + "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1: skipping shared memory + pin_memory. " + "Params remain as regular CPU tensors for OffloadExecutor.", + local_rank, + ) + # No shared memory creation, no pin_memory_in_place. + # OffloadExecutor will use unpinned H2D transfers (functional but slower). + del full_state_dict, grouped_params + gc.collect() + elif per_rank_shm: # EP > 1: each rank holds a unique expert shard. # Stagger writes so only one rank at a time allocates the mmap # file on /dev/shm (~43 GB). Without staggering, all 8 ranks From 9e6e73149a35ace06145ca85371f058efb4cd609 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Fri, 28 Aug 2026 23:15:40 +0800 Subject: [PATCH 07/54] fix(offload): preserve Parameter type in _fix_graph_device_placement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- magi_compiler/magi_backend/magi_backend.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index b8b1dad..7fe55d1 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -315,7 +315,10 @@ def _fix_graph_device_placement(self, module: torch.nn.Module): ev = node.meta.get('example_value') if ev is not None: if hasattr(ev, 'device') and str(ev.device) == 'cpu': - node.meta['example_value'] = ev.to(target_device) + new_ev = ev.to(target_device) + if isinstance(ev, torch.nn.Parameter): + new_ev = torch.nn.Parameter(new_ev, requires_grad=ev.requires_grad) + node.meta['example_value'] = new_ev needs_recompile = True cpu_fix_count += 1 elif isinstance(ev, (list, tuple)): @@ -323,7 +326,10 @@ def _fix_graph_device_placement(self, module: torch.nn.Module): any_fixed = False for item in ev: if hasattr(item, 'device') and str(item.device) == 'cpu': - fixed_list.append(item.to(target_device)) + new_item = item.to(target_device) + if isinstance(item, torch.nn.Parameter): + new_item = torch.nn.Parameter(new_item, requires_grad=item.requires_grad) + fixed_list.append(new_item) any_fixed = True cpu_fix_count += 1 else: From 3d570188278bc11ac7ac186d97910679f7505bfc Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 02:06:21 +0800 Subject: [PATCH 08/54] feat(offload): add timing instrumentation and incremental pin_memory - 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 --- magi_compiler/_api.py | 38 +++++++++++++++++++----- magi_compiler/offload/offload_warpper.py | 34 +++++++++++++++++++++ 2 files changed, 65 insertions(+), 7 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 1a989c6..0cc7ad1 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -573,13 +573,37 @@ def _force_cpu(t): skip_shm = os.environ.get("MAGI_OFFLOAD_SKIP_SHM", "0") == "1" if skip_shm: - magi_logger.info( - "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1: skipping shared memory + pin_memory. " - "Params remain as regular CPU tensors for OffloadExecutor.", - local_rank, - ) - # No shared memory creation, no pin_memory_in_place. - # OffloadExecutor will use unpinned H2D transfers (functional but slower). + pin_budget_gb = float(os.environ.get("MAGI_OFFLOAD_PIN_BUDGET_GB", "0")) + if pin_budget_gb > 0: + limit = int(pin_budget_gb * (1024 ** 3)) + params_with_size = [] + for name, tensor in full_state_dict.items(): + if tensor.device.type == "cpu": + params_with_size.append((name, tensor, tensor.numel() * tensor.element_size())) + params_with_size.sort(key=lambda x: x[2], reverse=True) + pinned_bytes = 0 + pinned_count = 0 + for name, tensor, size in params_with_size: + if pinned_bytes + size > limit: + continue + try: + pin_memory_in_place(tensor) + pinned_bytes += size + pinned_count += 1 + except RuntimeError: + break + total_bytes = sum(s for _, _, s in params_with_size) + magi_logger.info( + "[Rank %d] Pinned %d params (%.2f GB / %.2f GB total, budget=%.1f GB)", + local_rank, pinned_count, pinned_bytes / (1024**3), + total_bytes / (1024**3), pin_budget_gb, + ) + else: + magi_logger.info( + "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1, PIN_BUDGET=0: " + "params remain as unpinned CPU tensors.", + local_rank, + ) del full_state_dict, grouped_params gc.collect() elif per_rank_shm: diff --git a/magi_compiler/offload/offload_warpper.py b/magi_compiler/offload/offload_warpper.py index 4f2c788..cddd56a 100644 --- a/magi_compiler/offload/offload_warpper.py +++ b/magi_compiler/offload/offload_warpper.py @@ -14,6 +14,8 @@ import collections import operator +import os +import time from typing import Any, Dict import torch @@ -23,10 +25,14 @@ from magi_compiler.config import CompileConfig from magi_compiler.offload.profiler import OffloadProfiler from magi_compiler.offload.scheduler import OffloadRuntimeContext, SchedulerFactory +from magi_compiler.utils import magi_logger from magi_compiler.utils.nvtx import add_nvtx_event from ..magi_depyf.timeline import observe_lifecycle +_OFFLOAD_DEBUG = os.environ.get("MAGI_OFFLOAD_DEBUG", "0") == "1" + + class OffloadExecutor: def __init__(self, graph_module: GraphModule, compile_config: CompileConfig): @@ -153,13 +159,25 @@ def __call__(self, *args): ) need_profile = self.second_call + _step_h2d_ms = 0.0 + _step_compute_ms = 0.0 + _step_count = 0 + for node in self.graph_module.graph.nodes: if node.op == "placeholder": continue elif node.op == "call_module": + if _OFFLOAD_DEBUG: + torch.cuda.synchronize() + _t0 = time.perf_counter() + self.scheduler.prefetch(node.name, runtime_ctx) + if _OFFLOAD_DEBUG: + torch.cuda.synchronize() + _t1 = time.perf_counter() + if need_profile: if torch.distributed.is_initialized(): torch.distributed.barrier() @@ -172,6 +190,13 @@ def __call__(self, *args): env[node] = getattr(self.graph_module, node.target)(*s_args, **s_kwargs) del s_args, s_kwargs + if _OFFLOAD_DEBUG: + torch.cuda.synchronize() + _t2 = time.perf_counter() + _step_h2d_ms += (_t1 - _t0) * 1000 + _step_compute_ms += (_t2 - _t1) * 1000 + _step_count += 1 + if need_profile: if torch.distributed.is_initialized(): torch.distributed.barrier() @@ -189,6 +214,15 @@ def __call__(self, *args): env[node] = node.target(*f_args, **f_kwargs) elif node.op == "output": + if _OFFLOAD_DEBUG and _step_count > 0: + _rank = int(os.environ.get("RANK", "0")) + if _rank == 0: + magi_logger.info( + "[offload-timing] submods=%d h2d=%.1fms compute=%.1fms total=%.1fms", + _step_count, _step_h2d_ms, _step_compute_ms, + _step_h2d_ms + _step_compute_ms, + ) + if self.second_call: self._finalize_warmup() self.second_call = False From 772eadf226f6881659002b9c4739728ee38dec96 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 02:59:18 +0800 Subject: [PATCH 09/54] fix(offload): stagger pin_memory across ranks to avoid OOM --- magi_compiler/_api.py | 52 ++++++++++++++++++++++++------------------- 1 file changed, 29 insertions(+), 23 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 0cc7ad1..5c481f3 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -575,29 +575,35 @@ def _force_cpu(t): if skip_shm: pin_budget_gb = float(os.environ.get("MAGI_OFFLOAD_PIN_BUDGET_GB", "0")) if pin_budget_gb > 0: - limit = int(pin_budget_gb * (1024 ** 3)) - params_with_size = [] - for name, tensor in full_state_dict.items(): - if tensor.device.type == "cpu": - params_with_size.append((name, tensor, tensor.numel() * tensor.element_size())) - params_with_size.sort(key=lambda x: x[2], reverse=True) - pinned_bytes = 0 - pinned_count = 0 - for name, tensor, size in params_with_size: - if pinned_bytes + size > limit: - continue - try: - pin_memory_in_place(tensor) - pinned_bytes += size - pinned_count += 1 - except RuntimeError: - break - total_bytes = sum(s for _, _, s in params_with_size) - magi_logger.info( - "[Rank %d] Pinned %d params (%.2f GB / %.2f GB total, budget=%.1f GB)", - local_rank, pinned_count, pinned_bytes / (1024**3), - total_bytes / (1024**3), pin_budget_gb, - ) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + for turn in range(world_size): + if local_rank == turn: + limit = int(pin_budget_gb * (1024 ** 3)) + params_with_size = [] + for name, tensor in full_state_dict.items(): + if tensor.device.type == "cpu": + params_with_size.append((name, tensor, tensor.numel() * tensor.element_size())) + params_with_size.sort(key=lambda x: x[2], reverse=True) + pinned_bytes = 0 + pinned_count = 0 + for name, tensor, size in params_with_size: + if pinned_bytes + size > limit: + continue + try: + pin_memory_in_place(tensor) + pinned_bytes += size + pinned_count += 1 + except RuntimeError as e: + magi_logger.warning("[Rank %d] pin failed at %.2f GB: %s", local_rank, pinned_bytes / (1024**3), e) + break + total_bytes = sum(s for _, _, s in params_with_size) + magi_logger.info( + "[Rank %d] Pinned %d params (%.2f GB / %.2f GB total, budget=%.1f GB)", + local_rank, pinned_count, pinned_bytes / (1024**3), + total_bytes / (1024**3), pin_budget_gb, + ) + if dist.is_initialized(): + dist.barrier() else: magi_logger.info( "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1, PIN_BUDGET=0: " From 6b2efa47fc6e1d071caeaa6fe91fe6be04931448 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 15:00:08 +0800 Subject: [PATCH 10/54] feat(offload): add per-rank pin_memory timing and RLIMIT_MEMLOCK logging Log start/end time for each rank during staggered pin, plus OS memlock limit. Helps debug slow NFS page faults during cudaHostRegister. --- magi_compiler/_api.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 5c481f3..dbc946c 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -575,9 +575,20 @@ def _force_cpu(t): if skip_shm: pin_budget_gb = float(os.environ.get("MAGI_OFFLOAD_PIN_BUDGET_GB", "0")) if pin_budget_gb > 0: + import time as _time + import resource as _resource world_size = dist.get_world_size() if dist.is_initialized() else 1 + _soft, _hard = _resource.getrlimit(_resource.RLIMIT_MEMLOCK) + magi_logger.info( + "[Rank %d] pin_memory: world=%d budget=%.1f GB, RLIMIT_MEMLOCK soft=%s hard=%s", + local_rank, world_size, pin_budget_gb, + "unlimited" if _soft == _resource.RLIM_INFINITY else f"{_soft/(1024**3):.1f}GB", + "unlimited" if _hard == _resource.RLIM_INFINITY else f"{_hard/(1024**3):.1f}GB", + ) for turn in range(world_size): if local_rank == turn: + _t_pin_start = _time.perf_counter() + magi_logger.info("[Rank %d] pin_memory START (turn %d/%d)", local_rank, turn, world_size) limit = int(pin_budget_gb * (1024 ** 3)) params_with_size = [] for name, tensor in full_state_dict.items(): @@ -597,11 +608,15 @@ def _force_cpu(t): magi_logger.warning("[Rank %d] pin failed at %.2f GB: %s", local_rank, pinned_bytes / (1024**3), e) break total_bytes = sum(s for _, _, s in params_with_size) + _t_pin_end = _time.perf_counter() magi_logger.info( - "[Rank %d] Pinned %d params (%.2f GB / %.2f GB total, budget=%.1f GB)", - local_rank, pinned_count, pinned_bytes / (1024**3), + "[Rank %d] pin_memory DONE (turn %d/%d) %.1fs: %d params (%.2f GB / %.2f GB total, budget=%.1f GB)", + local_rank, turn, world_size, _t_pin_end - _t_pin_start, + pinned_count, pinned_bytes / (1024**3), total_bytes / (1024**3), pin_budget_gb, ) + else: + magi_logger.info("[Rank %d] pin_memory WAIT (turn %d/%d, rank %d pinning)", local_rank, turn, world_size, turn) if dist.is_initialized(): dist.barrier() else: From a198f151126e07be9aec831890dcbeed372a7baa Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 15:06:57 +0800 Subject: [PATCH 11/54] feat(offload): parallel-wave pin_memory with auto-detect concurrency MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- magi_compiler/_api.py | 56 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 47 insertions(+), 9 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index dbc946c..3214a68 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -578,17 +578,55 @@ def _force_cpu(t): import time as _time import resource as _resource world_size = dist.get_world_size() if dist.is_initialized() else 1 + _soft, _hard = _resource.getrlimit(_resource.RLIMIT_MEMLOCK) + + per_rank_bytes = sum( + t.numel() * t.element_size() + for t in full_state_dict.values() + if t.device.type == "cpu" + ) + per_rank_gb = per_rank_bytes / (1024 ** 3) + + _pin_concurrency_env = os.environ.get("MAGI_OFFLOAD_PIN_CONCURRENCY", "") + if _pin_concurrency_env: + pin_concurrency = int(_pin_concurrency_env) + else: + try: + with open("/proc/meminfo") as _f: + for _line in _f: + if _line.startswith("MemTotal:"): + total_ram_kb = int(_line.split()[1]) + break + else: + total_ram_kb = 512 * 1024 * 1024 + total_ram_gb = total_ram_kb / (1024 * 1024) + except Exception: + total_ram_gb = 512.0 + safe_ram_gb = total_ram_gb / 2 + pin_concurrency = max(1, int(safe_ram_gb / per_rank_gb)) if per_rank_gb > 0 else world_size + pin_concurrency = min(pin_concurrency, world_size) + + num_waves = (world_size + pin_concurrency - 1) // pin_concurrency magi_logger.info( - "[Rank %d] pin_memory: world=%d budget=%.1f GB, RLIMIT_MEMLOCK soft=%s hard=%s", - local_rank, world_size, pin_budget_gb, + "[Rank %d] pin_memory: world=%d budget=%.1f GB, per_rank=%.2f GB, " + "concurrency=%d (waves=%d), RLIMIT_MEMLOCK soft=%s hard=%s", + local_rank, world_size, pin_budget_gb, per_rank_gb, + pin_concurrency, num_waves, "unlimited" if _soft == _resource.RLIM_INFINITY else f"{_soft/(1024**3):.1f}GB", "unlimited" if _hard == _resource.RLIM_INFINITY else f"{_hard/(1024**3):.1f}GB", ) - for turn in range(world_size): - if local_rank == turn: + + my_wave = local_rank // pin_concurrency + for wave in range(num_waves): + if wave == my_wave: _t_pin_start = _time.perf_counter() - magi_logger.info("[Rank %d] pin_memory START (turn %d/%d)", local_rank, turn, world_size) + magi_logger.info( + "[Rank %d] pin_memory START (wave %d/%d, ranks %d-%d)", + local_rank, wave + 1, num_waves, + wave * pin_concurrency, + min((wave + 1) * pin_concurrency, world_size) - 1, + ) limit = int(pin_budget_gb * (1024 ** 3)) params_with_size = [] for name, tensor in full_state_dict.items(): @@ -610,13 +648,13 @@ def _force_cpu(t): total_bytes = sum(s for _, _, s in params_with_size) _t_pin_end = _time.perf_counter() magi_logger.info( - "[Rank %d] pin_memory DONE (turn %d/%d) %.1fs: %d params (%.2f GB / %.2f GB total, budget=%.1f GB)", - local_rank, turn, world_size, _t_pin_end - _t_pin_start, + "[Rank %d] pin_memory DONE (wave %d/%d) %.1fs: %d params " + "(%.2f GB / %.2f GB, budget=%.1f GB)", + local_rank, wave + 1, num_waves, + _t_pin_end - _t_pin_start, pinned_count, pinned_bytes / (1024**3), total_bytes / (1024**3), pin_budget_gb, ) - else: - magi_logger.info("[Rank %d] pin_memory WAIT (turn %d/%d, rank %d pinning)", local_rank, turn, world_size, turn) if dist.is_initialized(): dist.barrier() else: From 25537493c05cddaeace6da786120440ee30f9206 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 15:29:19 +0800 Subject: [PATCH 12/54] cleanup: remove _OFFLOAD_DEBUG instrumentation from OffloadExecutor 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. --- magi_compiler/offload/offload_warpper.py | 29 ------------------------ 1 file changed, 29 deletions(-) diff --git a/magi_compiler/offload/offload_warpper.py b/magi_compiler/offload/offload_warpper.py index cddd56a..008dbce 100644 --- a/magi_compiler/offload/offload_warpper.py +++ b/magi_compiler/offload/offload_warpper.py @@ -30,7 +30,6 @@ from ..magi_depyf.timeline import observe_lifecycle -_OFFLOAD_DEBUG = os.environ.get("MAGI_OFFLOAD_DEBUG", "0") == "1" @@ -159,25 +158,13 @@ def __call__(self, *args): ) need_profile = self.second_call - _step_h2d_ms = 0.0 - _step_compute_ms = 0.0 - _step_count = 0 - for node in self.graph_module.graph.nodes: if node.op == "placeholder": continue elif node.op == "call_module": - if _OFFLOAD_DEBUG: - torch.cuda.synchronize() - _t0 = time.perf_counter() - self.scheduler.prefetch(node.name, runtime_ctx) - if _OFFLOAD_DEBUG: - torch.cuda.synchronize() - _t1 = time.perf_counter() - if need_profile: if torch.distributed.is_initialized(): torch.distributed.barrier() @@ -190,13 +177,6 @@ def __call__(self, *args): env[node] = getattr(self.graph_module, node.target)(*s_args, **s_kwargs) del s_args, s_kwargs - if _OFFLOAD_DEBUG: - torch.cuda.synchronize() - _t2 = time.perf_counter() - _step_h2d_ms += (_t1 - _t0) * 1000 - _step_compute_ms += (_t2 - _t1) * 1000 - _step_count += 1 - if need_profile: if torch.distributed.is_initialized(): torch.distributed.barrier() @@ -214,15 +194,6 @@ def __call__(self, *args): env[node] = node.target(*f_args, **f_kwargs) elif node.op == "output": - if _OFFLOAD_DEBUG and _step_count > 0: - _rank = int(os.environ.get("RANK", "0")) - if _rank == 0: - magi_logger.info( - "[offload-timing] submods=%d h2d=%.1fms compute=%.1fms total=%.1fms", - _step_count, _step_h2d_ms, _step_compute_ms, - _step_h2d_ms + _step_compute_ms, - ) - if self.second_call: self._finalize_warmup() self.second_call = False From a91c9725f5663370751a3991337381d4a526d5e9 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 15:40:39 +0800 Subject: [PATCH 13/54] refactor: clean up PR - remove dead code, extract helper, fix imports - 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 --- magi_compiler/_api.py | 186 +++++++++++---------- magi_compiler/magi_backend/magi_backend.py | 7 +- magi_compiler/offload/offload_warpper.py | 5 - 3 files changed, 99 insertions(+), 99 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 3214a68..933191b 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -500,6 +500,101 @@ def _check_dynamic_arg_dims(inferred_dims: dict[str, int | list[int]], target_fu assert base_k in inspect.signature(target_func).parameters, f"Argument {base_k} (from {k}) not found in {target_func}" + +def _staggered_pin_memory( + full_state_dict: dict[str, torch.Tensor], + local_rank: int, + pin_budget_gb: float, +): + """Pin CPU tensors in coordinated waves to avoid host OOM. + + Automatically determines how many ranks can pin concurrently based on + host RAM and per-rank parameter size. Override with env var + ``MAGI_OFFLOAD_PIN_CONCURRENCY``. + """ + import resource + import time + + world_size = dist.get_world_size() if dist.is_initialized() else 1 + soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) + + per_rank_bytes = sum( + t.numel() * t.element_size() + for t in full_state_dict.values() + if t.device.type == "cpu" + ) + per_rank_gb = per_rank_bytes / (1024 ** 3) + + concurrency_env = os.environ.get("MAGI_OFFLOAD_PIN_CONCURRENCY", "") + if concurrency_env: + pin_concurrency = int(concurrency_env) + else: + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal:"): + total_ram_kb = int(line.split()[1]) + break + else: + total_ram_kb = 512 * 1024 * 1024 + total_ram_gb = total_ram_kb / (1024 * 1024) + except Exception: + total_ram_gb = 512.0 + safe_ram_gb = total_ram_gb / 2 + pin_concurrency = max(1, int(safe_ram_gb / per_rank_gb)) if per_rank_gb > 0 else world_size + pin_concurrency = min(pin_concurrency, world_size) + + num_waves = (world_size + pin_concurrency - 1) // pin_concurrency + magi_logger.info( + "[Rank %d] pin_memory: world=%d per_rank=%.2f GB, " + "concurrency=%d (waves=%d), RLIMIT_MEMLOCK soft=%s hard=%s", + local_rank, world_size, per_rank_gb, + pin_concurrency, num_waves, + "unlimited" if soft == resource.RLIM_INFINITY else f"{soft / (1024 ** 3):.1f}GB", + "unlimited" if hard == resource.RLIM_INFINITY else f"{hard / (1024 ** 3):.1f}GB", + ) + + limit_bytes = int(pin_budget_gb * (1024 ** 3)) + my_wave = local_rank // pin_concurrency + for wave in range(num_waves): + if wave == my_wave: + t0 = time.perf_counter() + params = [ + (name, tensor, tensor.numel() * tensor.element_size()) + for name, tensor in full_state_dict.items() + if tensor.device.type == "cpu" + ] + params.sort(key=lambda x: x[2], reverse=True) + + pinned_bytes = 0 + pinned_count = 0 + for name, tensor, size in params: + if pinned_bytes + size > limit_bytes: + continue + try: + pin_memory_in_place(tensor) + pinned_bytes += size + pinned_count += 1 + except RuntimeError as e: + magi_logger.warning( + "[Rank %d] pin failed at %.2f GB: %s", + local_rank, pinned_bytes / (1024 ** 3), e, + ) + break + + total_bytes = sum(s for _, _, s in params) + magi_logger.info( + "[Rank %d] pin_memory DONE (wave %d/%d) %.1fs: " + "%d params (%.2f / %.2f GB, budget=%.1f GB)", + local_rank, wave + 1, num_waves, + time.perf_counter() - t0, + pinned_count, pinned_bytes / (1024 ** 3), + total_bytes / (1024 ** 3), pin_budget_gb, + ) + if dist.is_initialized(): + dist.barrier() + + def _patch_cpu_offload_apply(cls: type[nn.Module]): magi_logger.info(f"Enabling CPU offload for {cls}") _orig_apply = cls._apply @@ -575,88 +670,7 @@ def _force_cpu(t): if skip_shm: pin_budget_gb = float(os.environ.get("MAGI_OFFLOAD_PIN_BUDGET_GB", "0")) if pin_budget_gb > 0: - import time as _time - import resource as _resource - world_size = dist.get_world_size() if dist.is_initialized() else 1 - - _soft, _hard = _resource.getrlimit(_resource.RLIMIT_MEMLOCK) - - per_rank_bytes = sum( - t.numel() * t.element_size() - for t in full_state_dict.values() - if t.device.type == "cpu" - ) - per_rank_gb = per_rank_bytes / (1024 ** 3) - - _pin_concurrency_env = os.environ.get("MAGI_OFFLOAD_PIN_CONCURRENCY", "") - if _pin_concurrency_env: - pin_concurrency = int(_pin_concurrency_env) - else: - try: - with open("/proc/meminfo") as _f: - for _line in _f: - if _line.startswith("MemTotal:"): - total_ram_kb = int(_line.split()[1]) - break - else: - total_ram_kb = 512 * 1024 * 1024 - total_ram_gb = total_ram_kb / (1024 * 1024) - except Exception: - total_ram_gb = 512.0 - safe_ram_gb = total_ram_gb / 2 - pin_concurrency = max(1, int(safe_ram_gb / per_rank_gb)) if per_rank_gb > 0 else world_size - pin_concurrency = min(pin_concurrency, world_size) - - num_waves = (world_size + pin_concurrency - 1) // pin_concurrency - magi_logger.info( - "[Rank %d] pin_memory: world=%d budget=%.1f GB, per_rank=%.2f GB, " - "concurrency=%d (waves=%d), RLIMIT_MEMLOCK soft=%s hard=%s", - local_rank, world_size, pin_budget_gb, per_rank_gb, - pin_concurrency, num_waves, - "unlimited" if _soft == _resource.RLIM_INFINITY else f"{_soft/(1024**3):.1f}GB", - "unlimited" if _hard == _resource.RLIM_INFINITY else f"{_hard/(1024**3):.1f}GB", - ) - - my_wave = local_rank // pin_concurrency - for wave in range(num_waves): - if wave == my_wave: - _t_pin_start = _time.perf_counter() - magi_logger.info( - "[Rank %d] pin_memory START (wave %d/%d, ranks %d-%d)", - local_rank, wave + 1, num_waves, - wave * pin_concurrency, - min((wave + 1) * pin_concurrency, world_size) - 1, - ) - limit = int(pin_budget_gb * (1024 ** 3)) - params_with_size = [] - for name, tensor in full_state_dict.items(): - if tensor.device.type == "cpu": - params_with_size.append((name, tensor, tensor.numel() * tensor.element_size())) - params_with_size.sort(key=lambda x: x[2], reverse=True) - pinned_bytes = 0 - pinned_count = 0 - for name, tensor, size in params_with_size: - if pinned_bytes + size > limit: - continue - try: - pin_memory_in_place(tensor) - pinned_bytes += size - pinned_count += 1 - except RuntimeError as e: - magi_logger.warning("[Rank %d] pin failed at %.2f GB: %s", local_rank, pinned_bytes / (1024**3), e) - break - total_bytes = sum(s for _, _, s in params_with_size) - _t_pin_end = _time.perf_counter() - magi_logger.info( - "[Rank %d] pin_memory DONE (wave %d/%d) %.1fs: %d params " - "(%.2f GB / %.2f GB, budget=%.1f GB)", - local_rank, wave + 1, num_waves, - _t_pin_end - _t_pin_start, - pinned_count, pinned_bytes / (1024**3), - total_bytes / (1024**3), pin_budget_gb, - ) - if dist.is_initialized(): - dist.barrier() + _staggered_pin_memory(full_state_dict, local_rank, pin_budget_gb) else: magi_logger.info( "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1, PIN_BUDGET=0: " @@ -785,12 +799,4 @@ def offload(obj): return {k: offload(v) for k, v in obj.items()} if isinstance(obj, (list, tuple)): return type(obj)(offload(i) for i in obj) - if hasattr(obj, "__dict__"): - import copy - obj = copy.copy(obj) - for attr_name in list(vars(obj)): - attr_val = getattr(obj, attr_name) - if isinstance(attr_val, torch.Tensor) and attr_val.is_cuda: - setattr(obj, attr_name, attr_val.cpu()) - return obj return obj diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 06ee445..658bdbe 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -15,6 +15,7 @@ import ast import dataclasses import pprint +import os import time from collections.abc import Callable from contextlib import contextmanager @@ -345,10 +346,8 @@ def _fix_graph_device_placement(self, module: torch.nn.Module): needs_recompile = True if needs_recompile: - import os as _os - if _os.environ.get('RANK', '0') == '0': - from magi_compiler.utils import magi_logger - magi_logger.info(f'[fix_device] fixed {cpu_fix_count} CPU example_values to cuda:{target_device}') + if os.environ.get('RANK', '0') == '0': + magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device) module.recompile() @observe_lifecycle("piecewise_compile") diff --git a/magi_compiler/offload/offload_warpper.py b/magi_compiler/offload/offload_warpper.py index 008dbce..4f2c788 100644 --- a/magi_compiler/offload/offload_warpper.py +++ b/magi_compiler/offload/offload_warpper.py @@ -14,8 +14,6 @@ import collections import operator -import os -import time from typing import Any, Dict import torch @@ -25,14 +23,11 @@ from magi_compiler.config import CompileConfig from magi_compiler.offload.profiler import OffloadProfiler from magi_compiler.offload.scheduler import OffloadRuntimeContext, SchedulerFactory -from magi_compiler.utils import magi_logger from magi_compiler.utils.nvtx import add_nvtx_event from ..magi_depyf.timeline import observe_lifecycle - - class OffloadExecutor: def __init__(self, graph_module: GraphModule, compile_config: CompileConfig): self.graph_module = graph_module From 4964e90419053f35ef714c6ced770d2586c08251 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 15:54:19 +0800 Subject: [PATCH 14/54] refactor: extract SHM helpers, eliminate duplication in _patch_cpu_offload_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 --- magi_compiler/_api.py | 146 +++++++++++++++++++----------------------- 1 file changed, 66 insertions(+), 80 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 933191b..f0a47a8 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -501,6 +501,55 @@ def _check_dynamic_arg_dims(inferred_dims: dict[str, int | list[int]], target_fu +def _shm_path(cls_name: str, dtype: torch.dtype, rank: int | None = None) -> str: + """Build the /dev/shm path for a shared weight file.""" + dtype_str = str(dtype).split(".")[-1] + suffix = f"_rank{rank}" if rank is not None else "" + return f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}{suffix}.bin" + + +def _pack_params_flat(flat: torch.Tensor, param_list: list[tuple[str, torch.Tensor]]) -> None: + """Copy a list of named tensors into a contiguous flat buffer.""" + offset = 0 + for _, tensor in param_list: + numel = tensor.numel() + flat[offset : offset + numel].copy_(tensor.view(-1)) + offset += numel + + +def _split_flat_to_params( + flat: torch.Tensor, + param_list: list[tuple[str, torch.Tensor]], +) -> dict[str, torch.Tensor]: + """Return views into *flat* shaped like the original parameters.""" + out: dict[str, torch.Tensor] = {} + offset = 0 + for name, orig in param_list: + numel = orig.numel() + view = flat[offset : offset + numel].view(orig.shape) + if orig.requires_grad: + view.requires_grad_(True) + out[name] = view + offset += numel + return out + + +def _create_shm_tensor( + shm_path: str, + param_list: list[tuple[str, torch.Tensor]], + dtype: torch.dtype, +) -> torch.Tensor: + """Create a shared-memory mmap file, pack *param_list* into it, return the giant tensor.""" + total_numel = sum(t.numel() for _, t in param_list) + elem_size = torch.empty(0, dtype=dtype).element_size() + with open(shm_path, "wb") as f: + f.truncate(total_numel * elem_size) + giant = torch.from_file(shm_path, shared=True, size=total_numel, dtype=dtype, device="cpu") + _pack_params_flat(giant, param_list) + return giant + + + def _staggered_pin_memory( full_state_dict: dict[str, torch.Tensor], local_rank: int, @@ -680,103 +729,40 @@ def _force_cpu(t): del full_state_dict, grouped_params gc.collect() elif per_rank_shm: - # EP > 1: each rank holds a unique expert shard. - # Stagger writes so only one rank at a time allocates the mmap - # file on /dev/shm (~43 GB). Without staggering, all 8 ranks - # would simultaneously need old_params + shm_file = ~87 GB each, - # totalling ~700 GB and exceeding the 512 Gi container limit. + # EP > 1: each rank writes its own mmap (staggered to cap peak host memory). + cls_name = self.__class__.__name__ world_size = dist.get_world_size() for turn in range(world_size): if local_rank == turn: for dtype, param_list in grouped_params.items(): - dtype_str = str(dtype).split(".")[-1] - cls_name = self.__class__.__name__ - shm_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}_rank{local_rank}.bin" - total_numel = sum(t.numel() for _, t in param_list) - elem_size = torch.empty(0, dtype=dtype).element_size() - - with open(shm_path, "wb") as f: - f.truncate(total_numel * elem_size) - - giant = torch.from_file(shm_path, shared=True, size=total_numel, dtype=dtype, device="cpu") - offset = 0 - for _, tensor in param_list: - numel = tensor.numel() - giant[offset : offset + numel].copy_(tensor.view(-1)) - offset += numel - + path = _shm_path(cls_name, dtype, rank=local_rank) + giant = _create_shm_tensor(path, param_list, dtype) pin_memory_in_place(giant) self._magi_giant_buffers.append(giant) - - offset = 0 - for name, original_tensor in param_list: - numel = original_tensor.numel() - shared_param = giant[offset : offset + numel].view(original_tensor.shape) - if original_tensor.requires_grad: - shared_param.requires_grad_(True) - shared_state_dict[name] = shared_param - offset += numel - - if os.path.exists(shm_path): - os.remove(shm_path) - + shared_state_dict.update(_split_flat_to_params(giant, param_list)) + if os.path.exists(path): + os.remove(path) self.load_state_dict(shared_state_dict, assign=True) del full_state_dict, grouped_params gc.collect() - dist.barrier() - else: - # EP <= 1: all ranks have identical weights; rank 0 writes, all read. + # EP <= 1: rank 0 writes mmap, all ranks share the same pages. + cls_name = self.__class__.__name__ dist.barrier() for dtype, param_list in grouped_params.items(): - dtype_str = str(dtype).split(".")[-1] - cls_name = self.__class__.__name__ - shared_bin_path = f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{cls_name}.bin" + path = _shm_path(cls_name, dtype) total_numel = sum(t.numel() for _, t in param_list) - if local_rank == 0: - flat_buffer = torch.zeros(total_numel, dtype=dtype) - offset = 0 - for _, tensor in param_list: - numel = tensor.numel() - flat_buffer[offset : offset + numel].copy_(tensor.view(-1)) - offset += numel - - if dtype == torch.bfloat16: - flat_buffer.view(torch.int16).numpy().tofile(shared_bin_path) - elif dtype.itemsize == 1 and dtype.is_floating_point: - flat_buffer.view(torch.uint8).numpy().tofile(shared_bin_path) - else: - flat_buffer.numpy().tofile(shared_bin_path) - - del flat_buffer - gc.collect() - + _create_shm_tensor(path, param_list, dtype) dist.barrier() - - giant_shared_tensor = torch.from_file( - shared_bin_path, shared=True, size=total_numel, dtype=dtype, device="cpu" - ) - self._magi_giant_buffers.append(giant_shared_tensor) - - pin_memory_in_place(giant_shared_tensor) - - offset = 0 - for name, original_tensor in param_list: - numel = original_tensor.numel() - shared_param = giant_shared_tensor[offset : offset + numel].view(original_tensor.shape) - - if original_tensor.requires_grad: - shared_param.requires_grad_(True) - - shared_state_dict[name] = shared_param - offset += numel - + giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") + pin_memory_in_place(giant) + self._magi_giant_buffers.append(giant) + shared_state_dict.update(_split_flat_to_params(giant, param_list)) dist.barrier() - if local_rank == 0 and os.path.exists(shared_bin_path): - os.remove(shared_bin_path) - + if local_rank == 0 and os.path.exists(path): + os.remove(path) self.load_state_dict(shared_state_dict, assign=True) else: From e42d77efa6e9718d81c0be3c564cad4f2a5980d7 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 16:03:17 +0800 Subject: [PATCH 15/54] refactor: unify SHM branches into _materialize_shm_weights() - 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 --- magi_compiler/_api.py | 97 ++++++++++++++++++++++++------------------- 1 file changed, 55 insertions(+), 42 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index f0a47a8..3f7eeca 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -550,6 +550,56 @@ def _create_shm_tensor( +def _materialize_shm_weights( + module: nn.Module, + grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], + local_rank: int, + per_rank: bool, +) -> None: + """Replace module params with pinned shared-memory tensors. + + per_rank=True (EP > 1): each rank writes its own mmap, staggered. + per_rank=False (EP <= 1): rank 0 writes, all ranks share pages. + """ + cls_name = module.__class__.__name__ + shared_state: dict[str, torch.Tensor] = {} + buffers: list[torch.Tensor] = [] + + if per_rank: + world_size = dist.get_world_size() + for turn in range(world_size): + if local_rank == turn: + for dtype, param_list in grouped_params.items(): + path = _shm_path(cls_name, dtype, rank=local_rank) + giant = _create_shm_tensor(path, param_list, dtype) + pin_memory_in_place(giant) + buffers.append(giant) + shared_state.update(_split_flat_to_params(giant, param_list)) + if os.path.exists(path): + os.remove(path) + dist.barrier() + else: + dist.barrier() + for dtype, param_list in grouped_params.items(): + path = _shm_path(cls_name, dtype) + total_numel = sum(t.numel() for _, t in param_list) + if local_rank == 0: + _create_shm_tensor(path, param_list, dtype) + dist.barrier() + giant = torch.from_file( + path, shared=True, size=total_numel, dtype=dtype, device="cpu", + ) + pin_memory_in_place(giant) + buffers.append(giant) + shared_state.update(_split_flat_to_params(giant, param_list)) + dist.barrier() + if local_rank == 0 and os.path.exists(path): + os.remove(path) + + module._magi_giant_buffers = buffers + module.load_state_dict(shared_state, assign=True) + + def _staggered_pin_memory( full_state_dict: dict[str, torch.Tensor], local_rank: int, @@ -710,13 +760,9 @@ def _force_cpu(t): grouped_params[dt] = [] grouped_params[dt].append((name, tensor)) - shared_state_dict = {} - self._magi_giant_buffers = [] - - per_rank_shm = ep_size > 1 - skip_shm = os.environ.get("MAGI_OFFLOAD_SKIP_SHM", "0") == "1" if skip_shm: + self._magi_giant_buffers = [] pin_budget_gb = float(os.environ.get("MAGI_OFFLOAD_PIN_BUDGET_GB", "0")) if pin_budget_gb > 0: _staggered_pin_memory(full_state_dict, local_rank, pin_budget_gb) @@ -726,44 +772,11 @@ def _force_cpu(t): "params remain as unpinned CPU tensors.", local_rank, ) - del full_state_dict, grouped_params - gc.collect() - elif per_rank_shm: - # EP > 1: each rank writes its own mmap (staggered to cap peak host memory). - cls_name = self.__class__.__name__ - world_size = dist.get_world_size() - for turn in range(world_size): - if local_rank == turn: - for dtype, param_list in grouped_params.items(): - path = _shm_path(cls_name, dtype, rank=local_rank) - giant = _create_shm_tensor(path, param_list, dtype) - pin_memory_in_place(giant) - self._magi_giant_buffers.append(giant) - shared_state_dict.update(_split_flat_to_params(giant, param_list)) - if os.path.exists(path): - os.remove(path) - self.load_state_dict(shared_state_dict, assign=True) - del full_state_dict, grouped_params - gc.collect() - dist.barrier() else: - # EP <= 1: rank 0 writes mmap, all ranks share the same pages. - cls_name = self.__class__.__name__ - dist.barrier() - for dtype, param_list in grouped_params.items(): - path = _shm_path(cls_name, dtype) - total_numel = sum(t.numel() for _, t in param_list) - if local_rank == 0: - _create_shm_tensor(path, param_list, dtype) - dist.barrier() - giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") - pin_memory_in_place(giant) - self._magi_giant_buffers.append(giant) - shared_state_dict.update(_split_flat_to_params(giant, param_list)) - dist.barrier() - if local_rank == 0 and os.path.exists(path): - os.remove(path) - self.load_state_dict(shared_state_dict, assign=True) + _materialize_shm_weights(self, grouped_params, local_rank, per_rank=(ep_size > 1)) + + del full_state_dict, grouped_params + gc.collect() else: From 62880d1e33e39fcfbf4bf92597a337ea8130395a Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 16:29:00 +0800 Subject: [PATCH 16/54] style: auto-format (black + isort) --- magi_compiler/_api.py | 69 +++++++------------- magi_compiler/magi_backend/magi_backend.py | 2 +- tests/feature_tests/test_ep_shared_memory.py | 9 +-- 3 files changed, 27 insertions(+), 53 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 3f7eeca..0c8577a 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -500,7 +500,6 @@ def _check_dynamic_arg_dims(inferred_dims: dict[str, int | list[int]], target_fu assert base_k in inspect.signature(target_func).parameters, f"Argument {base_k} (from {k}) not found in {target_func}" - def _shm_path(cls_name: str, dtype: torch.dtype, rank: int | None = None) -> str: """Build the /dev/shm path for a shared weight file.""" dtype_str = str(dtype).split(".")[-1] @@ -517,10 +516,7 @@ def _pack_params_flat(flat: torch.Tensor, param_list: list[tuple[str, torch.Tens offset += numel -def _split_flat_to_params( - flat: torch.Tensor, - param_list: list[tuple[str, torch.Tensor]], -) -> dict[str, torch.Tensor]: +def _split_flat_to_params(flat: torch.Tensor, param_list: list[tuple[str, torch.Tensor]]) -> dict[str, torch.Tensor]: """Return views into *flat* shaped like the original parameters.""" out: dict[str, torch.Tensor] = {} offset = 0 @@ -534,11 +530,7 @@ def _split_flat_to_params( return out -def _create_shm_tensor( - shm_path: str, - param_list: list[tuple[str, torch.Tensor]], - dtype: torch.dtype, -) -> torch.Tensor: +def _create_shm_tensor(shm_path: str, param_list: list[tuple[str, torch.Tensor]], dtype: torch.dtype) -> torch.Tensor: """Create a shared-memory mmap file, pack *param_list* into it, return the giant tensor.""" total_numel = sum(t.numel() for _, t in param_list) elem_size = torch.empty(0, dtype=dtype).element_size() @@ -549,12 +541,8 @@ def _create_shm_tensor( return giant - def _materialize_shm_weights( - module: nn.Module, - grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], - local_rank: int, - per_rank: bool, + module: nn.Module, grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], local_rank: int, per_rank: bool ) -> None: """Replace module params with pinned shared-memory tensors. @@ -586,9 +574,7 @@ def _materialize_shm_weights( if local_rank == 0: _create_shm_tensor(path, param_list, dtype) dist.barrier() - giant = torch.from_file( - path, shared=True, size=total_numel, dtype=dtype, device="cpu", - ) + giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") pin_memory_in_place(giant) buffers.append(giant) shared_state.update(_split_flat_to_params(giant, param_list)) @@ -600,11 +586,7 @@ def _materialize_shm_weights( module.load_state_dict(shared_state, assign=True) -def _staggered_pin_memory( - full_state_dict: dict[str, torch.Tensor], - local_rank: int, - pin_budget_gb: float, -): +def _staggered_pin_memory(full_state_dict: dict[str, torch.Tensor], local_rank: int, pin_budget_gb: float): """Pin CPU tensors in coordinated waves to avoid host OOM. Automatically determines how many ranks can pin concurrently based on @@ -617,12 +599,8 @@ def _staggered_pin_memory( world_size = dist.get_world_size() if dist.is_initialized() else 1 soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) - per_rank_bytes = sum( - t.numel() * t.element_size() - for t in full_state_dict.values() - if t.device.type == "cpu" - ) - per_rank_gb = per_rank_bytes / (1024 ** 3) + per_rank_bytes = sum(t.numel() * t.element_size() for t in full_state_dict.values() if t.device.type == "cpu") + per_rank_gb = per_rank_bytes / (1024**3) concurrency_env = os.environ.get("MAGI_OFFLOAD_PIN_CONCURRENCY", "") if concurrency_env: @@ -645,15 +623,17 @@ def _staggered_pin_memory( num_waves = (world_size + pin_concurrency - 1) // pin_concurrency magi_logger.info( - "[Rank %d] pin_memory: world=%d per_rank=%.2f GB, " - "concurrency=%d (waves=%d), RLIMIT_MEMLOCK soft=%s hard=%s", - local_rank, world_size, per_rank_gb, - pin_concurrency, num_waves, + "[Rank %d] pin_memory: world=%d per_rank=%.2f GB, " "concurrency=%d (waves=%d), RLIMIT_MEMLOCK soft=%s hard=%s", + local_rank, + world_size, + per_rank_gb, + pin_concurrency, + num_waves, "unlimited" if soft == resource.RLIM_INFINITY else f"{soft / (1024 ** 3):.1f}GB", "unlimited" if hard == resource.RLIM_INFINITY else f"{hard / (1024 ** 3):.1f}GB", ) - limit_bytes = int(pin_budget_gb * (1024 ** 3)) + limit_bytes = int(pin_budget_gb * (1024**3)) my_wave = local_rank // pin_concurrency for wave in range(num_waves): if wave == my_wave: @@ -675,20 +655,20 @@ def _staggered_pin_memory( pinned_bytes += size pinned_count += 1 except RuntimeError as e: - magi_logger.warning( - "[Rank %d] pin failed at %.2f GB: %s", - local_rank, pinned_bytes / (1024 ** 3), e, - ) + magi_logger.warning("[Rank %d] pin failed at %.2f GB: %s", local_rank, pinned_bytes / (1024**3), e) break total_bytes = sum(s for _, _, s in params) magi_logger.info( - "[Rank %d] pin_memory DONE (wave %d/%d) %.1fs: " - "%d params (%.2f / %.2f GB, budget=%.1f GB)", - local_rank, wave + 1, num_waves, + "[Rank %d] pin_memory DONE (wave %d/%d) %.1fs: " "%d params (%.2f / %.2f GB, budget=%.1f GB)", + local_rank, + wave + 1, + num_waves, time.perf_counter() - t0, - pinned_count, pinned_bytes / (1024 ** 3), - total_bytes / (1024 ** 3), pin_budget_gb, + pinned_count, + pinned_bytes / (1024**3), + total_bytes / (1024**3), + pin_budget_gb, ) if dist.is_initialized(): dist.barrier() @@ -768,8 +748,7 @@ def _force_cpu(t): _staggered_pin_memory(full_state_dict, local_rank, pin_budget_gb) else: magi_logger.info( - "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1, PIN_BUDGET=0: " - "params remain as unpinned CPU tensors.", + "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1, PIN_BUDGET=0: " "params remain as unpinned CPU tensors.", local_rank, ) else: diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 658bdbe..0b00f87 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -14,8 +14,8 @@ import ast import dataclasses -import pprint import os +import pprint import time from collections.abc import Callable from contextlib import contextmanager diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 670cdd2..505f21b 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -38,10 +38,8 @@ import torch.multiprocessing as mp import torch.nn as nn - _skip_no_dist = pytest.mark.skipif( - not hasattr(dist, "is_gloo_available") or not dist.is_gloo_available(), - reason="requires gloo backend", + not hasattr(dist, "is_gloo_available") or not dist.is_gloo_available(), reason="requires gloo backend" ) @@ -209,10 +207,7 @@ def test_ep_fix_preserves_per_rank_shards(): r1 = torch.load(f"{result_file}_1.pt", weights_only=True) assert r0["matches_own"], "rank 0 should keep its own weights" - assert r1["matches_own"], ( - "FIX FAILED: rank 1 should keep its own weights when EP > 1, " - "but they were overwritten." - ) + assert r1["matches_own"], "FIX FAILED: rank 1 should keep its own weights when EP > 1, " "but they were overwritten." assert not torch.equal(r0["weight"], r1["weight"]), ( "With EP > 1, each rank should have DIFFERENT expert weights. " "If they're equal, the per-rank shm path did not work correctly." From 6dc57fb6b1db2fe1d95a23b73f826f44ac9ce422 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sat, 29 Aug 2026 18:46:13 +0800 Subject: [PATCH 17/54] fix(offload): rewrite .to(device('cpu')) in FX graph for CPU-offload tracing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- magi_compiler/magi_backend/magi_backend.py | 36 +++++ .../feature_tests/test_fix_to_cpu_in_graph.py | 132 ++++++++++++++++++ 2 files changed, 168 insertions(+) create mode 100644 tests/feature_tests/test_fix_to_cpu_in_graph.py diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 0b00f87..22bdd54 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -313,6 +313,42 @@ def _fix_graph_device_placement(self, module: torch.nn.Module): node.update_kwarg('device', target_device) needs_recompile = True + # Fix .to(device('cpu')) calls baked during CPU-offload tracing. + # Without _deep_cuda, Dynamo traces with CPU tensors and specialises + # .to(x.device) as .to(device('cpu')). After we move example_values + # to CUDA below, these hardcoded .to(cpu) nodes create CUDA-vs-CPU + # mismatches in PiecewiseCompileInterpreter.run(). + def _is_cpu_device(val): + if isinstance(val, torch.device): + return val.type == 'cpu' + if isinstance(val, str): + return val == 'cpu' + return False + + for node in module.graph.nodes: + if node.op == 'call_method' and node.target == 'to': + new_args = list(node.args) + changed = False + for i, arg in enumerate(new_args): + if _is_cpu_device(arg): + new_args[i] = torch.device('cuda', target_device) + changed = True + if changed: + node.args = tuple(new_args) + needs_recompile = True + if 'device' in node.kwargs and _is_cpu_device(node.kwargs['device']): + node.update_kwarg('device', torch.device('cuda', target_device)) + needs_recompile = True + + if node.op == 'call_function' and 'device' in node.kwargs: + kdev = node.kwargs['device'] + is_already_handled = node.target in factory_functions or ( + hasattr(node.target, '__name__') and node.target.__name__ in ['empty', 'zeros', 'ones', 'full'] + ) + if not is_already_handled and _is_cpu_device(kdev): + node.update_kwarg('device', torch.device('cuda', target_device)) + needs_recompile = True + # Fix ALL nodes with CPU example_values — not just get_attr/placeholder. # model_cpu_offload traces with CPU FakeTensors; Inductor autotuning # creates benchmark tensors on the example_value device, and the diff --git a/tests/feature_tests/test_fix_to_cpu_in_graph.py b/tests/feature_tests/test_fix_to_cpu_in_graph.py new file mode 100644 index 0000000..2a4af51 --- /dev/null +++ b/tests/feature_tests/test_fix_to_cpu_in_graph.py @@ -0,0 +1,132 @@ +"""Test: _fix_graph_device_placement rewrites .to(device('cpu')) in FX graphs. + +Root cause (commit e0c7277): _deep_cuda was removed, so Dynamo traces with CPU +tensors. .to(x.device) gets specialised to .to(device('cpu')) as a literal +constant in the FX graph. _fix_graph_device_placement must rewrite these nodes +to .to(device('cuda')) alongside the existing example_value metadata fix. +""" + +import pytest +import torch +import torch.nn as nn +import torch.fx as fx + + +def _build_graph_with_to_cpu(): + """Build an FX graph that mirrors the offload-traced pattern: + + mapping_on_cpu = mapping.to(device('cpu')) + out = x.index_select(0, mapping_on_cpu) + """ + graph = fx.Graph() + x = graph.placeholder("x") + mapping = graph.placeholder("mapping") + to_node = graph.call_method("to", args=(mapping, torch.device("cpu"))) + idx_node = graph.call_method("index_select", args=(x, 0, to_node)) + graph.output(idx_node) + + gm = fx.GraphModule(nn.Module(), graph) + + x.meta["example_value"] = torch.randn(8, 32, dtype=torch.bfloat16) + mapping.meta["example_value"] = torch.randperm(8) + to_node.meta["example_value"] = torch.randperm(8) + idx_node.meta["example_value"] = torch.randn(8, 32, dtype=torch.bfloat16) + + return gm + + +def _is_cpu_device(val): + if isinstance(val, torch.device): + return val.type == "cpu" + if isinstance(val, str): + return val == "cpu" + return False + + +def _apply_metadata_only_fix(gm, target_device=0): + for node in gm.graph.nodes: + ev = node.meta.get("example_value") + if ev is not None and hasattr(ev, "device") and str(ev.device) == "cpu": + node.meta["example_value"] = ev.to(target_device) + gm.recompile() + + +def _apply_full_fix(gm, target_device=0): + for node in gm.graph.nodes: + if node.op == "call_method" and node.target == "to": + new_args = list(node.args) + changed = False + for i, arg in enumerate(new_args): + if _is_cpu_device(arg): + new_args[i] = torch.device("cuda", target_device) + changed = True + if changed: + node.args = tuple(new_args) + if "device" in node.kwargs and _is_cpu_device(node.kwargs["device"]): + node.update_kwarg("device", torch.device("cuda", target_device)) + _apply_metadata_only_fix(gm, target_device) + + +def _run_with_fake_tensors(gm, cuda_device=0): + from torch._subclasses.fake_tensor import FakeTensorMode + + x_real = torch.randn(8, 32, dtype=torch.bfloat16, device=f"cuda:{cuda_device}") + mapping_real = torch.randperm(8, device=f"cuda:{cuda_device}") + + with FakeTensorMode() as fm: + x_fake = fm.from_tensor(x_real) + mapping_fake = fm.from_tensor(mapping_real) + with torch.no_grad(): + return gm(x_fake, mapping_fake) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestFixToCpuInGraph: + + def test_graph_has_to_cpu_node(self): + gm = _build_graph_with_to_cpu() + found = any( + isinstance(arg, torch.device) and arg.type == "cpu" + for node in gm.graph.nodes + if node.op == "call_method" and node.target == "to" + for arg in node.args + ) + assert found, "Graph should contain .to(device('cpu'))" + + def test_metadata_only_fix_fails(self): + """Fixing only example_values leaves .to(cpu) -> device mismatch.""" + gm = _build_graph_with_to_cpu() + _apply_metadata_only_fix(gm) + with pytest.raises(RuntimeError, match=r"[Dd]evice|FakeTensor"): + _run_with_fake_tensors(gm) + + def test_full_fix_succeeds(self): + """Rewriting .to(cpu) -> .to(cuda) + metadata fix -> no error.""" + gm = _build_graph_with_to_cpu() + _apply_full_fix(gm) + out = _run_with_fake_tensors(gm) + assert str(out.device).startswith("cuda") + + def test_no_residual_to_cpu(self): + gm = _build_graph_with_to_cpu() + _apply_full_fix(gm) + for node in gm.graph.nodes: + if node.op == "call_method" and node.target == "to": + for arg in node.args: + assert not _is_cpu_device(arg), f"Residual .to(cpu): {node}" + + def test_to_dtype_untouched(self): + """Rewrite must NOT affect .to(dtype) calls.""" + graph = fx.Graph() + x = graph.placeholder("x") + to_bf16 = graph.call_method("to", args=(x, torch.bfloat16)) + graph.output(to_bf16) + gm = fx.GraphModule(nn.Module(), graph) + x.meta["example_value"] = torch.randn(4, 8) + to_bf16.meta["example_value"] = torch.randn(4, 8, dtype=torch.bfloat16) + + _apply_full_fix(gm) + + for node in gm.graph.nodes: + if node.op == "call_method" and node.target == "to": + assert node.args[1] is torch.bfloat16 From 0b4f6bd0285ce418c9049bae331b98e30412e1d8 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Sun, 30 Aug 2026 02:34:22 +0800 Subject: [PATCH 18/54] style: auto-format test_fix_to_cpu_in_graph.py --- .../feature_tests/test_fix_to_cpu_in_graph.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/tests/feature_tests/test_fix_to_cpu_in_graph.py b/tests/feature_tests/test_fix_to_cpu_in_graph.py index 2a4af51..326d7ed 100644 --- a/tests/feature_tests/test_fix_to_cpu_in_graph.py +++ b/tests/feature_tests/test_fix_to_cpu_in_graph.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Test: _fix_graph_device_placement rewrites .to(device('cpu')) in FX graphs. Root cause (commit e0c7277): _deep_cuda was removed, so Dynamo traces with CPU @@ -8,15 +22,15 @@ import pytest import torch -import torch.nn as nn import torch.fx as fx +import torch.nn as nn def _build_graph_with_to_cpu(): """Build an FX graph that mirrors the offload-traced pattern: - mapping_on_cpu = mapping.to(device('cpu')) - out = x.index_select(0, mapping_on_cpu) + mapping_on_cpu = mapping.to(device('cpu')) + out = x.index_select(0, mapping_on_cpu) """ graph = fx.Graph() x = graph.placeholder("x") @@ -82,7 +96,6 @@ def _run_with_fake_tensors(gm, cuda_device=0): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") class TestFixToCpuInGraph: - def test_graph_has_to_cpu_node(self): gm = _build_graph_with_to_cpu() found = any( From f87f3cf57c64d521b5c7d6dc9bc1498333610eb3 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 14:06:20 +0800 Subject: [PATCH 19/54] refactor: simplify _fix_graph_device_placement - Remove redundant factory_functions block; the unified call_function handler with _is_cpu_device already covers all device kwargs - Extract _move_to_device helper to deduplicate CPU->CUDA tensor logic - Simplify to single return value with identity check (is not) - Remove redundant rank-0 guard around magi_logger.info (already default) --- magi_compiler/magi_backend/magi_backend.py | 65 +++++++--------------- 1 file changed, 19 insertions(+), 46 deletions(-) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 22bdd54..ffccb33 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -288,31 +288,6 @@ def _fix_graph_device_placement(self, module: torch.nn.Module): needs_recompile = False target_device = torch.cuda.current_device() - factory_functions = [ - torch.empty, - torch.zeros, - torch.ones, - torch.full, - torch.rand, - torch.randn, - torch.arange, - torch.tensor, - torch.ops.aten.empty.memory_format, - ] - - for node in module.graph.nodes: - if node.op == 'call_function': - is_factory = node.target in factory_functions or ( - hasattr(node.target, '__name__') and node.target.__name__ in ['empty', 'zeros', 'ones', 'full'] - ) - - if is_factory: - if 'device' in node.kwargs: - current_dev = node.kwargs['device'] - if str(current_dev) == 'cpu' or current_dev == torch.device('cpu'): - node.update_kwarg('device', target_device) - needs_recompile = True - # Fix .to(device('cpu')) calls baked during CPU-offload tracing. # Without _deep_cuda, Dynamo traces with CPU tensors and specialises # .to(x.device) as .to(device('cpu')). After we move example_values @@ -341,11 +316,7 @@ def _is_cpu_device(val): needs_recompile = True if node.op == 'call_function' and 'device' in node.kwargs: - kdev = node.kwargs['device'] - is_already_handled = node.target in factory_functions or ( - hasattr(node.target, '__name__') and node.target.__name__ in ['empty', 'zeros', 'ones', 'full'] - ) - if not is_already_handled and _is_cpu_device(kdev): + if _is_cpu_device(node.kwargs['device']): node.update_kwarg('device', torch.device('cuda', target_device)) needs_recompile = True @@ -353,37 +324,39 @@ def _is_cpu_device(val): # model_cpu_offload traces with CPU FakeTensors; Inductor autotuning # creates benchmark tensors on the example_value device, and the # codegen may pick CPU-specific paths if it sees CPU metadata. + def _move_to_device(val): + if hasattr(val, 'device') and str(val.device) == 'cpu': + new_val = val.to(target_device) + if isinstance(val, torch.nn.Parameter): + new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad) + return new_val + return val + cpu_fix_count = 0 for node in module.graph.nodes: ev = node.meta.get('example_value') if ev is not None: - if hasattr(ev, 'device') and str(ev.device) == 'cpu': - new_ev = ev.to(target_device) - if isinstance(ev, torch.nn.Parameter): - new_ev = torch.nn.Parameter(new_ev, requires_grad=ev.requires_grad) - node.meta['example_value'] = new_ev - needs_recompile = True - cpu_fix_count += 1 + if hasattr(ev, 'device'): + new_ev = _move_to_device(ev) + if new_ev is not ev: + node.meta['example_value'] = new_ev + needs_recompile = True + cpu_fix_count += 1 elif isinstance(ev, (list, tuple)): fixed_list = [] any_fixed = False for item in ev: - if hasattr(item, 'device') and str(item.device) == 'cpu': - new_item = item.to(target_device) - if isinstance(item, torch.nn.Parameter): - new_item = torch.nn.Parameter(new_item, requires_grad=item.requires_grad) - fixed_list.append(new_item) + new_item = _move_to_device(item) + fixed_list.append(new_item) + if new_item is not item: any_fixed = True cpu_fix_count += 1 - else: - fixed_list.append(item) if any_fixed: node.meta['example_value'] = type(ev)(fixed_list) needs_recompile = True if needs_recompile: - if os.environ.get('RANK', '0') == '0': - magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device) + magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device) module.recompile() @observe_lifecycle("piecewise_compile") From 6551e0f1daca24fb37398578134edb26cd228cfe Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 14:08:17 +0800 Subject: [PATCH 20/54] style: reorder call_function before call_method in _fix_graph_device_placement --- magi_compiler/magi_backend/magi_backend.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index ffccb33..8b519eb 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -301,6 +301,11 @@ def _is_cpu_device(val): return False for node in module.graph.nodes: + if node.op == 'call_function' and 'device' in node.kwargs: + if _is_cpu_device(node.kwargs['device']): + node.update_kwarg('device', torch.device('cuda', target_device)) + needs_recompile = True + if node.op == 'call_method' and node.target == 'to': new_args = list(node.args) changed = False @@ -315,11 +320,6 @@ def _is_cpu_device(val): node.update_kwarg('device', torch.device('cuda', target_device)) needs_recompile = True - if node.op == 'call_function' and 'device' in node.kwargs: - if _is_cpu_device(node.kwargs['device']): - node.update_kwarg('device', torch.device('cuda', target_device)) - needs_recompile = True - # Fix ALL nodes with CPU example_values — not just get_attr/placeholder. # model_cpu_offload traces with CPU FakeTensors; Inductor autotuning # creates benchmark tensors on the example_value device, and the From aeb7edcff13a251150a6f55e928cd668b66a12f0 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 15:44:26 +0800 Subject: [PATCH 21/54] perf: streaming shm materialize to halve CPU memory peak Replace batch copy-all-then-replace pattern with streaming copy-and-replace per parameter. This keeps peak RSS near 1x model size instead of 2x during _materialize_shm_weights. Changes: - Add _assign_param, _stream_copy_and_replace, _create_empty_shm - Rewrite _materialize_shm_weights to use streaming - Remove per-rank staggering loop (no longer needed at 1x peak) - Release full_state_dict before materialize in caller - Add test_shm_memory_peak.py with 4 tests (peak, streaming, correctness, speed) --- magi_compiler/_api.py | 81 +++-- tests/feature_tests/test_shm_memory_peak.py | 332 ++++++++++++++++++++ 2 files changed, 395 insertions(+), 18 deletions(-) create mode 100644 tests/feature_tests/test_shm_memory_peak.py diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 0c8577a..66c8f40 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -530,6 +530,18 @@ def _split_flat_to_params(flat: torch.Tensor, param_list: list[tuple[str, torch. return out +def _assign_param(module: nn.Module, dotted_name: str, new_tensor: torch.Tensor) -> None: + """Replace a single parameter/buffer in *module* by its dotted path.""" + parts = dotted_name.rsplit(".", 1) + parent = module.get_submodule(parts[0]) if len(parts) == 2 else module + attr = parts[-1] + old = getattr(parent, attr) + if isinstance(old, nn.Parameter): + parent.register_parameter(attr, nn.Parameter(new_tensor, requires_grad=new_tensor.requires_grad)) + else: + setattr(parent, attr, new_tensor) + + def _create_shm_tensor(shm_path: str, param_list: list[tuple[str, torch.Tensor]], dtype: torch.dtype) -> torch.Tensor: """Create a shared-memory mmap file, pack *param_list* into it, return the giant tensor.""" total_numel = sum(t.numel() for _, t in param_list) @@ -541,49 +553,81 @@ def _create_shm_tensor(shm_path: str, param_list: list[tuple[str, torch.Tensor]] return giant +def _stream_copy_and_replace( + module: nn.Module, + giant: torch.Tensor, + param_list: list[tuple[str, torch.Tensor]], +) -> None: + """Copy each param into *giant*, replace in module immediately. + + By replacing before moving to the next param, only one param's worth + of duplication exists at any moment (peak ≈ 1× instead of 2×). + """ + offset = 0 + for i, (name, tensor) in enumerate(param_list): + numel = tensor.numel() + giant[offset : offset + numel].copy_(tensor.view(-1)) + view = giant[offset : offset + numel].view(tensor.shape) + if tensor.requires_grad: + view.requires_grad_(True) + _assign_param(module, name, view) + param_list[i] = (name, view) + offset += numel + + +def _create_empty_shm(shm_path: str, total_numel: int, dtype: torch.dtype) -> torch.Tensor: + """Create an empty mmap file and return the mapped tensor.""" + elem_size = torch.empty(0, dtype=dtype).element_size() + with open(shm_path, "wb") as f: + f.truncate(total_numel * elem_size) + return torch.from_file(shm_path, shared=True, size=total_numel, dtype=dtype, device="cpu") + + def _materialize_shm_weights( module: nn.Module, grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], local_rank: int, per_rank: bool ) -> None: """Replace module params with pinned shared-memory tensors. - per_rank=True (EP > 1): each rank writes its own mmap, staggered. + Uses streaming copy-and-replace so only one parameter is duplicated + at a time, keeping peak RSS near 1× model size instead of 2×. + + per_rank=True (EP > 1): each rank writes its own mmap concurrently. per_rank=False (EP <= 1): rank 0 writes, all ranks share pages. """ cls_name = module.__class__.__name__ - shared_state: dict[str, torch.Tensor] = {} buffers: list[torch.Tensor] = [] if per_rank: - world_size = dist.get_world_size() - for turn in range(world_size): - if local_rank == turn: - for dtype, param_list in grouped_params.items(): - path = _shm_path(cls_name, dtype, rank=local_rank) - giant = _create_shm_tensor(path, param_list, dtype) - pin_memory_in_place(giant) - buffers.append(giant) - shared_state.update(_split_flat_to_params(giant, param_list)) - if os.path.exists(path): - os.remove(path) - dist.barrier() + for dtype, param_list in grouped_params.items(): + path = _shm_path(cls_name, dtype, rank=local_rank) + total_numel = sum(t.numel() for _, t in param_list) + giant = _create_empty_shm(path, total_numel, dtype) + _stream_copy_and_replace(module, giant, param_list) + pin_memory_in_place(giant) + buffers.append(giant) + if os.path.exists(path): + os.remove(path) + dist.barrier() else: dist.barrier() for dtype, param_list in grouped_params.items(): path = _shm_path(cls_name, dtype) total_numel = sum(t.numel() for _, t in param_list) if local_rank == 0: - _create_shm_tensor(path, param_list, dtype) + giant = _create_empty_shm(path, total_numel, dtype) + _stream_copy_and_replace(module, giant, param_list) dist.barrier() - giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") + if local_rank != 0: + giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") + _stream_copy_and_replace(module, giant, param_list) pin_memory_in_place(giant) buffers.append(giant) - shared_state.update(_split_flat_to_params(giant, param_list)) dist.barrier() if local_rank == 0 and os.path.exists(path): os.remove(path) module._magi_giant_buffers = buffers - module.load_state_dict(shared_state, assign=True) + gc.collect() def _staggered_pin_memory(full_state_dict: dict[str, torch.Tensor], local_rank: int, pin_budget_gb: float): @@ -752,6 +796,7 @@ def _force_cpu(t): local_rank, ) else: + full_state_dict = None _materialize_shm_weights(self, grouped_params, local_rank, per_rank=(ep_size > 1)) del full_state_dict, grouped_params diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py new file mode 100644 index 0000000..49254c3 --- /dev/null +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Memory-peak tests for shared-memory weight materialization. + +Background +---------- +``_materialize_shm_weights`` copies ALL parameters into an mmap file, then +calls ``load_state_dict(assign=True)`` to replace them. During the copy the +original params (RssAnon) and the new mmap pages (RssFile) coexist, pushing +the process to ~2× model size. + +The streaming alternative replaces each parameter immediately after copying, +so only one parameter's worth of duplication exists at any moment. + +Measurement +----------- +Each test runs in a **subprocess** (clean VmHWM baseline). +VmHWM (RSS high-water mark from ``/proc/self/status``) tracks the growth +contributed by the mmap copy: + +- **batch**: VmHWM growth ≈ model_size (old params + mmap ≈ 2×) +- **streaming**: VmHWM growth ≈ model_size / num_params (≪ 0.3×) + +No distributed / CUDA required. +""" + +import gc +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.nn as nn + +_IS_LINUX = os.path.exists("/proc/self/status") +_skip_no_procfs = pytest.mark.skipif( + not _IS_LINUX, reason="requires /proc/self/status for VmHWM" +) + +PARAM_MB = 512 +NUM_PARAMS = 4 + + +def _read_vm(key: str = "VmHWM") -> float: + """Read a VmXxx field from /proc/self/status (MB).""" + with open("/proc/self/status") as f: + for line in f: + if line.startswith(key + ":"): + return int(line.split()[1]) / 1024 + raise RuntimeError(f"{key} not found") + + +class HeavyModule(nn.Module): + def __init__(self, numel_per_param: int, num_params: int = NUM_PARAMS, + dtype: torch.dtype = torch.bfloat16): + super().__init__() + for i in range(num_params): + self.register_parameter( + f"w{i}", nn.Parameter(torch.randn(numel_per_param, dtype=dtype)) + ) + + def forward(self, x): + return x + + +# ── batch (current code pattern) ──────────────────────────── + +def _batch_materialize(module: nn.Module, shm_dir: str) -> None: + full_state_dict = module.state_dict() + grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} + for name, tensor in full_state_dict.items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + + shared_state: dict[str, torch.Tensor] = {} + buffers: list[torch.Tensor] = [] + + for dtype, param_list in grouped.items(): + total_numel = sum(t.numel() for _, t in param_list) + elem_size = torch.empty(0, dtype=dtype).element_size() + path = os.path.join(shm_dir, f"batch_{dtype}.bin") + with open(path, "wb") as f: + f.truncate(total_numel * elem_size) + giant = torch.from_file( + path, shared=True, size=total_numel, dtype=dtype, device="cpu" + ) + offset = 0 + for _, tensor in param_list: + n = tensor.numel() + giant[offset : offset + n].copy_(tensor.view(-1)) + offset += n + offset = 0 + for name, orig in param_list: + n = orig.numel() + view = giant[offset : offset + n].view(orig.shape) + if orig.requires_grad: + view.requires_grad_(True) + shared_state[name] = view + offset += n + buffers.append(giant) + if os.path.exists(path): + os.remove(path) + + module.load_state_dict(shared_state, assign=True) + module._buffers_ref = buffers + + +# ── streaming (fix) ───────────────────────────────────────── + +def _assign_param(module: nn.Module, dotted_name: str, + new_tensor: torch.Tensor) -> None: + parts = dotted_name.rsplit(".", 1) + parent = module.get_submodule(parts[0]) if len(parts) == 2 else module + attr = parts[-1] + old = getattr(parent, attr) + if isinstance(old, nn.Parameter): + parent.register_parameter( + attr, nn.Parameter(new_tensor, requires_grad=new_tensor.requires_grad) + ) + else: + setattr(parent, attr, new_tensor) + + +def _streaming_materialize(module: nn.Module, shm_dir: str) -> None: + full_state_dict = module.state_dict() + grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} + for name, tensor in full_state_dict.items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + del full_state_dict + gc.collect() + + buffers: list[torch.Tensor] = [] + for dtype, param_list in grouped.items(): + total_numel = sum(t.numel() for _, t in param_list) + elem_size = torch.empty(0, dtype=dtype).element_size() + path = os.path.join(shm_dir, f"stream_{dtype}.bin") + with open(path, "wb") as f: + f.truncate(total_numel * elem_size) + giant = torch.from_file( + path, shared=True, size=total_numel, dtype=dtype, device="cpu" + ) + offset = 0 + for i, (name, tensor) in enumerate(param_list): + n = tensor.numel() + giant[offset : offset + n].copy_(tensor.view(-1)) + view = giant[offset : offset + n].view(tensor.shape) + if tensor.requires_grad: + view.requires_grad_(True) + _assign_param(module, name, view) + param_list[i] = (name, view) + offset += n + buffers.append(giant) + if os.path.exists(path): + os.remove(path) + + module._buffers_ref = buffers + gc.collect() + + +# ── subprocess workers ────────────────────────────────────── + +def _worker(result_dict, param_mb, materialize_fn): + elem_bytes = 2 # bf16 + numel_per = param_mb * 1024 * 1024 // (elem_bytes * NUM_PARAMS) + model = HeavyModule(numel_per, NUM_PARAMS) + gc.collect() + + hwm_before = _read_vm("VmHWM") + with tempfile.TemporaryDirectory() as d: + materialize_fn(model, d) + gc.collect() + hwm_after = _read_vm("VmHWM") + + result_dict["hwm_before"] = hwm_before + result_dict["hwm_after"] = hwm_after + result_dict["param_mb"] = param_mb + + +def _run_in_subprocess(materialize_fn, param_mb): + ctx = mp.get_context("fork") + mgr = ctx.Manager() + result = mgr.dict() + p = ctx.Process(target=_worker, args=(result, param_mb, materialize_fn)) + p.start() + p.join(timeout=120) + assert p.exitcode == 0, f"subprocess exited with code {p.exitcode}" + return dict(result) + + +# ── tests ─────────────────────────────────────────────────── + +@_skip_no_procfs +def test_batch_materialize_has_high_peak(): + """BUG REPRO: batch materialize adds ~1× model size as mmap overhead. + + At peak: old params (RssAnon) + mmap copy (RssFile) ≈ 2× model size. + VmHWM growth measures the mmap portion, expected > 0.5× model_size. + """ + r = _run_in_subprocess(_batch_materialize, PARAM_MB) + growth = r["hwm_after"] - r["hwm_before"] + pm = r["param_mb"] + + print(f"\n[batch] hwm_before={r['hwm_before']:.0f} MB, " + f"hwm_after={r['hwm_after']:.0f} MB, " + f"growth={growth:.0f} MB, model_size={pm} MB, " + f"ratio={growth / pm:.2f}x") + + assert growth > pm * 0.5, ( + f"Expected mmap overhead > {pm * 0.5:.0f} MB (0.5× model) " + f"but got {growth:.0f} MB ({growth / pm:.2f}×). " + f"The 2× peak may have been optimized away." + ) + + +@_skip_no_procfs +def test_streaming_materialize_low_peak(): + """FIX VERIFIED: streaming avoids the mmap overhead peak. + + By replacing each param immediately, only ~1/N of the model is ever + duplicated. VmHWM growth should be well under 0.3× model size. + """ + r = _run_in_subprocess(_streaming_materialize, PARAM_MB) + growth = r["hwm_after"] - r["hwm_before"] + pm = r["param_mb"] + + print(f"\n[streaming] hwm_before={r['hwm_before']:.0f} MB, " + f"hwm_after={r['hwm_after']:.0f} MB, " + f"growth={growth:.0f} MB, model_size={pm} MB, " + f"ratio={growth / pm:.2f}x") + + assert growth < pm * 0.3, ( + f"Expected mmap overhead < {pm * 0.3:.0f} MB (0.3× model) " + f"but got {growth:.0f} MB ({growth / pm:.2f}×). " + f"Streaming fix did not reduce peak." + ) + + +@_skip_no_procfs +def test_streaming_preserves_weights(): + """Streaming must produce identical weights to batch.""" + numel = 1024 + torch.manual_seed(42) + model_a = HeavyModule(numel, NUM_PARAMS) + torch.manual_seed(42) + model_b = HeavyModule(numel, NUM_PARAMS) + + with tempfile.TemporaryDirectory() as d1, tempfile.TemporaryDirectory() as d2: + _batch_materialize(model_a, d1) + _streaming_materialize(model_b, d2) + + for name in model_a.state_dict(): + assert torch.equal( + model_a.state_dict()[name], model_b.state_dict()[name] + ), f"Mismatch on '{name}'" + + +# ── speed ─────────────────────────────────────────────────── + +import time + + +def _speed_worker(result_dict, param_mb, num_params, materialize_fn, repeats): + elem_bytes = 2 + numel_per = param_mb * 1024 * 1024 // (elem_bytes * num_params) + times = [] + for _ in range(repeats): + torch.manual_seed(42) + model = HeavyModule(numel_per, num_params) + gc.collect() + with tempfile.TemporaryDirectory() as d: + t0 = time.perf_counter() + materialize_fn(model, d) + times.append(time.perf_counter() - t0) + del model + gc.collect() + result_dict["times"] = times + result_dict["avg"] = sum(times) / len(times) + + +def _run_speed_subprocess(materialize_fn, param_mb, num_params=NUM_PARAMS, + repeats=3): + ctx = mp.get_context("fork") + mgr = ctx.Manager() + result = mgr.dict() + p = ctx.Process( + target=_speed_worker, + args=(result, param_mb, num_params, materialize_fn, repeats), + ) + p.start() + p.join(timeout=300) + assert p.exitcode == 0, f"subprocess exited with code {p.exitcode}" + return dict(result) + + +@_skip_no_procfs +def test_streaming_not_slower_than_batch(): + """Streaming must not be significantly slower than batch. + + Allows up to 1.20x slowdown to account for per-param register_parameter + overhead. In practice streaming is often faster on large models because + it avoids the final load_state_dict bulk copy. + """ + mb = PARAM_MB + max_slowdown = 1.20 + + r_batch = _run_speed_subprocess(_batch_materialize, mb) + r_stream = _run_speed_subprocess(_streaming_materialize, mb) + + ratio = r_stream["avg"] / r_batch["avg"] + print(f"\n[speed] model={mb} MB, num_params={NUM_PARAMS}" + f" batch={r_batch['avg']:.3f}s" + f" streaming={r_stream['avg']:.3f}s" + f" ratio={ratio:.2f}x") + + assert ratio < max_slowdown, ( + f"Streaming is {ratio:.2f}x slower than batch " + f"(limit {max_slowdown}x). " + f"batch={r_batch['times']}, stream={r_stream['times']}" + ) From bb695f6f487284c00452d93215920ebe16931f2a Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 16:52:55 +0800 Subject: [PATCH 22/54] refactor: remove SKIP_SHM/PIN_BUDGET, keep only SHM materialize path - Delete _staggered_pin_memory and its env vars (MAGI_OFFLOAD_SKIP_SHM, MAGI_OFFLOAD_PIN_BUDGET_GB, MAGI_OFFLOAD_PIN_CONCURRENCY) - Always use _materialize_shm_weights for CPU offload weight management - Add host_memory.py utility for standardized CPU memory peak tracking - Instrument key offload milestones with fmt_host_mem() logging --- magi_compiler/_api.py | 108 ++--------------------------- magi_compiler/utils/host_memory.py | 62 +++++++++++++++++ 2 files changed, 68 insertions(+), 102 deletions(-) create mode 100644 magi_compiler/utils/host_memory.py diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 66c8f40..d5cdf74 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -29,6 +29,7 @@ from magi_compiler.cuda.cudart import pin_memory_in_place from magi_compiler.magi_backend.magi_compiler_base import MagiCompileState from magi_compiler.utils import compilation_counter, envs, magi_logger +from magi_compiler.utils.host_memory import fmt_host_mem from magi_compiler.utils.compile_time_monitor import CompileMonitor from .config import CompileConfig, CompileMode @@ -630,94 +631,6 @@ def _materialize_shm_weights( gc.collect() -def _staggered_pin_memory(full_state_dict: dict[str, torch.Tensor], local_rank: int, pin_budget_gb: float): - """Pin CPU tensors in coordinated waves to avoid host OOM. - - Automatically determines how many ranks can pin concurrently based on - host RAM and per-rank parameter size. Override with env var - ``MAGI_OFFLOAD_PIN_CONCURRENCY``. - """ - import resource - import time - - world_size = dist.get_world_size() if dist.is_initialized() else 1 - soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) - - per_rank_bytes = sum(t.numel() * t.element_size() for t in full_state_dict.values() if t.device.type == "cpu") - per_rank_gb = per_rank_bytes / (1024**3) - - concurrency_env = os.environ.get("MAGI_OFFLOAD_PIN_CONCURRENCY", "") - if concurrency_env: - pin_concurrency = int(concurrency_env) - else: - try: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemTotal:"): - total_ram_kb = int(line.split()[1]) - break - else: - total_ram_kb = 512 * 1024 * 1024 - total_ram_gb = total_ram_kb / (1024 * 1024) - except Exception: - total_ram_gb = 512.0 - safe_ram_gb = total_ram_gb / 2 - pin_concurrency = max(1, int(safe_ram_gb / per_rank_gb)) if per_rank_gb > 0 else world_size - pin_concurrency = min(pin_concurrency, world_size) - - num_waves = (world_size + pin_concurrency - 1) // pin_concurrency - magi_logger.info( - "[Rank %d] pin_memory: world=%d per_rank=%.2f GB, " "concurrency=%d (waves=%d), RLIMIT_MEMLOCK soft=%s hard=%s", - local_rank, - world_size, - per_rank_gb, - pin_concurrency, - num_waves, - "unlimited" if soft == resource.RLIM_INFINITY else f"{soft / (1024 ** 3):.1f}GB", - "unlimited" if hard == resource.RLIM_INFINITY else f"{hard / (1024 ** 3):.1f}GB", - ) - - limit_bytes = int(pin_budget_gb * (1024**3)) - my_wave = local_rank // pin_concurrency - for wave in range(num_waves): - if wave == my_wave: - t0 = time.perf_counter() - params = [ - (name, tensor, tensor.numel() * tensor.element_size()) - for name, tensor in full_state_dict.items() - if tensor.device.type == "cpu" - ] - params.sort(key=lambda x: x[2], reverse=True) - - pinned_bytes = 0 - pinned_count = 0 - for name, tensor, size in params: - if pinned_bytes + size > limit_bytes: - continue - try: - pin_memory_in_place(tensor) - pinned_bytes += size - pinned_count += 1 - except RuntimeError as e: - magi_logger.warning("[Rank %d] pin failed at %.2f GB: %s", local_rank, pinned_bytes / (1024**3), e) - break - - total_bytes = sum(s for _, _, s in params) - magi_logger.info( - "[Rank %d] pin_memory DONE (wave %d/%d) %.1fs: " "%d params (%.2f / %.2f GB, budget=%.1f GB)", - local_rank, - wave + 1, - num_waves, - time.perf_counter() - t0, - pinned_count, - pinned_bytes / (1024**3), - total_bytes / (1024**3), - pin_budget_gb, - ) - if dist.is_initialized(): - dist.barrier() - - def _patch_cpu_offload_apply(cls: type[nn.Module]): magi_logger.info(f"Enabling CPU offload for {cls}") _orig_apply = cls._apply @@ -769,6 +682,7 @@ def _force_cpu(t): return fn(t).cpu() _orig_apply(self, _force_cpu) + magi_logger.info('[offload] after _force_cpu: %s', fmt_host_mem()) # create shared memory tensors for all parameters/buffers on CPU ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", os.environ.get("EP_SIZE", "1"))) @@ -784,23 +698,13 @@ def _force_cpu(t): grouped_params[dt] = [] grouped_params[dt].append((name, tensor)) - skip_shm = os.environ.get("MAGI_OFFLOAD_SKIP_SHM", "0") == "1" - if skip_shm: - self._magi_giant_buffers = [] - pin_budget_gb = float(os.environ.get("MAGI_OFFLOAD_PIN_BUDGET_GB", "0")) - if pin_budget_gb > 0: - _staggered_pin_memory(full_state_dict, local_rank, pin_budget_gb) - else: - magi_logger.info( - "[Rank %d] MAGI_OFFLOAD_SKIP_SHM=1, PIN_BUDGET=0: " "params remain as unpinned CPU tensors.", - local_rank, - ) - else: - full_state_dict = None - _materialize_shm_weights(self, grouped_params, local_rank, per_rank=(ep_size > 1)) + full_state_dict = None + _materialize_shm_weights(self, grouped_params, local_rank, per_rank=(ep_size > 1)) + magi_logger.info('[offload] after SHM materialize: %s', fmt_host_mem()) del full_state_dict, grouped_params gc.collect() + magi_logger.info('[offload] after gc.collect: %s', fmt_host_mem()) else: diff --git a/magi_compiler/utils/host_memory.py b/magi_compiler/utils/host_memory.py new file mode 100644 index 0000000..29ca21d --- /dev/null +++ b/magi_compiler/utils/host_memory.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +"""Lightweight host (CPU) memory introspection for Linux.""" + +from __future__ import annotations + +import os + + +def get_host_mem_gb() -> dict[str, float]: + """Read /proc/self/status and return key memory metrics in GiB. + + Returns a dict with keys: + vm_peak – peak virtual memory (VmPeak) + vm_rss – current resident set (VmRSS) + vm_hwm – high-water mark of RSS (VmHWM) + rss_anon – anonymous (heap/stack) resident pages (RssAnon) + rss_file – file-backed (mmap) resident pages (RssFile) + rss_shmem – shared-memory resident pages (RssShmem) + Missing fields default to 0. + """ + fields = { + "VmPeak": "vm_peak", + "VmRSS": "vm_rss", + "VmHWM": "vm_hwm", + "RssAnon": "rss_anon", + "RssFile": "rss_file", + "RssShmem": "rss_shmem", + } + result: dict[str, float] = {v: 0.0 for v in fields.values()} + try: + with open("/proc/self/status") as f: + for line in f: + key = line.split(":")[0] + if key in fields: + kb = int(line.split()[1]) + result[fields[key]] = kb / (1024 * 1024) + except (OSError, ValueError): + pass + return result + + +def get_total_ram_gb() -> float: + """Return total host RAM in GiB from /proc/meminfo.""" + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal:"): + return int(line.split()[1]) / (1024 * 1024) + except (OSError, ValueError): + pass + return 0.0 + + +def fmt_host_mem(mem: dict[str, float] | None = None) -> str: + """One-line human-readable summary of current host memory.""" + if mem is None: + mem = get_host_mem_gb() + return ( + f"VmHWM={mem['vm_hwm']:.1f}G " + f"VmRSS={mem['vm_rss']:.1f}G " + f"(anon={mem['rss_anon']:.1f}G file={mem['rss_file']:.1f}G shm={mem['rss_shmem']:.1f}G)" + ) From 5fcc59f613a430382a80b5e2dfb3dbd990ceda2c Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 17:03:12 +0800 Subject: [PATCH 23/54] style: apply pre-commit formatting --- magi_compiler/_api.py | 8 +-- magi_compiler/magi_backend/magi_backend.py | 1 - magi_compiler/utils/host_memory.py | 1 - tests/feature_tests/test_shm_memory_peak.py | 72 +++++++++------------ 4 files changed, 34 insertions(+), 48 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index d5cdf74..5b53ed0 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -29,8 +29,8 @@ from magi_compiler.cuda.cudart import pin_memory_in_place from magi_compiler.magi_backend.magi_compiler_base import MagiCompileState from magi_compiler.utils import compilation_counter, envs, magi_logger -from magi_compiler.utils.host_memory import fmt_host_mem from magi_compiler.utils.compile_time_monitor import CompileMonitor +from magi_compiler.utils.host_memory import fmt_host_mem from .config import CompileConfig, CompileMode @@ -554,11 +554,7 @@ def _create_shm_tensor(shm_path: str, param_list: list[tuple[str, torch.Tensor]] return giant -def _stream_copy_and_replace( - module: nn.Module, - giant: torch.Tensor, - param_list: list[tuple[str, torch.Tensor]], -) -> None: +def _stream_copy_and_replace(module: nn.Module, giant: torch.Tensor, param_list: list[tuple[str, torch.Tensor]]) -> None: """Copy each param into *giant*, replace in module immediately. By replacing before moving to the next param, only one param's worth diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 8b519eb..83ad49f 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -14,7 +14,6 @@ import ast import dataclasses -import os import pprint import time from collections.abc import Callable diff --git a/magi_compiler/utils/host_memory.py b/magi_compiler/utils/host_memory.py index 29ca21d..d61bf89 100644 --- a/magi_compiler/utils/host_memory.py +++ b/magi_compiler/utils/host_memory.py @@ -3,7 +3,6 @@ from __future__ import annotations -import os def get_host_mem_gb() -> dict[str, float]: diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index 49254c3..feb6467 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -47,9 +47,7 @@ import torch.nn as nn _IS_LINUX = os.path.exists("/proc/self/status") -_skip_no_procfs = pytest.mark.skipif( - not _IS_LINUX, reason="requires /proc/self/status for VmHWM" -) +_skip_no_procfs = pytest.mark.skipif(not _IS_LINUX, reason="requires /proc/self/status for VmHWM") PARAM_MB = 512 NUM_PARAMS = 4 @@ -65,13 +63,10 @@ def _read_vm(key: str = "VmHWM") -> float: class HeavyModule(nn.Module): - def __init__(self, numel_per_param: int, num_params: int = NUM_PARAMS, - dtype: torch.dtype = torch.bfloat16): + def __init__(self, numel_per_param: int, num_params: int = NUM_PARAMS, dtype: torch.dtype = torch.bfloat16): super().__init__() for i in range(num_params): - self.register_parameter( - f"w{i}", nn.Parameter(torch.randn(numel_per_param, dtype=dtype)) - ) + self.register_parameter(f"w{i}", nn.Parameter(torch.randn(numel_per_param, dtype=dtype))) def forward(self, x): return x @@ -79,6 +74,7 @@ def forward(self, x): # ── batch (current code pattern) ──────────────────────────── + def _batch_materialize(module: nn.Module, shm_dir: str) -> None: full_state_dict = module.state_dict() grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} @@ -94,9 +90,7 @@ def _batch_materialize(module: nn.Module, shm_dir: str) -> None: path = os.path.join(shm_dir, f"batch_{dtype}.bin") with open(path, "wb") as f: f.truncate(total_numel * elem_size) - giant = torch.from_file( - path, shared=True, size=total_numel, dtype=dtype, device="cpu" - ) + giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") offset = 0 for _, tensor in param_list: n = tensor.numel() @@ -120,16 +114,14 @@ def _batch_materialize(module: nn.Module, shm_dir: str) -> None: # ── streaming (fix) ───────────────────────────────────────── -def _assign_param(module: nn.Module, dotted_name: str, - new_tensor: torch.Tensor) -> None: + +def _assign_param(module: nn.Module, dotted_name: str, new_tensor: torch.Tensor) -> None: parts = dotted_name.rsplit(".", 1) parent = module.get_submodule(parts[0]) if len(parts) == 2 else module attr = parts[-1] old = getattr(parent, attr) if isinstance(old, nn.Parameter): - parent.register_parameter( - attr, nn.Parameter(new_tensor, requires_grad=new_tensor.requires_grad) - ) + parent.register_parameter(attr, nn.Parameter(new_tensor, requires_grad=new_tensor.requires_grad)) else: setattr(parent, attr, new_tensor) @@ -149,9 +141,7 @@ def _streaming_materialize(module: nn.Module, shm_dir: str) -> None: path = os.path.join(shm_dir, f"stream_{dtype}.bin") with open(path, "wb") as f: f.truncate(total_numel * elem_size) - giant = torch.from_file( - path, shared=True, size=total_numel, dtype=dtype, device="cpu" - ) + giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") offset = 0 for i, (name, tensor) in enumerate(param_list): n = tensor.numel() @@ -172,6 +162,7 @@ def _streaming_materialize(module: nn.Module, shm_dir: str) -> None: # ── subprocess workers ────────────────────────────────────── + def _worker(result_dict, param_mb, materialize_fn): elem_bytes = 2 # bf16 numel_per = param_mb * 1024 * 1024 // (elem_bytes * NUM_PARAMS) @@ -202,6 +193,7 @@ def _run_in_subprocess(materialize_fn, param_mb): # ── tests ─────────────────────────────────────────────────── + @_skip_no_procfs def test_batch_materialize_has_high_peak(): """BUG REPRO: batch materialize adds ~1× model size as mmap overhead. @@ -213,10 +205,12 @@ def test_batch_materialize_has_high_peak(): growth = r["hwm_after"] - r["hwm_before"] pm = r["param_mb"] - print(f"\n[batch] hwm_before={r['hwm_before']:.0f} MB, " - f"hwm_after={r['hwm_after']:.0f} MB, " - f"growth={growth:.0f} MB, model_size={pm} MB, " - f"ratio={growth / pm:.2f}x") + print( + f"\n[batch] hwm_before={r['hwm_before']:.0f} MB, " + f"hwm_after={r['hwm_after']:.0f} MB, " + f"growth={growth:.0f} MB, model_size={pm} MB, " + f"ratio={growth / pm:.2f}x" + ) assert growth > pm * 0.5, ( f"Expected mmap overhead > {pm * 0.5:.0f} MB (0.5× model) " @@ -236,10 +230,12 @@ def test_streaming_materialize_low_peak(): growth = r["hwm_after"] - r["hwm_before"] pm = r["param_mb"] - print(f"\n[streaming] hwm_before={r['hwm_before']:.0f} MB, " - f"hwm_after={r['hwm_after']:.0f} MB, " - f"growth={growth:.0f} MB, model_size={pm} MB, " - f"ratio={growth / pm:.2f}x") + print( + f"\n[streaming] hwm_before={r['hwm_before']:.0f} MB, " + f"hwm_after={r['hwm_after']:.0f} MB, " + f"growth={growth:.0f} MB, model_size={pm} MB, " + f"ratio={growth / pm:.2f}x" + ) assert growth < pm * 0.3, ( f"Expected mmap overhead < {pm * 0.3:.0f} MB (0.3× model) " @@ -262,9 +258,7 @@ def test_streaming_preserves_weights(): _streaming_materialize(model_b, d2) for name in model_a.state_dict(): - assert torch.equal( - model_a.state_dict()[name], model_b.state_dict()[name] - ), f"Mismatch on '{name}'" + assert torch.equal(model_a.state_dict()[name], model_b.state_dict()[name]), f"Mismatch on '{name}'" # ── speed ─────────────────────────────────────────────────── @@ -290,15 +284,11 @@ def _speed_worker(result_dict, param_mb, num_params, materialize_fn, repeats): result_dict["avg"] = sum(times) / len(times) -def _run_speed_subprocess(materialize_fn, param_mb, num_params=NUM_PARAMS, - repeats=3): +def _run_speed_subprocess(materialize_fn, param_mb, num_params=NUM_PARAMS, repeats=3): ctx = mp.get_context("fork") mgr = ctx.Manager() result = mgr.dict() - p = ctx.Process( - target=_speed_worker, - args=(result, param_mb, num_params, materialize_fn, repeats), - ) + p = ctx.Process(target=_speed_worker, args=(result, param_mb, num_params, materialize_fn, repeats)) p.start() p.join(timeout=300) assert p.exitcode == 0, f"subprocess exited with code {p.exitcode}" @@ -320,10 +310,12 @@ def test_streaming_not_slower_than_batch(): r_stream = _run_speed_subprocess(_streaming_materialize, mb) ratio = r_stream["avg"] / r_batch["avg"] - print(f"\n[speed] model={mb} MB, num_params={NUM_PARAMS}" - f" batch={r_batch['avg']:.3f}s" - f" streaming={r_stream['avg']:.3f}s" - f" ratio={ratio:.2f}x") + print( + f"\n[speed] model={mb} MB, num_params={NUM_PARAMS}" + f" batch={r_batch['avg']:.3f}s" + f" streaming={r_stream['avg']:.3f}s" + f" ratio={ratio:.2f}x" + ) assert ratio < max_slowdown, ( f"Streaming is {ratio:.2f}x slower than batch " From f458d62f7ca38036f715d14262778d123952ce10 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 17:11:01 +0800 Subject: [PATCH 24/54] style: remove extra blank line in host_memory.py --- magi_compiler/utils/host_memory.py | 1 - 1 file changed, 1 deletion(-) diff --git a/magi_compiler/utils/host_memory.py b/magi_compiler/utils/host_memory.py index d61bf89..e54e419 100644 --- a/magi_compiler/utils/host_memory.py +++ b/magi_compiler/utils/host_memory.py @@ -4,7 +4,6 @@ from __future__ import annotations - def get_host_mem_gb() -> dict[str, float]: """Read /proc/self/status and return key memory metrics in GiB. From c1bdb90d45ccf6e66368466ebba0455bb2f6a57c Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 17:27:05 +0800 Subject: [PATCH 25/54] refactor: parse ep_size from MAGI_COMPILE_TOPOLOGY_KEY 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. --- magi_compiler/_api.py | 4 ++-- magi_compiler/config.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 5b53ed0..7296f37 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -32,7 +32,7 @@ from magi_compiler.utils.compile_time_monitor import CompileMonitor from magi_compiler.utils.host_memory import fmt_host_mem -from .config import CompileConfig, CompileMode +from .config import CompileConfig, CompileMode, get_topology_dim # ============================================================================= @@ -681,7 +681,7 @@ def _force_cpu(t): magi_logger.info('[offload] after _force_cpu: %s', fmt_host_mem()) # create shared memory tensors for all parameters/buffers on CPU - ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", os.environ.get("EP_SIZE", "1"))) + ep_size = get_topology_dim("ep") if dist.is_initialized(): local_rank = int(os.environ.get("LOCAL_RANK", 0)) full_state_dict = self.state_dict() diff --git a/magi_compiler/config.py b/magi_compiler/config.py index e226786..28adb37 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -429,6 +429,22 @@ def _get_parallel_topology() -> str: return f"ws{torch.distributed.get_world_size()}" +def get_topology_dim(dim: str, default: int = 1) -> int: + """Extract a parallel dimension size from ``MAGI_COMPILE_TOPOLOGY_KEY``. + + The key is a ``_``-joined string like ``cp8_dp1_ep8_tp1`` set by the + host framework's ParallelStateManager. Returns *default* if the key + is absent or the dimension is not present. + """ + import re + + topo = os.environ.get("MAGI_COMPILE_TOPOLOGY_KEY", "") + if not topo: + return default + m = re.search(rf"(?:^|_){re.escape(dim)}(\d+)", topo) + return int(m.group(1)) if m else default + + def model_rank_dir_name(model_idx: int, model_tag: str | None) -> str: """Directory name: ``model_{idx}[_{tag}]_rank_{rank}_{topology}``. From 0ac1ea0df430ced8f6bdb1a87c4969bfe99972ab Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 19:15:28 +0800 Subject: [PATCH 26/54] test: use MAGI_COMPILE_TOPOLOGY_KEY in EP shared-memory test 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. --- tests/feature_tests/test_ep_shared_memory.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 505f21b..20d7e5c 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -140,14 +140,15 @@ def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_fil os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = "29502" os.environ["LOCAL_RANK"] = str(rank) - os.environ["ENGINE_CONFIG__EP_SIZE"] = str(world_size) + os.environ["MAGI_COMPILE_TOPOLOGY_KEY"] = f"cp1_dp1_ep{world_size}_tp1" dist.init_process_group("gloo", rank=rank, world_size=world_size) torch.manual_seed(seed_per_rank[rank]) model = FakeExpertBlock(num_experts=4, dim=8) original_weight = model.expert_weight.data.clone() - ep_size = int(os.environ.get("ENGINE_CONFIG__EP_SIZE", "1")) + from magi_compiler.config import get_topology_dim + ep_size = get_topology_dim("ep") _shm_write_read(model, local_rank=rank, shared_dir=shared_dir, per_rank=(ep_size > 1)) weight_after = model.state_dict()["expert_weight"] From cec5cd48a45284561a27b00d0451a044fdca199b Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 19:31:54 +0800 Subject: [PATCH 27/54] test: remove unnecessary gloo skipif guard gloo is always available in standard PyTorch installations. --- tests/feature_tests/test_ep_shared_memory.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 20d7e5c..6381778 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -38,11 +38,6 @@ import torch.multiprocessing as mp import torch.nn as nn -_skip_no_dist = pytest.mark.skipif( - not hasattr(dist, "is_gloo_available") or not dist.is_gloo_available(), reason="requires gloo backend" -) - - class FakeExpertBlock(nn.Module): """Tiny module simulating an EP-sharded expert block.""" @@ -163,7 +158,6 @@ def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_fil # ─────────────────────────────────────────────────────────────── -@_skip_no_dist def test_shared_memory_overwrites_ep_shards(): """ BUG REPRO: with original shared-memory logic, rank 1's expert weights @@ -191,7 +185,6 @@ def test_shared_memory_overwrites_ep_shards(): ) -@_skip_no_dist def test_ep_fix_preserves_per_rank_shards(): """ FIX VERIFIED: with EP_SIZE > 1, per-rank shared-memory files ensure each From fa1ce068cc97e02427ba12228ef09da2964ddb55 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 20:26:00 +0800 Subject: [PATCH 28/54] refactor(test): import production helpers instead of reimplementing - 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) --- tests/feature_tests/test_ep_shared_memory.py | 91 +++++++++----------- tests/feature_tests/test_shm_memory_peak.py | 91 +++++++------------- 2 files changed, 73 insertions(+), 109 deletions(-) diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 6381778..2e20cd7 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -32,12 +32,18 @@ import os import tempfile -import pytest import torch import torch.distributed as dist import torch.multiprocessing as mp import torch.nn as nn +from magi_compiler._api import ( + _create_empty_shm, + _pack_params_flat, + _split_flat_to_params, + _stream_copy_and_replace, +) + class FakeExpertBlock(nn.Module): """Tiny module simulating an EP-sharded expert block.""" @@ -49,60 +55,49 @@ def forward(self, x): return x @ self.expert_weight.T +def _group_params(module: nn.Module) -> dict[torch.dtype, list[tuple[str, torch.Tensor]]]: + grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} + for name, tensor in module.state_dict().items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + return grouped + + def _shm_write_read(module, local_rank, shared_dir, per_rank): - """ - Core shared-memory logic extracted from _patch_cpu_offload_apply. + """Mirrors _materialize_shm_weights using production helpers. - per_rank=False → original buggy path (rank 0 writes, all read same file). - per_rank=True → fixed path (each rank writes its own file). + per_rank=False → original buggy path (rank 0 writes, all read same file + via load_state_dict — every rank gets rank 0's weights). + per_rank=True → fixed path (each rank writes its own file, uses + _stream_copy_and_replace to keep its own weights). """ - full_state_dict = module.state_dict() - grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} - for name, tensor in full_state_dict.items(): - dt = tensor.dtype - grouped.setdefault(dt, []).append((name, tensor)) - - writer_rank = local_rank if per_rank else 0 - shared_state = {} - - for dtype, param_list in grouped.items(): - dtype_str = str(dtype).split(".")[-1] - suffix = f"_rank{local_rank}" if per_rank else "" - shared_path = os.path.join(shared_dir, f"shared_{dtype_str}{suffix}.bin") - total_numel = sum(t.numel() for _, t in param_list) - - if local_rank == writer_rank: - flat = torch.zeros(total_numel, dtype=dtype) - off = 0 - for _, t in param_list: - n = t.numel() - flat[off : off + n].copy_(t.view(-1)) - off += n - if dtype == torch.bfloat16: - flat.view(torch.int16).numpy().tofile(shared_path) - else: - flat.numpy().tofile(shared_path) - del flat - gc.collect() - - dist.barrier() - - giant = torch.from_file(shared_path, shared=True, size=total_numel, dtype=dtype, device="cpu") - off = 0 - for name, orig in param_list: - n = orig.numel() - shared_state[name] = giant[off : off + n].view(orig.shape) - off += n - - dist.barrier() - if per_rank: + grouped = _group_params(module) + + if per_rank: + for dtype, param_list in grouped.items(): + suffix = f"_rank{local_rank}" + shared_path = os.path.join(shared_dir, f"shared_{str(dtype).split('.')[-1]}{suffix}.bin") + total_numel = sum(t.numel() for _, t in param_list) + giant = _create_empty_shm(shared_path, total_numel, dtype) + _stream_copy_and_replace(module, giant, param_list) + dist.barrier() if os.path.exists(shared_path): os.remove(shared_path) - else: + else: + shared_state = {} + for dtype, param_list in grouped.items(): + shared_path = os.path.join(shared_dir, f"shared_{str(dtype).split('.')[-1]}.bin") + total_numel = sum(t.numel() for _, t in param_list) + if local_rank == 0: + giant = _create_empty_shm(shared_path, total_numel, dtype) + _pack_params_flat(giant, param_list) + dist.barrier() + if local_rank != 0: + giant = torch.from_file(shared_path, shared=True, size=total_numel, dtype=dtype, device="cpu") + shared_state.update(_split_flat_to_params(giant, param_list)) + dist.barrier() if local_rank == 0 and os.path.exists(shared_path): os.remove(shared_path) - - module.load_state_dict(shared_state, assign=True) + module.load_state_dict(shared_state, assign=True) # ─────────────────────────────────────────────────────────────── diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index feb6467..e7ea585 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -42,12 +42,15 @@ import os import tempfile -import pytest import torch import torch.nn as nn -_IS_LINUX = os.path.exists("/proc/self/status") -_skip_no_procfs = pytest.mark.skipif(not _IS_LINUX, reason="requires /proc/self/status for VmHWM") +from magi_compiler._api import ( + _assign_param, + _create_empty_shm, + _pack_params_flat, + _stream_copy_and_replace, +) PARAM_MB = 512 NUM_PARAMS = 4 @@ -75,35 +78,32 @@ def forward(self, x): # ── batch (current code pattern) ──────────────────────────── -def _batch_materialize(module: nn.Module, shm_dir: str) -> None: - full_state_dict = module.state_dict() +def _group_params(module: nn.Module) -> dict[torch.dtype, list[tuple[str, torch.Tensor]]]: + """Group module params by dtype (shared by both batch and streaming paths).""" grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} - for name, tensor in full_state_dict.items(): + for name, tensor in module.state_dict().items(): grouped.setdefault(tensor.dtype, []).append((name, tensor)) + return grouped + + +def _batch_materialize(module: nn.Module, shm_dir: str) -> None: + """Old batch approach: copy ALL params into mmap, then load_state_dict. + Intentionally reimplemented here (NOT imported) because this is the + BUGGY baseline we want to prove has ~2x peak. Production code no + longer uses this pattern. + """ + grouped = _group_params(module) shared_state: dict[str, torch.Tensor] = {} buffers: list[torch.Tensor] = [] for dtype, param_list in grouped.items(): total_numel = sum(t.numel() for _, t in param_list) - elem_size = torch.empty(0, dtype=dtype).element_size() path = os.path.join(shm_dir, f"batch_{dtype}.bin") - with open(path, "wb") as f: - f.truncate(total_numel * elem_size) - giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") - offset = 0 - for _, tensor in param_list: - n = tensor.numel() - giant[offset : offset + n].copy_(tensor.view(-1)) - offset += n - offset = 0 - for name, orig in param_list: - n = orig.numel() - view = giant[offset : offset + n].view(orig.shape) - if orig.requires_grad: - view.requires_grad_(True) - shared_state[name] = view - offset += n + giant = _create_empty_shm(path, total_numel, dtype) + _pack_params_flat(giant, param_list) + from magi_compiler._api import _split_flat_to_params + shared_state.update(_split_flat_to_params(giant, param_list)) buffers.append(giant) if os.path.exists(path): os.remove(path) @@ -115,43 +115,16 @@ def _batch_materialize(module: nn.Module, shm_dir: str) -> None: # ── streaming (fix) ───────────────────────────────────────── -def _assign_param(module: nn.Module, dotted_name: str, new_tensor: torch.Tensor) -> None: - parts = dotted_name.rsplit(".", 1) - parent = module.get_submodule(parts[0]) if len(parts) == 2 else module - attr = parts[-1] - old = getattr(parent, attr) - if isinstance(old, nn.Parameter): - parent.register_parameter(attr, nn.Parameter(new_tensor, requires_grad=new_tensor.requires_grad)) - else: - setattr(parent, attr, new_tensor) - - def _streaming_materialize(module: nn.Module, shm_dir: str) -> None: - full_state_dict = module.state_dict() - grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} - for name, tensor in full_state_dict.items(): - grouped.setdefault(tensor.dtype, []).append((name, tensor)) - del full_state_dict - gc.collect() - + """Uses the REAL production functions to test actual behavior.""" + grouped = _group_params(module) buffers: list[torch.Tensor] = [] + for dtype, param_list in grouped.items(): total_numel = sum(t.numel() for _, t in param_list) - elem_size = torch.empty(0, dtype=dtype).element_size() path = os.path.join(shm_dir, f"stream_{dtype}.bin") - with open(path, "wb") as f: - f.truncate(total_numel * elem_size) - giant = torch.from_file(path, shared=True, size=total_numel, dtype=dtype, device="cpu") - offset = 0 - for i, (name, tensor) in enumerate(param_list): - n = tensor.numel() - giant[offset : offset + n].copy_(tensor.view(-1)) - view = giant[offset : offset + n].view(tensor.shape) - if tensor.requires_grad: - view.requires_grad_(True) - _assign_param(module, name, view) - param_list[i] = (name, view) - offset += n + giant = _create_empty_shm(path, total_numel, dtype) + _stream_copy_and_replace(module, giant, param_list) buffers.append(giant) if os.path.exists(path): os.remove(path) @@ -194,7 +167,6 @@ def _run_in_subprocess(materialize_fn, param_mb): # ── tests ─────────────────────────────────────────────────── -@_skip_no_procfs def test_batch_materialize_has_high_peak(): """BUG REPRO: batch materialize adds ~1× model size as mmap overhead. @@ -219,7 +191,6 @@ def test_batch_materialize_has_high_peak(): ) -@_skip_no_procfs def test_streaming_materialize_low_peak(): """FIX VERIFIED: streaming avoids the mmap overhead peak. @@ -244,7 +215,6 @@ def test_streaming_materialize_low_peak(): ) -@_skip_no_procfs def test_streaming_preserves_weights(): """Streaming must produce identical weights to batch.""" numel = 1024 @@ -295,16 +265,15 @@ def _run_speed_subprocess(materialize_fn, param_mb, num_params=NUM_PARAMS, repea return dict(result) -@_skip_no_procfs def test_streaming_not_slower_than_batch(): """Streaming must not be significantly slower than batch. - Allows up to 1.20x slowdown to account for per-param register_parameter + Allows up to 1.50x slowdown to account for per-param register_parameter overhead. In practice streaming is often faster on large models because it avoids the final load_state_dict bulk copy. """ mb = PARAM_MB - max_slowdown = 1.20 + max_slowdown = 1.50 r_batch = _run_speed_subprocess(_batch_materialize, mb) r_stream = _run_speed_subprocess(_streaming_materialize, mb) From d2ecc7241426a5610d5af00283743be578b5edf5 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 20:39:16 +0800 Subject: [PATCH 29/54] test: tighten memory peak thresholds (batch >0.8x, streaming <0.2x) --- tests/feature_tests/test_shm_memory_peak.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index e7ea585..a4b66ee 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -171,7 +171,7 @@ def test_batch_materialize_has_high_peak(): """BUG REPRO: batch materialize adds ~1× model size as mmap overhead. At peak: old params (RssAnon) + mmap copy (RssFile) ≈ 2× model size. - VmHWM growth measures the mmap portion, expected > 0.5× model_size. + VmHWM growth measures the mmap portion, expected > 0.8× model_size. """ r = _run_in_subprocess(_batch_materialize, PARAM_MB) growth = r["hwm_after"] - r["hwm_before"] @@ -184,8 +184,8 @@ def test_batch_materialize_has_high_peak(): f"ratio={growth / pm:.2f}x" ) - assert growth > pm * 0.5, ( - f"Expected mmap overhead > {pm * 0.5:.0f} MB (0.5× model) " + assert growth > pm * 0.8, ( + f"Expected mmap overhead > {pm * 0.8:.0f} MB (0.8× model) " f"but got {growth:.0f} MB ({growth / pm:.2f}×). " f"The 2× peak may have been optimized away." ) @@ -195,7 +195,7 @@ def test_streaming_materialize_low_peak(): """FIX VERIFIED: streaming avoids the mmap overhead peak. By replacing each param immediately, only ~1/N of the model is ever - duplicated. VmHWM growth should be well under 0.3× model size. + duplicated. VmHWM growth should be well under 0.2× model size. """ r = _run_in_subprocess(_streaming_materialize, PARAM_MB) growth = r["hwm_after"] - r["hwm_before"] @@ -208,8 +208,8 @@ def test_streaming_materialize_low_peak(): f"ratio={growth / pm:.2f}x" ) - assert growth < pm * 0.3, ( - f"Expected mmap overhead < {pm * 0.3:.0f} MB (0.3× model) " + assert growth < pm * 0.2, ( + f"Expected mmap overhead < {pm * 0.2:.0f} MB (0.2× model) " f"but got {growth:.0f} MB ({growth / pm:.2f}×). " f"Streaming fix did not reduce peak." ) From 410b9c393154d6a2a5d2d634c30d990252279ba1 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 20:43:35 +0800 Subject: [PATCH 30/54] refactor: _fix_graph_device_placement as @staticmethod, recursive _move_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 --- magi_compiler/magi_backend/magi_backend.py | 143 ++++++++++----------- 1 file changed, 65 insertions(+), 78 deletions(-) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 83ad49f..b81c190 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -279,84 +279,71 @@ def __init__( # extra_traceback is attribute of torch.fx.Interpreter, when it is True, it annoyingly dumps the torch.fx.Graph on errors. self.extra_traceback = False - def _fix_graph_device_placement(self, module: torch.nn.Module): - for name, child in module.named_children(): - self._fix_graph_device_placement(child) - - if isinstance(module, torch.fx.GraphModule): - needs_recompile = False - target_device = torch.cuda.current_device() - - # Fix .to(device('cpu')) calls baked during CPU-offload tracing. - # Without _deep_cuda, Dynamo traces with CPU tensors and specialises - # .to(x.device) as .to(device('cpu')). After we move example_values - # to CUDA below, these hardcoded .to(cpu) nodes create CUDA-vs-CPU - # mismatches in PiecewiseCompileInterpreter.run(). - def _is_cpu_device(val): - if isinstance(val, torch.device): - return val.type == 'cpu' - if isinstance(val, str): - return val == 'cpu' - return False - - for node in module.graph.nodes: - if node.op == 'call_function' and 'device' in node.kwargs: - if _is_cpu_device(node.kwargs['device']): - node.update_kwarg('device', torch.device('cuda', target_device)) - needs_recompile = True - - if node.op == 'call_method' and node.target == 'to': - new_args = list(node.args) - changed = False - for i, arg in enumerate(new_args): - if _is_cpu_device(arg): - new_args[i] = torch.device('cuda', target_device) - changed = True - if changed: - node.args = tuple(new_args) - needs_recompile = True - if 'device' in node.kwargs and _is_cpu_device(node.kwargs['device']): - node.update_kwarg('device', torch.device('cuda', target_device)) - needs_recompile = True - - # Fix ALL nodes with CPU example_values — not just get_attr/placeholder. - # model_cpu_offload traces with CPU FakeTensors; Inductor autotuning - # creates benchmark tensors on the example_value device, and the - # codegen may pick CPU-specific paths if it sees CPU metadata. - def _move_to_device(val): - if hasattr(val, 'device') and str(val.device) == 'cpu': - new_val = val.to(target_device) - if isinstance(val, torch.nn.Parameter): - new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad) - return new_val - return val - - cpu_fix_count = 0 - for node in module.graph.nodes: - ev = node.meta.get('example_value') - if ev is not None: - if hasattr(ev, 'device'): - new_ev = _move_to_device(ev) - if new_ev is not ev: - node.meta['example_value'] = new_ev - needs_recompile = True - cpu_fix_count += 1 - elif isinstance(ev, (list, tuple)): - fixed_list = [] - any_fixed = False - for item in ev: - new_item = _move_to_device(item) - fixed_list.append(new_item) - if new_item is not item: - any_fixed = True - cpu_fix_count += 1 - if any_fixed: - node.meta['example_value'] = type(ev)(fixed_list) - needs_recompile = True - - if needs_recompile: - magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device) - module.recompile() + @staticmethod + def _fix_graph_device_placement(module: torch.nn.Module): + for _, child in module.named_children(): + PiecewiseCompileInterpreter._fix_graph_device_placement(child) + + if not isinstance(module, torch.fx.GraphModule): + return + + needs_recompile = False + target_device = torch.cuda.current_device() + + def _is_cpu_device(val): + if isinstance(val, torch.device): + return val.type == 'cpu' + return isinstance(val, str) and val == 'cpu' + + # --- Rewrite hardcoded .to(cpu) / factory(device='cpu') nodes --- + + for node in module.graph.nodes: + if node.op == 'call_function' and 'device' in node.kwargs: + if _is_cpu_device(node.kwargs['device']): + node.update_kwarg('device', torch.device('cuda', target_device)) + needs_recompile = True + + if node.op == 'call_method' and node.target == 'to': + new_args = list(node.args) + changed = False + for i, arg in enumerate(new_args): + if _is_cpu_device(arg): + new_args[i] = torch.device('cuda', target_device) + changed = True + if changed: + node.args = tuple(new_args) + needs_recompile = True + if 'device' in node.kwargs and _is_cpu_device(node.kwargs['device']): + node.update_kwarg('device', torch.device('cuda', target_device)) + needs_recompile = True + + # --- Fix CPU example_values (recursive for nested list/tuple) --- + + def _move_to_device(val): + if isinstance(val, (list, tuple)): + items = [_move_to_device(v) for v in val] + return type(val)(items) if any(n is not o for n, o in zip(items, val)) else val + if hasattr(val, 'device') and str(val.device) == 'cpu': + new_val = val.to(target_device) + if isinstance(val, torch.nn.Parameter): + new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad) + return new_val + return val + + cpu_fix_count = 0 + for node in module.graph.nodes: + ev = node.meta.get('example_value') + if ev is None: + continue + new_ev = _move_to_device(ev) + if new_ev is not ev: + node.meta['example_value'] = new_ev + needs_recompile = True + cpu_fix_count += 1 + + if needs_recompile: + magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device) + module.recompile() @observe_lifecycle("piecewise_compile") def run(self, *args): From b98a4c1c09f4ca5ec54b14905e0d8f386e03b61d Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 20:50:19 +0800 Subject: [PATCH 31/54] refactor: promote _device_is_cpu/_recursive_to_device to class staticmethods MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- magi_compiler/magi_backend/magi_backend.py | 45 ++++++++++++---------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index b81c190..619b526 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -279,6 +279,26 @@ def __init__( # extra_traceback is attribute of torch.fx.Interpreter, when it is True, it annoyingly dumps the torch.fx.Graph on errors. self.extra_traceback = False + @staticmethod + def _device_is_cpu(val) -> bool: + """Check whether *val* represents a CPU device (torch.device or str).""" + if isinstance(val, torch.device): + return val.type == 'cpu' + return isinstance(val, str) and val == 'cpu' + + @staticmethod + def _recursive_to_device(val, target_device: int): + """Recursively move tensor-like *val* (or nested list/tuple) to *target_device*.""" + if isinstance(val, (list, tuple)): + items = [PiecewiseCompileInterpreter._recursive_to_device(v, target_device) for v in val] + return type(val)(items) if any(n is not o for n, o in zip(items, val)) else val + if hasattr(val, 'device') and str(val.device) == 'cpu': + new_val = val.to(target_device) + if isinstance(val, torch.nn.Parameter): + new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad) + return new_val + return val + @staticmethod def _fix_graph_device_placement(module: torch.nn.Module): for _, child in module.named_children(): @@ -289,17 +309,13 @@ def _fix_graph_device_placement(module: torch.nn.Module): needs_recompile = False target_device = torch.cuda.current_device() - - def _is_cpu_device(val): - if isinstance(val, torch.device): - return val.type == 'cpu' - return isinstance(val, str) and val == 'cpu' + _device_is_cpu = PiecewiseCompileInterpreter._device_is_cpu # --- Rewrite hardcoded .to(cpu) / factory(device='cpu') nodes --- for node in module.graph.nodes: if node.op == 'call_function' and 'device' in node.kwargs: - if _is_cpu_device(node.kwargs['device']): + if _device_is_cpu(node.kwargs['device']): node.update_kwarg('device', torch.device('cuda', target_device)) needs_recompile = True @@ -307,35 +323,24 @@ def _is_cpu_device(val): new_args = list(node.args) changed = False for i, arg in enumerate(new_args): - if _is_cpu_device(arg): + if _device_is_cpu(arg): new_args[i] = torch.device('cuda', target_device) changed = True if changed: node.args = tuple(new_args) needs_recompile = True - if 'device' in node.kwargs and _is_cpu_device(node.kwargs['device']): + if 'device' in node.kwargs and _device_is_cpu(node.kwargs['device']): node.update_kwarg('device', torch.device('cuda', target_device)) needs_recompile = True # --- Fix CPU example_values (recursive for nested list/tuple) --- - def _move_to_device(val): - if isinstance(val, (list, tuple)): - items = [_move_to_device(v) for v in val] - return type(val)(items) if any(n is not o for n, o in zip(items, val)) else val - if hasattr(val, 'device') and str(val.device) == 'cpu': - new_val = val.to(target_device) - if isinstance(val, torch.nn.Parameter): - new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad) - return new_val - return val - cpu_fix_count = 0 for node in module.graph.nodes: ev = node.meta.get('example_value') if ev is None: continue - new_ev = _move_to_device(ev) + new_ev = PiecewiseCompileInterpreter._recursive_to_device(ev, target_device) if new_ev is not ev: node.meta['example_value'] = new_ev needs_recompile = True From 1f1709e677f1ce7f57911967b0fd90a4a953e3f3 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 20:52:47 +0800 Subject: [PATCH 32/54] test: use production staticmethods instead of reimplementing helpers 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. --- .../feature_tests/test_fix_to_cpu_in_graph.py | 37 +++++-------------- 1 file changed, 9 insertions(+), 28 deletions(-) diff --git a/tests/feature_tests/test_fix_to_cpu_in_graph.py b/tests/feature_tests/test_fix_to_cpu_in_graph.py index 326d7ed..788300b 100644 --- a/tests/feature_tests/test_fix_to_cpu_in_graph.py +++ b/tests/feature_tests/test_fix_to_cpu_in_graph.py @@ -25,6 +25,11 @@ import torch.fx as fx import torch.nn as nn +from magi_compiler.magi_backend.magi_backend import PiecewiseCompileInterpreter + +_fix = PiecewiseCompileInterpreter._fix_graph_device_placement +_device_is_cpu = PiecewiseCompileInterpreter._device_is_cpu + def _build_graph_with_to_cpu(): """Build an FX graph that mirrors the offload-traced pattern: @@ -49,14 +54,6 @@ def _build_graph_with_to_cpu(): return gm -def _is_cpu_device(val): - if isinstance(val, torch.device): - return val.type == "cpu" - if isinstance(val, str): - return val == "cpu" - return False - - def _apply_metadata_only_fix(gm, target_device=0): for node in gm.graph.nodes: ev = node.meta.get("example_value") @@ -65,22 +62,6 @@ def _apply_metadata_only_fix(gm, target_device=0): gm.recompile() -def _apply_full_fix(gm, target_device=0): - for node in gm.graph.nodes: - if node.op == "call_method" and node.target == "to": - new_args = list(node.args) - changed = False - for i, arg in enumerate(new_args): - if _is_cpu_device(arg): - new_args[i] = torch.device("cuda", target_device) - changed = True - if changed: - node.args = tuple(new_args) - if "device" in node.kwargs and _is_cpu_device(node.kwargs["device"]): - node.update_kwarg("device", torch.device("cuda", target_device)) - _apply_metadata_only_fix(gm, target_device) - - def _run_with_fake_tensors(gm, cuda_device=0): from torch._subclasses.fake_tensor import FakeTensorMode @@ -116,17 +97,17 @@ def test_metadata_only_fix_fails(self): def test_full_fix_succeeds(self): """Rewriting .to(cpu) -> .to(cuda) + metadata fix -> no error.""" gm = _build_graph_with_to_cpu() - _apply_full_fix(gm) + _fix(gm) out = _run_with_fake_tensors(gm) assert str(out.device).startswith("cuda") def test_no_residual_to_cpu(self): gm = _build_graph_with_to_cpu() - _apply_full_fix(gm) + _fix(gm) for node in gm.graph.nodes: if node.op == "call_method" and node.target == "to": for arg in node.args: - assert not _is_cpu_device(arg), f"Residual .to(cpu): {node}" + assert not _device_is_cpu(arg), f"Residual .to(cpu): {node}" def test_to_dtype_untouched(self): """Rewrite must NOT affect .to(dtype) calls.""" @@ -138,7 +119,7 @@ def test_to_dtype_untouched(self): x.meta["example_value"] = torch.randn(4, 8) to_bf16.meta["example_value"] = torch.randn(4, 8, dtype=torch.bfloat16) - _apply_full_fix(gm) + _fix(gm) for node in gm.graph.nodes: if node.op == "call_method" and node.target == "to": From f516ba4b3d93280000c9ac2f3535839d27a3b4fb Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 20:55:11 +0800 Subject: [PATCH 33/54] refactor: move device-fix helpers to module-level pure functions _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. --- magi_compiler/magi_backend/magi_backend.py | 140 +++++++++--------- .../feature_tests/test_fix_to_cpu_in_graph.py | 8 +- 2 files changed, 72 insertions(+), 76 deletions(-) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 619b526..8a487d2 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -252,6 +252,73 @@ def _maybe_store_cache_entry( return True +def _device_is_cpu(val) -> bool: + """Check whether *val* represents a CPU device (torch.device or str).""" + if isinstance(val, torch.device): + return val.type == 'cpu' + return isinstance(val, str) and val == 'cpu' + + +def _recursive_to_device(val, target_device: int): + """Recursively move tensor-like *val* (or nested list/tuple) to *target_device*.""" + if isinstance(val, (list, tuple)): + items = [_recursive_to_device(v, target_device) for v in val] + return type(val)(items) if any(n is not o for n, o in zip(items, val)) else val + if hasattr(val, 'device') and str(val.device) == 'cpu': + new_val = val.to(target_device) + if isinstance(val, torch.nn.Parameter): + new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad) + return new_val + return val + + +def fix_graph_device_placement(module: torch.nn.Module): + """Rewrite CPU device refs and example_values to CUDA in an FX graph.""" + for _, child in module.named_children(): + fix_graph_device_placement(child) + + if not isinstance(module, torch.fx.GraphModule): + return + + needs_recompile = False + target_device = torch.cuda.current_device() + + for node in module.graph.nodes: + if node.op == 'call_function' and 'device' in node.kwargs: + if _device_is_cpu(node.kwargs['device']): + node.update_kwarg('device', torch.device('cuda', target_device)) + needs_recompile = True + + if node.op == 'call_method' and node.target == 'to': + new_args = list(node.args) + changed = False + for i, arg in enumerate(new_args): + if _device_is_cpu(arg): + new_args[i] = torch.device('cuda', target_device) + changed = True + if changed: + node.args = tuple(new_args) + needs_recompile = True + if 'device' in node.kwargs and _device_is_cpu(node.kwargs['device']): + node.update_kwarg('device', torch.device('cuda', target_device)) + needs_recompile = True + + cpu_fix_count = 0 + for node in module.graph.nodes: + ev = node.meta.get('example_value') + if ev is None: + continue + new_ev = _recursive_to_device(ev, target_device) + if new_ev is not ev: + node.meta['example_value'] = new_ev + needs_recompile = True + cpu_fix_count += 1 + + if needs_recompile: + magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device) + module.recompile() + + class PiecewiseCompileInterpreter(torch.fx.Interpreter): """ Code adapted from `torch.fx.passes.shape_prop.ShapeProp`. @@ -279,82 +346,11 @@ def __init__( # extra_traceback is attribute of torch.fx.Interpreter, when it is True, it annoyingly dumps the torch.fx.Graph on errors. self.extra_traceback = False - @staticmethod - def _device_is_cpu(val) -> bool: - """Check whether *val* represents a CPU device (torch.device or str).""" - if isinstance(val, torch.device): - return val.type == 'cpu' - return isinstance(val, str) and val == 'cpu' - - @staticmethod - def _recursive_to_device(val, target_device: int): - """Recursively move tensor-like *val* (or nested list/tuple) to *target_device*.""" - if isinstance(val, (list, tuple)): - items = [PiecewiseCompileInterpreter._recursive_to_device(v, target_device) for v in val] - return type(val)(items) if any(n is not o for n, o in zip(items, val)) else val - if hasattr(val, 'device') and str(val.device) == 'cpu': - new_val = val.to(target_device) - if isinstance(val, torch.nn.Parameter): - new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad) - return new_val - return val - - @staticmethod - def _fix_graph_device_placement(module: torch.nn.Module): - for _, child in module.named_children(): - PiecewiseCompileInterpreter._fix_graph_device_placement(child) - - if not isinstance(module, torch.fx.GraphModule): - return - - needs_recompile = False - target_device = torch.cuda.current_device() - _device_is_cpu = PiecewiseCompileInterpreter._device_is_cpu - - # --- Rewrite hardcoded .to(cpu) / factory(device='cpu') nodes --- - - for node in module.graph.nodes: - if node.op == 'call_function' and 'device' in node.kwargs: - if _device_is_cpu(node.kwargs['device']): - node.update_kwarg('device', torch.device('cuda', target_device)) - needs_recompile = True - - if node.op == 'call_method' and node.target == 'to': - new_args = list(node.args) - changed = False - for i, arg in enumerate(new_args): - if _device_is_cpu(arg): - new_args[i] = torch.device('cuda', target_device) - changed = True - if changed: - node.args = tuple(new_args) - needs_recompile = True - if 'device' in node.kwargs and _device_is_cpu(node.kwargs['device']): - node.update_kwarg('device', torch.device('cuda', target_device)) - needs_recompile = True - - # --- Fix CPU example_values (recursive for nested list/tuple) --- - - cpu_fix_count = 0 - for node in module.graph.nodes: - ev = node.meta.get('example_value') - if ev is None: - continue - new_ev = PiecewiseCompileInterpreter._recursive_to_device(ev, target_device) - if new_ev is not ev: - node.meta['example_value'] = new_ev - needs_recompile = True - cpu_fix_count += 1 - - if needs_recompile: - magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device) - module.recompile() - @observe_lifecycle("piecewise_compile") def run(self, *args): fake_args = self._build_fake_args(args) if self.compile_config.offload_config.model_cpu_offload: - self._fix_graph_device_placement(self.module) + fix_graph_device_placement(self.module) for i, arg in enumerate(fake_args): if isinstance(arg, torch.Tensor): fake_args[i] = arg.cuda() diff --git a/tests/feature_tests/test_fix_to_cpu_in_graph.py b/tests/feature_tests/test_fix_to_cpu_in_graph.py index 788300b..abad031 100644 --- a/tests/feature_tests/test_fix_to_cpu_in_graph.py +++ b/tests/feature_tests/test_fix_to_cpu_in_graph.py @@ -25,10 +25,10 @@ import torch.fx as fx import torch.nn as nn -from magi_compiler.magi_backend.magi_backend import PiecewiseCompileInterpreter - -_fix = PiecewiseCompileInterpreter._fix_graph_device_placement -_device_is_cpu = PiecewiseCompileInterpreter._device_is_cpu +from magi_compiler.magi_backend.magi_backend import ( + _device_is_cpu, + fix_graph_device_placement as _fix, +) def _build_graph_with_to_cpu(): From 0de84025dc68f757530d31c9407e34f471f85617 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 21:05:51 +0800 Subject: [PATCH 34/54] test: call production _materialize_shm_weights directly in EP 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. --- tests/feature_tests/test_ep_shared_memory.py | 109 +++++++------------ 1 file changed, 37 insertions(+), 72 deletions(-) diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 2e20cd7..ed3e617 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -15,34 +15,29 @@ """ Regression tests for the EP shared-memory weight corruption bug. -Root cause: _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's data -on every rank destroyed expert weight diversity and produced garbled output. +Root cause: _materialize_shm_weights with per_rank=False creates a single +shared-memory file from local_rank=0 and has ALL ranks map it. With expert +parallelism (EP > 1), each rank holds a different expert shard; sharing one +mmap means the last writer's data overwrites everyone else's views. -Fix: when EP_SIZE > 1, each rank writes its OWN shared-memory file so that -expert shards are preserved. When EP_SIZE <= 1, the original rank-0-writes -all-read scheme is safe (weights are identical across ranks). +Fix: when ep_size > 1, the caller passes per_rank=True so each rank writes +its OWN shared-memory file, preserving expert shard diversity. These tests use torch.multiprocessing.spawn with 2 workers and the gloo backend to reproduce the exact multi-rank scenario without a real model. """ -import gc import os import tempfile +from unittest.mock import patch import torch import torch.distributed as dist import torch.multiprocessing as mp import torch.nn as nn -from magi_compiler._api import ( - _create_empty_shm, - _pack_params_flat, - _split_flat_to_params, - _stream_copy_and_replace, -) +from magi_compiler._api import _materialize_shm_weights + class FakeExpertBlock(nn.Module): """Tiny module simulating an EP-sharded expert block.""" @@ -57,47 +52,19 @@ def forward(self, x): def _group_params(module: nn.Module) -> dict[torch.dtype, list[tuple[str, torch.Tensor]]]: grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} - for name, tensor in module.state_dict().items(): - grouped.setdefault(tensor.dtype, []).append((name, tensor)) + for name, param in module.named_parameters(): + grouped.setdefault(param.dtype, []).append((name, param.data)) return grouped -def _shm_write_read(module, local_rank, shared_dir, per_rank): - """Mirrors _materialize_shm_weights using production helpers. - - per_rank=False → original buggy path (rank 0 writes, all read same file - via load_state_dict — every rank gets rank 0's weights). - per_rank=True → fixed path (each rank writes its own file, uses - _stream_copy_and_replace to keep its own weights). - """ - grouped = _group_params(module) - - if per_rank: - for dtype, param_list in grouped.items(): - suffix = f"_rank{local_rank}" - shared_path = os.path.join(shared_dir, f"shared_{str(dtype).split('.')[-1]}{suffix}.bin") - total_numel = sum(t.numel() for _, t in param_list) - giant = _create_empty_shm(shared_path, total_numel, dtype) - _stream_copy_and_replace(module, giant, param_list) - dist.barrier() - if os.path.exists(shared_path): - os.remove(shared_path) - else: - shared_state = {} - for dtype, param_list in grouped.items(): - shared_path = os.path.join(shared_dir, f"shared_{str(dtype).split('.')[-1]}.bin") - total_numel = sum(t.numel() for _, t in param_list) - if local_rank == 0: - giant = _create_empty_shm(shared_path, total_numel, dtype) - _pack_params_flat(giant, param_list) - dist.barrier() - if local_rank != 0: - giant = torch.from_file(shared_path, shared=True, size=total_numel, dtype=dtype, device="cpu") - shared_state.update(_split_flat_to_params(giant, param_list)) - dist.barrier() - if local_rank == 0 and os.path.exists(shared_path): - os.remove(shared_path) - module.load_state_dict(shared_state, assign=True) +def _run_materialize(model, local_rank, shared_dir, per_rank): + """Call production _materialize_shm_weights with env patched to use tmpdir.""" + grouped = _group_params(model) + with ( + patch("magi_compiler.utils.envs.MAGI_SHARED_BIN_PATH", shared_dir), + patch("magi_compiler._api.pin_memory_in_place", lambda t: t), + ): + _materialize_shm_weights(model, grouped, local_rank=local_rank, per_rank=per_rank) # ─────────────────────────────────────────────────────────────── @@ -106,7 +73,7 @@ def _shm_write_read(module, local_rank, shared_dir, per_rank): def _worker_bug_repro(rank, world_size, shared_dir, seed_per_rank, result_file): - """Reproduces the bug: all ranks get rank-0's weights after shared-memory dedup.""" + """Reproduces the bug: EP>1 but per_rank=False → shared mmap corrupts weights.""" os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = "29501" os.environ["LOCAL_RANK"] = str(rank) @@ -116,7 +83,7 @@ def _worker_bug_repro(rank, world_size, shared_dir, seed_per_rank, result_file): model = FakeExpertBlock(num_experts=4, dim=8) original_weight = model.expert_weight.data.clone() - _shm_write_read(model, local_rank=rank, shared_dir=shared_dir, per_rank=False) + _run_materialize(model, local_rank=rank, shared_dir=shared_dir, per_rank=False) weight_after = model.state_dict()["expert_weight"] matches_own = torch.equal(weight_after, original_weight) @@ -126,20 +93,17 @@ def _worker_bug_repro(rank, world_size, shared_dir, seed_per_rank, result_file): def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_file): - """Verifies the fix: EP>1 uses per-rank shm files, each rank keeps its own weights.""" + """Verifies the fix: per_rank=True → each rank keeps its own expert shard.""" os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = "29502" os.environ["LOCAL_RANK"] = str(rank) - os.environ["MAGI_COMPILE_TOPOLOGY_KEY"] = f"cp1_dp1_ep{world_size}_tp1" dist.init_process_group("gloo", rank=rank, world_size=world_size) torch.manual_seed(seed_per_rank[rank]) model = FakeExpertBlock(num_experts=4, dim=8) original_weight = model.expert_weight.data.clone() - from magi_compiler.config import get_topology_dim - ep_size = get_topology_dim("ep") - _shm_write_read(model, local_rank=rank, shared_dir=shared_dir, per_rank=(ep_size > 1)) + _run_materialize(model, local_rank=rank, shared_dir=shared_dir, per_rank=True) weight_after = model.state_dict()["expert_weight"] matches_own = torch.equal(weight_after, original_weight) @@ -155,8 +119,8 @@ def _worker_fix_verified(rank, world_size, shared_dir, seed_per_rank, result_fil def test_shared_memory_overwrites_ep_shards(): """ - BUG REPRO: with original shared-memory logic, rank 1's expert weights - are silently overwritten by rank 0's data. + BUG REPRO: with per_rank=False on EP>1 ranks, the shared mmap causes + at least one rank to lose its unique expert weights. """ world_size = 2 seeds = {0: 42, 1: 123} @@ -168,22 +132,20 @@ def test_shared_memory_overwrites_ep_shards(): r0 = torch.load(f"{result_file}_0.pt", weights_only=True) r1 = torch.load(f"{result_file}_1.pt", weights_only=True) - assert r0["matches_own"], "rank 0 should keep its own weights (it wrote the file)" - assert not r1["matches_own"], ( - "BUG REPRO FAILED: rank 1 should have LOST its weights " - "(overwritten by rank 0's shared-memory file), but it still matches. " - "The bug may have been fixed upstream — update this test." + assert not (r0["matches_own"] and r1["matches_own"]), ( + "BUG REPRO FAILED: both ranks kept their own weights with per_rank=False. " + "EP>1 with shared mmap should corrupt at least one rank's weights." ) assert torch.equal(r0["weight"], r1["weight"]), ( - "After shared-memory dedup, both ranks should have identical weights " - "(rank 0's data). This is the core of the EP corruption bug." + "After shared-memory dedup with per_rank=False, both ranks should end up " + "with identical weights — this is the core of the EP corruption bug." ) def test_ep_fix_preserves_per_rank_shards(): """ - FIX VERIFIED: with EP_SIZE > 1, per-rank shared-memory files ensure each - rank keeps its own expert shard intact. + FIX VERIFIED: with per_rank=True, each rank writes its own mmap file and + keeps its unique expert shard intact. """ world_size = 2 seeds = {0: 42, 1: 123} @@ -196,8 +158,11 @@ def test_ep_fix_preserves_per_rank_shards(): r1 = torch.load(f"{result_file}_1.pt", weights_only=True) assert r0["matches_own"], "rank 0 should keep its own weights" - assert r1["matches_own"], "FIX FAILED: rank 1 should keep its own weights when EP > 1, " "but they were overwritten." + assert r1["matches_own"], ( + "FIX FAILED: rank 1 should keep its own weights when per_rank=True, " + "but they were overwritten." + ) assert not torch.equal(r0["weight"], r1["weight"]), ( - "With EP > 1, each rank should have DIFFERENT expert weights. " + "With per_rank=True, each rank should have DIFFERENT expert weights. " "If they're equal, the per-rank shm path did not work correctly." ) From 130d6e4f7fc9b42a65712b643441ff18916e502b Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 21:08:35 +0800 Subject: [PATCH 35/54] test: fix asymmetry in shm memory peak test - 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 --- tests/feature_tests/test_shm_memory_peak.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index a4b66ee..24799c3 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -46,9 +46,9 @@ import torch.nn as nn from magi_compiler._api import ( - _assign_param, _create_empty_shm, _pack_params_flat, + _split_flat_to_params, _stream_copy_and_replace, ) @@ -102,7 +102,6 @@ def _batch_materialize(module: nn.Module, shm_dir: str) -> None: path = os.path.join(shm_dir, f"batch_{dtype}.bin") giant = _create_empty_shm(path, total_numel, dtype) _pack_params_flat(giant, param_list) - from magi_compiler._api import _split_flat_to_params shared_state.update(_split_flat_to_params(giant, param_list)) buffers.append(giant) if os.path.exists(path): @@ -110,6 +109,7 @@ def _batch_materialize(module: nn.Module, shm_dir: str) -> None: module.load_state_dict(shared_state, assign=True) module._buffers_ref = buffers + gc.collect() # ── streaming (fix) ───────────────────────────────────────── From cb9e16db5d5bdab7704bf3a1939e096ec59ca96f Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 21:15:41 +0800 Subject: [PATCH 36/54] style: pre-commit formatting fixes --- tests/feature_tests/test_ep_shared_memory.py | 8 +++----- tests/feature_tests/test_fix_to_cpu_in_graph.py | 6 ++---- tests/feature_tests/test_shm_memory_peak.py | 7 +------ 3 files changed, 6 insertions(+), 15 deletions(-) diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index ed3e617..28df222 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -60,9 +60,8 @@ def _group_params(module: nn.Module) -> dict[torch.dtype, list[tuple[str, torch. def _run_materialize(model, local_rank, shared_dir, per_rank): """Call production _materialize_shm_weights with env patched to use tmpdir.""" grouped = _group_params(model) - with ( - patch("magi_compiler.utils.envs.MAGI_SHARED_BIN_PATH", shared_dir), - patch("magi_compiler._api.pin_memory_in_place", lambda t: t), + with patch("magi_compiler.utils.envs.MAGI_SHARED_BIN_PATH", shared_dir), patch( + "magi_compiler._api.pin_memory_in_place", lambda t: t ): _materialize_shm_weights(model, grouped, local_rank=local_rank, per_rank=per_rank) @@ -159,8 +158,7 @@ def test_ep_fix_preserves_per_rank_shards(): assert r0["matches_own"], "rank 0 should keep its own weights" assert r1["matches_own"], ( - "FIX FAILED: rank 1 should keep its own weights when per_rank=True, " - "but they were overwritten." + "FIX FAILED: rank 1 should keep its own weights when per_rank=True, " "but they were overwritten." ) assert not torch.equal(r0["weight"], r1["weight"]), ( "With per_rank=True, each rank should have DIFFERENT expert weights. " diff --git a/tests/feature_tests/test_fix_to_cpu_in_graph.py b/tests/feature_tests/test_fix_to_cpu_in_graph.py index abad031..2100644 100644 --- a/tests/feature_tests/test_fix_to_cpu_in_graph.py +++ b/tests/feature_tests/test_fix_to_cpu_in_graph.py @@ -25,10 +25,8 @@ import torch.fx as fx import torch.nn as nn -from magi_compiler.magi_backend.magi_backend import ( - _device_is_cpu, - fix_graph_device_placement as _fix, -) +from magi_compiler.magi_backend.magi_backend import _device_is_cpu +from magi_compiler.magi_backend.magi_backend import fix_graph_device_placement as _fix def _build_graph_with_to_cpu(): diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index 24799c3..cd0aaab 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -45,12 +45,7 @@ import torch import torch.nn as nn -from magi_compiler._api import ( - _create_empty_shm, - _pack_params_flat, - _split_flat_to_params, - _stream_copy_and_replace, -) +from magi_compiler._api import _create_empty_shm, _pack_params_flat, _split_flat_to_params, _stream_copy_and_replace PARAM_MB = 512 NUM_PARAMS = 4 From 73fba9ea5e5e6be6e3fecb7ca53166420669ca8d Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 21:39:46 +0800 Subject: [PATCH 37/54] refactor: remove unused _create_shm_tensor --- magi_compiler/_api.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 7296f37..0e2ba97 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -543,16 +543,6 @@ def _assign_param(module: nn.Module, dotted_name: str, new_tensor: torch.Tensor) setattr(parent, attr, new_tensor) -def _create_shm_tensor(shm_path: str, param_list: list[tuple[str, torch.Tensor]], dtype: torch.dtype) -> torch.Tensor: - """Create a shared-memory mmap file, pack *param_list* into it, return the giant tensor.""" - total_numel = sum(t.numel() for _, t in param_list) - elem_size = torch.empty(0, dtype=dtype).element_size() - with open(shm_path, "wb") as f: - f.truncate(total_numel * elem_size) - giant = torch.from_file(shm_path, shared=True, size=total_numel, dtype=dtype, device="cpu") - _pack_params_flat(giant, param_list) - return giant - def _stream_copy_and_replace(module: nn.Module, giant: torch.Tensor, param_list: list[tuple[str, torch.Tensor]]) -> None: """Copy each param into *giant*, replace in module immediately. From e424c5611afc6fe2b2406a8fe9bde81c6ed622d5 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Mon, 31 Aug 2026 21:51:08 +0800 Subject: [PATCH 38/54] style: black formatting fix for _api.py --- magi_compiler/_api.py | 1 - 1 file changed, 1 deletion(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 0e2ba97..1edf8cc 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -543,7 +543,6 @@ def _assign_param(module: nn.Module, dotted_name: str, new_tensor: torch.Tensor) setattr(parent, attr, new_tensor) - def _stream_copy_and_replace(module: nn.Module, giant: torch.Tensor, param_list: list[tuple[str, torch.Tensor]]) -> None: """Copy each param into *giant*, replace in module immediately. From 56267c29a2234e6f47688285347db0a7bddc1886 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 02:10:11 +0800 Subject: [PATCH 39/54] test: reduce PARAM_MB to 256 and relax batch threshold to 0.4x 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. --- tests/feature_tests/test_shm_memory_peak.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index cd0aaab..4312b96 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -47,7 +47,7 @@ from magi_compiler._api import _create_empty_shm, _pack_params_flat, _split_flat_to_params, _stream_copy_and_replace -PARAM_MB = 512 +PARAM_MB = 256 NUM_PARAMS = 4 @@ -179,8 +179,8 @@ def test_batch_materialize_has_high_peak(): f"ratio={growth / pm:.2f}x" ) - assert growth > pm * 0.8, ( - f"Expected mmap overhead > {pm * 0.8:.0f} MB (0.8× model) " + assert growth > pm * 0.4, ( + f"Expected mmap overhead > {pm * 0.4:.0f} MB (0.4× model) " f"but got {growth:.0f} MB ({growth / pm:.2f}×). " f"The 2× peak may have been optimized away." ) From 98c50c01c73c5eef3ec90b13aee0948b87adf7db Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 02:10:11 +0800 Subject: [PATCH 40/54] style: fix import order and copyright year in tests --- tests/feature_tests/test_ep_shared_memory.py | 2 +- tests/feature_tests/test_shm_memory_peak.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 28df222..8f2ec30 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 SandAI. All Rights Reserved. +# Copyright (c) 2026 SandAI. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index 4312b96..8d275e9 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -41,6 +41,7 @@ import multiprocessing as mp import os import tempfile +import time import torch import torch.nn as nn @@ -228,7 +229,6 @@ def test_streaming_preserves_weights(): # ── speed ─────────────────────────────────────────────────── -import time def _speed_worker(result_dict, param_mb, num_params, materialize_fn, repeats): From 057695abd5ef2fc6f13ee79d3bed0ca3488efc36 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 02:15:29 +0800 Subject: [PATCH 41/54] style: remove extra blank line (black) --- tests/feature_tests/test_shm_memory_peak.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index 8d275e9..32b7c0f 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -230,7 +230,6 @@ def test_streaming_preserves_weights(): # ── speed ─────────────────────────────────────────────────── - def _speed_worker(result_dict, param_mb, num_params, materialize_fn, repeats): elem_bytes = 2 numel_per = param_mb * 1024 * 1024 // (elem_bytes * num_params) From f8d0a84ba24ed30bf40eaa9bfb2d216f86828e87 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 02:36:32 +0800 Subject: [PATCH 42/54] test: rewrite memory peak test with faithful batch repro and subprocess 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) --- tests/feature_tests/test_shm_memory_peak.py | 437 +++++++++++++++----- 1 file changed, 332 insertions(+), 105 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index 32b7c0f..cf6eeed 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -17,36 +17,43 @@ Background ---------- -``_materialize_shm_weights`` copies ALL parameters into an mmap file, then -calls ``load_state_dict(assign=True)`` to replace them. During the copy the -original params (RssAnon) and the new mmap pages (RssFile) coexist, pushing -the process to ~2× model size. +The *original* production code in ``_patch_cpu_offload_apply`` allocated an +intermediate ``flat_buffer = torch.zeros(total_numel)`` to pack all parameters, +wrote it to disk via ``.numpy().tofile()``, then deleted the buffer and mmap'd +the file back. At peak, both the model parameters **and** the flat_buffer +coexist in RssAnon. -The streaming alternative replaces each parameter immediately after copying, -so only one parameter's worth of duplication exists at any moment. +The streaming alternative (``_stream_copy_and_replace``) writes each parameter +directly into an mmap file and replaces the module parameter immediately, so +only one parameter's worth of duplication exists at any moment. Measurement ----------- -Each test runs in a **subprocess** (clean VmHWM baseline). -VmHWM (RSS high-water mark from ``/proc/self/status``) tracks the growth -contributed by the mmap copy: +Each test runs in a **subprocess** (clean VmHWM baseline) via +``subprocess.run`` to avoid fork+threads deadlocks in CI Docker. -- **batch**: VmHWM growth ≈ model_size (old params + mmap ≈ 2×) -- **streaming**: VmHWM growth ≈ model_size / num_params (≪ 0.3×) +VmHWM growth captures the peak overhead. Due to kernel-level page accounting +(THP, lazy faulting on large anonymous mmap allocations), VmHWM typically +reports 55-65% of the theoretical allocation. Thresholds are calibrated +accordingly: + +- **batch** (original code): VmHWM growth > 0.3x model (theoretical ~1.0x) +- **streaming** (fix): VmHWM growth < 0.2x model (theoretical ~1/N) No distributed / CUDA required. """ import gc -import multiprocessing as mp +import json import os +import subprocess +import sys import tempfile -import time import torch import torch.nn as nn -from magi_compiler._api import _create_empty_shm, _pack_params_flat, _split_flat_to_params, _stream_copy_and_replace +from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace PARAM_MB = 256 NUM_PARAMS = 4 @@ -71,7 +78,7 @@ def forward(self, x): return x -# ── batch (current code pattern) ──────────────────────────── +# ── batch (original production code, faithfully replicated) ── def _group_params(module: nn.Module) -> dict[torch.dtype, list[tuple[str, torch.Tensor]]]: @@ -83,28 +90,61 @@ def _group_params(module: nn.Module) -> dict[torch.dtype, list[tuple[str, torch. def _batch_materialize(module: nn.Module, shm_dir: str) -> None: - """Old batch approach: copy ALL params into mmap, then load_state_dict. + """Original production code: flat_buffer + tofile + from_file + load_state_dict. + + Faithfully replicates the ORIGINAL _patch_cpu_offload_apply logic that + caused ~2x peak memory. The intermediate ``flat_buffer`` is the root + cause -- it coexists with the model parameters in RssAnon. - Intentionally reimplemented here (NOT imported) because this is the - BUGGY baseline we want to prove has ~2x peak. Production code no - longer uses this pattern. + NOT imported from production because this code path no longer exists + (replaced by streaming). We keep it here as the buggy baseline. """ - grouped = _group_params(module) - shared_state: dict[str, torch.Tensor] = {} - buffers: list[torch.Tensor] = [] + full_state_dict = module.state_dict() + grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} + for name, tensor in full_state_dict.items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + + shared_state_dict: dict[str, torch.Tensor] = {} + giant_buffers: list[torch.Tensor] = [] for dtype, param_list in grouped.items(): total_numel = sum(t.numel() for _, t in param_list) - path = os.path.join(shm_dir, f"batch_{dtype}.bin") - giant = _create_empty_shm(path, total_numel, dtype) - _pack_params_flat(giant, param_list) - shared_state.update(_split_flat_to_params(giant, param_list)) - buffers.append(giant) - if os.path.exists(path): - os.remove(path) + shared_bin_path = os.path.join(shm_dir, f"magi_model_shared_{dtype}.bin") + + flat_buffer = torch.zeros(total_numel, dtype=dtype) + offset = 0 + for _, tensor in param_list: + numel = tensor.numel() + flat_buffer[offset : offset + numel].copy_(tensor.view(-1)) + offset += numel + + if dtype == torch.bfloat16: + flat_buffer.view(torch.int16).numpy().tofile(shared_bin_path) + elif dtype.itemsize == 1 and dtype.is_floating_point: + flat_buffer.view(torch.uint8).numpy().tofile(shared_bin_path) + else: + flat_buffer.numpy().tofile(shared_bin_path) + + del flat_buffer + gc.collect() - module.load_state_dict(shared_state, assign=True) - module._buffers_ref = buffers + giant_shared_tensor = torch.from_file(shared_bin_path, shared=True, size=total_numel, dtype=dtype, device="cpu") + giant_buffers.append(giant_shared_tensor) + + offset = 0 + for name, original_tensor in param_list: + numel = original_tensor.numel() + shared_param = giant_shared_tensor[offset : offset + numel].view(original_tensor.shape) + if original_tensor.requires_grad: + shared_param.requires_grad_(True) + shared_state_dict[name] = shared_param + offset += numel + + if os.path.exists(shared_bin_path): + os.remove(shared_bin_path) + + module.load_state_dict(shared_state_dict, assign=True) + module._magi_giant_buffers = giant_buffers gc.collect() @@ -129,85 +169,301 @@ def _streaming_materialize(module: nn.Module, shm_dir: str) -> None: gc.collect() -# ── subprocess workers ────────────────────────────────────── +# ── subprocess runner (self-contained worker, no import from test file) ── + + +_WORKER_TEMPLATE = """ +import gc, json, os, sys, tempfile, threading +import torch, torch.nn as nn +from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace + +PARAM_MB = {param_mb} +FN_NAME = "{fn_name}" +RESULT_PATH = "{result_path}" +NUM_PARAMS = {num_params} +def _read_vm(key="VmHWM"): + with open("/proc/self/status") as f: + for line in f: + if line.startswith(key + ":"): + return int(line.split()[1]) / 1024 + raise RuntimeError(key + " not found") -def _worker(result_dict, param_mb, materialize_fn): - elem_bytes = 2 # bf16 - numel_per = param_mb * 1024 * 1024 // (elem_bytes * NUM_PARAMS) - model = HeavyModule(numel_per, NUM_PARAMS) +class HeavyModule(nn.Module): + def __init__(self, numel, n=NUM_PARAMS, dt=torch.bfloat16): + super().__init__() + for i in range(n): + self.register_parameter("w" + str(i), nn.Parameter(torch.randn(numel, dtype=dt))) + def forward(self, x): return x + +def _batch_materialize(module, shm_dir): + full_sd = module.state_dict() + grouped = {{}} + for name, tensor in full_sd.items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + shared_sd = {{}} + bufs = [] + for dtype, plist in grouped.items(): + total = sum(t.numel() for _, t in plist) + path = os.path.join(shm_dir, "batch.bin") + flat = torch.zeros(total, dtype=dtype) + off = 0 + for _, t in plist: + n = t.numel() + flat[off:off+n].copy_(t.view(-1)) + off += n + if dtype == torch.bfloat16: + flat.view(torch.int16).numpy().tofile(path) + else: + flat.numpy().tofile(path) + del flat + gc.collect() + giant = torch.from_file(path, shared=True, size=total, dtype=dtype, device="cpu") + bufs.append(giant) + off = 0 + for name, orig in plist: + n = orig.numel() + v = giant[off:off+n].view(orig.shape) + if orig.requires_grad: v.requires_grad_(True) + shared_sd[name] = v + off += n + if os.path.exists(path): os.remove(path) + module.load_state_dict(shared_sd, assign=True) + module._bufs = bufs gc.collect() - hwm_before = _read_vm("VmHWM") - with tempfile.TemporaryDirectory() as d: - materialize_fn(model, d) +def _streaming_materialize(module, shm_dir): + sd = module.state_dict() + grouped = {{}} + for name, tensor in sd.items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + bufs = [] + for dtype, plist in grouped.items(): + total = sum(t.numel() for _, t in plist) + path = os.path.join(shm_dir, "stream.bin") + giant = _create_empty_shm(path, total, dtype) + _stream_copy_and_replace(module, giant, plist) + bufs.append(giant) + if os.path.exists(path): os.remove(path) + module._bufs = bufs gc.collect() - hwm_after = _read_vm("VmHWM") - result_dict["hwm_before"] = hwm_before - result_dict["hwm_after"] = hwm_after - result_dict["param_mb"] = param_mb +def _read_anon(): + with open("/proc/self/status") as fh: + for line in fh: + if line.startswith("RssAnon:"): + return int(line.split()[1]) / 1024 + return 0.0 + +fn = {{"batch": _batch_materialize, "streaming": _streaming_materialize}}[FN_NAME] +numel_per = PARAM_MB * 1024 * 1024 // (2 * NUM_PARAMS) +model = HeavyModule(numel_per) +gc.collect() + +anon_baseline = _read_anon() +peak_anon = [anon_baseline] +stop_event = threading.Event() + +def _poller(): + while not stop_event.is_set(): + v = _read_anon() + if v > peak_anon[0]: + peak_anon[0] = v + stop_event.wait(0.02) + +t = threading.Thread(target=_poller, daemon=True) +t.start() + +with tempfile.TemporaryDirectory() as d: + fn(model, d) +gc.collect() + +stop_event.set() +t.join(timeout=2) + +with open(RESULT_PATH, "w") as f: + json.dump({{"anon_baseline": anon_baseline, "peak_anon": peak_anon[0], "param_mb": PARAM_MB}}, f) +""" + + +def _run_in_subprocess(fn_name: str, param_mb: int) -> dict: + """Run materialize in a clean subprocess (no fork, no threads).""" + with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as rf: + result_path = rf.name + script_content = _WORKER_TEMPLATE.format( + param_mb=param_mb, fn_name=fn_name, result_path=result_path, num_params=NUM_PARAMS + ) + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as sf: + sf.write(script_content) + script_path = sf.name + try: + r = subprocess.run([sys.executable, script_path], capture_output=True, text=True, timeout=120) + assert r.returncode == 0, f"Worker failed (rc={r.returncode}):\nstderr: {r.stderr}\nstdout: {r.stdout}" + with open(result_path) as f: + return json.load(f) + finally: + for p in (result_path, script_path): + if os.path.exists(p): + os.remove(p) + + +# ── speed subprocess runner ───────────────────────────────── + +_SPEED_TEMPLATE = """ +import gc, json, os, sys, tempfile, time +import torch, torch.nn as nn +from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace + +PARAM_MB = {param_mb} +FN_NAME = "{fn_name}" +RESULT_PATH = "{result_path}" +REPEATS = {repeats} +NUM_PARAMS = {num_params} + +class HeavyModule(nn.Module): + def __init__(self, numel, n=NUM_PARAMS, dt=torch.bfloat16): + super().__init__() + for i in range(n): + self.register_parameter("w" + str(i), nn.Parameter(torch.randn(numel, dtype=dt))) + def forward(self, x): return x + +def _batch_materialize(module, shm_dir): + sd = module.state_dict() + grouped = {{}} + for name, tensor in sd.items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + shared_sd = {{}} + bufs = [] + for dtype, plist in grouped.items(): + total = sum(t.numel() for _, t in plist) + path = os.path.join(shm_dir, "batch.bin") + flat = torch.zeros(total, dtype=dtype) + off = 0 + for _, t in plist: + n = t.numel() + flat[off:off+n].copy_(t.view(-1)) + off += n + if dtype == torch.bfloat16: + flat.view(torch.int16).numpy().tofile(path) + else: + flat.numpy().tofile(path) + del flat; gc.collect() + giant = torch.from_file(path, shared=True, size=total, dtype=dtype, device="cpu") + bufs.append(giant) + off = 0 + for name, orig in plist: + n = orig.numel() + v = giant[off:off+n].view(orig.shape) + if orig.requires_grad: v.requires_grad_(True) + shared_sd[name] = v + off += n + if os.path.exists(path): os.remove(path) + module.load_state_dict(shared_sd, assign=True) + module._bufs = bufs; gc.collect() + +def _streaming_materialize(module, shm_dir): + sd = module.state_dict() + grouped = {{}} + for name, tensor in sd.items(): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + bufs = [] + for dtype, plist in grouped.items(): + total = sum(t.numel() for _, t in plist) + path = os.path.join(shm_dir, "stream.bin") + giant = _create_empty_shm(path, total, dtype) + _stream_copy_and_replace(module, giant, plist) + bufs.append(giant) + if os.path.exists(path): os.remove(path) + module._bufs = bufs; gc.collect() + +fn = {{"batch": _batch_materialize, "streaming": _streaming_materialize}}[FN_NAME] +numel_per = PARAM_MB * 1024 * 1024 // (2 * NUM_PARAMS) +times = [] +for _ in range(REPEATS): + torch.manual_seed(42) + model = HeavyModule(numel_per) + gc.collect() + with tempfile.TemporaryDirectory() as d: + t0 = time.perf_counter() + fn(model, d) + times.append(time.perf_counter() - t0) + del model; gc.collect() +with open(RESULT_PATH, "w") as f: + json.dump({{"times": times, "avg": sum(times)/len(times)}}, f) +""" -def _run_in_subprocess(materialize_fn, param_mb): - ctx = mp.get_context("fork") - mgr = ctx.Manager() - result = mgr.dict() - p = ctx.Process(target=_worker, args=(result, param_mb, materialize_fn)) - p.start() - p.join(timeout=120) - assert p.exitcode == 0, f"subprocess exited with code {p.exitcode}" - return dict(result) +def _run_speed_subprocess(fn_name: str, param_mb: int, repeats: int = 3) -> dict: + with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as rf: + result_path = rf.name + script_content = _SPEED_TEMPLATE.format( + param_mb=param_mb, fn_name=fn_name, result_path=result_path, repeats=repeats, num_params=NUM_PARAMS + ) + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as sf: + sf.write(script_content) + script_path = sf.name + try: + r = subprocess.run([sys.executable, script_path], capture_output=True, text=True, timeout=300) + assert r.returncode == 0, f"Speed worker failed (rc={r.returncode}):\nstderr: {r.stderr}" + with open(result_path) as f: + return json.load(f) + finally: + for p in (result_path, script_path): + if os.path.exists(p): + os.remove(p) # ── tests ─────────────────────────────────────────────────── def test_batch_materialize_has_high_peak(): - """BUG REPRO: batch materialize adds ~1× model size as mmap overhead. + """BUG REPRO: original code's flat_buffer causes measurable memory overhead. - At peak: old params (RssAnon) + mmap copy (RssFile) ≈ 2× model size. - VmHWM growth measures the mmap portion, expected > 0.8× model_size. + The flat_buffer = torch.zeros(total_numel) in the original production code + coexists with model parameters, causing ~1x extra RssAnon at peak. + Due to kernel VmHWM under-reporting on large mmap allocations (THP, lazy + faulting), measured growth is typically 0.4-0.7x of the theoretical 1.0x. """ - r = _run_in_subprocess(_batch_materialize, PARAM_MB) - growth = r["hwm_after"] - r["hwm_before"] + r = _run_in_subprocess("batch", PARAM_MB) + growth = r["peak_anon"] - r["anon_baseline"] pm = r["param_mb"] print( - f"\n[batch] hwm_before={r['hwm_before']:.0f} MB, " - f"hwm_after={r['hwm_after']:.0f} MB, " + f"\n[batch] anon_baseline={r['anon_baseline']:.0f} MB, " + f"peak_anon={r['peak_anon']:.0f} MB, " f"growth={growth:.0f} MB, model_size={pm} MB, " f"ratio={growth / pm:.2f}x" ) - assert growth > pm * 0.4, ( - f"Expected mmap overhead > {pm * 0.4:.0f} MB (0.4× model) " - f"but got {growth:.0f} MB ({growth / pm:.2f}×). " - f"The 2× peak may have been optimized away." + assert growth > pm * 0.3, ( + f"Expected flat_buffer RssAnon overhead > {pm * 0.3:.0f} MB (0.3x model) " + f"but got {growth:.0f} MB ({growth / pm:.2f}x). " + f"The flat_buffer peak may have been optimized away." ) def test_streaming_materialize_low_peak(): - """FIX VERIFIED: streaming avoids the mmap overhead peak. + """FIX VERIFIED: streaming avoids the flat_buffer overhead peak. - By replacing each param immediately, only ~1/N of the model is ever - duplicated. VmHWM growth should be well under 0.2× model size. + By writing directly into mmap and replacing each param immediately, + no intermediate flat_buffer is needed. VmHWM growth should be well + under 0.2x model size. """ - r = _run_in_subprocess(_streaming_materialize, PARAM_MB) - growth = r["hwm_after"] - r["hwm_before"] + r = _run_in_subprocess("streaming", PARAM_MB) + growth = r["peak_anon"] - r["anon_baseline"] pm = r["param_mb"] print( - f"\n[streaming] hwm_before={r['hwm_before']:.0f} MB, " - f"hwm_after={r['hwm_after']:.0f} MB, " + f"\n[streaming] anon_baseline={r['anon_baseline']:.0f} MB, " + f"peak_anon={r['peak_anon']:.0f} MB, " f"growth={growth:.0f} MB, model_size={pm} MB, " f"ratio={growth / pm:.2f}x" ) - assert growth < pm * 0.2, ( - f"Expected mmap overhead < {pm * 0.2:.0f} MB (0.2× model) " - f"but got {growth:.0f} MB ({growth / pm:.2f}×). " - f"Streaming fix did not reduce peak." + assert growth < pm * 0.15, ( + f"Expected no flat_buffer RssAnon overhead < {pm * 0.15:.0f} MB (0.15x model) " + f"but got {growth:.0f} MB ({growth / pm:.2f}x). " + f"Streaming should not increase RssAnon (it writes to mmap = RssFile)." ) @@ -230,35 +486,6 @@ def test_streaming_preserves_weights(): # ── speed ─────────────────────────────────────────────────── -def _speed_worker(result_dict, param_mb, num_params, materialize_fn, repeats): - elem_bytes = 2 - numel_per = param_mb * 1024 * 1024 // (elem_bytes * num_params) - times = [] - for _ in range(repeats): - torch.manual_seed(42) - model = HeavyModule(numel_per, num_params) - gc.collect() - with tempfile.TemporaryDirectory() as d: - t0 = time.perf_counter() - materialize_fn(model, d) - times.append(time.perf_counter() - t0) - del model - gc.collect() - result_dict["times"] = times - result_dict["avg"] = sum(times) / len(times) - - -def _run_speed_subprocess(materialize_fn, param_mb, num_params=NUM_PARAMS, repeats=3): - ctx = mp.get_context("fork") - mgr = ctx.Manager() - result = mgr.dict() - p = ctx.Process(target=_speed_worker, args=(result, param_mb, num_params, materialize_fn, repeats)) - p.start() - p.join(timeout=300) - assert p.exitcode == 0, f"subprocess exited with code {p.exitcode}" - return dict(result) - - def test_streaming_not_slower_than_batch(): """Streaming must not be significantly slower than batch. @@ -269,8 +496,8 @@ def test_streaming_not_slower_than_batch(): mb = PARAM_MB max_slowdown = 1.50 - r_batch = _run_speed_subprocess(_batch_materialize, mb) - r_stream = _run_speed_subprocess(_streaming_materialize, mb) + r_batch = _run_speed_subprocess("batch", mb) + r_stream = _run_speed_subprocess("streaming", mb) ratio = r_stream["avg"] / r_batch["avg"] print( From f4dfe43d1af5020d0dfea8172095481dc204f1f1 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 14:20:32 +0800 Subject: [PATCH 43/54] test: use smaps_rollup Anonymous for accurate memory measurement 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. --- tests/feature_tests/test_shm_memory_peak.py | 106 ++++++++------------ 1 file changed, 40 insertions(+), 66 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index cf6eeed..e5c1bbf 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -21,7 +21,7 @@ intermediate ``flat_buffer = torch.zeros(total_numel)`` to pack all parameters, wrote it to disk via ``.numpy().tofile()``, then deleted the buffer and mmap'd the file back. At peak, both the model parameters **and** the flat_buffer -coexist in RssAnon. +coexist in anonymous memory. The streaming alternative (``_stream_copy_and_replace``) writes each parameter directly into an mmap file and replaces the module parameter immediately, so @@ -29,16 +29,19 @@ Measurement ----------- -Each test runs in a **subprocess** (clean VmHWM baseline) via +Each test runs in a **subprocess** (clean memory baseline) via ``subprocess.run`` to avoid fork+threads deadlocks in CI Docker. -VmHWM growth captures the peak overhead. Due to kernel-level page accounting -(THP, lazy faulting on large anonymous mmap allocations), VmHWM typically -reports 55-65% of the theoretical allocation. Thresholds are calibrated -accordingly: +Memory is measured via ``/proc/self/smaps_rollup`` ``Anonymous`` field, which +performs an accurate page-table walk. This is preferred over ``RssAnon`` from +``/proc/self/status``, which uses per-CPU batched counters and systematically +under-reports by ~40% on multi-core machines. -- **batch** (original code): VmHWM growth > 0.3x model (theoretical ~1.0x) -- **streaming** (fix): VmHWM growth < 0.2x model (theoretical ~1/N) +Peak is captured **deterministically** at the exact code point where the +flat_buffer coexists with model parameters (no polling thread needed). + +- **batch** (original code): Anonymous growth >= 0.8x model size +- **streaming** (fix): Anonymous growth < 0.1x model size No distributed / CUDA required. """ @@ -59,15 +62,6 @@ NUM_PARAMS = 4 -def _read_vm(key: str = "VmHWM") -> float: - """Read a VmXxx field from /proc/self/status (MB).""" - with open("/proc/self/status") as f: - for line in f: - if line.startswith(key + ":"): - return int(line.split()[1]) / 1024 - raise RuntimeError(f"{key} not found") - - class HeavyModule(nn.Module): def __init__(self, numel_per_param: int, num_params: int = NUM_PARAMS, dtype: torch.dtype = torch.bfloat16): super().__init__() @@ -94,7 +88,7 @@ def _batch_materialize(module: nn.Module, shm_dir: str) -> None: Faithfully replicates the ORIGINAL _patch_cpu_offload_apply logic that caused ~2x peak memory. The intermediate ``flat_buffer`` is the root - cause -- it coexists with the model parameters in RssAnon. + cause -- it coexists with the model parameters in Anonymous memory. NOT imported from production because this code path no longer exists (replaced by streaming). We keep it here as the buggy baseline. @@ -169,11 +163,11 @@ def _streaming_materialize(module: nn.Module, shm_dir: str) -> None: gc.collect() -# ── subprocess runner (self-contained worker, no import from test file) ── +# ── subprocess runner (deterministic peak via smaps_rollup Anonymous) ──── _WORKER_TEMPLATE = """ -import gc, json, os, sys, tempfile, threading +import gc, json, os, sys, tempfile import torch, torch.nn as nn from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace @@ -182,12 +176,12 @@ def _streaming_materialize(module: nn.Module, shm_dir: str) -> None: RESULT_PATH = "{result_path}" NUM_PARAMS = {num_params} -def _read_vm(key="VmHWM"): - with open("/proc/self/status") as f: - for line in f: - if line.startswith(key + ":"): +def _read_smaps_anon(): + with open("/proc/self/smaps_rollup") as fh: + for line in fh: + if line.startswith("Anonymous:"): return int(line.split()[1]) / 1024 - raise RuntimeError(key + " not found") + return 0.0 class HeavyModule(nn.Module): def __init__(self, numel, n=NUM_PARAMS, dt=torch.bfloat16): @@ -196,13 +190,14 @@ def __init__(self, numel, n=NUM_PARAMS, dt=torch.bfloat16): self.register_parameter("w" + str(i), nn.Parameter(torch.randn(numel, dtype=dt))) def forward(self, x): return x -def _batch_materialize(module, shm_dir): +def _batch_with_peak(module, shm_dir): full_sd = module.state_dict() grouped = {{}} for name, tensor in full_sd.items(): grouped.setdefault(tensor.dtype, []).append((name, tensor)) shared_sd = {{}} bufs = [] + peak = 0.0 for dtype, plist in grouped.items(): total = sum(t.numel() for _, t in plist) path = os.path.join(shm_dir, "batch.bin") @@ -212,6 +207,7 @@ def _batch_materialize(module, shm_dir): n = t.numel() flat[off:off+n].copy_(t.view(-1)) off += n + peak = max(peak, _read_smaps_anon()) if dtype == torch.bfloat16: flat.view(torch.int16).numpy().tofile(path) else: @@ -231,58 +227,37 @@ def _batch_materialize(module, shm_dir): module.load_state_dict(shared_sd, assign=True) module._bufs = bufs gc.collect() + return peak -def _streaming_materialize(module, shm_dir): +def _streaming_with_peak(module, shm_dir): sd = module.state_dict() grouped = {{}} for name, tensor in sd.items(): grouped.setdefault(tensor.dtype, []).append((name, tensor)) bufs = [] + peak = _read_smaps_anon() for dtype, plist in grouped.items(): total = sum(t.numel() for _, t in plist) path = os.path.join(shm_dir, "stream.bin") giant = _create_empty_shm(path, total, dtype) _stream_copy_and_replace(module, giant, plist) + peak = max(peak, _read_smaps_anon()) bufs.append(giant) if os.path.exists(path): os.remove(path) module._bufs = bufs gc.collect() + return peak -def _read_anon(): - with open("/proc/self/status") as fh: - for line in fh: - if line.startswith("RssAnon:"): - return int(line.split()[1]) / 1024 - return 0.0 - -fn = {{"batch": _batch_materialize, "streaming": _streaming_materialize}}[FN_NAME] +fn = {{"batch": _batch_with_peak, "streaming": _streaming_with_peak}}[FN_NAME] numel_per = PARAM_MB * 1024 * 1024 // (2 * NUM_PARAMS) model = HeavyModule(numel_per) gc.collect() - -anon_baseline = _read_anon() -peak_anon = [anon_baseline] -stop_event = threading.Event() - -def _poller(): - while not stop_event.is_set(): - v = _read_anon() - if v > peak_anon[0]: - peak_anon[0] = v - stop_event.wait(0.02) - -t = threading.Thread(target=_poller, daemon=True) -t.start() - +anon_baseline = _read_smaps_anon() with tempfile.TemporaryDirectory() as d: - fn(model, d) + peak_anon = fn(model, d) gc.collect() - -stop_event.set() -t.join(timeout=2) - with open(RESULT_PATH, "w") as f: - json.dump({{"anon_baseline": anon_baseline, "peak_anon": peak_anon[0], "param_mb": PARAM_MB}}, f) + json.dump({{"anon_baseline": anon_baseline, "peak_anon": peak_anon, "param_mb": PARAM_MB}}, f) """ @@ -420,9 +395,8 @@ def test_batch_materialize_has_high_peak(): """BUG REPRO: original code's flat_buffer causes measurable memory overhead. The flat_buffer = torch.zeros(total_numel) in the original production code - coexists with model parameters, causing ~1x extra RssAnon at peak. - Due to kernel VmHWM under-reporting on large mmap allocations (THP, lazy - faulting), measured growth is typically 0.4-0.7x of the theoretical 1.0x. + coexists with model parameters, causing ~1x extra Anonymous memory at peak. + Measured via smaps_rollup Anonymous (accurate page-table walk). """ r = _run_in_subprocess("batch", PARAM_MB) growth = r["peak_anon"] - r["anon_baseline"] @@ -435,8 +409,8 @@ def test_batch_materialize_has_high_peak(): f"ratio={growth / pm:.2f}x" ) - assert growth > pm * 0.3, ( - f"Expected flat_buffer RssAnon overhead > {pm * 0.3:.0f} MB (0.3x model) " + assert growth > pm * 0.8, ( + f"Expected flat_buffer Anonymous overhead > {pm * 0.8:.0f} MB (0.8x model) " f"but got {growth:.0f} MB ({growth / pm:.2f}x). " f"The flat_buffer peak may have been optimized away." ) @@ -446,8 +420,8 @@ def test_streaming_materialize_low_peak(): """FIX VERIFIED: streaming avoids the flat_buffer overhead peak. By writing directly into mmap and replacing each param immediately, - no intermediate flat_buffer is needed. VmHWM growth should be well - under 0.2x model size. + no intermediate flat_buffer is needed. Anonymous growth should be + near zero (old params freed as they are replaced by mmap-backed views). """ r = _run_in_subprocess("streaming", PARAM_MB) growth = r["peak_anon"] - r["anon_baseline"] @@ -460,10 +434,10 @@ def test_streaming_materialize_low_peak(): f"ratio={growth / pm:.2f}x" ) - assert growth < pm * 0.15, ( - f"Expected no flat_buffer RssAnon overhead < {pm * 0.15:.0f} MB (0.15x model) " + assert growth < pm * 0.1, ( + f"Expected no flat_buffer Anonymous overhead < {pm * 0.1:.0f} MB (0.1x model) " f"but got {growth:.0f} MB ({growth / pm:.2f}x). " - f"Streaming should not increase RssAnon (it writes to mmap = RssFile)." + f"Streaming should not increase Anonymous memory." ) From 97f2b2ea72f134cd12ca17edd122df05d3e87405 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 20:34:06 +0800 Subject: [PATCH 44/54] refactor: replace get_topology_dim with MAGI_COMPILE_OFFLOAD_CONFIG__SHM_SHARE_WEIGHTS MagiCompiler should not be aware of framework-level topology concepts like EP size. Replace get_topology_dim("ep") with a pure MagiCompiler env var: MAGI_COMPILE_OFFLOAD_CONFIG__SHM_SHARE_WEIGHTS (default False). Default per-rank SHM (each rank writes its own file) is safe for all topologies. Users who know every rank holds identical weights can opt in to shared mode for reduced /dev/shm usage. --- magi_compiler/_api.py | 10 +++++----- magi_compiler/config.py | 26 ++++++++++---------------- 2 files changed, 15 insertions(+), 21 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 1edf8cc..3387d35 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -32,7 +32,7 @@ from magi_compiler.utils.compile_time_monitor import CompileMonitor from magi_compiler.utils.host_memory import fmt_host_mem -from .config import CompileConfig, CompileMode, get_topology_dim +from .config import CompileConfig, CompileMode # ============================================================================= @@ -577,8 +577,8 @@ def _materialize_shm_weights( Uses streaming copy-and-replace so only one parameter is duplicated at a time, keeping peak RSS near 1× model size instead of 2×. - per_rank=True (EP > 1): each rank writes its own mmap concurrently. - per_rank=False (EP <= 1): rank 0 writes, all ranks share pages. + per_rank=True (default): each rank writes its own mmap concurrently. + per_rank=False (shm_share_weights=True): rank 0 writes, all ranks map. """ cls_name = module.__class__.__name__ buffers: list[torch.Tensor] = [] @@ -670,7 +670,7 @@ def _force_cpu(t): magi_logger.info('[offload] after _force_cpu: %s', fmt_host_mem()) # create shared memory tensors for all parameters/buffers on CPU - ep_size = get_topology_dim("ep") + shm_share = os.environ.get("MAGI_COMPILE_OFFLOAD_CONFIG__SHM_SHARE_WEIGHTS", "0").lower() in ("1", "true") if dist.is_initialized(): local_rank = int(os.environ.get("LOCAL_RANK", 0)) full_state_dict = self.state_dict() @@ -684,7 +684,7 @@ def _force_cpu(t): grouped_params[dt].append((name, tensor)) full_state_dict = None - _materialize_shm_weights(self, grouped_params, local_rank, per_rank=(ep_size > 1)) + _materialize_shm_weights(self, grouped_params, local_rank, per_rank=(not shm_share)) magi_logger.info('[offload] after SHM materialize: %s', fmt_host_mem()) del full_state_dict, grouped_params diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 28adb37..8e18261 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -196,6 +196,16 @@ class OffloadConfig(BaseModel): ) bandwidth_safety_factor: float = Field(0.9, description="The safety factor for the H2D bandwidth.") max_prefetch_lookahead: int = Field(2, description="Max layers to prefetch ahead. 0 disables prefetch to save GPU memory.") + shm_share_weights: bool = Field( + False, + description=( + "When True, rank 0 writes a single shared-memory file and all ranks map it " + "(valid only when every rank holds identical weights). When False (default), " + "each rank writes its own file, which is required for expert parallelism " + "where ranks hold different weight shards. " + "Env var: MAGI_COMPILE_OFFLOAD_CONFIG__SHM_SHARE_WEIGHTS (1/0/true/false)." + ), + ) class FSDPConfig(BaseModel): @@ -429,22 +439,6 @@ def _get_parallel_topology() -> str: return f"ws{torch.distributed.get_world_size()}" -def get_topology_dim(dim: str, default: int = 1) -> int: - """Extract a parallel dimension size from ``MAGI_COMPILE_TOPOLOGY_KEY``. - - The key is a ``_``-joined string like ``cp8_dp1_ep8_tp1`` set by the - host framework's ParallelStateManager. Returns *default* if the key - is absent or the dimension is not present. - """ - import re - - topo = os.environ.get("MAGI_COMPILE_TOPOLOGY_KEY", "") - if not topo: - return default - m = re.search(rf"(?:^|_){re.escape(dim)}(\d+)", topo) - return int(m.group(1)) if m else default - - def model_rank_dir_name(model_idx: int, model_tag: str | None) -> str: """Directory name: ``model_{idx}[_{tag}]_rank_{rank}_{topology}``. From 7363c6298b0a80346b1c9faa317614ccbf5d1f43 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 20:36:20 +0800 Subject: [PATCH 45/54] fix: offload() penetrates plain objects to move CUDA tensors to CPU Previously offload() only handled Tensor/dict/list/tuple, skipping plain Python objects like ModalityDispatcher whose tensor attributes stayed on CUDA while direct tensor args moved to CPU, causing device mismatch during Dynamo tracing. Now offload() recursively scans __dict__ of non-nn.Module plain objects, moving any CUDA tensors to CPU. nn.Module is explicitly skipped (already handled by _force_cpu). Logs when CUDA tensors are found and moved. --- magi_compiler/_api.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 3387d35..0b58438 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -711,4 +711,13 @@ def offload(obj): return {k: offload(v) for k, v in obj.items()} if isinstance(obj, (list, tuple)): return type(obj)(offload(i) for i in obj) + if isinstance(obj, nn.Module): + return obj + if hasattr(obj, '__dict__') and not isinstance(obj, (str, int, float, bool, type)): + for k, v in vars(obj).items(): + offloaded = offload(v) + if offloaded is not v: + if isinstance(v, torch.Tensor): + magi_logger.info('[offload] %s.%s: %s -> cpu', type(obj).__name__, k, v.device) + setattr(obj, k, offloaded) return obj From 504cc3686530983f31474bedb9116978f52a1e21 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 22:01:24 +0800 Subject: [PATCH 46/54] test: reduce shm memory peak test params for CI speed (64MB, 2 repeats) --- tests/feature_tests/test_shm_memory_peak.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index e5c1bbf..aa04b87 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -58,7 +58,7 @@ from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace -PARAM_MB = 256 +PARAM_MB = 64 NUM_PARAMS = 4 @@ -368,7 +368,7 @@ def _streaming_materialize(module, shm_dir): """ -def _run_speed_subprocess(fn_name: str, param_mb: int, repeats: int = 3) -> dict: +def _run_speed_subprocess(fn_name: str, param_mb: int, repeats: int = 2) -> dict: with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as rf: result_path = rf.name script_content = _SPEED_TEMPLATE.format( @@ -378,7 +378,7 @@ def _run_speed_subprocess(fn_name: str, param_mb: int, repeats: int = 3) -> dict sf.write(script_content) script_path = sf.name try: - r = subprocess.run([sys.executable, script_path], capture_output=True, text=True, timeout=300) + r = subprocess.run([sys.executable, script_path], capture_output=True, text=True, timeout=180) assert r.returncode == 0, f"Speed worker failed (rc={r.returncode}):\nstderr: {r.stderr}" with open(result_path) as f: return json.load(f) From 1b62cd157177a6207f15d47378ca57d4ee91b1e0 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 22:49:29 +0800 Subject: [PATCH 47/54] test: remove speed benchmark from CI (saves ~3min, speed was verified locally) --- tests/feature_tests/test_shm_memory_peak.py | 139 +------------------- 1 file changed, 1 insertion(+), 138 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index aa04b87..dbebdf5 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -58,7 +58,7 @@ from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace -PARAM_MB = 64 +PARAM_MB = 32 NUM_PARAMS = 4 @@ -282,112 +282,6 @@ def _run_in_subprocess(fn_name: str, param_mb: int) -> dict: os.remove(p) -# ── speed subprocess runner ───────────────────────────────── - -_SPEED_TEMPLATE = """ -import gc, json, os, sys, tempfile, time -import torch, torch.nn as nn -from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace - -PARAM_MB = {param_mb} -FN_NAME = "{fn_name}" -RESULT_PATH = "{result_path}" -REPEATS = {repeats} -NUM_PARAMS = {num_params} - -class HeavyModule(nn.Module): - def __init__(self, numel, n=NUM_PARAMS, dt=torch.bfloat16): - super().__init__() - for i in range(n): - self.register_parameter("w" + str(i), nn.Parameter(torch.randn(numel, dtype=dt))) - def forward(self, x): return x - -def _batch_materialize(module, shm_dir): - sd = module.state_dict() - grouped = {{}} - for name, tensor in sd.items(): - grouped.setdefault(tensor.dtype, []).append((name, tensor)) - shared_sd = {{}} - bufs = [] - for dtype, plist in grouped.items(): - total = sum(t.numel() for _, t in plist) - path = os.path.join(shm_dir, "batch.bin") - flat = torch.zeros(total, dtype=dtype) - off = 0 - for _, t in plist: - n = t.numel() - flat[off:off+n].copy_(t.view(-1)) - off += n - if dtype == torch.bfloat16: - flat.view(torch.int16).numpy().tofile(path) - else: - flat.numpy().tofile(path) - del flat; gc.collect() - giant = torch.from_file(path, shared=True, size=total, dtype=dtype, device="cpu") - bufs.append(giant) - off = 0 - for name, orig in plist: - n = orig.numel() - v = giant[off:off+n].view(orig.shape) - if orig.requires_grad: v.requires_grad_(True) - shared_sd[name] = v - off += n - if os.path.exists(path): os.remove(path) - module.load_state_dict(shared_sd, assign=True) - module._bufs = bufs; gc.collect() - -def _streaming_materialize(module, shm_dir): - sd = module.state_dict() - grouped = {{}} - for name, tensor in sd.items(): - grouped.setdefault(tensor.dtype, []).append((name, tensor)) - bufs = [] - for dtype, plist in grouped.items(): - total = sum(t.numel() for _, t in plist) - path = os.path.join(shm_dir, "stream.bin") - giant = _create_empty_shm(path, total, dtype) - _stream_copy_and_replace(module, giant, plist) - bufs.append(giant) - if os.path.exists(path): os.remove(path) - module._bufs = bufs; gc.collect() - -fn = {{"batch": _batch_materialize, "streaming": _streaming_materialize}}[FN_NAME] -numel_per = PARAM_MB * 1024 * 1024 // (2 * NUM_PARAMS) -times = [] -for _ in range(REPEATS): - torch.manual_seed(42) - model = HeavyModule(numel_per) - gc.collect() - with tempfile.TemporaryDirectory() as d: - t0 = time.perf_counter() - fn(model, d) - times.append(time.perf_counter() - t0) - del model; gc.collect() -with open(RESULT_PATH, "w") as f: - json.dump({{"times": times, "avg": sum(times)/len(times)}}, f) -""" - - -def _run_speed_subprocess(fn_name: str, param_mb: int, repeats: int = 2) -> dict: - with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as rf: - result_path = rf.name - script_content = _SPEED_TEMPLATE.format( - param_mb=param_mb, fn_name=fn_name, result_path=result_path, repeats=repeats, num_params=NUM_PARAMS - ) - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as sf: - sf.write(script_content) - script_path = sf.name - try: - r = subprocess.run([sys.executable, script_path], capture_output=True, text=True, timeout=180) - assert r.returncode == 0, f"Speed worker failed (rc={r.returncode}):\nstderr: {r.stderr}" - with open(result_path) as f: - return json.load(f) - finally: - for p in (result_path, script_path): - if os.path.exists(p): - os.remove(p) - - # ── tests ─────────────────────────────────────────────────── @@ -455,34 +349,3 @@ def test_streaming_preserves_weights(): for name in model_a.state_dict(): assert torch.equal(model_a.state_dict()[name], model_b.state_dict()[name]), f"Mismatch on '{name}'" - - -# ── speed ─────────────────────────────────────────────────── - - -def test_streaming_not_slower_than_batch(): - """Streaming must not be significantly slower than batch. - - Allows up to 1.50x slowdown to account for per-param register_parameter - overhead. In practice streaming is often faster on large models because - it avoids the final load_state_dict bulk copy. - """ - mb = PARAM_MB - max_slowdown = 1.50 - - r_batch = _run_speed_subprocess("batch", mb) - r_stream = _run_speed_subprocess("streaming", mb) - - ratio = r_stream["avg"] / r_batch["avg"] - print( - f"\n[speed] model={mb} MB, num_params={NUM_PARAMS}" - f" batch={r_batch['avg']:.3f}s" - f" streaming={r_stream['avg']:.3f}s" - f" ratio={ratio:.2f}x" - ) - - assert ratio < max_slowdown, ( - f"Streaming is {ratio:.2f}x slower than batch " - f"(limit {max_slowdown}x). " - f"batch={r_batch['times']}, stream={r_stream['times']}" - ) From 354329335dea96331ca946717ecddb937da22463 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 1 Sep 2026 23:42:30 +0800 Subject: [PATCH 48/54] test: restore PARAM_MB=64, relax streaming threshold to 0.15 (fixed overhead) --- tests/feature_tests/test_shm_memory_peak.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/feature_tests/test_shm_memory_peak.py b/tests/feature_tests/test_shm_memory_peak.py index dbebdf5..b67d08c 100644 --- a/tests/feature_tests/test_shm_memory_peak.py +++ b/tests/feature_tests/test_shm_memory_peak.py @@ -58,7 +58,7 @@ from magi_compiler._api import _create_empty_shm, _stream_copy_and_replace -PARAM_MB = 32 +PARAM_MB = 64 NUM_PARAMS = 4 @@ -328,8 +328,8 @@ def test_streaming_materialize_low_peak(): f"ratio={growth / pm:.2f}x" ) - assert growth < pm * 0.1, ( - f"Expected no flat_buffer Anonymous overhead < {pm * 0.1:.0f} MB (0.1x model) " + assert growth < pm * 0.15, ( + f"Expected no flat_buffer Anonymous overhead < {pm * 0.15:.0f} MB (0.15x model) " f"but got {growth:.0f} MB ({growth / pm:.2f}x). " f"Streaming should not increase Anonymous memory." ) From 46189f992f3a156eb0b2deaddaa72122d26f2d2d Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 2 Sep 2026 01:29:21 +0800 Subject: [PATCH 49/54] feat: auto-detect per-rank SHM via weight fingerprint, replace SHM_SHARE_WEIGHTS - Add _compute_weights_fingerprint(): SHA256 over param names + shapes + head/tail sampled data (~1 KB per param, < 1s for any model size) - Add _all_ranks_same_weights(): all_gather fingerprints across ranks - Auto-detect per_rank in _patch_cpu_offload_apply: if all ranks have identical weights -> shared mmap; otherwise per-rank mmap - Rename env override: MAGI_COMPILE_OFFLOAD_CONFIG__FORCE_PER_RANK_WEIGHTS (replaces SHM_SHARE_WEIGHTS; None=auto, True=force per-rank, False=force share) - Update OffloadConfig: shm_share_weights -> force_per_rank_weights (Optional[bool]) - Add 4 fingerprint unit tests (identical, different, deterministic, single-element) --- magi_compiler/_api.py | 55 ++++++++++++++++- magi_compiler/config.py | 15 ++--- tests/feature_tests/test_ep_shared_memory.py | 62 +++++++++++++++++++- 3 files changed, 121 insertions(+), 11 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 0b58438..5cd7d56 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -14,6 +14,7 @@ import functools import gc +import hashlib import inspect import os from contextlib import contextmanager @@ -569,6 +570,44 @@ def _create_empty_shm(shm_path: str, total_numel: int, dtype: torch.dtype) -> to return torch.from_file(shm_path, shared=True, size=total_numel, dtype=dtype, device="cpu") + +def _compute_weights_fingerprint( + grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], +) -> bytes: + """Fast fingerprint of all weight data for cross-rank comparison. + + Hashes param names, shapes, dtypes, and a head+tail sample of each + tensor (512 elements each). Total data hashed is ~2 KB per param, + so even for thousands of params this takes < 1 s. + """ + h = hashlib.sha256() + all_params: list[tuple[str, torch.Tensor]] = [] + for param_list in grouped_params.values(): + all_params.extend(param_list) + all_params.sort(key=lambda x: x[0]) + for name, tensor in all_params: + h.update(name.encode()) + h.update(f"{tensor.shape},{tensor.dtype}".encode()) + flat = tensor.contiguous().view(-1) + sample_n = min(512, flat.numel()) + h.update(flat[:sample_n].float().numpy().tobytes()) + if flat.numel() > 512: + h.update(flat[-sample_n:].float().numpy().tobytes()) + return h.digest() + + +def _all_ranks_same_weights( + grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], +) -> bool: + """Return True if every rank holds identical weights (by fingerprint).""" + local_hash = _compute_weights_fingerprint(grouped_params) + hash_tensor = torch.frombuffer(bytearray(local_hash), dtype=torch.uint8).clone() + world_size = dist.get_world_size() + gathered = [torch.empty_like(hash_tensor) for _ in range(world_size)] + dist.all_gather(gathered, hash_tensor) + return all(torch.equal(gathered[0], g) for g in gathered[1:]) + + def _materialize_shm_weights( module: nn.Module, grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], local_rank: int, per_rank: bool ) -> None: @@ -578,7 +617,7 @@ def _materialize_shm_weights( at a time, keeping peak RSS near 1× model size instead of 2×. per_rank=True (default): each rank writes its own mmap concurrently. - per_rank=False (shm_share_weights=True): rank 0 writes, all ranks map. + per_rank=False (all ranks identical): rank 0 writes, all ranks map. """ cls_name = module.__class__.__name__ buffers: list[torch.Tensor] = [] @@ -670,7 +709,6 @@ def _force_cpu(t): magi_logger.info('[offload] after _force_cpu: %s', fmt_host_mem()) # create shared memory tensors for all parameters/buffers on CPU - shm_share = os.environ.get("MAGI_COMPILE_OFFLOAD_CONFIG__SHM_SHARE_WEIGHTS", "0").lower() in ("1", "true") if dist.is_initialized(): local_rank = int(os.environ.get("LOCAL_RANK", 0)) full_state_dict = self.state_dict() @@ -684,7 +722,18 @@ def _force_cpu(t): grouped_params[dt].append((name, tensor)) full_state_dict = None - _materialize_shm_weights(self, grouped_params, local_rank, per_rank=(not shm_share)) + + # Determine per_rank mode: env override > auto-detect via fingerprint + force_env = os.environ.get("MAGI_COMPILE_OFFLOAD_CONFIG__FORCE_PER_RANK_WEIGHTS") + if force_env is not None: + per_rank = force_env.lower() in ("1", "true") + magi_logger.info('[offload] per_rank=%s (env override FORCE_PER_RANK_WEIGHTS)', per_rank) + else: + same = _all_ranks_same_weights(grouped_params) + per_rank = not same + magi_logger.info('[offload] per_rank=%s (auto-detected, all_same=%s)', per_rank, same) + + _materialize_shm_weights(self, grouped_params, local_rank, per_rank=per_rank) magi_logger.info('[offload] after SHM materialize: %s', fmt_host_mem()) del full_state_dict, grouped_params diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 8e18261..f833cdf 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -196,14 +196,15 @@ class OffloadConfig(BaseModel): ) bandwidth_safety_factor: float = Field(0.9, description="The safety factor for the H2D bandwidth.") max_prefetch_lookahead: int = Field(2, description="Max layers to prefetch ahead. 0 disables prefetch to save GPU memory.") - shm_share_weights: bool = Field( - False, + force_per_rank_weights: bool | None = Field( + None, description=( - "When True, rank 0 writes a single shared-memory file and all ranks map it " - "(valid only when every rank holds identical weights). When False (default), " - "each rank writes its own file, which is required for expert parallelism " - "where ranks hold different weight shards. " - "Env var: MAGI_COMPILE_OFFLOAD_CONFIG__SHM_SHARE_WEIGHTS (1/0/true/false)." + "Override for per-rank shared memory mode. When None (default), " + "MagiCompiler auto-detects by comparing weight fingerprints across " + "ranks: if all ranks hold identical weights, a single shared mmap is " + "used; otherwise each rank writes its own file. Set to True to force " + "per-rank mode (e.g. expert parallelism), or False to force sharing. " + "Env var: MAGI_COMPILE_OFFLOAD_CONFIG__FORCE_PER_RANK_WEIGHTS (1/0/true/false)." ), ) diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index 8f2ec30..f3348f3 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -36,7 +36,7 @@ import torch.multiprocessing as mp import torch.nn as nn -from magi_compiler._api import _materialize_shm_weights +from magi_compiler._api import _compute_weights_fingerprint, _materialize_shm_weights class FakeExpertBlock(nn.Module): @@ -164,3 +164,63 @@ def test_ep_fix_preserves_per_rank_shards(): "With per_rank=True, each rank should have DIFFERENT expert weights. " "If they're equal, the per-rank shm path did not work correctly." ) + +# ─────────────────────────────────────────────────────────────── +# Fingerprint tests +# ─────────────────────────────────────────────────────────────── + + +def _group_params_for_fp(module): + grouped = {} + for name, param in module.named_parameters(): + grouped.setdefault(param.dtype, []).append((name, param.data)) + return grouped + + +def test_fingerprint_identical_models(): + """Two models with the same seed produce the same fingerprint.""" + torch.manual_seed(42) + m1 = FakeExpertBlock(num_experts=4, dim=8) + torch.manual_seed(42) + m2 = FakeExpertBlock(num_experts=4, dim=8) + + fp1 = _compute_weights_fingerprint(_group_params_for_fp(m1)) + fp2 = _compute_weights_fingerprint(_group_params_for_fp(m2)) + assert fp1 == fp2, "Identical models should have identical fingerprints" + + +def test_fingerprint_different_models(): + """Two models with different seeds produce different fingerprints.""" + torch.manual_seed(42) + m1 = FakeExpertBlock(num_experts=4, dim=8) + torch.manual_seed(123) + m2 = FakeExpertBlock(num_experts=4, dim=8) + + fp1 = _compute_weights_fingerprint(_group_params_for_fp(m1)) + fp2 = _compute_weights_fingerprint(_group_params_for_fp(m2)) + assert fp1 != fp2, "Models with different weights should have different fingerprints" + + +def test_fingerprint_is_deterministic(): + """Calling fingerprint twice on the same model gives the same result.""" + torch.manual_seed(42) + m = FakeExpertBlock(num_experts=4, dim=8) + g = _group_params_for_fp(m) + + fp1 = _compute_weights_fingerprint(g) + fp2 = _compute_weights_fingerprint(g) + assert fp1 == fp2, "Fingerprint should be deterministic" + + +def test_fingerprint_detects_single_element_change(): + """Changing one element in one parameter changes the fingerprint.""" + torch.manual_seed(42) + m1 = FakeExpertBlock(num_experts=4, dim=8) + torch.manual_seed(42) + m2 = FakeExpertBlock(num_experts=4, dim=8) + m2.expert_weight.data[0, 0] += 1.0 + + fp1 = _compute_weights_fingerprint(_group_params_for_fp(m1)) + fp2 = _compute_weights_fingerprint(_group_params_for_fp(m2)) + assert fp1 != fp2, "Single element change should be detected" + From c47e6b75b0bdf50285176bc2cf67fbd428004dab Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 2 Sep 2026 01:37:16 +0800 Subject: [PATCH 50/54] style: format with black --- magi_compiler/_api.py | 9 ++------- tests/feature_tests/test_ep_shared_memory.py | 2 +- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 5cd7d56..0745765 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -570,10 +570,7 @@ def _create_empty_shm(shm_path: str, total_numel: int, dtype: torch.dtype) -> to return torch.from_file(shm_path, shared=True, size=total_numel, dtype=dtype, device="cpu") - -def _compute_weights_fingerprint( - grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], -) -> bytes: +def _compute_weights_fingerprint(grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]]) -> bytes: """Fast fingerprint of all weight data for cross-rank comparison. Hashes param names, shapes, dtypes, and a head+tail sample of each @@ -596,9 +593,7 @@ def _compute_weights_fingerprint( return h.digest() -def _all_ranks_same_weights( - grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]], -) -> bool: +def _all_ranks_same_weights(grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]]) -> bool: """Return True if every rank holds identical weights (by fingerprint).""" local_hash = _compute_weights_fingerprint(grouped_params) hash_tensor = torch.frombuffer(bytearray(local_hash), dtype=torch.uint8).clone() diff --git a/tests/feature_tests/test_ep_shared_memory.py b/tests/feature_tests/test_ep_shared_memory.py index f3348f3..65c1bf3 100644 --- a/tests/feature_tests/test_ep_shared_memory.py +++ b/tests/feature_tests/test_ep_shared_memory.py @@ -165,6 +165,7 @@ def test_ep_fix_preserves_per_rank_shards(): "If they're equal, the per-rank shm path did not work correctly." ) + # ─────────────────────────────────────────────────────────────── # Fingerprint tests # ─────────────────────────────────────────────────────────────── @@ -223,4 +224,3 @@ def test_fingerprint_detects_single_element_change(): fp1 = _compute_weights_fingerprint(_group_params_for_fp(m1)) fp2 = _compute_weights_fingerprint(_group_params_for_fp(m2)) assert fp1 != fp2, "Single element change should be detected" - From 64f8b0cadc6f3ee62b9a3c23849900cae772157f Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 2 Sep 2026 11:58:49 +0800 Subject: [PATCH 51/54] refactor: read force_per_rank_weights from CompileConfig instead of os.environ --- magi_compiler/_api.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 0745765..4e8b8bb 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -216,7 +216,7 @@ def _magi_compile_class( raise AttributeError(f"{cls.__name__} has no callable method '{method_name}'") if issubclass(cls, nn.Module) and conf.offload_config.model_cpu_offload: - _patch_cpu_offload_apply(cls) + _patch_cpu_offload_apply(cls, conf) old_init = cls.__init__ @@ -650,7 +650,7 @@ def _materialize_shm_weights( gc.collect() -def _patch_cpu_offload_apply(cls: type[nn.Module]): +def _patch_cpu_offload_apply(cls: type[nn.Module], conf: CompileConfig): magi_logger.info(f"Enabling CPU offload for {cls}") _orig_apply = cls._apply @@ -718,11 +718,11 @@ def _force_cpu(t): full_state_dict = None - # Determine per_rank mode: env override > auto-detect via fingerprint - force_env = os.environ.get("MAGI_COMPILE_OFFLOAD_CONFIG__FORCE_PER_RANK_WEIGHTS") - if force_env is not None: - per_rank = force_env.lower() in ("1", "true") - magi_logger.info('[offload] per_rank=%s (env override FORCE_PER_RANK_WEIGHTS)', per_rank) + # Determine per_rank mode: config override > auto-detect via fingerprint + force = conf.offload_config.force_per_rank_weights + if force is not None: + per_rank = force + magi_logger.info('[offload] per_rank=%s (config force_per_rank_weights)', per_rank) else: same = _all_ranks_same_weights(grouped_params) per_rank = not same From 931d985d09297298e277785c80317dae92a3d4d4 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 2 Sep 2026 17:00:55 +0800 Subject: [PATCH 52/54] fix: use gloo group for CPU all_gather in weight fingerprint (NCCL compat) Extract get_cpu_gloo_group() into utils/dist_utils.py, shared by _all_ranks_same_weights and _get_cost_sync_group. Fixes RuntimeError on NCCL-default process groups where CPU tensors are rejected. --- magi_compiler/_api.py | 4 ++- magi_compiler/profiling/runtime_estimator.py | 15 ++-------- magi_compiler/utils/dist_utils.py | 31 ++++++++++++++++++++ 3 files changed, 37 insertions(+), 13 deletions(-) create mode 100644 magi_compiler/utils/dist_utils.py diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 4e8b8bb..250ada8 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -595,11 +595,13 @@ def _compute_weights_fingerprint(grouped_params: dict[torch.dtype, list[tuple[st def _all_ranks_same_weights(grouped_params: dict[torch.dtype, list[tuple[str, torch.Tensor]]]) -> bool: """Return True if every rank holds identical weights (by fingerprint).""" + from magi_compiler.utils.dist_utils import get_cpu_gloo_group + local_hash = _compute_weights_fingerprint(grouped_params) hash_tensor = torch.frombuffer(bytearray(local_hash), dtype=torch.uint8).clone() world_size = dist.get_world_size() gathered = [torch.empty_like(hash_tensor) for _ in range(world_size)] - dist.all_gather(gathered, hash_tensor) + dist.all_gather(gathered, hash_tensor, group=get_cpu_gloo_group()) return all(torch.equal(gathered[0], g) for g in gathered[1:]) diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index 490ace1..1c3539e 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -59,8 +59,6 @@ # Dedicated GLOO (CPU) group for the cost sync, built once -- keeps it off the # NCCL process groups the forward uses (cannot desync weight-gather / CP comms). -_COST_SYNC_GROUP = "uninit" - def snode_issues_collective(snode: BaseSchedulerNode) -> bool: """ @@ -74,18 +72,11 @@ def snode_issues_collective(snode: BaseSchedulerNode) -> bool: return _extern_has_internal_collective(snode) + def _get_cost_sync_group(): - global _COST_SYNC_GROUP - import torch.distributed as dist + from magi_compiler.utils.dist_utils import get_cpu_gloo_group - if _COST_SYNC_GROUP != "uninit": - return _COST_SYNC_GROUP - try: - _COST_SYNC_GROUP = dist.new_group(backend="gloo") - except Exception as exc: # noqa: BLE001 - magi_logger.warning("cost-sync: gloo group unavailable (%s); using default group", exc) - _COST_SYNC_GROUP = None - return _COST_SYNC_GROUP + return get_cpu_gloo_group() @dataclasses.dataclass diff --git a/magi_compiler/utils/dist_utils.py b/magi_compiler/utils/dist_utils.py new file mode 100644 index 0000000..92b723f --- /dev/null +++ b/magi_compiler/utils/dist_utils.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +"""Distributed utilities for MagiCompiler.""" + +from __future__ import annotations + +import torch.distributed as dist + +from magi_compiler.utils.logger import magi_logger + +_CPU_GLOO_GROUP = "uninit" + + +def get_cpu_gloo_group() -> dist.ProcessGroup | None: + """Return a gloo process group for CPU-tensor collectives. + + The default process group is typically NCCL, which only supports CUDA + tensors. This helper lazily creates a gloo group so that CPU tensors + can participate in collectives like ``all_gather`` and ``all_reduce``. + + Returns ``None`` if gloo is unavailable (the caller should fall back + to the default group or skip the collective). + """ + global _CPU_GLOO_GROUP + if _CPU_GLOO_GROUP != "uninit": + return _CPU_GLOO_GROUP + try: + _CPU_GLOO_GROUP = dist.new_group(backend="gloo") + except Exception as exc: # noqa: BLE001 + magi_logger.warning("get_cpu_gloo_group: gloo unavailable (%s); returning None", exc) + _CPU_GLOO_GROUP = None + return _CPU_GLOO_GROUP From 9e81fd3933d19a3e337c2a4e6e0812854a5f93c4 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 2 Sep 2026 17:55:41 +0800 Subject: [PATCH 53/54] fix: handle None gloo group in _all_ranks_same_weights (safe fallback to per_rank=True) --- magi_compiler/_api.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 250ada8..de84bc6 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -597,11 +597,16 @@ def _all_ranks_same_weights(grouped_params: dict[torch.dtype, list[tuple[str, to """Return True if every rank holds identical weights (by fingerprint).""" from magi_compiler.utils.dist_utils import get_cpu_gloo_group + group = get_cpu_gloo_group() + if group is None: + magi_logger.warning('[offload] gloo group unavailable, assuming per_rank=True (safe default)') + return False + local_hash = _compute_weights_fingerprint(grouped_params) hash_tensor = torch.frombuffer(bytearray(local_hash), dtype=torch.uint8).clone() world_size = dist.get_world_size() gathered = [torch.empty_like(hash_tensor) for _ in range(world_size)] - dist.all_gather(gathered, hash_tensor, group=get_cpu_gloo_group()) + dist.all_gather(gathered, hash_tensor, group=group) return all(torch.equal(gathered[0], g) for g in gathered[1:]) From d5846032f645f015d09234d46f24223caf581c56 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 2 Sep 2026 18:05:29 +0800 Subject: [PATCH 54/54] style: fix black formatting in runtime_estimator.py --- magi_compiler/profiling/runtime_estimator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index 1c3539e..db3246a 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -60,6 +60,7 @@ # Dedicated GLOO (CPU) group for the cost sync, built once -- keeps it off the # NCCL process groups the forward uses (cannot desync weight-gather / CP comms). + def snode_issues_collective(snode: BaseSchedulerNode) -> bool: """ True if replaying / running this snode issues NCCL (collective AG, or a @@ -72,7 +73,6 @@ def snode_issues_collective(snode: BaseSchedulerNode) -> bool: return _extern_has_internal_collective(snode) - def _get_cost_sync_group(): from magi_compiler.utils.dist_utils import get_cpu_gloo_group