contrib: add MiMo-V2.5 (FP8 on Trn2)

contrib/models/MiMo-V2.5/perf_test/0_setup.sh

@@ -0,0 +1,113 @@
+#!/bin/bash
+# Setup for MiMo-V2.5 vLLM benchmarking on Trn2.
+#
+# Clones upstream vllm-project/vllm-neuron at release-0.5.0 and applies
+# vllm-neuron-patch.patch, which adds a runtime registration hook so the
+# contrib NeuronMiMoV2ForCausalLM is plugged into both NxDI's MODEL_TYPES
+# (under the key "mimov2") and vLLM's ModelRegistry (as
+# MiMoV2ForCausalLM) at vllm-neuron plugin init time.
+#
+# Then downloads XiaomiMiMo/MiMo-V2.5 from HuggingFace (FP8 blockwise, ~320 GB).
+set -e
+
+echo "=========================================="
+echo "Setup: vllm-neuron + MiMo-V2.5 weights"
+echo "=========================================="
+
+source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate
+
+# Resolve repo-relative paths up front — we cd into $HOME/vllm-neuron below,
+# after which $0's relative form would no longer resolve.
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+PATCH_FILE="$SCRIPT_DIR/vllm-neuron-patch.patch"
+CONTRIB_SRC="$(cd "$SCRIPT_DIR/.." && pwd)/src"
+
+echo ""
+echo "[1/2] Installing vllm-neuron (release-0.5.0) with the contrib registration patch..."
+
+if [ ! -d $HOME/vllm-neuron ]; then
+    git clone --branch release-0.5.0 https://github.com/vllm-project/vllm-neuron.git $HOME/vllm-neuron
+fi
+
+cd $HOME/vllm-neuron
+
+# Apply patch (idempotent via `git apply --check` first).
+if git apply --check "$PATCH_FILE" 2>/dev/null; then
+    git apply "$PATCH_FILE"
+    echo "  Applied $PATCH_FILE"
+else
+    echo "  Patch already applied or conflicts; continuing."
+fi
+
+pip install --extra-index-url=https://pip.repos.neuron.amazonaws.com -e .
+
+python3 -c "import vllm_neuron; print('vllm-neuron installed:', vllm_neuron.__file__)"
+
+echo ""
+echo "[2/2] Downloading MiMo-V2.5 FP8 weights from HuggingFace..."
+
+MIMO_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5}"
+if [ -d "$MIMO_PATH" ] && [ "$(ls "$MIMO_PATH"/*.safetensors 2>/dev/null | wc -l)" -gt 0 ]; then
+    echo "  MiMo-V2.5 weights already exist at $MIMO_PATH, skipping download"
+else
+    mkdir -p "$MIMO_PATH"
+    huggingface-cli download XiaomiMiMo/MiMo-V2.5 --local-dir "$MIMO_PATH" --max-workers 16
+    echo "  Download complete: $(du -sh $MIMO_PATH | cut -f1)"
+fi
+
+NEURON_FP8_PATH="${MIMO_PATH}-Neuron-FP8"
+COMPILED_PATH="/opt/dlami/nvme/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8_vllm"
+
+echo ""
+echo "========================================================================"
+echo "Next steps"
+echo "========================================================================"
+echo ""
+echo "1. Preprocess the FP8 checkpoint for Neuron (~16 min, ~15 GB peak RAM):"
+echo ""
+echo "     python $CONTRIB_SRC/conversion_script/preprocess_mimo_v2_5_fp8.py \\"
+echo "         --hf_model_path $MIMO_PATH \\"
+echo "         --save_path     $NEURON_FP8_PATH \\"
+echo "         --tp_degree 64"
+echo ""
+echo "2. Export the environment variables used by the smoke / bench scripts:"
+echo ""
+echo "     # --- Required ---"
+echo "     # Contrib package src (registers NeuronMiMoV2ForCausalLM with vllm-neuron)."
+echo "     export NXDI_CONTRIB_MIMO_V2_5_SRC=$CONTRIB_SRC"
+echo "     # vLLM's builtin arch validator only knows MiMoV2FlashForCausalLM, so the"
+echo "     # preprocess rewrites the checkpoint's config.json architectures to that"
+echo "     # name. Alias V2.5 src to the Flash env var so vllm-neuron's contrib hook"
+echo "     # registers mimov2flash -> our V2.5 NeuronMiMoV2ForCausalLM class."
+echo "     export NXDI_CONTRIB_MIMO_V2_FLASH_SRC=\"\$NXDI_CONTRIB_MIMO_V2_5_SRC\""
+echo "     # Preprocessed Neuron-FP8 checkpoint."
+echo "     export MIMO_V2_5_PATH=$NEURON_FP8_PATH"
+echo ""
+echo "     # --- Optional (recommended) ---"
+echo "     # vLLM compiles into <checkpoint>/neuron-compiled-artifacts/<hash>/ by"
+echo "     # default. Pin it to a persistent shared location so multiple configs"
+echo "     # don't collide and you can reuse the NEFF / sharded weights across runs."
+echo "     export NEURON_COMPILED_ARTIFACTS=$COMPILED_PATH"
+echo "     # NxDI's HLO/NEFF staging workdir (.hlo_module.pb etc). Default is"
+echo "     # /tmp/nxd_model/<compile-name>/; on Trn2 /tmp is wiped by the nightly"
+echo "     # reboot, and parallel compiles sharing the same basename silently"
+echo "     # overwrite each other's staged HLOs. Pin to a unique per-config"
+echo "     # directory under persistent storage."
+echo "     export BASE_COMPILE_WORK_DIR=/opt/dlami/nvme/tmp/nxd_model/\$(basename $COMPILED_PATH)"
+echo "     # First-time compile of V2.5's 256-expert MoE takes ~30 min (NEFF HLO +"
+echo "     # shard_checkpoint for 64 ranks). Extend vLLM's ready timeout."
+echo "     export VLLM_ENGINE_READY_TIMEOUT_S=7200"
+echo ""
+echo "3a. Run the one-shot benchmark (launches + benches + tears down):"
+echo ""
+echo "      bash $SCRIPT_DIR/bench_mimo_v2_5.sh"
+echo ""
+echo "3b. ...OR keep a server up and probe it manually:"
+echo ""
+echo "      # shell 1: server in foreground (Ctrl-C to stop)"
+echo "      bash $SCRIPT_DIR/start_vllm_server.sh"
+echo ""
+echo "      # shell 2: once 'Application startup complete.' prints,"
+echo "      bash $SCRIPT_DIR/sanity_check.sh"
+echo "      CONCURRENCY=16 NUM_PROMPTS=128 bash $SCRIPT_DIR/run_bench_single.sh"
+echo ""

contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+set -e
+
+# MiMo-V2.5 FP8 vLLM benchmark on Trn2. One-shot wrapper:
+#   launch server -> sanity check -> bench at c=1,16,32 -> stop server.
+#
+# This script composes three building blocks in perf_test/:
+#   start_vllm_server.sh  - server launch + env-var setup (backgrounded here)
+#   sanity_check.sh       - one-shot curl against the running server
+#   run_bench_single.sh   - one concurrency level of `vllm bench serve`
+#
+# Use those directly if you want to keep a long-running server and iterate
+# on bench parameters from another shell.
+#
+# Server recipe: TP=64, moe_tp=1/moe_ep=64, BS=32, continuous batching.
+# BS=32 is the smallest working batch size on the FP8 path (NxDI's TKG
+# path refuses Expert Parallelism with BS < num_experts/top_k = 256/8 = 32).
+# BS=1 single-stream latency demos are not currently supported on V2.5 FP8.
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+PORT="${PORT:-8000}"
+RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5}"
+CONFIG_NAME="bs32_tp64_moetp1_ep64"
+
+mkdir -p "$RESULTS_DIR"
+
+# Wait for vLLM server to be ready. First-time compile of the 256-expert
+# MoE model takes ~30 min and can stretch past 2 h under contention, so
+# poll for up to 2 h.
+wait_for_server() {
+    echo "  Waiting for vLLM server on port $PORT (up to 2 h for first compile)..."
+    local interval=10
+    local max_attempts=720
+    local start=$SECONDS
+    for i in $(seq 1 $max_attempts); do
+        if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
+            echo "  Server ready after $((SECONDS - start))s."
+            return 0
+        fi
+        if [ $((i % 6)) -eq 0 ]; then
+            echo "    ...still waiting ($((SECONDS - start))s elapsed)"
+        fi
+        sleep $interval
+    done
+    echo "  ERROR: Server did not start within $((max_attempts * interval))s"
+    return 1
+}
+
+stop_server() {
+    echo "  Stopping vLLM server..."
+    pkill -f "vllm.entrypoints.openai.api_server" 2>/dev/null || true
+    sleep 5
+}
+
+echo "=========================================="
+echo "MiMo-V2.5 FP8 Performance Benchmark"
+echo "=========================================="
+echo "Port:    $PORT"
+echo "Results: $RESULTS_DIR"
+echo ""
+
+# Start the server in the background. start_vllm_server.sh handles all the
+# env vars (MODEL_PATH, NEURON_COMPILED_ARTIFACTS, BASE_COMPILE_WORK_DIR,
+# contrib src registration, etc.) and execs `python3 -m vllm...`.
+bash "$SCRIPT_DIR/start_vllm_server.sh" &
+SERVER_PID=$!
+trap stop_server EXIT
+
+wait_for_server
+
+# One-shot sanity check (curl the chat endpoint).
+PORT="$PORT" bash "$SCRIPT_DIR/sanity_check.sh" || true
+
+# Three concurrency levels. run_bench_single.sh reads knobs from the
+# environment; see its header for all the options.
+PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \
+    CONCURRENCY=1  NUM_PROMPTS=16  bash "$SCRIPT_DIR/run_bench_single.sh"
+PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \
+    CONCURRENCY=16 NUM_PROMPTS=128 bash "$SCRIPT_DIR/run_bench_single.sh"
+PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \
+    CONCURRENCY=32 NUM_PROMPTS=128 bash "$SCRIPT_DIR/run_bench_single.sh"
+
+echo "=========================================="
+echo "MiMo-V2.5 FP8 benchmark complete!"
+echo "Results saved to: $RESULTS_DIR"
+echo "=========================================="
+ls -la "$RESULTS_DIR"