mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
d0f7e4f69c
- requirements-neuron.txt: requirements.txt + dev deps without tilelang / tile-kernels (pure-torch drop-ins are vendored), torch 2.14 + torch-neuronx 2.14 from the Neuron repository. - build_vendor.sh: nothing to build for neuron; setup.py drops tilelang / tile-kernels from install_requires on neuron. - setup.py and get_platform(): an exported DEEPINFER_ACCELERATOR wins over device detection, so images build where no accelerator is visible. - docker/neuron/Dockerfile (+ per-Dockerfile ignore): on the native PyTorch Neuron DLC public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314- sdk2.33.0-ubuntu26.04; build.sh with TORCH_DEVICE_BACKEND_AUTOLOAD=0 and an import self-check in Neuron CPU mode. - README: Trainium in supported hardware and a "DeepInfer on AWS Trainium" section (English + folded Chinese); the pip torch-neuronx wheel needs glibc >= 2.43, so on Ubuntu 24.04 layer the venv on the DLAMI env. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
32 lines
1.8 KiB
Docker
32 lines
1.8 KiB
Docker
# 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"]
|