mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
760e3cadc0
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
tilelang_ops interface tests (basic-mode stubs reported as skips via
nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
process group, device primitive checks, end-to-end vs golden on CPU and
on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
the port depends on or works around (fp8 strided-view indexing, uint64,
1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
70 lines
3.6 KiB
Bash
Executable File
70 lines
3.6 KiB
Bash
Executable File
#!/bin/bash
|
|
# One-time setup of a fresh Trn2 (Neuron DLAMI, SDK 2.33+) for the DeepInfer Neuron port.
|
|
#
|
|
# git clone -b neuron https://github.com/whn09/DeepInfer.git && cd DeepInfer
|
|
# bash tests/neuron/setup_trn2.sh # WORK=~/di by default
|
|
#
|
|
# Creates $WORK/venv (layered on the DLAMI's read-only torch-neuronx venv), builds the C++ extension,
|
|
# downloads the official reference implementation + tokenizer, generates the tiny checkpoint and the
|
|
# CPU golden. Idempotent: finished steps are skipped. Host CPU only -- never touches a NeuronCore.
|
|
set -euo pipefail
|
|
|
|
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
|
WORK="${WORK:-$HOME/di}"
|
|
OPT_VENV="${OPT_VENV:-/opt/aws_neuronx_venv_torch_neuronx_2_14_0}"
|
|
HF_REPO="deepseek-ai/DeepSeek-V4.1-Flash-1017"
|
|
mkdir -p "$WORK"
|
|
log() { echo "[setup $(date +%H:%M:%S)] $*"; }
|
|
|
|
# 1. venv layered on the read-only DLAMI torch-neuronx venv
|
|
if [ ! -x "$WORK/venv/bin/python" ]; then
|
|
log "venv -> $WORK/venv (layered on $OPT_VENV)"
|
|
PY="$(readlink -f "$OPT_VENV/bin/python")"
|
|
"$PY" -m venv "$WORK/venv"
|
|
SP="$("$WORK/venv/bin/python" -c 'import site; print(site.getsitepackages()[0])')"
|
|
echo "import site; site.addsitedir('$OPT_VENV/lib/python$("$PY" -c 'import sys; print(f"{sys.version_info[0]}.{sys.version_info[1]}")')/site-packages')" > "$SP/zz_opt_neuron.pth"
|
|
# neuronx-cc and the neuron tools live in the DLAMI venv
|
|
echo "export PATH=\"\$PATH:$OPT_VENV/bin\"" >> "$WORK/venv/bin/activate"
|
|
fi
|
|
# shellcheck disable=SC1091
|
|
source "$WORK/venv/bin/activate"
|
|
python -c "import torch, torch_neuronx; print('torch', torch.__version__, 'torch-neuronx', torch_neuronx.__version__)"
|
|
which neuronx-cc >/dev/null || { echo "neuronx-cc not on PATH"; exit 1; }
|
|
|
|
# 2. Python deps: requirements.txt minus torch / tilelang / tile-kernels (pure-torch shims are vendored)
|
|
log "pip deps"
|
|
pip install -q einops "huggingface_hub<1" ninja orjson pillow posix_ipc pybind11 pyzmq requests scipy sympy tabulate \
|
|
tokenizers tqdm "transformers==4.50.0" xgrammar pytest psutil safetensors
|
|
|
|
# 3. C++ extension (host-only, -DWITH_NEURON)
|
|
if ! ls "$REPO"/deepinfer/_so/neuron-*/_C*.so >/dev/null 2>&1; then
|
|
log "build C++ extension"
|
|
(cd "$REPO" && DEEPINFER_ACCELERATOR=neuron MAX_JOBS=16 bash ./scripts/build.sh build > "$WORK/build.log" 2>&1) \
|
|
|| { tail -30 "$WORK/build.log"; exit 1; }
|
|
fi
|
|
|
|
# 4. official reference implementation + tokenizer + real config (small files only, no weights)
|
|
REF="$WORK/ref_dsv41"
|
|
if [ ! -f "$REF/inference/model.py" ]; then
|
|
log "download reference implementation from $HF_REPO"
|
|
mkdir -p "$REF/inference"
|
|
for f in inference/model.py inference/kernel.py inference/convert.py inference/engram.py inference/generate.py \
|
|
inference/vision.py inference/image_processor.py inference/config.json tokenizer.json tokenizer_config.json config.json; do
|
|
curl -fsSL "https://huggingface.co/$HF_REPO/resolve/main/$f" -o "$REF/$f"
|
|
done
|
|
fi
|
|
|
|
# 5. tiny checkpoint (~1.6 GB) and the CPU golden from the official reference
|
|
if [ ! -f "$WORK/tiny_ckpt/config.json" ]; then
|
|
log "tiny checkpoint -> $WORK/tiny_ckpt"
|
|
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python "$REPO/examples/neuron/make_tiny_checkpoint.py" \
|
|
--out "$WORK/tiny_ckpt" --tokenizer-dir "$REF" --real-config "$REF/config.json"
|
|
fi
|
|
if [ ! -f "$WORK/golden/long.pt" ]; then
|
|
log "golden (official reference on CPU) -> $WORK/golden"
|
|
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I "$REPO/examples/neuron/reference_golden.py" \
|
|
--hf-ckpt "$WORK/tiny_ckpt" --ref-dir "$REF/inference" --converted "$WORK/tiny_ckpt_ref_mp1" \
|
|
--out "$WORK/golden" --threads 32
|
|
fi
|
|
log "done. Next: bash tests/neuron/run_all.sh"
|