From cc90926a1c6247baafecdfcf243d94c1d1235666 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 21 Oct 2026 14:57:51 +0800 Subject: [PATCH] Neuron: full DeepSeek-V4.1-Flash loads and prefills on 64 NeuronCores (EP=64) Device-memory fixes found bringing up the released checkpoint (16.45 GiB of weights per 24 GB core): - Engram: the Neuron EngramBuffer keeps only this rank's SF rows, so skip the all-gather of the replicated SF table (~3 GB per layer, allocated on the device -> "Failed to allocate resource"). - gather_tl: index_select / advanced indexing of an int8 table allocates several times the table on the device (+4.9 GiB for the 1.3 GB embedding); gather through F.embedding on an int32 / fp16 view instead (bit-exact, no extra memory). - deep_gemm FP8/FP4 GEMMs: dequantize the weight in N chunks of 32M elements and one expert at a time instead of whole projections / all local experts in fp32 (OOM on a fragmented pool); bit-identical results. - transfer_to_neuron: import torch.utils._device explicitly (missing when torch_neuronx is not auto-loaded). examples/neuron/offline_inference_neuron.py: offline inference in basic mode (text prompts, no MTP / EPLB, fixed KV pool, max_model_len >= 2048 * 8 for the candidate-block filter), one process per NeuronCore. Co-Authored-By: Claude Opus 5.5 --- .../layers/tilelang_ops/neuron/gather.py | 17 +++- .../models/v4_modules/engram.py | 10 +- deepinfer/neuron_ops/transfer_to_neuron.py | 1 + .../vendor/deep_ep/buffers/engram.py | 8 +- deepinfer/neuron_ops/vendor/deep_gemm/gemm.py | 36 +++++-- examples/neuron/offline_inference_neuron.py | 98 +++++++++++++++++++ 6 files changed, 155 insertions(+), 15 deletions(-) create mode 100644 examples/neuron/offline_inference_neuron.py diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py index d573d44..02fa4d0 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py @@ -5,6 +5,21 @@ from collections.abc import Callable import torch +def _gather_rows(x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """``x[index]`` for a byte table ``x [n, D]``. + + On the neuron device, index_select / advanced indexing of an int8 table allocates several times the table + (+4.9 GiB for the 1.3 GB V4.1 embedding, which ran the full model out of HBM), while F.embedding through a + 4- or 2-byte view copies only the selected rows and is bit-exact (verified on arbitrary bytes, NaN patterns + included). + """ + if x.is_contiguous(): + for elem, dtype in ((4, torch.int32), (2, torch.float16)): + if x.shape[1] % elem == 0 and x.storage_offset() % elem == 0: + return torch.nn.functional.embedding(index, x.view(dtype)).view(torch.int8) + return torch.index_select(x, 0, index) + + def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: int) -> Callable[..., None]: """Kernel args: ``x [n, D]`` (row stride ``x_stride``), ``index [m]``, ``y [m, D]`` contiguous. @@ -21,7 +36,7 @@ def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: in return valid = (index >= 0) & (index < n) safe = torch.where(valid, index, torch.zeros_like(index)) - rows = torch.index_select(x[:, :D], 0, safe) + rows = _gather_rows(x[:, :D], safe) y.copy_(torch.where(valid.view(m, 1), rows, torch.zeros((), dtype=rows.dtype, device=rows.device))) return _gather diff --git a/deepinfer/model_executor/models/v4_modules/engram.py b/deepinfer/model_executor/models/v4_modules/engram.py index 1f40794..0f30312 100644 --- a/deepinfer/model_executor/models/v4_modules/engram.py +++ b/deepinfer/model_executor/models/v4_modules/engram.py @@ -96,12 +96,12 @@ class NgramEmbedding(nn.Module): def gather_sf(self) -> torch.Tensor: assert self.sf is not None local_sf = _pack_ue8m0_to_int(self.sf) - sf = local_sf.new_empty(self.engram_group.size() * self.num_local_embeddings, local_sf.size(1)) if PLATFORM.is_neuron: - # Gather raw bytes: gloo (host mode) has no int16/int32 collectives for this dtype. - dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.contiguous().view(torch.uint8), group=self.engram_group) - else: - dist.all_gather_into_tensor(sf, local_sf, group=self.engram_group) + # The Neuron EngramBuffer keeps only this rank's SF rows (it serves its own table shard), so skip the + # all-gather: the replicated table is ~3 GB per layer for the released checkpoint, allocated on the device. + return local_sf.contiguous() + sf = local_sf.new_empty(self.engram_group.size() * self.num_local_embeddings, local_sf.size(1)) + dist.all_gather_into_tensor(sf, local_sf, group=self.engram_group) return sf def release_storage(self) -> None: diff --git a/deepinfer/neuron_ops/transfer_to_neuron.py b/deepinfer/neuron_ops/transfer_to_neuron.py index ee48d44..fd6a22e 100644 --- a/deepinfer/neuron_ops/transfer_to_neuron.py +++ b/deepinfer/neuron_ops/transfer_to_neuron.py @@ -24,6 +24,7 @@ import types from typing import Any import torch +import torch.utils._device # not imported by `import torch` alone (e.g. TORCH_DEVICE_BACKEND_AUTOLOAD=0) _REAL_DEVICE = torch.device _TARGET: torch.device | None = None diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py index 8ddc2fa..1444d3e 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py @@ -176,8 +176,12 @@ class EngramBuffer(BufferBase): 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() - assert sf.size(0) == self.num_ranks * num_entries and sf.size(1) == self.num_sf_packs + 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()) diff --git a/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py b/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py index c49a5aa..812c68d 100644 --- a/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py +++ b/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py @@ -56,6 +56,23 @@ def _dequant_operand( return dequantize(data, sf, gran_k=recipe[2]) +# Dequantized weight rows per chunk (fp32 elements). The device is a 24 GB NeuronCore that already holds the +# packed weights: dequantizing a whole projection (or all local experts) to fp32 at once ran the full model out +# of HBM. Chunking along N keeps every output element's K-reduction intact, so results are unchanged. +_CHUNK_ELEMS = 1 << 25 + + +def _row_chunks(n: int, k: int) -> list[tuple[int, int]]: + step = max(1, _CHUNK_ELEMS // max(1, k)) + return [(n0, min(n, n0 + step)) for n0 in range(0, n, step)] + + +def _required_sf(sf, mn, k, recipe, is_sfa, num_groups=None, disable_ue8m0_cast=False): + return transform_sf_into_required_layout( + sf, mn, k, recipe, num_groups=num_groups, is_sfa=is_sfa, disable_ue8m0_cast=disable_ue8m0_cast + ) + + # --------------------------------------------------------------------------- # # BF16 # --------------------------------------------------------------------------- # @@ -177,8 +194,12 @@ def _fp8_fp4_gemm( return full_recipe = get_recipe(recipe, recipe_a, recipe_b) a_f32 = _dequant_operand(a, sfa, m, k, full_recipe, True, disable_ue8m0_cast=disable_ue8m0_cast) - b_f32 = _dequant_operand(b, sfb, n, k, full_recipe, False, disable_ue8m0_cast=disable_ue8m0_cast) - _write_output(a_f32 @ b_f32.transpose(0, 1), d, c, alpha) + sfb_req = _required_sf(sfb, n, k, full_recipe, False, disable_ue8m0_cast=disable_ue8m0_cast) + c_is_d = c is not None and _is_alias(c, d) + for n0, n1 in _row_chunks(n, k): + b_f32 = dequantize(b[n0:n1], sfb_req[n0:n1], gran_k=full_recipe[2]) + c_chunk = None if c is None else (d[:, n0:n1] if c_is_d else c[..., n0:n1]) + _write_output(a_f32 @ b_f32.transpose(0, 1), d[:, n0:n1], c_chunk, alpha) def fp8_fp4_gemm_nt( @@ -324,9 +345,7 @@ def m_grouped_fp8_fp4_gemm_nt_contiguous( return full_recipe = get_recipe(recipe, recipe_a, recipe_b) a_f32 = _dequant_operand(a, sfa, total_m, k, full_recipe, True, disable_ue8m0_cast=disable_ue8m0_cast) - b_f32 = _dequant_operand( - b, sfb, n, k, full_recipe, False, num_groups=num_groups, disable_ue8m0_cast=disable_ue8m0_cast - ) + sfb_req = _required_sf(sfb, n, k, full_recipe, False, num_groups=num_groups, disable_ue8m0_cast=disable_ue8m0_cast) spans = _psum_group_spans(grouped_layout, get_mk_alignment_for_contiguous_layout(), total_m) assert len(spans) == num_groups, f"{len(spans)=}, {num_groups=}" @@ -334,8 +353,11 @@ def m_grouped_fp8_fp4_gemm_nt_contiguous( for group_idx, (start, end) in enumerate(spans): if end <= start: continue - out = a_f32[start:end] @ b_f32[group_idx].transpose(0, 1) - d[start:end].copy_(out.to(d.dtype)) + # One expert (and one N chunk) at a time: never all local experts dequantized at once. + for n0, n1 in _row_chunks(n, k): + b_f32 = dequantize(b[group_idx, n0:n1], sfb_req[group_idx, n0:n1], gran_k=full_recipe[2]) + out = a_f32[start:end] @ b_f32.transpose(0, 1) + d[start:end, n0:n1].copy_(out.to(d.dtype)) def m_grouped_fp8_fp4_gemm_nn_contiguous( diff --git a/examples/neuron/offline_inference_neuron.py b/examples/neuron/offline_inference_neuron.py new file mode 100644 index 0000000..3047e5f --- /dev/null +++ b/examples/neuron/offline_inference_neuron.py @@ -0,0 +1,98 @@ +"""Offline inference with the full DeepSeek-V4.1-Flash on one trn2.48xlarge (Neuron basic mode). + +Same flow as examples/offline_inference.py (launch_dp -> one engine process per DP rank, each owning one +NeuronCore), with the settings basic mode needs: + * text-only prompts (VL kernels are stubbed in basic mode); + * no DSpark/MTP drafting (DEEPINFER_USE_NEXTN_DECODE=0) and no EPLB; + * a fixed KV-block pool instead of the engine's memory-profiling prefill (a max_num_batched_tokens dummy + prefill does not fit next to the eager pure-torch kernels); + * small batches and few output tokens: eager execution compiles every new op/shape on the first run. + + export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0 + python examples/neuron/offline_inference_neuron.py --model_name /mnt/nvme/models/DeepSeek-V4.1-Flash-1017 --ep 64 +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +from pathlib import Path + +EXAMPLES = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(EXAMPLES)) + +from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env, setup_cuda_visible_devices # noqa: E402 +from _launch import launch_dp # noqa: E402 + +PROMPTS = [ + "问题：法国的首都是哪里？\n答案：巴黎\n\n问题：一年有几个月？\n答案：12 个月\n\n问题：中国的首都是哪里？\n答案：", + "Question: What is the largest planet in the solar system?\nAnswer: Jupiter\n\n" + "Question: How many legs does a spider have?\nAnswer: Eight\n\n" + "Question: Who wrote Romeo and Juliet?\nAnswer:", +] + +parser = argparse.ArgumentParser(description="V4.1 offline inference on Neuron (basic mode)") +parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME, help="HF model id or local HF checkpoint directory") +parser.add_argument("--ep", default=64, type=int) +parser.add_argument("--max_tokens", default=16, type=int) +parser.add_argument("--num_kv_blocks", default=176, type=int, help="fixed KV pool per rank; must hold max_model_len tokens") +parser.add_argument("--max_model_len", default=16384, type=int, help="the candidate-block filter needs >= candidate_topk_blocks * candidate_block_size (2048 * 8)") +parser.add_argument("--max_num_batched_tokens", default=1024, type=int) +args = parser.parse_args() +dp_size = args.ep + + +def _fixed_kv_pool(num_blocks: int) -> None: + """Skip WorkerBase.profile_num_available_blocks (a max_num_batched_tokens dummy prefill).""" + from deepinfer.worker.base import WorkerBase + + def profile_num_available_blocks(self): # noqa: ANN001, ANN202 + self.model_runner.set_block_size(self.cache_config.block_size) + self.cache_config.num_gpu_blocks = num_blocks + return num_blocks + + WorkerBase.profile_num_available_blocks = profile_num_available_blocks + + +def main(dp_size: int, dp_rank: int) -> None: + os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0" + os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron") + setup_common_env(model_name=args.model_name, sampler_force_topk=4096, set_override_experts=False, enable_eplb=False) + # Pick this rank's NeuronCore before anything initializes deepinfer.PLATFORM (importing deepinfer.worker does): + # PlatformNeuron claims the core from CUDA_VISIBLE_DEVICES once, at init. + from deep_cluster.hardware import get_actual_num_gpus + + setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus()) + _fixed_kv_pool(args.num_kv_blocks) + + llm = create_llm( + args.model_name, + dp_rank, + ep_size=args.ep, + cp_size=1, + env={"DEEPINFER_LOGGING_INTERVAL_STEP": "1"}, + data_parallel_size=dp_size, + max_num_batched_tokens=args.max_num_batched_tokens, + max_model_len=args.max_model_len, + max_num_seqs=len(PROMPTS), + gpu_memory_utilization=0.85, + lookahead_tokens=0, + ) + + from deepinfer.sampling_params import SamplingParams + + params = [SamplingParams(temperature=0.0, max_tokens=args.max_tokens) for _ in PROMPTS] + outputs = llm.generate(prompts=PROMPTS, sampling_params=params, use_tqdm=False) + for i, output in enumerate(outputs): + text = llm.llm_engine.tokenizer.decode([x.token_id for x in output]) + logps = [x.token_logprobs for x in output] + ppl = math.exp(-sum(logps) / len(logps)) if logps else float("nan") + if dp_rank == 0: + print(f"DP {dp_rank}; INPUT: {PROMPTS[i]!r}; OUTPUT: {text!r}; output_tokens={len(output)}; decode_ppl={ppl}", flush=True) + llm.close() + + +if __name__ == "__main__": + launch_dp(main, dp_size)