From d44b0c25da4eb1e73bb1cf8639d52116149870e0 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 10 Sep 2026 14:49:21 +0800 Subject: [PATCH] =?UTF-8?q?fix(offload):=20skip=20shmem=20for=20per=5Frank?= =?UTF-8?q?=20weights=20=E2=80=94=20use=20plain=20pinned=20tensor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When per_rank=True (EP scenario), each rank has unique weights and never shares them with other ranks. The mmap file was created only to be read by the same rank and immediately deleted — pure overhead. Replace _create_empty_shm + os.remove with a plain torch.empty, keeping the streaming copy-and-replace and pin_memory_in_place unchanged. --- magi_compiler/_api.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 033a020..e73fd49 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -618,22 +618,20 @@ 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 (default): each rank writes its own mmap concurrently. - per_rank=False (all ranks identical): rank 0 writes, all ranks map. + per_rank=True (default): each rank gets a plain pinned tensor (no mmap/shmem + needed because no cross-rank sharing occurs). + per_rank=False (all ranks identical): rank 0 writes an mmap, all ranks map. """ cls_name = module.__class__.__name__ buffers: list[torch.Tensor] = [] if per_rank: 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) + giant = torch.empty(total_numel, dtype=dtype, device="cpu") _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()