diff --git a/README.md b/README.md index 7143caa..9ee8e58 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,7 @@ Supported hardware: - NVIDIA SM100 GPUs - Huawei Ascend 950 NPUs +- AWS Trainium2 (native torch-neuronx, basic mode — see [DeepInfer on AWS Trainium](#deepinfer-on-aws-trainium)) Required for a complete inference service (not included in this release): @@ -50,13 +51,14 @@ Required for a complete inference service (not included in this release): - Python 3.12 - CUDA 13 (NVIDIA only) - CANN 9.2 (Ascend only) +- Neuron SDK 2.33 with torch-neuronx 2.14 (Trainium only) ### Build from source Before building DeepInfer, please create and activate an empty Python 3.12 environment. ``` bash -# Use `ascend` instead for the Ascend platform +# Use `ascend` instead for the Ascend platform, `neuron` for AWS Trainium # Keep it exported in the environment where you run DeepInfer export DEEPINFER_ACCELERATOR=nvidia ./build.sh @@ -74,6 +76,74 @@ python examples/offline_inference.py See [examples](examples/README.md) for details. +## DeepInfer on AWS Trainium + +DeepInfer runs on Trainium2 through the native PyTorch `neuron` device (torch-neuronx 2.14, no torch-xla), +structured like the Ascend port: `PlatformNeuron`, the `tilelang_ops/neuron` kernel backend and pure-torch +drop-ins for DeepGEMM / FlashMLA / DeepEP / DeepSelect / TileKernels in `deepinfer/neuron_ops/vendor`. +Design notes and conventions: [deepinfer/neuron_ops/NEURON_PORT.md](deepinfer/neuron_ops/NEURON_PORT.md). + +**Status — basic mode.** Eager execution; DSpark/MTP, zero-overhead scheduling, CP, EPLB, YOCO bounded replay and +VL are off. EP > 1 runs on real neuron collectives. A tiny checkpoint in the real FP8/FP4 layout matches the +official reference implementation within its own noise floor, on CPU and on a NeuronCore +([examples/neuron](examples/neuron/README.md)). Performance work (static graphs, NKI kernels) and the remaining +features are tracked in the [roadmap](https://github.com/whn09/DeepInfer/issues/19). + +On a Neuron DLAMI (trn2.48xlarge, SDK 2.33): + +``` bash +bash tests/neuron/setup_trn2.sh # venv on the DLAMI torch-neuronx env, build, tiny checkpoint + golden +bash tests/neuron/run_all.sh cpu # host-only tests +CORES=0-7 bash tests/neuron/run_all.sh device +``` + +or build by hand in a venv layered on the DLAMI's torch-neuronx environment (the `torch-neuronx` wheel on the +Neuron pip repository needs glibc >= 2.43, newer than Ubuntu 24.04's; the DLAMI ships a compatible build): + +``` bash +python3.12 -m venv venv && echo "import site; site.addsitedir('/opt/aws_neuronx_venv_torch_neuronx_2_14_0/lib/python3.12/site-packages')" \ + > venv/lib/python3.12/site-packages/zz_neuron.pth && source venv/bin/activate +export PATH=$PATH:/opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin # neuronx-cc +export DEEPINFER_ACCELERATOR=neuron +./build.sh # requirements-neuron.txt (torch / torch-neuronx already satisfied) + C++ build +python examples/neuron/offline_inference_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 --ep 64 +``` + +
+中文 + +DeepInfer 通过原生 PyTorch 的 `neuron` 设备(torch-neuronx 2.14,不依赖 torch-xla)在 Trainium2 上运行,结构与 +Ascend 移植一致:`PlatformNeuron`、`tilelang_ops/neuron` kernel 后端,以及 `deepinfer/neuron_ops/vendor` 中 +DeepGEMM / FlashMLA / DeepEP / DeepSelect / TileKernels 的纯 torch 实现。设计说明与约定见 +[deepinfer/neuron_ops/NEURON_PORT.md](deepinfer/neuron_ops/NEURON_PORT.md)。 + +**现状:basic mode。** 逐算子即时执行;DSpark/MTP、零开销调度、CP、EPLB、YOCO 有界重放和 VL 均关闭。EP > 1 使用 +真实的 neuron 集合通信。按真实 FP8/FP4 布局生成的小 checkpoint 在 CPU 和 NeuronCore 上都与官方参考实现一致(误差在 +参考实现自身的噪声范围内,见 [examples/neuron](examples/neuron/README.md))。性能优化(静态图、NKI kernel)和其余特性 +在 [roadmap](https://github.com/whn09/DeepInfer/issues/19) 中跟踪。 + +在 Neuron DLAMI(trn2.48xlarge,SDK 2.33)上: + +``` bash +bash tests/neuron/setup_trn2.sh # 基于 DLAMI 的 torch-neuronx 环境建 venv、编译、生成小 checkpoint 和真值 +bash tests/neuron/run_all.sh cpu # 只用 host 的测试 +CORES=0-7 bash tests/neuron/run_all.sh device +``` + +或者在叠加于 DLAMI torch-neuronx 环境之上的 venv 里手动构建(Neuron pip 源上的 `torch-neuronx` wheel 需要 +glibc >= 2.43,比 Ubuntu 24.04 的版本新;DLAMI 自带的是兼容的构建): + +``` bash +python3.12 -m venv venv && echo "import site; site.addsitedir('/opt/aws_neuronx_venv_torch_neuronx_2_14_0/lib/python3.12/site-packages')" \ + > venv/lib/python3.12/site-packages/zz_neuron.pth && source venv/bin/activate +export PATH=$PATH:/opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin # neuronx-cc +export DEEPINFER_ACCELERATOR=neuron +./build.sh # requirements-neuron.txt(torch / torch-neuronx 已满足)+ C++ 编译 +python examples/neuron/offline_inference_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 --ep 64 +``` + +
+ ## Performance benchmarks | Prefill | Decode | diff --git a/deepinfer/platform.py b/deepinfer/platform.py index 745e736..b977b80 100644 --- a/deepinfer/platform.py +++ b/deepinfer/platform.py @@ -285,6 +285,7 @@ class PlatformNeuron(Platform): def get_platform() -> Platform: from deep_cluster.hardware import accelerator_kind - kind = accelerator_kind() + # An explicitly exported DEEPINFER_ACCELERATOR wins (e.g. Neuron CPU mode on a host without /dev/neuron*). + kind = os.environ.get("DEEPINFER_ACCELERATOR") or accelerator_kind() platform_class = {"ascend": PlatformNPU, "neuron": PlatformNeuron}.get(kind, PlatformGPU) return platform_class() diff --git a/docker/neuron/Dockerfile b/docker/neuron/Dockerfile new file mode 100644 index 0000000..bc01407 --- /dev/null +++ b/docker/neuron/Dockerfile @@ -0,0 +1,31 @@ +# DeepInfer on AWS Trainium (Neuron basic mode), on top of the native PyTorch Neuron DLC. +# +# The base image ships torch 2.14 + torch-neuronx 2.14 (no torch-xla), neuronx-cc, nki and the Neuron tools on +# Ubuntu 26.04 / Python 3.14 (Neuron DLCs are Ubuntu 26.04-only from SDK 2.33; the torch-neuronx wheel needs +# glibc >= 2.43). The Neuron driver (aws-neuronx-dkms) must be installed on the host. +# +# docker build -f docker/neuron/Dockerfile -t deepinfer-neuron . +# docker run --rm -it --privileged --ipc=host --network=host \ +# -v /mnt/nvme/models:/models deepinfer-neuron \ +# python examples/neuron/offline_inference_neuron.py --model_name /models/DeepSeek-V4.1-Flash-1017 --ep 64 +ARG BASE_IMAGE=public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314-sdk2.33.0-ubuntu26.04 +FROM ${BASE_IMAGE} + +ENV DEEPINFER_ACCELERATOR=neuron \ + PIP_NO_CACHE_DIR=1 \ + PIP_ROOT_USER_ACTION=ignore + +WORKDIR /opt/deepinfer +COPY . . + +# requirements-neuron.txt (torch / torch-neuronx are already satisfied by the base image), no vendor build +# (Neuron uses the pure-torch drop-ins in deepinfer/neuron_ops/vendor), then the host-only C++ extension. +# No NeuronCore is visible during `docker build`: keep torch from auto-loading torch_neuronx (it queries the device +# when it initializes); compiling the extension only needs the torch headers. The self-check runs in Neuron CPU mode. +RUN TORCH_DEVICE_BACKEND_AUTOLOAD=0 MAX_JOBS=16 ./build.sh \ + && TORCH_DEVICE_BACKEND_AUTOLOAD=0 DEEPINFER_NEURON_DEVICE=cpu python -c "import deepinfer, torch; deepinfer.PLATFORM; \ +from deepinfer._C import engine; from deepinfer.model_executor.layers.tilelang_ops import ops; \ +print('deepinfer-neuron OK:', type(deepinfer.PLATFORM).__name__, 'torch', torch.__version__)" + +ENV PYTHONPATH=/opt/deepinfer +CMD ["/bin/bash"] diff --git a/docker/neuron/Dockerfile.dockerignore b/docker/neuron/Dockerfile.dockerignore new file mode 100644 index 0000000..daed45d --- /dev/null +++ b/docker/neuron/Dockerfile.dockerignore @@ -0,0 +1,8 @@ +# Applies to docker/neuron/Dockerfile only (BuildKit per-Dockerfile ignore file). +.git +build +deepinfer/_so +**/__pycache__ +*.pyc +# Kernel sources for CUDA / Ascend; Neuron images use the vendored pure-torch packages instead. +vendor diff --git a/requirements-neuron.txt b/requirements-neuron.txt new file mode 100644 index 0000000..465785f --- /dev/null +++ b/requirements-neuron.txt @@ -0,0 +1,46 @@ +# AWS Trainium (Neuron SDK 2.33+, native torch-neuronx -- no torch-xla). +# Same dependencies as requirements.txt + requirements-dev.txt, except: +# * tilelang / tile-kernels: not needed -- Neuron uses the pure-torch drop-ins vendored in +# deepinfer/neuron_ops/vendor (deep_gemm, flash_mla, deep_ep, deep_select, tile_kernels, tilelang); +# * torch comes with torch-neuronx (torch 2.14), not the CUDA/NPU builds. +# On a Neuron DLAMI the torch-neuronx stack is already in /opt/aws_neuronx_venv_torch_neuronx_2_14_0; +# tests/neuron/setup_trn2.sh layers a venv on it instead of reinstalling. +--extra-index-url https://pip.repos.neuron.amazonaws.com +--find-links https://download.pytorch.org/whl/cpu/torch/ +torch==2.14.0 +torch-neuronx==2.14.* +neuronx-cc==2.* +nki>=0.6 + +einops==0.8.2 +huggingface_hub==0.36.2 +ninja==1.13.0 +numpy==2.3.5 +orjson==3.11.7 +pillow==11.3.0 +posix_ipc==1.3.2 +psutil==7.2.2 +pybind11==3.0.1 +pyzmq==27.1.0 +requests==2.32.5 +safetensors==0.8.0 +scipy==1.17.0 +sympy==1.14.0 +tabulate==0.9.0 +tokenizers==0.21.4 +tqdm==4.67.3 +transformers==4.50.0 +xgrammar==0.2.2 + +# development / tests +clang-format==22.1.3 +execnet==2.1.2 +pybind11-stubgen==2.5.5 +pytest==8.4.2 +pytest-xdist==3.8.0 +ruff==0.15.0 +tree-sitter==0.26.0 +tree-sitter-cpp==0.23.4 +matplotlib==3.11.2 +fastapi==0.128.0 +uvicorn==0.40.0 diff --git a/scripts/build_vendor.sh b/scripts/build_vendor.sh index 0c616a4..6c57cb7 100755 --- a/scripts/build_vendor.sh +++ b/scripts/build_vendor.sh @@ -15,8 +15,13 @@ case "${DEEPINFER_ACCELERATOR:-nvidia}" in PACKAGES=(DeepSelect FlashMLA DeepEP-Ascend DeepGEMM-Ascend) INSTALL=(deep_select flash_mla deep_ep deep_gemm) ;; + neuron) + # Neuron uses the pure-torch drop-ins vendored in deepinfer/neuron_ops/vendor; nothing to build. + echo "✅ Neuron: vendored kernel packages are pure torch (deepinfer/neuron_ops/vendor), nothing to build" + exit 0 + ;; *) - echo "Unsupported DEEPINFER_ACCELERATOR: '${DEEPINFER_ACCELERATOR}' (expect nvidia or ascend)" >&2 + echo "Unsupported DEEPINFER_ACCELERATOR: '${DEEPINFER_ACCELERATOR}' (expect nvidia, ascend or neuron)" >&2 exit 1 ;; esac diff --git a/setup.py b/setup.py index d58496f..b6ec01d 100644 --- a/setup.py +++ b/setup.py @@ -9,8 +9,11 @@ from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CppExtension, C from deep_cluster.hardware import accelerator_kind -IS_NPU = accelerator_kind() == "ascend" -IS_NEURON = accelerator_kind() == "neuron" +# An explicitly exported DEEPINFER_ACCELERATOR (build.sh requires it) wins over device detection, so that images +# can be built where the accelerator is not visible (e.g. `docker build`). +_ACCELERATOR = os.environ.get("DEEPINFER_ACCELERATOR") or accelerator_kind() +IS_NPU = _ACCELERATOR == "ascend" +IS_NEURON = _ACCELERATOR == "neuron" ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) @@ -96,7 +99,11 @@ def get_requirements() -> list[str]: """Get Python package dependencies from requirements.txt (comments and blank lines stripped).""" with open(get_path("requirements.txt")) as f: lines = (line.split("#", 1)[0].strip() for line in f) - return [line for line in lines if line] + reqs = [line for line in lines if line] + if IS_NEURON: + # Neuron vendors pure-torch tilelang / tile_kernels (deepinfer/neuron_ops/vendor); see requirements-neuron.txt. + reqs = [r for r in reqs if not r.startswith(("tilelang", "tile-kernels"))] + return reqs setuptools.setup(