coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron engram shared tables: one file for the whole group (#28)

Browse Source 
The name came from the parent pid, which CP followers do not share with their
driver (they are its children): followers created a second table, never
unlinked (622 GB of /dev/shm leaked over five runs, then a host OOM), and each
side read zeros for the other's shards. Rank 0 now publishes the path through
the distributed store, and opening asserts the size so it can never create a
file.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-24 07:08:37 +08:00
parent d44cf253c5
commit 2836e365ea
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py vendored
View File

@ -208,12 +208,18 @@ class EngramBuffer(BufferBase):
def _shared(self, name: str, rows: int, row_bytes: int) -> torch.Tensor:
"""A ``[rows, row_bytes]`` uint8 tensor mapping one /dev/shm file shared by the group's ranks."""
# Ranks of one job share their launcher: its pid names the job.
path = f"/dev/shm/deepinfer_engram_{os.getppid()}_{name}"
# One name for the whole group: rank 0 publishes its pid through the distributed store. (The parent pid is not
# shared: CP followers are children of their engine's driver, not of the launcher.)
store = dist.distributed_c10d._get_default_store()
key = f"deepinfer_engram_shm_{name}"
if self.rank_idx == 0:
path = f"/dev/shm/deepinfer_engram_{os.getpid()}_{name}"
with open(path, "wb") as f:
f.truncate(rows * row_bytes)
store.set(key, path)
comm.barrier(self)
path = store.get(key).decode()
assert os.path.getsize(path) == rows * row_bytes, path # opening must never create a second table
t = torch.from_file(path, shared=True, size=rows * row_bytes, dtype=torch.uint8).view(rows, row_bytes)
comm.barrier(self)
if self.rank_idx == 0: