mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-12 09:23:54 +08:00
2836e365ea
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>
349 lines
18 KiB
Python
349 lines
18 KiB
Python
"""Mirrors `deep_ep.buffers.engram` (`EngramBuffer`) of DeepEP-Ascend.
|
|
|
|
What Engram needs from DeepEP: the n-gram embedding tables of every Engram layer are
|
|
**sharded over the Engram-parallel ranks** (`models/v4_modules/engram.py:107` keeps only
|
|
`num_local_embeddings` rows per rank), and every step each rank must read arbitrary rows of
|
|
the *global* table for its own tokens. Upstream serves those reads out of RDMA-registered
|
|
storage; this port exchanges the row indices and the rows themselves with all-to-all.
|
|
|
|
Index space (confirmed by `tests/engram/test_engram.py:80-95`): the reference global table is
|
|
`all_gather(local_storage)`, so a global index `g` of layer `l` lives on rank
|
|
`g // num_entries_per_layer[l]`, row `g % num_entries_per_layer[l]`.
|
|
|
|
Layout contract (`csrc/buffers/engram.hpp:139-200`):
|
|
* `write(storages, sfs)`: `storages[l]` is this rank's `[num_entries_l, hidden]` shard;
|
|
`sfs[l]` is the **replicated** `[num_ranks * num_entries_l, num_sf_packs]` table of which
|
|
only this rank's shard is kept.
|
|
* `fetch(indices)`: `indices` is int32 `[num_layers, num_tokens, num_entries_per_token]`;
|
|
the returned per-layer hook yields `[num_tokens, num_entries_per_token * hidden]`
|
|
(and, in FP8 mode, `[num_tokens, num_entries_per_token * num_sf_packs]`, column-major
|
|
with stride `(1, num_tokens)` when `use_tma_aligned_col_major_sf=True`).
|
|
|
|
Basic mode: the request counts are host values, so `fetch` host-syncs once per layer.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import os
|
|
from typing import Callable, List, Optional, Tuple
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from .. import comm
|
|
from .base import BufferBase
|
|
from ..utils.indexing import byte_rows, empty_rows, put_rows, take_rows
|
|
from ..utils.math import align
|
|
|
|
# NEURON: row-major scales (the TMA column-major layout only adds transposes; small int16 ones crash neuronx-cc,
|
|
# issue #9). Consumers index scales logically.
|
|
_COL_MAJOR_SF = False
|
|
|
|
|
|
def _accept_legacy_storage_size(init):
|
|
@functools.wraps(init)
|
|
def wrapper(self, *args, **kwargs):
|
|
if 'num_rdma_storage_size' in kwargs:
|
|
if 'num_rdma_storage_bytes' in kwargs:
|
|
raise TypeError('Specify only one of num_rdma_storage_bytes and num_rdma_storage_size')
|
|
kwargs['num_rdma_storage_bytes'] = kwargs.pop('num_rdma_storage_size')
|
|
return init(self, *args, **kwargs)
|
|
return wrapper
|
|
|
|
|
|
|
|
def _take_rows_to(table: torch.Tensor, index: torch.Tensor, device: torch.device) -> torch.Tensor:
|
|
"""Gather rows where the table lives (Engram tables default to host memory, which the GPU reads
|
|
through UVA), then move the result to the compute device."""
|
|
return take_rows(table, index.to(table.device)).to(device)
|
|
|
|
class EngramBuffer(BufferBase):
|
|
"""Distributed n-gram embedding lookup over a sharded table."""
|
|
|
|
barrier = comm.barrier
|
|
get_comm_stream = staticmethod(comm.get_comm_stream)
|
|
get_physical_domain_size = comm.get_physical_domain_size
|
|
get_logical_domain_size = comm.get_logical_domain_size
|
|
|
|
@_accept_legacy_storage_size
|
|
def __init__(self,
|
|
group: dist.ProcessGroup,
|
|
num_gpu_bytes: int,
|
|
num_rdma_storage_bytes: int,
|
|
use_cpu_rdma_storage: bool = False,
|
|
num_allocated_qps: int = 0,
|
|
qp_depth: int = 0,
|
|
restrict_rd_atomic: bool = False,
|
|
allow_hybrid_mode: bool = True,
|
|
sl_idx: Optional[int] = None,
|
|
num_gpu_timeout_secs: int = 100,
|
|
explicitly_destroy: bool = False,
|
|
*, use_huge1g_page: bool = True):
|
|
"""Create the buffer. QP / service-level / CPU-storage knobs are accepted and ignored.
|
|
|
|
Unlike Ascend (which rejects `use_cpu_rdma_storage=True`), this port accepts it:
|
|
DeepInfer passes it whenever `DEEPINFER_ENGRAM_FORCE_GPU_STORAGE` is off
|
|
(`parallel_utils/userbuffer.py:147`). There is no RDMA storage here; the table shard
|
|
simply stays on whatever device `write()` is given.
|
|
"""
|
|
alignment = comm.get_num_allocation_alignment()
|
|
assert num_gpu_bytes > 0 and num_gpu_bytes % alignment == 0
|
|
assert num_rdma_storage_bytes > 0 and num_rdma_storage_bytes % alignment == 0
|
|
self.group = group
|
|
self.rank_idx = group.rank()
|
|
self.num_ranks = group.size()
|
|
self.num_gpu_bytes = num_gpu_bytes
|
|
self.num_bytes = num_gpu_bytes
|
|
self.num_rdma_storage_bytes = num_rdma_storage_bytes
|
|
self.use_cpu_rdma_storage = use_cpu_rdma_storage
|
|
self.allow_hybrid_mode = allow_hybrid_mode
|
|
self.num_allocated_qps = 0
|
|
super().__init__(explicitly_destroy)
|
|
|
|
self.num_entries_per_layer: List[int] = []
|
|
self.hidden = 0
|
|
self.num_max_tokens = 0
|
|
self.num_entries_per_token = 0
|
|
self.num_sf_packs = 0
|
|
self.dtype: Optional[torch.dtype] = None
|
|
self._storages: List[torch.Tensor] = []
|
|
self._sf_shards: List[torch.Tensor] = []
|
|
|
|
def destroy(self) -> None:
|
|
super().destroy()
|
|
self._storages = []
|
|
self._sf_shards = []
|
|
|
|
# --- Sizing ---------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def get_storage_size_hint(group: dist.ProcessGroup,
|
|
num_entries_per_layer: List[int], hidden: int,
|
|
num_max_tokens: int, num_entries_per_token: int,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
num_sf_packs: int = 0) -> Tuple[int, int]:
|
|
"""Allocation-aligned (receive + SF, per-rank storage) sizes; same formula as Ascend."""
|
|
assert num_entries_per_layer and all(num_entries > 0 for num_entries in num_entries_per_layer)
|
|
assert hidden > 0 and num_max_tokens > 0 and num_entries_per_token > 0 and num_sf_packs >= 0
|
|
assert dtype in (torch.bfloat16, torch.float8_e4m3fn)
|
|
assert (dtype == torch.float8_e4m3fn) == (num_sf_packs > 0)
|
|
alignment = comm.get_num_allocation_alignment()
|
|
num_recv_bytes = len(num_entries_per_layer) * num_max_tokens * num_entries_per_token * hidden * dtype.itemsize
|
|
num_sf_bytes = sum(num_entries_per_layer) * num_sf_packs * torch.int16.itemsize
|
|
num_storage_bytes = sum(num_entries_per_layer) * hidden * dtype.itemsize
|
|
return align(num_recv_bytes + num_sf_bytes, alignment), align(num_storage_bytes, alignment)
|
|
|
|
@staticmethod
|
|
def get_theoretical_config(group: dist.ProcessGroup,
|
|
num_layers: int,
|
|
num_max_tokens: int,
|
|
num_entries_per_token: int,
|
|
num_max_read_requests_per_qp: int = 16) -> Tuple[int, int, bool]:
|
|
"""`(num_allocated_qps, qp_depth, restrict_rd_atomic)`; no QPs on Neuron."""
|
|
return 0, 0, False
|
|
|
|
def set_config(self, num_entries_per_layer: List[int], hidden: int,
|
|
num_max_tokens: int, num_entries_per_token: int,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
num_sf_packs: int = 0) -> None:
|
|
"""Configure the local shard geometry and the maximum fetch geometry."""
|
|
assert not self.destroyed
|
|
self.get_storage_size_hint(self.group, num_entries_per_layer, hidden,
|
|
num_max_tokens, num_entries_per_token, dtype, num_sf_packs)
|
|
self.num_entries_per_layer = list(num_entries_per_layer)
|
|
self.hidden = hidden
|
|
self.num_max_tokens = num_max_tokens
|
|
self.num_entries_per_token = num_entries_per_token
|
|
self.num_sf_packs = num_sf_packs
|
|
self.dtype = dtype
|
|
|
|
# --- Table upload ---------------------------------------------------------------
|
|
|
|
def write(self, storages: List[torch.Tensor],
|
|
sfs: Optional[List[torch.Tensor]] = None) -> None:
|
|
"""Copy this rank's table shard (and its slice of the replicated SF table) in."""
|
|
assert not self.destroyed
|
|
num_layers = len(self.num_entries_per_layer)
|
|
assert num_layers > 0, 'set_config must be called before write'
|
|
assert len(storages) == num_layers
|
|
assert (sfs is not None) == (self.num_sf_packs > 0)
|
|
|
|
for layer_idx, storage in enumerate(storages):
|
|
assert storage.dim() == 2 and storage.is_contiguous()
|
|
assert storage.size(0) == self.num_entries_per_layer[layer_idx]
|
|
assert storage.size(1) == self.hidden
|
|
assert storage.dtype == self.dtype
|
|
if self._use_shared_tables(storages, sfs):
|
|
self._write_shared(storages, sfs)
|
|
return
|
|
self._storages = [storage.clone() for storage in storages]
|
|
|
|
self._sf_shards = []
|
|
if sfs is not None:
|
|
for layer_idx, sf in enumerate(sfs):
|
|
num_entries = self.num_entries_per_layer[layer_idx]
|
|
assert sf.dim() == 2 and sf.is_contiguous() and sf.size(1) == self.num_sf_packs
|
|
if sf.size(0) == num_entries:
|
|
# Already this rank's slice (DeepInfer on Neuron skips the replicated-table all-gather).
|
|
self._sf_shards.append(sf.clone())
|
|
continue
|
|
assert sf.size(0) == self.num_ranks * num_entries
|
|
begin = self.rank_idx * num_entries
|
|
self._sf_shards.append(sf[begin:begin + num_entries].clone())
|
|
|
|
comm.barrier(self)
|
|
|
|
# --- Single-host shared tables (Neuron) -------------------------------------------
|
|
|
|
def _use_shared_tables(self, storages: List[torch.Tensor], sfs: Optional[List[torch.Tensor]]) -> bool:
|
|
"""NEURON: every rank of a trn2 instance shares one host, so the whole table can sit in shared host memory and
|
|
each rank reads its own rows: no all-to-all, no host sync on split sizes (DEEPINFER_NEURON_ENGRAM_SHARED=0 off)."""
|
|
return (
|
|
os.environ.get("DEEPINFER_NEURON_ENGRAM_SHARED", "1") == "1"
|
|
and self.num_ranks > 1
|
|
and all(t.device.type == "cpu" for t in storages)
|
|
and (sfs is None or all(sf.size(0) == n for sf, n in zip(sfs, self.num_entries_per_layer)))
|
|
)
|
|
|
|
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."""
|
|
# 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:
|
|
os.unlink(path) # the mappings stay valid; the memory goes away with the last process
|
|
return t
|
|
|
|
def _write_shared(self, storages: List[torch.Tensor], sfs: Optional[List[torch.Tensor]]) -> None:
|
|
self._storages, self._sf_shards = [], []
|
|
for layer_idx, storage in enumerate(storages):
|
|
n = self.num_entries_per_layer[layer_idx]
|
|
begin = self.rank_idx * n
|
|
table = self._shared(f"t{layer_idx}", self.num_ranks * n, self.hidden * storage.element_size())
|
|
table[begin:begin + n] = byte_rows(storage)
|
|
self._storages.append(table)
|
|
if sfs is not None:
|
|
sf = sfs[layer_idx]
|
|
sf_table = self._shared(f"s{layer_idx}", self.num_ranks * n, sf.size(1) * sf.element_size())
|
|
sf_table[begin:begin + n] = byte_rows(sf.contiguous())
|
|
self._sf_shards.append((sf_table, sf.dtype))
|
|
self._shared_tables = True
|
|
comm.barrier(self)
|
|
|
|
def _fetch_layer_shared(self, layer_idx: int, indices: torch.Tensor) -> torch.Tensor | tuple:
|
|
device = indices.device
|
|
num_tokens, num_entries_per_token = indices.shape
|
|
flat = indices.reshape(-1).to(device="cpu", dtype=torch.int64)
|
|
flat = flat.clamp(0, self._storages[layer_idx].size(0) - 1)
|
|
rows = self._storages[layer_idx].index_select(0, flat).view(self.dtype)
|
|
fetched = rows.reshape(num_tokens, num_entries_per_token * self.hidden).to(device)
|
|
if not self._sf_shards:
|
|
return fetched
|
|
sf_table, sf_dtype = self._sf_shards[layer_idx]
|
|
sf_rows = sf_table.index_select(0, flat).view(sf_dtype)
|
|
return fetched, sf_rows.reshape(num_tokens, num_entries_per_token * self.num_sf_packs).to(device)
|
|
|
|
# --- Fetch ----------------------------------------------------------------------
|
|
|
|
def _fetch_layer(self, layer_idx: int, indices: torch.Tensor,
|
|
use_tma_aligned_col_major_sf: bool):
|
|
"""Gather `indices` ([num_tokens, num_entries_per_token], global ids) of one layer."""
|
|
device = indices.device
|
|
num_tokens, num_entries_per_token = indices.shape
|
|
num_entries = self.num_entries_per_layer[layer_idx]
|
|
storage = self._storages[layer_idx]
|
|
sf_shard = self._sf_shards[layer_idx] if self._sf_shards else None
|
|
num_requests = num_tokens * num_entries_per_token
|
|
flat = indices.reshape(-1).to(torch.int64)
|
|
owner = torch.div(flat, num_entries, rounding_mode='floor')
|
|
local = flat - owner * num_entries
|
|
|
|
if self.num_ranks == 1:
|
|
rows = _take_rows_to(storage, local, device)
|
|
sf_rows = _take_rows_to(sf_shard, local, device) if sf_shard is not None else None
|
|
else:
|
|
# Group the requests by owner rank (stable: ascending request position).
|
|
order = torch.argsort(owner, stable=True)
|
|
# NEURON: host sync — all-to-all split sizes must be host values.
|
|
send_counts = torch.zeros(self.num_ranks, dtype=torch.int64, device=device)
|
|
send_counts.index_add_(0, owner, torch.ones_like(owner))
|
|
send_counts = send_counts.tolist()
|
|
recv_counts_t = torch.empty(self.num_ranks, dtype=torch.int64, device=device)
|
|
comm.all_to_all_single(recv_counts_t,
|
|
torch.tensor(send_counts, dtype=torch.int64, device=device),
|
|
None, None, self.group)
|
|
recv_counts = recv_counts_t.tolist()
|
|
num_served = int(sum(recv_counts))
|
|
|
|
send_idx = local[order]
|
|
served_idx = torch.empty((num_served,), dtype=torch.int64, device=device)
|
|
comm.all_to_all_single(served_idx, send_idx, recv_counts, send_counts, self.group)
|
|
served_idx = served_idx.clamp(0, max(num_entries - 1, 0))
|
|
|
|
served_rows = _take_rows_to(storage, served_idx, device).contiguous()
|
|
answer_rows = empty_rows(num_requests, self.hidden, storage.dtype, device)
|
|
# Split sizes count rows: `all_to_all_single` splits along dim 0.
|
|
comm.all_to_all_single(byte_rows(answer_rows), byte_rows(served_rows),
|
|
send_counts, recv_counts, self.group)
|
|
rows = empty_rows(num_requests, self.hidden, storage.dtype, device)
|
|
put_rows(rows, order, answer_rows)
|
|
|
|
sf_rows = None
|
|
if sf_shard is not None:
|
|
served_sf = _take_rows_to(sf_shard, served_idx, device).contiguous()
|
|
answer_sf = empty_rows(num_requests, self.num_sf_packs, sf_shard.dtype, device)
|
|
comm.all_to_all_single(byte_rows(answer_sf), byte_rows(served_sf),
|
|
send_counts, recv_counts, self.group)
|
|
sf_rows = torch.empty_like(answer_sf)
|
|
sf_rows[order] = answer_sf
|
|
|
|
fetched = rows.view(num_tokens, num_entries_per_token * self.hidden)
|
|
if sf_rows is None:
|
|
return fetched
|
|
sf_flat = sf_rows.view(num_tokens, num_entries_per_token * self.num_sf_packs)
|
|
if use_tma_aligned_col_major_sf and _COL_MAJOR_SF:
|
|
# Column-major view: stride (1, num_tokens), no token-dimension padding
|
|
# (`as_strided` in csrc/buffers/engram.hpp:177).
|
|
num_sf_cols = sf_flat.size(1)
|
|
fetched_sf = torch.empty(num_sf_cols * num_tokens, dtype=sf_flat.dtype,
|
|
device=device).as_strided((num_tokens, num_sf_cols),
|
|
(1, num_tokens))
|
|
fetched_sf.copy_(sf_flat)
|
|
else:
|
|
fetched_sf = sf_flat
|
|
return fetched, fetched_sf
|
|
|
|
# The tables live in host memory: the gather (and its host-sized all-to-all) runs eagerly, outside the decode
|
|
# graph; the graph resumes with the fetched rows.
|
|
@torch.compiler.disable
|
|
def fetch(self, indices: torch.Tensor, num_qps: int = 0,
|
|
use_tma_aligned_col_major_sf: bool = True) -> List[Callable]:
|
|
"""Fetch all layers; returns one result hook per layer (work happens eagerly here)."""
|
|
assert not self.destroyed
|
|
assert indices.dim() == 3 and indices.is_contiguous()
|
|
assert indices.dtype == torch.int32
|
|
num_layers, num_tokens, num_entries_per_token = indices.shape
|
|
assert num_layers == len(self.num_entries_per_layer)
|
|
assert num_tokens <= self.num_max_tokens
|
|
assert num_entries_per_token == self.num_entries_per_token
|
|
assert self._storages, 'write() must be called before fetch()'
|
|
|
|
if getattr(self, "_shared_tables", False):
|
|
results = [self._fetch_layer_shared(layer_idx, indices[layer_idx]) for layer_idx in range(num_layers)]
|
|
else:
|
|
results = [self._fetch_layer(layer_idx, indices[layer_idx], use_tma_aligned_col_major_sf)
|
|
for layer_idx in range(num_layers)]
|
|
return [(lambda result=result: result) for result in results]
|