[Fix] Skip shmem for per-rank weights — use plain pinned tensor

[cennn]

🗂️ PR Category

• 🚀 Optimization (performance, memory, etc.)

📝 Description

When per_rank=True (EP scenario, each rank has unique weights), the current code creates an mmap shared-memory file via _create_empty_shm, writes into it, then immediately deletes it — the file is never read by another rank. This is pure filesystem overhead.

Change: Replace _create_empty_shm + os.remove with a plain torch.empty(), keeping the streaming _stream_copy_and_replace and pin_memory_in_place unchanged.

The per_rank=False path (all ranks share identical weights) is untouched — it still uses mmap so rank 0 writes once and all other ranks map the same file.

Before

path = _shm_path(cls_name, dtype, rank=local_rank)
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)

After

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)

[jiahy0825]

LGTM