From 1a3e80ece1b41f6e0eeb48af43b44e6dda1bda9f Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Wed, 1 Jul 2026 16:14:37 +0000 Subject: [PATCH 01/25] Add SGLang disaggregated P/D inference full-overlay Dockerfiles and tooling Introduce single full-overlay Dockerfiles for SGLang disaggregated prefill/decode inference that merge the RCCL, MoRI, and NIXL/Mooncake KV-transfer layers into one build (no base-image chaining), plus the supporting run and benchmark scripts. Key changes: - New Dockerfiles: `sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile` and the `*.oci-rdma62.*` variant (rdma-core v62 baked in for OCI-CX7 hosts). - `run.sh` native launcher with in-container rank-ordered node-IP discovery (`ip_rendezvous.py`), so `IPADDRS`/`SGLANG_NODE_IPS` need not be forwarded. - `parse_to_csv.py` rewritten for comprehensive metric extraction (best-throughput iteration) into the madengine perf-CSV schema. - Resilient benchmark sweep in `benchmark_xPyD.sh` (fail-fast + point retries) writing to the madengine-expected perf CSV path. - Updated `sglang_disagg_mori_io_ep.sh` / `sglang_disagg_server.sh` and README. Co-authored-by: Ilia Kosarev --- ...l_overlay.oci-rdma62.ubuntu.amd.Dockerfile | 270 ++++++++++++++ ...ference_full_overlay.ubuntu.amd.Dockerfile | 250 +++++++++++++ scripts/sglang_disagg/README.MD | 53 ++- scripts/sglang_disagg/benchmark_xPyD.sh | 184 +++++++--- scripts/sglang_disagg/ip_rendezvous.py | 151 ++++++++ scripts/sglang_disagg/parse_to_csv.py | 331 +++++++++++------- scripts/sglang_disagg/run.sh | 217 ++++++++++++ .../sglang_disagg/sglang_disagg_mori_io_ep.sh | 103 +++--- scripts/sglang_disagg/sglang_disagg_server.sh | 30 +- 9 files changed, 1352 insertions(+), 237 deletions(-) create mode 100644 docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile create mode 100644 docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile create mode 100755 scripts/sglang_disagg/ip_rendezvous.py create mode 100755 scripts/sglang_disagg/run.sh diff --git a/docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile b/docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile new file mode 100644 index 0000000..7d5ae50 --- /dev/null +++ b/docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile @@ -0,0 +1,270 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# SGLang Disaggregated P/D — MERGED full overlay (RCCL + MoRI + KV-transfer) +# +# Single-Dockerfile equivalent of the three chained overlays: +# base sglang -> RCCL overlay (+smifix) -> MoRI overlay -> KV-transfer (RIXL) +# Built in one `docker build` step so madengine's single-dockerfile build path +# produces the whole stack at once (no base_docker chaining between overlays). +# +# Pins (override via --build-arg): +# RCCL : ROCm/rocm-systems develop @ RCCL_COMMIT (default 78e8ba0) + smifix +# MoRI : ROCm/mori @ MORI_COMMIT (default a14e6992, includes #366) +# NIXL : ROCm/RIXL @ RIXL_COMMIT (default f33a5599) + ROCm/ucx @ UCX_COMMIT +# (the AMD NIXL implementation; exposed to SGLang as `nixl` via alias). +# NOTE: this is NOT ai-dynamo/nixl. The rocm KV transport for +# KV_TRANSFER_BACKEND=nixl is ROCm/RIXL; recipe mirrors the proven +# scripts/kvcache_transfer_bench/Dockerfile. +############################################################################### +ARG BASE_DOCKER=lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x +FROM $BASE_DOCKER + +SHELL ["/bin/bash", "-o", "pipefail", "-c"] +USER root + +############################################################################### +# 1) RCCL overlay — rebuild RCCL from source so the RCCL under test wins over +# the base image's librccl. (mirrors sglang_disagg_inference_rccl_overlay) +############################################################################### +ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git +ARG RCCL_BRANCH=develop +ARG RCCL_COMMIT=78e8ba0 +ARG RCCL_INSTALL_DIR=/opt/rccl +ARG BUILD_GPU_TARGETS=gfx942 + +ENV RCCL_HOME=${RCCL_INSTALL_DIR} +# Prepend the overlay RCCL so it wins over the base image's librccl. +ENV LD_LIBRARY_PATH=${RCCL_INSTALL_DIR}/lib:${LD_LIBRARY_PATH} + +RUN mkdir -p "${RCCL_INSTALL_DIR}" +WORKDIR /opt + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch && \ + rm -rf /var/lib/apt/lists/* + +RUN if [[ -n "${RCCL_COMMIT}" ]]; then \ + git clone "${RCCL_REPO}" /tmp/rccl; \ + else \ + git clone --depth 1 --branch "${RCCL_BRANCH}" "${RCCL_REPO}" /tmp/rccl; \ + fi && \ + cd /tmp/rccl && \ + if [[ -n "${RCCL_BRANCH}" ]]; then git checkout "${RCCL_BRANCH}" || true; fi && \ + if [[ -n "${RCCL_COMMIT}" ]]; then git checkout "${RCCL_COMMIT}"; fi && \ + if [[ -d projects/rccl ]]; then \ + cd projects/rccl && git submodule update --init --recursive; \ + echo "/tmp/rccl/projects/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + else \ + git submodule update --init --recursive; \ + echo "/tmp/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + fi + +RUN set -e && \ + BLD_RCCL_HOME=$(cat /tmp/BLD_RCCL_HOME.txt) && \ + cd "${BLD_RCCL_HOME}" && \ + ./install.sh --amdgpu_targets="${BUILD_GPU_TARGETS}" --prefix="${RCCL_INSTALL_DIR}" + +RUN rm -rf /tmp/rccl /tmp/BLD_RCCL_HOME.txt + +# Re-add the rocm_smi dependency the rocm-systems RCCL build drops (smifix). +# Newer rocm-systems RCCL builds librccl WITHOUT a DT_NEEDED on librocm_smi64.so; +# torch's libtorch_hip.so relies on librccl to transitively pull in rocm_smi, so +# without this `import torch` dies with "undefined symbol: rsmi_init". +RUN set -e; \ + command -v patchelf >/dev/null 2>&1 || pip install --no-cache-dir patchelf >/dev/null 2>&1 \ + || { apt-get -o Acquire::Retries=5 update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends patchelf && rm -rf /var/lib/apt/lists/*; }; \ + librccl="$(readlink -f ${RCCL_INSTALL_DIR}/lib/librccl.so)"; \ + smi_soname="$(basename "$(ls /opt/rocm*/lib/librocm_smi64.so.[0-9]* 2>/dev/null | grep -E 'librocm_smi64\.so\.[0-9]+$' | head -1)")"; \ + test -n "${smi_soname}" || { echo "ROCM_SMI_SONAME_NOT_FOUND"; exit 1; }; \ + if readelf -d "${librccl}" 2>/dev/null | grep -q "NEEDED.*${smi_soname}"; then \ + echo "RCCL_SMI_NEEDED_ALREADY_PRESENT ${smi_soname}"; \ + else \ + patchelf --add-needed "${smi_soname}" "${librccl}" && echo "RCCL_SMI_NEEDED_ADDED ${smi_soname} -> ${librccl}"; \ + fi + +# Sanity: overlay librccl present AND torch imports with overlay librccl first. +RUN set -e; \ + test -e "${RCCL_INSTALL_DIR}/lib/librccl.so" || { echo "RCCL_OVERLAY_MISSING"; exit 1; }; \ + echo "RCCL_OVERLAY_OK $(ls -l ${RCCL_INSTALL_DIR}/lib/librccl.so*)"; \ + python3 -c "import torch; print('RCCL_OVERLAY_TORCH_OK', torch.__version__)" + +RUN pip list 2>/dev/null | grep -iE "sglang|torch" || true + +############################################################################### +# 2) MoRI overlay — substitutable MoE EP all-to-all (dispatch/combine, IBGDA). +# Default ROCm/mori @ a14e6992 includes #366 (internode decode hang fix). +############################################################################### +ARG MORI_REPO=https://github.com/ROCm/mori.git +ARG MORI_BRANCH=main +ARG MORI_COMMIT=a14e6992ffa95478e83127fe2672afff2840856f +ARG MORI_WHEEL_URL= +ARG MORI_GPU_ARCHS=gfx942 +ARG MORI_VERSION=1.2.0 +ARG MORI_SRC_DIR=/sgl-workspace/mori + +ENV MORI_GPU_ARCHS=${MORI_GPU_ARCHS} \ + MORI_SKIP_PRECOMPILE=1 \ + CMAKE_BUILD_TYPE=Release \ + SETUPTOOLS_SCM_PRETEND_VERSION=${MORI_VERSION} + +WORKDIR /sgl-workspace + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch && \ + rm -rf /var/lib/apt/lists/* + +RUN set -e; \ + if [[ -n "${MORI_WHEEL_URL}" ]]; then \ + echo "[mori-overlay] installing prebuilt wheel: ${MORI_WHEEL_URL}"; \ + pip install --no-cache-dir --force-reinstall "${MORI_WHEEL_URL}"; \ + else \ + echo "[mori-overlay] source build ${MORI_REPO}@${MORI_BRANCH}${MORI_COMMIT:+ (${MORI_COMMIT})} archs=${MORI_GPU_ARCHS}"; \ + rm -rf "${MORI_SRC_DIR}"; \ + if [[ -n "${MORI_COMMIT}" ]]; then \ + git clone "${MORI_REPO}" "${MORI_SRC_DIR}"; \ + cd "${MORI_SRC_DIR}" && git checkout "${MORI_COMMIT}"; \ + else \ + git clone --depth 1 --branch "${MORI_BRANCH}" "${MORI_REPO}" "${MORI_SRC_DIR}"; \ + fi; \ + cd "${MORI_SRC_DIR}" && git submodule update --init --recursive || true; \ + pip install --no-cache-dir --force-reinstall .; \ + fi + +# Sanity: MoRI must import and report the overlay version. +RUN python3 -c "import mori, importlib.metadata as m; \ +print('MORI_OVERLAY_OK', getattr(mori,'__version__', m.version('amd_mori')))" \ + || { echo 'MORI_OVERLAY_IMPORT_FAILED'; exit 1; } + +############################################################################### +# 3) KV-transfer overlay — ROCm UCX + ROCm/RIXL (the AMD "nixl" KV transport). +# Mirrors the proven recipe in scripts/kvcache_transfer_bench/Dockerfile. +# SGLang's KV_TRANSFER_BACKEND=nixl imports `nixl`; we alias rixl -> nixl. +############################################################################### +ARG ROCM_PATH=/opt/rocm +ARG KV_WORKSPACE=/sgl-workspace +ARG UCX_REPO=https://github.com/ROCm/ucx.git +ARG UCX_COMMIT=da3fac2a +ARG RIXL_REPO=https://github.com/ROCm/RIXL.git +ARG RIXL_COMMIT=f33a5599 + +ENV UCX_HOME=${KV_WORKSPACE}/ucx +ENV RIXL_HOME=${KV_WORKSPACE}/rixl +ENV PATH=${UCX_HOME}/bin:${PATH} +ENV LD_LIBRARY_PATH=${RIXL_HOME}/lib/x86_64-linux-gnu:${UCX_HOME}/lib:${LD_LIBRARY_PATH} + +WORKDIR ${KV_WORKSPACE} + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch autoconf automake libtool \ + cython3 libaio-dev libibverbs-dev librdmacm-dev libpci-dev \ + libgflags-dev libgoogle-glog-dev && \ + rm -rf /var/lib/apt/lists/* + +# meson/ninja/pybind11 build frontends for the RIXL meson build + wheel. +RUN pip install --no-cache-dir -U meson ninja pybind11 pyyaml build wheel + +# --- ROCm UCX --------------------------------------------------------------- +RUN set -e; \ + git clone "${UCX_REPO}" "${UCX_HOME}.src" && cd "${UCX_HOME}.src" && \ + git checkout "${UCX_COMMIT}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${UCX_HOME}" --enable-shared --disable-static \ + --disable-doxygen-doc --enable-optimizations --enable-devel-headers \ + --with-rocm="${ROCM_PATH}" --with-verbs --with-dm --enable-mt && \ + make -j"$(nproc)" && make install && \ + cd "${KV_WORKSPACE}" && rm -rf "${UCX_HOME}.src" + +# --- ROCm/RIXL (the AMD NIXL implementation) -------------------------------- +RUN set -e; \ + git clone "${RIXL_REPO}" "${RIXL_HOME}" && cd "${RIXL_HOME}" && \ + git checkout "${RIXL_COMMIT}" && (git submodule update --init --recursive || true); \ + meson setup build --prefix="${RIXL_HOME}" -Ducx_path="${UCX_HOME}" -Drocm_path="${ROCM_PATH}" && \ + cd build && ninja && ninja install + +# Install the meson-built RIXL python package into site-packages directly. +# The upstream contrib/build-wheel.sh requires uv + py3.12 + auditwheel, which +# the py3.10 sglang base lacks; `ninja install` already produced the +# cpython-310 bindings under the RIXL prefix, so copy them into site-packages +# and alias `nixl` -> `rixl` for SGLang's KV_TRANSFER_BACKEND=nixl import path. +RUN set -e; \ + echo "=== RIXL python install tree ==="; \ + find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | sed 's/^/PYDIR: /'; \ + pysp="$(find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | head -1)"; \ + test -n "${pysp}" || { echo "RIXL_PY_INSTALL_NOT_FOUND"; find "${RIXL_HOME}" -name '_bindings*.so' -o -name '*.py' | head; exit 1; }; \ + echo "RIXL_PY_SRC=${pysp}"; ls -la "${pysp}"; \ + SP="$(python3 -c 'import site; print(site.getsitepackages()[0])')"; \ + copied=0; \ + for d in "${pysp}"/*/; do \ + name="$(basename "$d")"; \ + case "$name" in *dist-info|__pycache__) continue;; esac; \ + rm -rf "${SP}/${name}"; cp -a "$d" "${SP}/${name}"; echo "INSTALLED_PY_PKG ${name} -> ${SP}/${name}"; copied=1; \ + done; \ + test "$copied" = 1 || { echo "NO_RIXL_PKG_COPIED"; exit 1; }; \ + echo "import rixl, sys; sys.modules['nixl'] = rixl" > "${SP}/nixl_alias.pth"; \ + echo "NIXL_ALIAS_WRITTEN ${SP}/nixl_alias.pth" + +# Sanity: nixl (== rixl) import must succeed; report version + agent symbol. +RUN set -e; \ + python3 -c "import nixl; print('NIXL_IMPORT_OK', getattr(nixl,'__file__','?'))"; \ + python3 -c "import rixl, importlib.metadata as m; print('RIXL_VERSION', m.version('rixl'))" || true; \ + python3 -c "from nixl._api import nixl_agent, nixl_agent_config; print('NIXL_AGENT_OK')" \ + || echo "NIXL_AGENT_API_DIFFERS (verify sglang nixl import path at runtime)" + +############################################################################### +# 4) Runtime python deps + Mooncake KV-transfer backend. +# These were formerly built/pip-installed by the launcher at job start +# (scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh); baked into the image so +# the launcher no longer mutates the runtime environment. +############################################################################### +# Runtime python deps formerly pip-installed by the launcher. +RUN pip install --no-cache-dir py-spy pyyaml pandas \ + && pip install --no-cache-dir --ignore-installed --force-reinstall flask + +# Mooncake KV-transfer backend (KV_TRANSFER_BACKEND=mooncake). Mirrors +# docker/sglang_disagg_inference_kvtransfer_overlay.ubuntu.amd.Dockerfile:62-74. +# Default: pip wheel pin; set MOONCAKE_COMMIT for a source build. +ARG MOONCAKE_VERSION=0.3.6.post1 +ARG MOONCAKE_REPO=https://github.com/kvcache-ai/Mooncake.git +ARG MOONCAKE_COMMIT= +RUN set -e; \ + if [[ -n "${MOONCAKE_COMMIT}" ]]; then \ + echo "[full-overlay] Mooncake source ${MOONCAKE_REPO}@${MOONCAKE_COMMIT}"; \ + git clone "${MOONCAKE_REPO}" /tmp/mooncake && cd /tmp/mooncake && \ + git checkout "${MOONCAKE_COMMIT}" && (git submodule update --init --recursive || true); \ + pip install --no-cache-dir --force-reinstall . && rm -rf /tmp/mooncake; \ + elif [[ -n "${MOONCAKE_VERSION}" ]]; then \ + echo "[full-overlay] Mooncake pip mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ + pip install --no-cache-dir --force-reinstall "mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ + else \ + echo "[full-overlay] Mooncake: keeping base image version"; \ + fi + +# Sanity: mooncake must import alongside nixl/rixl (non-fatal — module path may vary). +RUN python3 -c "import mooncake; print('MOONCAKE_OVERLAY_OK')" \ + || echo "MOONCAKE_IMPORT_DIFFERS (verify sglang mooncake import path at runtime)" + +############################################################################### +# 5) OCI cluster workaround: build + install rdma-core v62 from source. +# The OCI-CX7 host stack needs a newer libibverbs/librdmacm/libmlx5 than the +# base image ships. Formerly built at job start by the launcher; baked here so +# the runtime env is not mutated. This variant is OCI-only; the base overlay +# keeps the image-default rdma-core. +############################################################################### +ARG RDMA_VER=v62.0 +RUN set -e; \ + git clone --branch "${RDMA_VER}" --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ + cd /tmp/rdma-core && mkdir -p build && cd build && \ + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ + ninja && ninja install && ldconfig && \ + rm -rf /tmp/rdma-core + +# Sanity: rebuilt libibverbs present and python still imports (linker not corrupted). +RUN set -e; \ + ls -l /usr/lib/libibverbs.so* 2>/dev/null || ls -l /usr/lib/*/libibverbs.so* 2>/dev/null || true; \ + python3 -c "import torch; print(\"OCI_RDMA62_TORCH_OK\", torch.__version__)" diff --git a/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile new file mode 100644 index 0000000..5e1f3d4 --- /dev/null +++ b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile @@ -0,0 +1,250 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# SGLang Disaggregated P/D — MERGED full overlay (RCCL + MoRI + KV-transfer) +# +# Single-Dockerfile equivalent of the three chained overlays: +# base sglang -> RCCL overlay (+smifix) -> MoRI overlay -> KV-transfer (RIXL) +# Built in one `docker build` step so madengine's single-dockerfile build path +# produces the whole stack at once (no base_docker chaining between overlays). +# +# Pins (override via --build-arg): +# RCCL : ROCm/rocm-systems develop @ RCCL_COMMIT (default 78e8ba0) + smifix +# MoRI : ROCm/mori @ MORI_COMMIT (default a14e6992, includes #366) +# NIXL : ROCm/RIXL @ RIXL_COMMIT (default f33a5599) + ROCm/ucx @ UCX_COMMIT +# (the AMD NIXL implementation; exposed to SGLang as `nixl` via alias). +# NOTE: this is NOT ai-dynamo/nixl. The rocm KV transport for +# KV_TRANSFER_BACKEND=nixl is ROCm/RIXL; recipe mirrors the proven +# scripts/kvcache_transfer_bench/Dockerfile. +############################################################################### +ARG BASE_DOCKER=lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x +FROM $BASE_DOCKER + +SHELL ["/bin/bash", "-o", "pipefail", "-c"] +USER root + +############################################################################### +# 1) RCCL overlay — rebuild RCCL from source so the RCCL under test wins over +# the base image's librccl. (mirrors sglang_disagg_inference_rccl_overlay) +############################################################################### +ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git +ARG RCCL_BRANCH=develop +ARG RCCL_COMMIT=78e8ba0 +ARG RCCL_INSTALL_DIR=/opt/rccl +ARG BUILD_GPU_TARGETS=gfx942 + +ENV RCCL_HOME=${RCCL_INSTALL_DIR} +# Prepend the overlay RCCL so it wins over the base image's librccl. +ENV LD_LIBRARY_PATH=${RCCL_INSTALL_DIR}/lib:${LD_LIBRARY_PATH} + +RUN mkdir -p "${RCCL_INSTALL_DIR}" +WORKDIR /opt + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch && \ + rm -rf /var/lib/apt/lists/* + +RUN if [[ -n "${RCCL_COMMIT}" ]]; then \ + git clone "${RCCL_REPO}" /tmp/rccl; \ + else \ + git clone --depth 1 --branch "${RCCL_BRANCH}" "${RCCL_REPO}" /tmp/rccl; \ + fi && \ + cd /tmp/rccl && \ + if [[ -n "${RCCL_BRANCH}" ]]; then git checkout "${RCCL_BRANCH}" || true; fi && \ + if [[ -n "${RCCL_COMMIT}" ]]; then git checkout "${RCCL_COMMIT}"; fi && \ + if [[ -d projects/rccl ]]; then \ + cd projects/rccl && git submodule update --init --recursive; \ + echo "/tmp/rccl/projects/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + else \ + git submodule update --init --recursive; \ + echo "/tmp/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + fi + +RUN set -e && \ + BLD_RCCL_HOME=$(cat /tmp/BLD_RCCL_HOME.txt) && \ + cd "${BLD_RCCL_HOME}" && \ + ./install.sh --amdgpu_targets="${BUILD_GPU_TARGETS}" --prefix="${RCCL_INSTALL_DIR}" + +RUN rm -rf /tmp/rccl /tmp/BLD_RCCL_HOME.txt + +# Re-add the rocm_smi dependency the rocm-systems RCCL build drops (smifix). +# Newer rocm-systems RCCL builds librccl WITHOUT a DT_NEEDED on librocm_smi64.so; +# torch's libtorch_hip.so relies on librccl to transitively pull in rocm_smi, so +# without this `import torch` dies with "undefined symbol: rsmi_init". +RUN set -e; \ + command -v patchelf >/dev/null 2>&1 || pip install --no-cache-dir patchelf >/dev/null 2>&1 \ + || { apt-get -o Acquire::Retries=5 update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends patchelf && rm -rf /var/lib/apt/lists/*; }; \ + librccl="$(readlink -f ${RCCL_INSTALL_DIR}/lib/librccl.so)"; \ + smi_soname="$(basename "$(ls /opt/rocm*/lib/librocm_smi64.so.[0-9]* 2>/dev/null | grep -E 'librocm_smi64\.so\.[0-9]+$' | head -1)")"; \ + test -n "${smi_soname}" || { echo "ROCM_SMI_SONAME_NOT_FOUND"; exit 1; }; \ + if readelf -d "${librccl}" 2>/dev/null | grep -q "NEEDED.*${smi_soname}"; then \ + echo "RCCL_SMI_NEEDED_ALREADY_PRESENT ${smi_soname}"; \ + else \ + patchelf --add-needed "${smi_soname}" "${librccl}" && echo "RCCL_SMI_NEEDED_ADDED ${smi_soname} -> ${librccl}"; \ + fi + +# Sanity: overlay librccl present AND torch imports with overlay librccl first. +RUN set -e; \ + test -e "${RCCL_INSTALL_DIR}/lib/librccl.so" || { echo "RCCL_OVERLAY_MISSING"; exit 1; }; \ + echo "RCCL_OVERLAY_OK $(ls -l ${RCCL_INSTALL_DIR}/lib/librccl.so*)"; \ + python3 -c "import torch; print('RCCL_OVERLAY_TORCH_OK', torch.__version__)" + +RUN pip list 2>/dev/null | grep -iE "sglang|torch" || true + +############################################################################### +# 2) MoRI overlay — substitutable MoE EP all-to-all (dispatch/combine, IBGDA). +# Default ROCm/mori @ a14e6992 includes #366 (internode decode hang fix). +############################################################################### +ARG MORI_REPO=https://github.com/ROCm/mori.git +ARG MORI_BRANCH=main +ARG MORI_COMMIT=a14e6992ffa95478e83127fe2672afff2840856f +ARG MORI_WHEEL_URL= +ARG MORI_GPU_ARCHS=gfx942 +ARG MORI_VERSION=1.2.0 +ARG MORI_SRC_DIR=/sgl-workspace/mori + +ENV MORI_GPU_ARCHS=${MORI_GPU_ARCHS} \ + MORI_SKIP_PRECOMPILE=1 \ + CMAKE_BUILD_TYPE=Release \ + SETUPTOOLS_SCM_PRETEND_VERSION=${MORI_VERSION} + +WORKDIR /sgl-workspace + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch && \ + rm -rf /var/lib/apt/lists/* + +RUN set -e; \ + if [[ -n "${MORI_WHEEL_URL}" ]]; then \ + echo "[mori-overlay] installing prebuilt wheel: ${MORI_WHEEL_URL}"; \ + pip install --no-cache-dir --force-reinstall "${MORI_WHEEL_URL}"; \ + else \ + echo "[mori-overlay] source build ${MORI_REPO}@${MORI_BRANCH}${MORI_COMMIT:+ (${MORI_COMMIT})} archs=${MORI_GPU_ARCHS}"; \ + rm -rf "${MORI_SRC_DIR}"; \ + if [[ -n "${MORI_COMMIT}" ]]; then \ + git clone "${MORI_REPO}" "${MORI_SRC_DIR}"; \ + cd "${MORI_SRC_DIR}" && git checkout "${MORI_COMMIT}"; \ + else \ + git clone --depth 1 --branch "${MORI_BRANCH}" "${MORI_REPO}" "${MORI_SRC_DIR}"; \ + fi; \ + cd "${MORI_SRC_DIR}" && git submodule update --init --recursive || true; \ + pip install --no-cache-dir --force-reinstall .; \ + fi + +# Sanity: MoRI must import and report the overlay version. +RUN python3 -c "import mori, importlib.metadata as m; \ +print('MORI_OVERLAY_OK', getattr(mori,'__version__', m.version('amd_mori')))" \ + || { echo 'MORI_OVERLAY_IMPORT_FAILED'; exit 1; } + +############################################################################### +# 3) KV-transfer overlay — ROCm UCX + ROCm/RIXL (the AMD "nixl" KV transport). +# Mirrors the proven recipe in scripts/kvcache_transfer_bench/Dockerfile. +# SGLang's KV_TRANSFER_BACKEND=nixl imports `nixl`; we alias rixl -> nixl. +############################################################################### +ARG ROCM_PATH=/opt/rocm +ARG KV_WORKSPACE=/sgl-workspace +ARG UCX_REPO=https://github.com/ROCm/ucx.git +ARG UCX_COMMIT=da3fac2a +ARG RIXL_REPO=https://github.com/ROCm/RIXL.git +ARG RIXL_COMMIT=f33a5599 + +ENV UCX_HOME=${KV_WORKSPACE}/ucx +ENV RIXL_HOME=${KV_WORKSPACE}/rixl +ENV PATH=${UCX_HOME}/bin:${PATH} +ENV LD_LIBRARY_PATH=${RIXL_HOME}/lib/x86_64-linux-gnu:${UCX_HOME}/lib:${LD_LIBRARY_PATH} + +WORKDIR ${KV_WORKSPACE} + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch autoconf automake libtool \ + cython3 libaio-dev libibverbs-dev librdmacm-dev libpci-dev \ + libgflags-dev libgoogle-glog-dev && \ + rm -rf /var/lib/apt/lists/* + +# meson/ninja/pybind11 build frontends for the RIXL meson build + wheel. +RUN pip install --no-cache-dir -U meson ninja pybind11 pyyaml build wheel + +# --- ROCm UCX --------------------------------------------------------------- +RUN set -e; \ + git clone "${UCX_REPO}" "${UCX_HOME}.src" && cd "${UCX_HOME}.src" && \ + git checkout "${UCX_COMMIT}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${UCX_HOME}" --enable-shared --disable-static \ + --disable-doxygen-doc --enable-optimizations --enable-devel-headers \ + --with-rocm="${ROCM_PATH}" --with-verbs --with-dm --enable-mt && \ + make -j"$(nproc)" && make install && \ + cd "${KV_WORKSPACE}" && rm -rf "${UCX_HOME}.src" + +# --- ROCm/RIXL (the AMD NIXL implementation) -------------------------------- +RUN set -e; \ + git clone "${RIXL_REPO}" "${RIXL_HOME}" && cd "${RIXL_HOME}" && \ + git checkout "${RIXL_COMMIT}" && (git submodule update --init --recursive || true); \ + meson setup build --prefix="${RIXL_HOME}" -Ducx_path="${UCX_HOME}" -Drocm_path="${ROCM_PATH}" && \ + cd build && ninja && ninja install + +# Install the meson-built RIXL python package into site-packages directly. +# The upstream contrib/build-wheel.sh requires uv + py3.12 + auditwheel, which +# the py3.10 sglang base lacks; `ninja install` already produced the +# cpython-310 bindings under the RIXL prefix, so copy them into site-packages +# and alias `nixl` -> `rixl` for SGLang's KV_TRANSFER_BACKEND=nixl import path. +RUN set -e; \ + echo "=== RIXL python install tree ==="; \ + find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | sed 's/^/PYDIR: /'; \ + pysp="$(find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | head -1)"; \ + test -n "${pysp}" || { echo "RIXL_PY_INSTALL_NOT_FOUND"; find "${RIXL_HOME}" -name '_bindings*.so' -o -name '*.py' | head; exit 1; }; \ + echo "RIXL_PY_SRC=${pysp}"; ls -la "${pysp}"; \ + SP="$(python3 -c 'import site; print(site.getsitepackages()[0])')"; \ + copied=0; \ + for d in "${pysp}"/*/; do \ + name="$(basename "$d")"; \ + case "$name" in *dist-info|__pycache__) continue;; esac; \ + rm -rf "${SP}/${name}"; cp -a "$d" "${SP}/${name}"; echo "INSTALLED_PY_PKG ${name} -> ${SP}/${name}"; copied=1; \ + done; \ + test "$copied" = 1 || { echo "NO_RIXL_PKG_COPIED"; exit 1; }; \ + echo "import rixl, sys; sys.modules['nixl'] = rixl" > "${SP}/nixl_alias.pth"; \ + echo "NIXL_ALIAS_WRITTEN ${SP}/nixl_alias.pth" + +# Sanity: nixl (== rixl) import must succeed; report version + agent symbol. +RUN set -e; \ + python3 -c "import nixl; print('NIXL_IMPORT_OK', getattr(nixl,'__file__','?'))"; \ + python3 -c "import rixl, importlib.metadata as m; print('RIXL_VERSION', m.version('rixl'))" || true; \ + python3 -c "from nixl._api import nixl_agent, nixl_agent_config; print('NIXL_AGENT_OK')" \ + || echo "NIXL_AGENT_API_DIFFERS (verify sglang nixl import path at runtime)" + +############################################################################### +# 4) Runtime python deps + Mooncake KV-transfer backend. +# These were formerly built/pip-installed by the launcher at job start +# (scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh); baked into the image so +# the launcher no longer mutates the runtime environment. +############################################################################### +# Runtime python deps formerly pip-installed by the launcher. +RUN pip install --no-cache-dir py-spy pyyaml pandas \ + && pip install --no-cache-dir --ignore-installed --force-reinstall flask + +# Mooncake KV-transfer backend (KV_TRANSFER_BACKEND=mooncake). Mirrors +# docker/sglang_disagg_inference_kvtransfer_overlay.ubuntu.amd.Dockerfile:62-74. +# Default: pip wheel pin; set MOONCAKE_COMMIT for a source build. +ARG MOONCAKE_VERSION=0.3.6.post1 +ARG MOONCAKE_REPO=https://github.com/kvcache-ai/Mooncake.git +ARG MOONCAKE_COMMIT= +RUN set -e; \ + if [[ -n "${MOONCAKE_COMMIT}" ]]; then \ + echo "[full-overlay] Mooncake source ${MOONCAKE_REPO}@${MOONCAKE_COMMIT}"; \ + git clone "${MOONCAKE_REPO}" /tmp/mooncake && cd /tmp/mooncake && \ + git checkout "${MOONCAKE_COMMIT}" && (git submodule update --init --recursive || true); \ + pip install --no-cache-dir --force-reinstall . && rm -rf /tmp/mooncake; \ + elif [[ -n "${MOONCAKE_VERSION}" ]]; then \ + echo "[full-overlay] Mooncake pip mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ + pip install --no-cache-dir --force-reinstall "mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ + else \ + echo "[full-overlay] Mooncake: keeping base image version"; \ + fi + +# Sanity: mooncake must import alongside nixl/rixl (non-fatal — module path may vary). +RUN python3 -c "import mooncake; print('MOONCAKE_OVERLAY_OK')" \ + || echo "MOONCAKE_IMPORT_DIFFERS (verify sglang mooncake import path at runtime)" diff --git a/scripts/sglang_disagg/README.MD b/scripts/sglang_disagg/README.MD index 0540b19..20dbc0d 100644 --- a/scripts/sglang_disagg/README.MD +++ b/scripts/sglang_disagg/README.MD @@ -54,12 +54,58 @@ docker build --build-arg MORI_COMMIT= -t sglang_disagg_pd_image -f sglang_d | File | Description | |------|-------------| -| `run_xPyD_models.slurm` | SLURM script to launch docker containers on all nodes via `sbatch` | -| `sglang_disagg_mori_io_ep.sh` | Container entrypoint — starts prefill/decode servers, proxy, and benchmark | +| `run.sh` | **madengine entrypoint** (models.json `scripts` target). Bridges madengine's `sglang-disagg` launcher env to the launchers below | +| `run_xPyD_models.slurm` | Standalone SLURM script to launch docker containers on all nodes via `sbatch` | +| `sglang_disagg_mori_io_ep.sh` | Container entrypoint (MoRI EP) — starts prefill/decode servers, proxy, and benchmark | +| `sglang_disagg_server.sh` | Container entrypoint (Mooncake/NIXL) | | `models.yaml` | Model-specific CLI flags for all supported models | | `mori_ep_env.sh` | RDMA/NCCL/Gloo environment variables | | `benchmark_xPyD.sh` | Concurrency sweep benchmark using sglang bench_serving | | `benchmark_parser.py` | Log parser for CONCURRENCY benchmark logs | +| `parse_to_csv.py` | Log parser for CONCURRENCY benchmark logs (emits summary CSV and madengine perf CSV) | + +## Running via madengine (`madengine run`) + +This workload is registered in `models.json` (e.g. `sglang_disagg_deepseek-r1`) and +runs through madengine's built-in **`sglang-disagg`** launcher. The launcher runs one +container per node and exports the cluster topology; `run.sh` maps it to the launchers +above: + +| madengine env | bridged to | +|---------------|-----------| +| `SGLANG_NODE_RANK` | `NODE_RANK` (0=proxy, 1..xP=prefill, rest=decode) | +| `SGLANG_DISAGG_PREFILL_NODES` / `..._DECODE_NODES` | `xP` / `yD` | +| `SGLANG_NODE_IPS` | `IPADDRS` (+ `MASTER_ADDR` = rank-0 IP) | +| `SGLANG_TP_SIZE` | `IO_EP_TP_SIZE` / per-server `--tp-size` | + +Model + transport come from deployment `env_vars` (or `args`): `MODEL_NAME` (short key +in `models.yaml`), `MODEL_PATH` (mounted weights), `RUN_MORI` (1=MoRI EP default, +0=Mooncake/NIXL), `DP_MODE` (1=DP-attention EP for DeepSeek-V3/R1 inter-node). + +Deploy with an additional-context config carrying `slurm` + `distributed.launcher: +"sglang-disagg"` (+ `distributed.sglang_disagg.{prefill_nodes,decode_nodes}` for a +custom split) + `env_vars`; see the Confluence package for a ready 5-node EP16 manifest. + +## Docker overlays (substitutable component versions) + +The base image (`docker/sglang_disagg_inference.ubuntu.amd.Dockerfile`) bakes RCCL, +MoRI and Mooncake/NIXL. To test *specific* versions of these components without +rebuilding the base, use the merged full overlay, which layers RCCL, MoRI and the +KV-transfer (Mooncake/NIXL) steps on top of `BASE_DOCKER` in a single image +(`base -> rccl -> mori -> kvtransfer`): + +| Overlay Dockerfile | Substitutes | Key build args | +|--------------------|-------------|----------------| +| `docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile` | RCCL + MoRI + Mooncake/NIXL (merged) | `RCCL_REPO/RCCL_BRANCH/RCCL_COMMIT`, `MORI_REPO/MORI_BRANCH/MORI_COMMIT` or `MORI_WHEEL_URL`, `MOONCAKE_VERSION` / `NIXL_VERSION` | +| `docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile` | Same, OCI RDMA 6.2 variant | (as above) | + +> **RCCL note:** a newer rocm-systems RCCL build links `librccl` without +> the `librocm_smi64.so` `DT_NEEDED` the base carried. Since the overlay is first +> on `LD_LIBRARY_PATH`, this would break `import torch` +> (`libtorch_hip.so: undefined symbol: rsmi_init`). The overlay Dockerfile +> re-adds the dependency (`patchelf --add-needed`) and sanity-checks `import +> torch` with the overlay librccl resolved first, so the full +> `base -> rccl -> mori -> kvtransfer` chain stays importable. ## Quick Start @@ -72,6 +118,7 @@ export xP=1 export yD=1 export MODEL_NAME=Llama-3.1-8B-Instruct export RUN_MORI=1 # MoRI (default). Set RUN_MORI=0 for Mooncake (KV_TRANSFER_BACKEND=mooncake) +export KV_TRANSFER_BACKEND=mori # Validated/default. Use nixl only with an image that provides nixl._api. # num_nodes = xP + yD sbatch -N 2 -n 2 --nodelist= run_xPyD_models.slurm @@ -115,7 +162,7 @@ Logs are written to `${LOG_PATH}/${SLURM_JOB_ID}/`: Parse benchmark results: ```bash -python3 benchmark_parser.py /benchmark_XXX_CONCURRENCY.log +python3 parse_to_csv.py /benchmark_XXX_CONCURRENCY.log -o results.csv ``` Smoke test from the proxy node: diff --git a/scripts/sglang_disagg/benchmark_xPyD.sh b/scripts/sglang_disagg/benchmark_xPyD.sh index 05e6e3c..598b917 100644 --- a/scripts/sglang_disagg/benchmark_xPyD.sh +++ b/scripts/sglang_disagg/benchmark_xPyD.sh @@ -1,61 +1,139 @@ #!/bin/bash +set -uo pipefail timestamp=$(date "+%Y%m%d_%H%M%S") -LOG="/run_logs/${SLURM_JOB_ID}/benchmark_${SLURM_JOB_ID}_${timestamp}_xP${xP}_yD${yD}_$MODEL_NAME" -echo "==== Benchmark Serving Concurrency Sweep Test ${LOG} ===== " -echo "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a ${LOG}_CONCURRENCY.log >/dev/null -echo "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a ${LOG}_CONCURRENCY.log >/dev/null - -sleep 60 -echo "Test run:" | tee -a ${LOG}_CONCURRENCY.log >/dev/null -python3 -m sglang.bench_serving \ - --model $MODEL_PATH \ - --backend sglang \ - --host 127.0.0.1 \ - --port 2322 \ - --dataset-name random \ - --random-input 1024 \ - --random-output 1024\ - --random-range-ratio 1.0 \ - --max-concurrency 512 \ - --num-prompt 1024 \ - --pd-separated \ - 2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null -echo "" -CON="8 16 32 64 128 256 512" -# ISL/OSL combinations — override via BENCHMARK_COMBINATIONS env var (space-separated, e.g. "1024/1024 8192/1024") +RUN_LOG_JOB_ID="${SLURM_JOB_ID:-0}" +RUN_LOG_DIR="/run_logs/${RUN_LOG_JOB_ID}" +mkdir -p "$RUN_LOG_DIR" 2>/dev/null || true + +LOG="${RUN_LOG_DIR}/benchmark_${RUN_LOG_JOB_ID}_${timestamp}_xP${xP}_yD${yD}_${MODEL_NAME}" +LOG_FILE="${LOG}_CONCURRENCY.log" + +BENCHMARK_ITR="${BENCHMARK_ITR:-1}" +if ! [[ "$BENCHMARK_ITR" =~ ^[0-9]+$ ]] || [[ "$BENCHMARK_ITR" -lt 1 ]]; then + echo "ERROR: BENCHMARK_ITR must be a positive integer, got '${BENCHMARK_ITR}'" >&2 + exit 1 +fi + +CON="${BENCHMARK_CONCURRENCY_LEVELS:-8 16 32 64 128 256 512}" +# ISL/OSL combinations; override with e.g. BENCHMARK_COMBINATIONS="1024/1024 8192/1024". IFS=' ' read -ra COMBINATIONS <<< "${BENCHMARK_COMBINATIONS:-1024/1024 8192/1024}" -echo "Benchmarking iterations: $BENCHMARK_ITR" | tee -a ${LOG}_CONCURRENCY.log >/dev/null -for i in {1..$BENCHMARK_ITR}; do - sleep 60 - echo "RUNNING: the benchserving script for iter: $i" | tee -a ${LOG}_CONCURRENCY.log >/dev/null + +# --- Per-point resilience (kept from #328) ----------------------------------- +# A single transient gateway circuit-open / server recovery hiccup should not +# nuke a multi-hour sweep. Retry the whole point after a cooldown; a persistently +# failing point still fails the run (fail-fast preserved). +# Set BENCHMARK_POINT_RETRIES=0 to restore single-attempt behavior. +BENCHMARK_FAIL_FAST="${BENCHMARK_FAIL_FAST:-1}" +BENCHMARK_POINT_RETRIES="${BENCHMARK_POINT_RETRIES:-2}" +BENCHMARK_RETRY_COOLDOWN_SECONDS="${BENCHMARK_RETRY_COOLDOWN_SECONDS:-45}" + +log_msg() { + echo "$*" | tee -a "$LOG_FILE" >/dev/null +} + +run_serving_bench() { + local isl="$1" + local osl="$2" + local con="$3" + local prompts="$4" + + # Extra context on a separate INFO line; the parseable marker below stays in + # the exact shape parse_to_csv.py expects: "RUNNING: prompts isl X osl Y con Z". + log_msg "INFO: model=${MODEL_NAME} xP=${xP} yD=${yD} job=${RUN_LOG_JOB_ID} prompts=${prompts}" + log_msg "RUNNING: prompts isl ${isl} osl ${osl} con ${con}" + python3 -m sglang.bench_serving \ + --model "$MODEL_PATH" \ + --backend sglang \ + --host 127.0.0.1 \ + --port 2322 \ + --dataset-name random \ + --random-input "$isl" \ + --random-output "$osl" \ + --random-range-ratio 1.0 \ + --max-concurrency "$con" \ + --num-prompt "$prompts" \ + --pd-separated \ + 2>&1 | tee -a "$LOG_FILE" >/dev/null +} + +log_msg "==== Benchmark Serving Concurrency Sweep Test ${LOG} =====" +log_msg "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" +log_msg "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" +log_msg "Benchmarking iterations: ${BENCHMARK_ITR}" +log_msg "Benchmarking combinations: ${COMBINATIONS[*]}" +log_msg "Benchmarking concurrency levels: ${CON}" + +benchmark_failures=0 + +# Optional warmup / precheck. Runs BEFORE the "iter: 1" marker, so parse_to_csv.py +# (which ignores everything prior to that marker) never counts it as a result. +if [[ "${BENCHMARK_PRECHECK:-0}" == "1" ]]; then + sleep "${BENCHMARK_START_DELAY_SECONDS:-10}" + log_msg "Test run:" + if ! run_serving_bench \ + "${BENCHMARK_PRECHECK_ISL:-128}" \ + "${BENCHMARK_PRECHECK_OSL:-16}" \ + "${BENCHMARK_PRECHECK_CONCURRENCY:-1}" \ + "${BENCHMARK_PRECHECK_PROMPTS:-1}"; then + benchmark_failures=$((benchmark_failures + 1)) + log_msg "ERROR: benchmark precheck failed" + [[ "$BENCHMARK_FAIL_FAST" == "1" ]] && exit 1 + fi +fi + +for ((i = 1; i <= BENCHMARK_ITR; i++)); do + sleep "${BENCHMARK_ITERATION_DELAY_SECONDS:-60}" + log_msg "RUNNING: the benchserving script for iter: $i" for combo in "${COMBINATIONS[@]}"; do - IFS="/" read -r isl osl <<< "$combo" - for con in $CON; do - p_con=$(($con * 2)) - if [ "$p_con" -lt 16 ]; then - p_con=16 - fi - echo "RUNNING: prompts $prompts isl $isl osl $osl con $con" | tee -a ${LOG}_CONCURRENCY.log >/dev/null - python3 -m sglang.bench_serving \ - --model $MODEL_PATH \ - --backend sglang \ - --host 127.0.0.1 \ - --port 2322 \ - --dataset-name random \ - --random-input $isl \ - --random-output $osl \ - --random-range-ratio 1.0 \ - --max-concurrency $con \ - --num-prompt $p_con \ - --pd-separated \ - 2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null - - sleep 10 + IFS="/" read -r isl osl <<< "$combo" + for con in $CON; do + p_con=$((con * 2)) + if [[ "$p_con" -lt 16 ]]; then + p_con=16 + fi + total_attempts=$((BENCHMARK_POINT_RETRIES + 1)) + point_ok=0 + for ((attempt = 1; attempt <= total_attempts; attempt++)); do + if run_serving_bench "$isl" "$osl" "$con" "$p_con"; then + point_ok=1 + break + fi + if [[ "$attempt" -lt "$total_attempts" ]]; then + log_msg "WARN: bench_serving failed (attempt ${attempt}/${total_attempts}) for iter=${i} isl=${isl} osl=${osl} con=${con}; retrying after ${BENCHMARK_RETRY_COOLDOWN_SECONDS}s (transient circuit-open / server recovery)" + sleep "${BENCHMARK_RETRY_COOLDOWN_SECONDS}" + fi + done + if [[ "$point_ok" -ne 1 ]]; then + benchmark_failures=$((benchmark_failures + 1)) + log_msg "ERROR: bench_serving failed after ${total_attempts} attempt(s) for iter=${i} isl=${isl} osl=${osl} con=${con}" + [[ "$BENCHMARK_FAIL_FAST" == "1" ]] && break 3 + fi + + sleep "${BENCHMARK_BETWEEN_RUN_SECONDS:-10}" done done done -python3 parse_to_csv.py ${LOG}_CONCURRENCY.log -o ${LOG}_CONCURRENCY.csv \ - --perf-csv /run_logs/${SLURM_JOB_ID}/perf.csv \ - --model-name "${MODEL_NAME}" \ - 2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null + +log_msg "==== Benchmark Serving Concurrency End Time ${LOG} =====" +log_msg "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" +log_msg "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" + +# --- Finalization: aligned with develop (parse_to_csv.py + MAD_OUTPUT_CSV) ---- +# Write the madengine multiple_results CSV to the container cwd under the name +# madengine harvests (MAD_OUTPUT_CSV, e.g. perf_sglang-disagg-.csv); keep +# a copy under /run_logs for post-mortem inspection. +PERF_CSV="${MAD_OUTPUT_CSV:-perf_sglang-disagg-${MODEL_NAME}.csv}" +if ! python3 parse_to_csv.py "${LOG}_CONCURRENCY.log" -o "${LOG}_CONCURRENCY.csv" \ + --perf-csv "${PERF_CSV}" \ + --model-name "${MODEL_NAME}" \ + 2>&1 | tee -a "$LOG_FILE" >/dev/null; then + benchmark_failures=$((benchmark_failures + 1)) + log_msg "ERROR: parse_to_csv.py failed to produce ${PERF_CSV}" +fi +cp -f "${PERF_CSV}" "${RUN_LOG_DIR}/" 2>/dev/null || true + +if [[ "$benchmark_failures" -ne 0 ]]; then + log_msg "ERROR: benchmark completed with ${benchmark_failures} failure(s)" + exit 1 +fi diff --git a/scripts/sglang_disagg/ip_rendezvous.py b/scripts/sglang_disagg/ip_rendezvous.py new file mode 100755 index 0000000..9091dab --- /dev/null +++ b/scripts/sglang_disagg/ip_rendezvous.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +"""In-container node-IP rendezvous for disaggregated P/D launchers. + +madengine forwards MASTER_ADDR/NODE_RANK/NNODES into each container but NOT the +rank-ordered node IP list. This helper rebuilds that list with stdlib only: + + rank 0 -> tiny TCP server: collects {rank: ip} from peers, then broadcasts + the comma-separated, rank-ordered IP list back to everyone. + rank>0 -> connects to MASTER_ADDR, reports its own IP, then blocks until + rank 0 broadcasts the full list. + +On success the rank-ordered "ip0,ip1,..." list is printed to stdout (exit 0); +on failure an empty line is printed (exit 2), so callers can use +`out="$(ip_rendezvous.py ... || true)"` and test for emptiness. + +NOTE: this file is intentionally duplicated under scripts/sglang_disagg/ and +scripts/vllm_dissag/ (each model packages its own script dir into the image). +Keep both copies identical. + +Usage: + ip_rendezvous.py +Env: + IP_SYNC_TIMEOUT total rendezvous budget in seconds (default 1800). +""" +import json +import os +import socket +import sys +import time + + +def recv_line(conn): + buf = b"" + while b"\n" not in buf: + chunk = conn.recv(4096) + if not chunk: + break + buf += chunk + return buf.split(b"\n", 1)[0].decode("utf-8", "replace") if buf else None + + +def serve(nnodes, host_ip, port, token, deadline): + """rank-0 path: collect peer IPs, broadcast the rank-ordered list.""" + by_rank = {0: host_ip} + # rank -> connection. Workers reconnect while we wait for slow nodes, so we + # keep only the latest socket per rank (closing the stale one) to avoid + # broadcasting over half-open sockets. + conns = {} + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("0.0.0.0", port)) + srv.listen(max(8, nnodes + 2)) + srv.settimeout(1.0) + try: + while len(by_rank) < nnodes and time.time() < deadline: + try: + conn, _ = srv.accept() + except socket.timeout: + continue + conn.settimeout(5.0) + line = recv_line(conn) + if not line: + conn.close() + continue + try: + msg = json.loads(line) + except Exception: + conn.close() + continue + if msg.get("token") != token: + conn.close() + continue + wrank = int(msg.get("rank", -1)) + wip = str(msg.get("ip", "")).strip() + if 0 <= wrank < nnodes and wip: + old = conns.get(wrank) + if old is not None: + try: + old.close() + except Exception: + pass + by_rank[wrank] = wip + conns[wrank] = conn + else: + conn.close() + if len(by_rank) != nnodes: + return None + ipaddrs = ",".join(by_rank[i] for i in range(nnodes)) + payload = (ipaddrs + "\n").encode() + # Best-effort broadcast: a single dead/stale peer socket must not abort + # the whole rendezvous, so swallow per-conn send errors. + for conn in conns.values(): + try: + conn.sendall(payload) + except Exception: + pass + finally: + try: + conn.close() + except Exception: + pass + return ipaddrs + finally: + srv.close() + + +def report(rank, host_ip, master_addr, port, token, deadline): + """rank>0 path: report own IP, block on recv until rank-0 broadcasts.""" + payload = (json.dumps({"token": token, "rank": rank, "ip": host_ip}) + "\n").encode() + while time.time() < deadline: + try: + sock = socket.create_connection((master_addr, port), timeout=3.0) + sock.sendall(payload) + # rank-0 broadcasts once, only after every node has reported, which + # can be far longer than a few seconds when peers load the image at + # different times. Block on recv for the remaining deadline so we + # stay connected and never miss that one-shot broadcast. + remaining = max(1.0, deadline - time.time()) + sock.settimeout(remaining) + line = recv_line(sock) + sock.close() + if line: + return line.strip() + except Exception: + time.sleep(1.0) + return None + + +def main(argv): + rank = int(argv[1]) + nnodes = int(argv[2]) + host_ip = argv[3] + master_addr = argv[4] + port = int(argv[5]) + job_id = argv[6] + token = f"JOB{job_id}" + timeout_s = int(os.environ.get("IP_SYNC_TIMEOUT", "1800")) + deadline = time.time() + timeout_s + if rank == 0: + result = serve(nnodes, host_ip, port, token, deadline) + else: + result = report(rank, host_ip, master_addr, port, token, deadline) + if result: + print(result) + return 0 + print("") + return 2 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/scripts/sglang_disagg/parse_to_csv.py b/scripts/sglang_disagg/parse_to_csv.py index 8d4d65d..951d2b5 100644 --- a/scripts/sglang_disagg/parse_to_csv.py +++ b/scripts/sglang_disagg/parse_to_csv.py @@ -1,170 +1,247 @@ #!/usr/bin/env python3 """ Parse SGLang benchmark log file and save results to CSV. -Extracts: Concurrency, Input tokens, Output tokens, Total Token throughput (tok/s) -For each configuration, takes the MAX Total Token throughput across all iterations. + +Extracts a full metric set per (isl, osl, concurrency) configuration: + - Request throughput (req/s) + - Input / Output / Total token throughput (tok/s) + - Mean E2E latency (ms), Mean TTFT (ms), Mean ITL (ms) + +For each configuration the row from the iteration with the MAX total token +throughput is retained (all metrics come from that single best run, so they +stay mutually consistent). """ import re import csv from pathlib import Path -from typing import List, Dict, Tuple -from collections import defaultdict +from typing import Dict, Tuple + +# Metric key -> madengine metric name. Order defines CSV row order per config. +METRIC_FIELDS = [ + ("request_throughput_req_s", "request_throughput_req_s"), + ("input_token_throughput_tok_s", "input_token_throughput_tok_s"), + ("output_token_throughput_tok_s", "output_token_throughput_tok_s"), + ("total_token_throughput_tok_s", "total_token_throughput_tok_s"), + ("mean_e2e_latency_ms", "mean_e2e_latency_ms"), + ("mean_ttft_ms", "mean_ttft_ms"), + ("mean_itl_ms", "mean_itl_ms"), +] + + +def _extract(pattern: str, text: str): + """Return the first capture group as float (commas stripped), else None.""" + m = re.search(pattern, text) + return float(m.group(1).replace(",", "")) if m else None def parse_benchmark_log(log_file: str) -> Dict[Tuple[int, int, int], Dict]: - """Parse benchmark log file and extract results, keeping max throughput per configuration.""" - results = defaultdict(lambda: {'concurrency': None, 'input_tokens': None, - 'output_tokens': None, 'max_throughput': 0.0}) - - with open(log_file, 'r') as f: + """Parse benchmark log, keeping the best-throughput run per configuration.""" + results: Dict[Tuple[int, int, int], Dict] = {} + + with open(log_file, "r") as f: content = f.read() - - # Find the start of the first iteration (ignore warmup) - first_iter_match = re.search(r'RUNNING: the benchserving script for iter: 1', content) + + # Find the start of the first iteration (ignore warmup). + first_iter_match = re.search(r"RUNNING: the benchserving script for iter: 1", content) if not first_iter_match: print("Warning: No iteration 1 found. Processing entire file.") start_pos = 0 else: start_pos = first_iter_match.start() - - # Process only from first iteration onwards + content = content[start_pos:] - - # Split by benchmark result sections - sections = re.split(r'============ Serving Benchmark Result ============', content) - - current_input_seq_len = None - current_output_seq_len = None - current_concurrency = None - - for i, section in enumerate(sections[1:], 1): # Skip first empty section - # Look for configuration in previous sections (from RUNNING line) - if i > 1: - prev_section = sections[i-1] - - # Extract config: prompts isl osl con - config_match = re.search(r'RUNNING: prompts\s+isl\s+(\d+)\s+osl\s+(\d+)\s+con\s+(\d+)', prev_section) - if config_match: - current_input_seq_len = int(config_match.group(1)) - current_output_seq_len = int(config_match.group(2)) - current_concurrency = int(config_match.group(3)) - - # Extract Total token throughput (tok/s) from benchmark result section - throughput_match = re.search(r'Total token throughput \(tok/s\):\s+([\d.]+)', section) - throughput = float(throughput_match.group(1)) if throughput_match else None - - # Only process if we have a valid configuration from RUNNING line and throughput - if current_input_seq_len and current_output_seq_len and current_concurrency and throughput is not None: - config_key = (current_input_seq_len, current_output_seq_len, current_concurrency) - - # Update results for this configuration - # Always use values from RUNNING line (isl, osl, con) - results[config_key]['concurrency'] = current_concurrency - results[config_key]['input_tokens'] = current_input_seq_len - results[config_key]['output_tokens'] = current_output_seq_len - - # Keep the maximum throughput - if throughput > results[config_key]['max_throughput']: - results[config_key]['max_throughput'] = throughput - + + # Split by benchmark result sections; config lives in the preceding section. + sections = re.split(r"============ Serving Benchmark Result ============", content) + + current_isl = None + current_osl = None + current_con = None + + for i, section in enumerate(sections[1:], 1): # Skip first (pre-first-result) chunk + prev_section = sections[i - 1] + + # Config emitted by benchmark_xPyD.sh: "RUNNING: prompts isl X osl Y con Z" + config_match = re.search( + r"RUNNING: prompts\s+isl\s+(\d+)\s+osl\s+(\d+)\s+con\s+(\d+)", prev_section + ) + if config_match: + current_isl = int(config_match.group(1)) + current_osl = int(config_match.group(2)) + current_con = int(config_match.group(3)) + + if current_isl is None or current_osl is None or current_con is None: + continue + + total_throughput = _extract(r"Total token throughput \(tok/s\):\s+([\d,\.]+)", section) + if total_throughput is None: + continue + + row = { + "isl": current_isl, + "osl": current_osl, + "con": current_con, + "request_throughput_req_s": _extract( + r"Request throughput \(req/s\):\s+([\d,\.]+)", section + ), + "input_token_throughput_tok_s": _extract( + r"Input token throughput \(tok/s\):\s+([\d,\.]+)", section + ), + "output_token_throughput_tok_s": _extract( + r"Output token throughput \(tok/s\):\s+([\d,\.]+)", section + ), + "total_token_throughput_tok_s": total_throughput, + "mean_e2e_latency_ms": _extract(r"Mean E2E Latency \(ms\):\s+([\d,\.]+)", section), + "mean_ttft_ms": _extract(r"Mean TTFT \(ms\):\s+([\d,\.]+)", section), + "mean_itl_ms": _extract(r"Mean ITL \(ms\):\s+([\d,\.]+)", section), + } + + config_key = (current_isl, current_osl, current_con) + best = results.get(config_key) + if best is None or total_throughput > best["total_token_throughput_tok_s"]: + results[config_key] = row + return results +def _sorted_results(results: Dict[Tuple[int, int, int], Dict]): + # Sort by concurrency, then isl, then osl. + return [ + results[key] + for key in sorted(results.keys(), key=lambda k: (k[2], k[0], k[1])) + ] + + def save_to_csv(results: Dict[Tuple[int, int, int], Dict], output_file: str): - """Save results to CSV file with specified columns.""" + """Save a wide per-configuration summary CSV (all metrics).""" if not results: print("No results to save.") return - - # Define column order - fieldnames = ['Concurrency', 'Input tokens', 'Output tokens', 'Total Token throughput (tok/s)'] - - with open(output_file, 'w', newline='') as f: + + fieldnames = [ + "Concurrency", + "Input tokens", + "Output tokens", + "Request throughput (req/s)", + "Input token throughput (tok/s)", + "Output token throughput (tok/s)", + "Total Token throughput (tok/s)", + "Mean E2E Latency (ms)", + "Mean TTFT (ms)", + "Mean ITL (ms)", + ] + + def _fmt(value): + return f"{value:.2f}" if value is not None else "" + + with open(output_file, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() - - # Sort by concurrency, then input tokens, then output tokens - for (input_tokens, output_tokens, concurrency), data in sorted(results.items(), - key=lambda x: (x[0][2], x[0][0], x[0][1])): - row = { - 'Concurrency': data['concurrency'], - 'Input tokens': data['input_tokens'], - 'Output tokens': data['output_tokens'], - 'Total Token throughput (tok/s)': f"{data['max_throughput']:.2f}" - } - writer.writerow(row) - + for data in _sorted_results(results): + writer.writerow( + { + "Concurrency": data["con"], + "Input tokens": data["isl"], + "Output tokens": data["osl"], + "Request throughput (req/s)": _fmt(data["request_throughput_req_s"]), + "Input token throughput (tok/s)": _fmt(data["input_token_throughput_tok_s"]), + "Output token throughput (tok/s)": _fmt(data["output_token_throughput_tok_s"]), + "Total Token throughput (tok/s)": _fmt(data["total_token_throughput_tok_s"]), + "Mean E2E Latency (ms)": _fmt(data["mean_e2e_latency_ms"]), + "Mean TTFT (ms)": _fmt(data["mean_ttft_ms"]), + "Mean ITL (ms)": _fmt(data["mean_itl_ms"]), + } + ) + print(f"Saved {len(results)} benchmark configurations to {output_file}") def _get_run_metadata(pipeline: str = "sglang"): """Collect run metadata from environment variables.""" import os - xP = os.environ.get('xP', '1') - yD = os.environ.get('yD', '1') - dp_mode = os.environ.get('DP_MODE', '0') - run_mori = os.environ.get('RUN_MORI', '0') - gpus = os.environ.get('GPUS_PER_NODE', '8') + + xP = os.environ.get("xP", "1") + yD = os.environ.get("yD", "1") + dp_mode = os.environ.get("DP_MODE", "0") + run_mori = os.environ.get("RUN_MORI", "0") + gpus = os.environ.get("GPUS_PER_NODE", "8") # Determine backend tag - if dp_mode == '1': - backend = 'mori_dp' - elif run_mori == '1': - backend = 'mori_io' + if dp_mode == "1": + backend = "mori_dp" + elif run_mori == "1": + backend = "mori_io" else: - backend = 'mooncake' + backend = "mooncake" return { - 'pipeline': pipeline, - 'deployment_type': f'disagg_{xP}P{yD}D', - 'tags': f'{pipeline}_disagg,{backend}', - 'n_gpus': str(int(xP) * int(gpus) + int(yD) * int(gpus)), - 'nnodes': str(int(xP) + int(yD)), - 'gpus_per_node': gpus, - 'docker_image': os.environ.get('DOCKER_IMAGE_NAME', ''), - 'machine_name': os.environ.get('SLURM_JOB_NODELIST', ''), - 'launcher': 'slurm_multi', - 'gpu_architecture': 'gfx942', + "pipeline": pipeline, + "deployment_type": f"disagg_{xP}P{yD}D", + "tags": f"{pipeline}_disagg,{backend}", + "n_gpus": str(int(xP) * int(gpus) + int(yD) * int(gpus)), + "nnodes": str(int(xP) + int(yD)), + "gpus_per_node": gpus, + "docker_image": os.environ.get("DOCKER_IMAGE_NAME", ""), + "machine_name": os.environ.get("SLURM_JOB_NODELIST", ""), + "launcher": "slurm_multi", + "gpu_architecture": "gfx942", } -def save_perf_csv(results: Dict[Tuple[int, int, int], Dict], output_file: str, - model_name: str = "", pipeline: str = "sglang"): - """Save results in madengine perf.csv format.""" +def save_perf_csv( + results: Dict[Tuple[int, int, int], Dict], + output_file: str, + model_name: str = "", + pipeline: str = "sglang", +): + """Save results in madengine perf.csv format (one row per config x metric).""" if not results: print("No results to save to perf.csv.") return + import os + meta = _get_run_metadata(pipeline) + xP = os.environ.get("xP", "1") + yD = os.environ.get("yD", "1") fieldnames = [ - 'model', 'n_gpus', 'nnodes', 'gpus_per_node', 'training_precision', - 'pipeline', 'args', 'tags', 'docker_file', 'base_docker', 'docker_sha', - 'docker_image', 'git_commit', 'machine_name', 'deployment_type', 'launcher', - 'gpu_architecture', 'performance', 'metric', 'relative_change', 'status', - 'build_duration', 'test_duration', 'dataname', 'data_provider_type', - 'data_size', 'data_download_duration', 'build_number', - 'additional_docker_run_options', + "model", "n_gpus", "nnodes", "gpus_per_node", "training_precision", + "pipeline", "args", "tags", "docker_file", "base_docker", "docker_sha", + "docker_image", "git_commit", "machine_name", "deployment_type", "launcher", + "gpu_architecture", "performance", "metric", "relative_change", "status", + "build_duration", "test_duration", "dataname", "data_provider_type", + "data_size", "data_download_duration", "build_number", + "additional_docker_run_options", ] - with open(output_file, 'w', newline='') as f: + rows_written = 0 + with open(output_file, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() - for (input_tokens, output_tokens, concurrency), data in sorted( - results.items(), key=lambda x: (x[0][2], x[0][0], x[0][1]) - ): - row = { - 'model': model_name, - 'performance': f"{data['max_throughput']:.2f}", - 'metric': f"tok/s (isl={data['input_tokens']} osl={data['output_tokens']} con={data['concurrency']})", - 'status': 'SUCCESS', - } - row.update(meta) - writer.writerow(row) + for data in _sorted_results(results): + config_key = ( + f"{xP}p{yD}d_isl{data['isl']}_osl{data['osl']}_con{data['con']}" + ) + for source_key, metric_name in METRIC_FIELDS: + value = data.get(source_key) + if value is None: + continue + row = { + "model": config_key, + "performance": f"{value:.2f}", + "metric": metric_name, + "status": "SUCCESS", + } + row.update(meta) + # model_name (catalog key) is prepended by madengine; keep the + # per-config key in `model` so metrics stay disambiguated. + writer.writerow(row) + rows_written += 1 - print(f"Saved {len(results)} rows to perf.csv: {output_file}") + print(f"Saved {rows_written} rows to perf.csv: {output_file}") def main(): @@ -172,40 +249,42 @@ def main(): import sys import argparse - parser = argparse.ArgumentParser(description='Parse SGLang benchmark log file and save results to CSV') - parser.add_argument('log_file', type=str, help='Path to benchmark log file') - parser.add_argument('-o', '--output', type=str, help='Output CSV file name (default: _results.csv)') - parser.add_argument('--perf-csv', type=str, help='Also generate madengine perf.csv at this path') - parser.add_argument('--model-name', type=str, default='', help='Model name for perf.csv') + parser = argparse.ArgumentParser( + description="Parse SGLang benchmark log file and save results to CSV" + ) + parser.add_argument("log_file", type=str, help="Path to benchmark log file") + parser.add_argument( + "-o", "--output", type=str, + help="Output CSV file name (default: _results.csv)", + ) + parser.add_argument( + "--perf-csv", type=str, help="Also generate madengine perf.csv at this path" + ) + parser.add_argument("--model-name", type=str, default="", help="Model name for perf.csv") args = parser.parse_args() log_file = args.log_file - # Check if file exists if not Path(log_file).exists(): print(f"Error: Log file not found: {log_file}") sys.exit(1) print(f"Parsing log file: {log_file}") - # Parse the log file results = parse_benchmark_log(log_file) if not results: print("No benchmark results found in log file.") - return + sys.exit(1) - # Generate output filename if args.output: output_file = args.output else: - output_file = Path(log_file).stem + '_results.csv' + output_file = Path(log_file).stem + "_results.csv" - # Save to CSV save_to_csv(results, output_file) - # Save madengine perf.csv if requested if args.perf_csv: save_perf_csv(results, args.perf_csv, args.model_name) @@ -214,5 +293,5 @@ def main(): print(f" Output file: {output_file}") -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh new file mode 100755 index 0000000..fe90926 --- /dev/null +++ b/scripts/sglang_disagg/run.sh @@ -0,0 +1,217 @@ +#!/bin/bash +# ============================================================================= +# madengine-native entrypoint for SGLang Disaggregated P/D inference. +# +# This is the `scripts` target referenced by models.json (sglang_disagg_*). +# madengine's built-in `sglang-disagg` SLURM/K8s launcher +# (src/madengine/deployment/slurm.py::_generate_sglang_disagg_command) runs ONE +# container per node and exports the cluster topology as env vars. This script +# bridges those into the env contract of the proven MAD-private disagg +# launchers and execs the right one (MoRI EP, or Mooncake/NIXL). +# +# madengine -> proven env mapping: +# SGLANG_NODE_RANK -> NODE_RANK (0=proxy, 1..xP=prefill, rest=decode) +# SGLANG_DISAGG_PREFILL_NODES -> xP +# SGLANG_DISAGG_DECODE_NODES -> yD +# SGLANG_NODE_IPS -> IPADDRS (comma-separated, rank order) +# SGLANG_TP_SIZE -> IO_EP_TP_SIZE / per-server --tp-size +# MASTER_PORT -> MASTER_PORT +# +# Model identity + transport come from deployment env_vars / models.json args: +# MODEL_NAME short catalog key in models.yaml (e.g. DeepSeek-R1) [required] +# MODEL_PATH weights dir mounted into the container [required] +# RUN_MORI 1 = MoRI EP all-to-all (default), 0 = Mooncake/NIXL +# KV_TRANSFER_BACKEND prefill->decode KV transfer backend (mori/mooncake/nixl) +# DP_MODE 1 = DP-attention EP (DeepSeek-V3/R1 inter-node), 0 = TP +# ============================================================================= +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# --- parse optional CLI args (models.json "args"), env takes precedence ------ +while [[ $# -gt 0 ]]; do + case "$1" in + --model-name|--model_name|--model_repo) MODEL_NAME="${MODEL_NAME:-$2}"; shift 2;; + --model-path|--model_path) MODEL_PATH="${MODEL_PATH:-$2}"; shift 2;; + --run-mori) RUN_MORI="${RUN_MORI:-$2}"; shift 2;; + --dp-mode) DP_MODE="${DP_MODE:-$2}"; shift 2;; + --xp|--xP|--prefill-nodes) xP="${xP:-$2}"; shift 2;; + --yd|--yD|--decode-nodes) yD="${yD:-$2}"; shift 2;; + --ipaddrs|--node-ips) IPADDRS="${IPADDRS:-$2}"; shift 2;; + --master-port) MASTER_PORT="${MASTER_PORT:-$2}"; shift 2;; + --gpus-per-node) GPUS_PER_NODE="${GPUS_PER_NODE:-$2}"; shift 2;; + --kv-transfer-backend|--transfer-backend) KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-$2}"; shift 2;; + *) shift;; + esac +done + +# --- bridge madengine sglang-disagg topology -> proven launcher env --------- +export NODE_RANK="${NODE_RANK:-${SGLANG_NODE_RANK:-${SLURM_PROCID:-0}}}" +# Persist run.sh output to the shared /run_logs mount so failures are +# diagnosable from the submission node (madengine discards model stdout). +# The dir is created from inside the container (root); make it world-writable +# so an ordinary user outside the container can clean it up afterwards. Fall +# back to /tmp when /run_logs is absent/not writable, otherwise the tee in the +# process substitution can't open its file and we lose the ERROR output we +# need for the post-mortem. +RUN_LOG_DIR="/run_logs/${SLURM_JOB_ID:-local}" +if mkdir -p "${RUN_LOG_DIR}" 2>/dev/null && [[ -w "${RUN_LOG_DIR}" ]]; then + chmod 0777 /run_logs 2>/dev/null || true + chmod -R 0777 "${RUN_LOG_DIR}" 2>/dev/null || true +else + RUN_LOG_DIR="/tmp/run_logs/${SLURM_JOB_ID:-local}" + mkdir -p "${RUN_LOG_DIR}" 2>/dev/null || true + chmod -R 0777 "${RUN_LOG_DIR}" 2>/dev/null || true +fi +export RUN_LOG_DIR +exec > >(tee -a "${RUN_LOG_DIR}/run_sh_rank${NODE_RANK}.log") 2>&1 +export xP="${xP:-${SGLANG_DISAGG_PREFILL_NODES:-1}}" +export yD="${yD:-${SGLANG_DISAGG_DECODE_NODES:-1}}" +export NNODES="${NNODES:-${SGLANG_DISAGG_TOTAL_NODES:-${WORLD_SIZE:-$((xP + yD))}}}" +export MASTER_PORT="${MASTER_PORT:-23731}" +export IO_EP_TP_SIZE="${IO_EP_TP_SIZE:-${SGLANG_TP_SIZE:-8}}" +# madengine forwards MASTER_ADDR (rank-0 host) into the container; use it as the +# rendezvous address for in-container IP discovery below. +export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" + +# IPADDRS = comma-separated node IPs in rank order. It may be supplied +# explicitly via env/arg or SGLANG_NODE_IPS. madengine's container runner does +# NOT forward SGLANG_NODE_IPS into the container, so when it is absent we +# self-discover the node IPs in-container using only the forwarded +# MASTER_ADDR/NODE_RANK/NNODES, mirroring scripts/vllm_dissag/run.sh. +export IPADDRS="${IPADDRS:-${SGLANG_NODE_IPS:-}}" + +# Prefer the routable transport interface (NCCL_SOCKET_IFNAME, eth0 here); +# hostname -I order is not deterministic and may surface link-local addrs. +# NCCL_SOCKET_IFNAME may be a comma-separated list (see set_env_vars.sh), so +# take only the first interface before handing it to `ip addr show`. +_HOST_IFACE="${NCCL_SOCKET_IFNAME:-eth0}"; _HOST_IFACE="${_HOST_IFACE%%,*}" +HOST_IP="$(ip -4 -o addr show "${_HOST_IFACE}" 2>/dev/null | awk '{print $4}' | cut -d/ -f1 | head -n1)" +[[ -z "${HOST_IP}" ]] && HOST_IP="$(hostname -I | awk '{print $1}')" +# Derive the rendezvous port from the job id. Use a wide modulo so two +# concurrent jobs whose ids happen to share a rank-0 host are very unlikely to +# collide (a narrow 10000-wide range collides whenever ids differ by a multiple +# of 10000). The value is deterministic across nodes (depends only on the job +# id, which is identical on every node). +IP_SYNC_PORT="${IP_SYNC_PORT:-$((20000 + (${SLURM_JOB_ID:-0} % 40000)))}" + +# rank-0 acts as a tiny TCP rendezvous server: peers connect to MASTER_ADDR and +# report their host IP; rank-0 aggregates the rank-ordered list and broadcasts. +# The implementation lives in the sibling ip_rendezvous.py so it can be unit +# tested and is shared (duplicated, kept identical) with scripts/vllm_dissag/. +_tcp_discover_ipaddrs() { + python3 "$SCRIPT_DIR/ip_rendezvous.py" \ + "${NODE_RANK}" "${NNODES}" "${HOST_IP}" "${MASTER_ADDR}" "${IP_SYNC_PORT}" "${SLURM_JOB_ID:-0}" +} + +# Method 1: SLURM nodelist (only if scontrol/getent exist inside the container). +if [[ -z "${IPADDRS}" && -n "${SLURM_JOB_NODELIST:-}" ]] && command -v scontrol >/dev/null 2>&1 && command -v getent >/dev/null 2>&1; then + _ips=""; _cnt=0 + while IFS= read -r _n; do + [[ -z "${_n}" ]] && continue + _ip="$(getent ahostsv4 "${_n}" | awk 'NR==1{print $1}')" + [[ -n "${_ip}" ]] && { _ips="${_ips:+${_ips},}${_ip}"; _cnt=$((_cnt + 1)); } + done < <(scontrol show hostnames "${SLURM_JOB_NODELIST}") + [[ "${_cnt}" == "${NNODES}" && -n "${_ips}" ]] && IPADDRS="${_ips}" +fi + +# Method 2: TCP rendezvous via forwarded MASTER_ADDR (primary path). +if [[ -z "${IPADDRS}" ]]; then + _tcp="$(_tcp_discover_ipaddrs || true)" + [[ -n "${_tcp}" ]] && IPADDRS="${_tcp}" +fi + +if [[ -z "${IPADDRS}" ]]; then + echo "ERROR: unable to determine IPADDRS (node IP list);" \ + "MASTER_ADDR=${MASTER_ADDR} NNODES=${NNODES} NODE_RANK=${NODE_RANK}." >&2 + exit 1 +fi +export IPADDRS +# rank-0 IP is the master/load-balancer host. +export MASTER_ADDR="$(echo "$IPADDRS" | cut -d',' -f1)" + +# --- model + transport defaults -------------------------------------------- +export MODEL_NAME="${MODEL_NAME:-}" +export MODEL_PATH="${MODEL_PATH:-}" +export RUN_MORI="${RUN_MORI:-1}" +export DP_MODE="${DP_MODE:-1}" +# Default the KV transfer backend from RUN_MORI so the non-MoRI path does not +# silently inherit an invalid `mori` backend (mooncake/NIXL use mooncake). +if [[ "$RUN_MORI" == "1" ]]; then + export KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mori}" +else + export KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mooncake}" +fi + +# Helper scripts (socket_barrier.py, set_env_vars.sh, benchmark_xPyD.sh, ...) +# live alongside this script; point the Mooncake-path lookups at them. +export MOONCAKE_COOKBOOK_PATH="${MOONCAKE_COOKBOOK_PATH:-$SCRIPT_DIR}" +mkdir -p "${RUN_LOG_DIR}" 2>/dev/null || true + +if [[ -z "$MODEL_NAME" || -z "$MODEL_PATH" ]]; then + echo "ERROR: MODEL_NAME and MODEL_PATH must be set (via deployment env_vars or args)." >&2 + echo " MODEL_NAME='$MODEL_NAME' MODEL_PATH='$MODEL_PATH'" >&2 + exit 1 +fi + +# --- stage model weights if missing (rank-0 gated, NFS-shared MODEL_PATH) ----- +# run.sh historically assumed pre-staged weights. Stage from HF when MODEL_PATH +# is incomplete. Only NODE_RANK 0 downloads; peers wait on a sentinel. +# Treat weights as complete only if our own download sentinel is present, OR a +# config.json plus at least one weight shard exists (covers manually pre-staged +# dirs). A partial earlier download (config.json present, shards missing, no +# sentinel) must NOT be mistaken for a complete stage. +_weights_complete() { + [[ -f "${MODEL_PATH}/.stage_done" ]] && return 0 + [[ -f "${MODEL_PATH}/config.json" ]] || return 1 + compgen -G "${MODEL_PATH}/*.safetensors" >/dev/null 2>&1 && return 0 + compgen -G "${MODEL_PATH}/*.bin" >/dev/null 2>&1 && return 0 + return 1 +} +if ! _weights_complete; then + _MODEL_REPO="${MODEL_REPO:-}" + if [[ -z "${_MODEL_REPO}" ]]; then + case "${MODEL_NAME}" in + DeepSeek-R1) _MODEL_REPO="deepseek-ai/DeepSeek-R1-0528" ;; + Qwen3-Next-80B) _MODEL_REPO="Qwen/Qwen3-Next-80B-A3B-Instruct" ;; + *) echo "ERROR: weights missing at ${MODEL_PATH}; set MODEL_REPO for ${MODEL_NAME}" >&2; exit 1 ;; + esac + fi + _SENTINEL="${MODEL_PATH}/.stage_done" + if [[ "${NODE_RANK:-0}" == "0" ]]; then + echo "[stage_weights] NODE_RANK=0 downloading ${_MODEL_REPO} -> ${MODEL_PATH}" + mkdir -p "${MODEL_PATH}" + # huggingface-cli is deprecated and refuses to run on newer huggingface_hub + # (it prints "use `hf` instead" and exits non-zero); prefer the `hf` CLI and + # fall back to huggingface-cli only on older images that lack `hf`. + if command -v hf >/dev/null 2>&1; then + _HF_DL=(hf download) + else + _HF_DL=(huggingface-cli download) + fi + HF_TOKEN="${MAD_SECRETS_HFTOKEN:-${HF_TOKEN:-}}" "${_HF_DL[@]}" "${_MODEL_REPO}" --local-dir "${MODEL_PATH}" || { echo "[stage_weights] download failed" >&2; exit 1; } + touch "${_SENTINEL}" + else + echo "[stage_weights] NODE_RANK=${NODE_RANK} waiting for weights at ${MODEL_PATH}" + for _i in $(seq 1 720); do + [[ -f "${_SENTINEL}" && -f "${MODEL_PATH}/config.json" ]] && break + sleep 10 + done + [[ -f "${MODEL_PATH}/config.json" ]] || { echo "[stage_weights] timed out waiting for weights" >&2; exit 1; } + fi +fi + + +echo "==============================================================" +echo " SGLang Disaggregated (madengine bridge)" +echo " MODEL_NAME=$MODEL_NAME MODEL_PATH=$MODEL_PATH" +echo " NODE_RANK=$NODE_RANK xP=$xP yD=$yD TP=$IO_EP_TP_SIZE" +echo " RUN_MORI=$RUN_MORI DP_MODE=$DP_MODE KV_TRANSFER_BACKEND=$KV_TRANSFER_BACKEND" +echo " MASTER_ADDR=$MASTER_ADDR IPADDRS=$IPADDRS" +echo "==============================================================" + +if [[ "$RUN_MORI" == "1" ]]; then + exec bash "$SCRIPT_DIR/sglang_disagg_mori_io_ep.sh" +else + exec bash "$SCRIPT_DIR/sglang_disagg_server.sh" +fi diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index f224493..9b739c2 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -76,53 +76,10 @@ BARRIER_PORT="${BARRIER_PORT:-4342}" # Dependencies and Environment Setup # ============================================================================= -pip install py-spy -pip install --ignore-installed --force-reinstall flask -pip install pyyaml - - -host_ip=$(ip route get 1.1.1.1 | awk '/src/ {print $7}') -host_name=$(hostname) - -if [[ "$PARALLEL_MODE" != "dp" && "$PARALLEL_MODE" != "tp" ]]; then - echo "ERROR: PARALLEL_MODE must be 'dp' or 'tp' (got: ${PARALLEL_MODE})" - exit 1 -fi - -# ============================================================================= -# Parallelism Settings -# ============================================================================= - -# DP_MODE=0: CLI --tp-size is IO_EP_TP_SIZE (default 8) on every worker; PREFILL_EP_SIZE/DECODE_EP_SIZE -# still scale with xP/yD×GPUS_PER_NODE for MoRI env (not passed as CLI unless DP_MODE=1). -# DP_MODE=1: --tp-size scales with cluster; --dp-size/--ep-size on CLI (same total degree as Nnodes×GPUS_PER_NODE). +# === Model-Specific Configuration from YAML === +# Parse models.yaml into MODEL_* flags. (The runtime build layer was moved into +# the image: docker/sglang_disagg_inference_full_overlay*.Dockerfile.) GPUS_PER_NODE="${GPUS_PER_NODE:-8}" -GENERIC_TP_SIZE="${GENERIC_TP_SIZE:-8}" - -if [[ "$DP_MODE" == "1" ]]; then - PREFILL_TP_SIZE=$((xP * GPUS_PER_NODE)) - DECODE_TP_SIZE=$((yD * GPUS_PER_NODE)) -else - PREFILL_TP_SIZE="${GENERIC_TP_SIZE}" - DECODE_TP_SIZE="${GENERIC_TP_SIZE}" -fi - - -if [[ "$DP_MODE" == "1" ]]; then - PREFILL_EP_SIZE=$((xP * GPUS_PER_NODE)) - DECODE_EP_SIZE=$((yD * GPUS_PER_NODE)) - PREFILL_DP_SIZE=$((xP * GPUS_PER_NODE)) - DECODE_DP_SIZE=$((yD * GPUS_PER_NODE)) - export PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE -else - unset PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE 2>/dev/null || true -fi -export PREFILL_TP_SIZE DECODE_TP_SIZE - -# ============================================================================= -# Model-Specific Configuration from YAML -# ============================================================================= - MODELS_YAML="${MODELS_YAML:-${SCRIPT_DIR}/models.yaml}" if [[ ! -f "$MODELS_YAML" ]]; then @@ -175,6 +132,49 @@ for key, value in exports.items(): PY )" +host_ip=$(ip route get 1.1.1.1 | awk '/src/ {print $7}') +host_name=$(hostname) + +if [[ "$PARALLEL_MODE" != "dp" && "$PARALLEL_MODE" != "tp" ]]; then + echo "ERROR: PARALLEL_MODE must be 'dp' or 'tp' (got: ${PARALLEL_MODE})" + exit 1 +fi + +# ============================================================================= +# Parallelism Settings +# ============================================================================= + +# DP_MODE=0: CLI --tp-size is IO_EP_TP_SIZE (default 8) on every worker; PREFILL_EP_SIZE/DECODE_EP_SIZE +# still scale with xP/yD×GPUS_PER_NODE for MoRI env (not passed as CLI unless DP_MODE=1). +# DP_MODE=1: --tp-size scales with cluster; --dp-size/--ep-size on CLI (same total degree as Nnodes×GPUS_PER_NODE). +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +GENERIC_TP_SIZE="${GENERIC_TP_SIZE:-8}" + +if [[ "$DP_MODE" == "1" ]]; then + PREFILL_TP_SIZE=$((xP * GPUS_PER_NODE)) + DECODE_TP_SIZE=$((yD * GPUS_PER_NODE)) +else + PREFILL_TP_SIZE="${GENERIC_TP_SIZE}" + DECODE_TP_SIZE="${GENERIC_TP_SIZE}" +fi + + +if [[ "$DP_MODE" == "1" ]]; then + PREFILL_EP_SIZE=$((xP * GPUS_PER_NODE)) + DECODE_EP_SIZE=$((yD * GPUS_PER_NODE)) + PREFILL_DP_SIZE=$((xP * GPUS_PER_NODE)) + DECODE_DP_SIZE=$((yD * GPUS_PER_NODE)) + export PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE +else + unset PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE 2>/dev/null || true +fi +export PREFILL_TP_SIZE DECODE_TP_SIZE + +# ============================================================================= +# Model-Specific Configuration from YAML +# ============================================================================= + + PREFILL_MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_MODE_FLAGS} ${MODEL_PREFILL_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}" DECODE_MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_MODE_FLAGS} ${MODEL_DECODE_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}" echo "Using model-specific configuration for: $MODEL_NAME (mode=${PARALLEL_MODE})" @@ -184,7 +184,7 @@ export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG MODEL_EXPERIMENTAL_FLAGS # shellcheck disable=SC1091 source "${SCRIPT_DIR}/mori_ep_env.sh" -# KV transfer backend: default mori, switchable to mooncake (Mooncake). +# KV transfer backend: default mori, switchable to mooncake or nixl. # Kept out of models.yaml so model config is backend-agnostic. _TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mori}" PREFILL_MODEL_CONFIG+=" --disaggregation-transfer-backend ${_TRANSFER_BACKEND}" @@ -534,15 +534,17 @@ PY echo "" fi + benchmark_status=0 if [[ "${SKIP_BENCHMARK:-0}" != "1" ]] && [[ -n "${MOONCAKE_COOKBOOK_PATH:-}" ]]; then if [[ -f "${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" ]]; then echo "Running ${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" ( cd "${MOONCAKE_COOKBOOK_PATH}" || exit 1 bash benchmark_xPyD.sh - ) + ) || benchmark_status=$? else echo "WARN: benchmark_xPyD.sh not found under MOONCAKE_COOKBOOK_PATH=${MOONCAKE_COOKBOOK_PATH}" >&2 + benchmark_status=1 fi fi @@ -552,6 +554,11 @@ PY echo "Killing the co-located prefill server (pid=${_node0_prefill_pid})" kill "${_node0_prefill_pid}" + if [[ "${benchmark_status}" -ne 0 ]]; then + echo "ERROR: benchmark failed with status ${benchmark_status}" >&2 + exit "${benchmark_status}" + fi + elif [[ "$NODE_RANK" -ge 1 && "$NODE_RANK" -lt "$xP" ]]; then echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-default})" # NODE_RANK 0..xP-1 map directly to PREFILL_NODE_RANK 0..xP-1 (proxy co-located on NODE_RANK=0). diff --git a/scripts/sglang_disagg/sglang_disagg_server.sh b/scripts/sglang_disagg/sglang_disagg_server.sh index 261dcdd..3db53d7 100755 --- a/scripts/sglang_disagg/sglang_disagg_server.sh +++ b/scripts/sglang_disagg/sglang_disagg_server.sh @@ -14,6 +14,17 @@ MODEL_NAME="${MODEL_NAME:-}" xP="${xP:-1}" yD="${yD:-1}" IPADDRS="${IPADDRS:-localhost}" +KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mooncake}" + +# KV_TRANSFER_BACKEND is interpolated into an eval'd launch command below; restrict +# it to known-good values to avoid invalid backends and shell-token injection. +case "$KV_TRANSFER_BACKEND" in + mori|mooncake|nixl) ;; + *) + echo "ERROR: unsupported KV_TRANSFER_BACKEND='$KV_TRANSFER_BACKEND' (expected: mori|mooncake|nixl)" >&2 + exit 1 + ;; +esac # ============================================================================= # Dependencies and Environment Setup @@ -39,6 +50,7 @@ declare -A MODEL_PREFILL_CONFIGS=( ["Llama-3.1-405B-Instruct-FP8-KV"]="--tp-size 8" ["amd-Llama-3.3-70B-Instruct-FP8-KV"]="--tp-size 8" ["DeepSeek-V3"]="--tp-size 8" + ["DeepSeek-R1"]="--tp-size 8" ) declare -A MODEL_DECODE_CONFIGS=( @@ -48,6 +60,7 @@ declare -A MODEL_DECODE_CONFIGS=( ["Llama-3.1-405B-Instruct-FP8-KV"]="--tp-size 8" ["amd-Llama-3.3-70B-Instruct-FP8-KV"]="--tp-size 8" ["DeepSeek-V3"]="--tp-size 8" + ["DeepSeek-R1"]="--tp-size 8" ) # ============================================================================= @@ -142,9 +155,9 @@ if [ "$NODE_RANK" -eq 0 ]; then #wait until prefill nodes get ready until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do if [ $SECONDS -ge $TIMEOUT_SECONDS ]; then - echo "Awaited ${SECONDS} seconds. Timeout reached. Signal not found in prefill ${i} file" \ - | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null - + echo "FATAL: awaited ${SECONDS}s; readiness signal not found for prefill ${i} (${LOG_FILE}). Aborting before launching the router." \ + | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >&2 + exit 1 fi sleep $SLEEP_SECONDS SECONDS=$(( SECONDS + SLEEP_SECONDS)) @@ -156,8 +169,9 @@ if [ "$NODE_RANK" -eq 0 ]; then #wait until decode nodes get ready until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do if [ $SECONDS -ge $TIMEOUT_SECONDS ]; then - echo "Awaited ${SECONDS} seconds. Timeout reached. Signal not found in decode ${i} file" \ - | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null + echo "FATAL: awaited ${SECONDS}s; readiness signal not found for decode ${i} (${LOG_FILE}). Aborting before launching the router." \ + | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >&2 + exit 1 fi sleep $SLEEP_SECONDS SECONDS=$(( SECONDS + SLEEP_SECONDS)) @@ -193,6 +207,7 @@ if [ "$NODE_RANK" -eq 0 ]; then elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -le "$xP" ]; then echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-'default'})" echo "Using prefill config: $PREFILL_MODEL_CONFIG" + echo "Using KV transfer backend: ${KV_TRANSFER_BACKEND}" PREFILL_CMD="MC_TE_METRIC=true python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ @@ -202,7 +217,7 @@ elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -le "$xP" ]; then --port 2322 \ --stream-output \ --trust-remote-code \ - --disaggregation-transfer-backend mooncake" + --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" if [[ -n "$PREFILL_MODEL_CONFIG" ]]; then PREFILL_CMD="$PREFILL_CMD $PREFILL_MODEL_CONFIG" @@ -229,6 +244,7 @@ elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -le "$xP" ]; then else echo "${host_name}:${host_ip} is Decode Node (Model: ${MODEL_NAME:-'default'})" echo "Using decode config: $DECODE_MODEL_CONFIG" + echo "Using KV transfer backend: ${KV_TRANSFER_BACKEND}" DECODE_CMD="python3 -m sglang.launch_server \ --model-path ${MODEL_PATH} \ @@ -238,7 +254,7 @@ else --port 2322 \ --stream-output \ --trust-remote-code \ - --disaggregation-transfer-backend mooncake" + --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" if [[ -n "$DECODE_MODEL_CONFIG" ]]; then DECODE_CMD="$DECODE_CMD $DECODE_MODEL_CONFIG" From 348e3602ed10b610d12a6afea99a8f28520b760d Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Wed, 1 Jul 2026 16:23:14 +0000 Subject: [PATCH 02/25] Add Primus Megatron-LM scaleout Dockerfile and benchmark reporting scripts Add a Primus Megatron-LM multi-node scaleout training image (candidate RCCL built from source) and the benchmark scripts that drive it and convert training output into the madengine perf-CSV format. Key changes: - New overlay Dockerfile `primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile` on the `rocm/primus:v26.4` base (RCCL built from source + librccl smifix, optional rdma-core for Broadcom Thor2 bnxt_re). - New scaleout scripts under `scripts/primus_scaleout/megatron-lm/`: `run.sh`, `primus_megatron-lm_benchmark_setup.sh`, `primus_megatron-lm_benchmark_report.sh`, and `primus_megatron-lm_benchmark_report.py` (covers Llama-3.1 8B/70B/405B). Co-authored-by: Ilia Kosarev --- ...n_train_rccl_overlay.ubuntu.amd.Dockerfile | 288 +++++++++ .../primus_megatron-lm_benchmark_report.py | 149 +++++ .../primus_megatron-lm_benchmark_report.sh | 610 ++++++++++++++++++ .../primus_megatron-lm_benchmark_setup.sh | 41 ++ scripts/primus_scaleout/megatron-lm/run.sh | 143 ++++ 5 files changed, 1231 insertions(+) create mode 100644 docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile create mode 100644 scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py create mode 100755 scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh create mode 100755 scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh create mode 100755 scripts/primus_scaleout/megatron-lm/run.sh diff --git a/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile new file mode 100644 index 0000000..7ff95b7 --- /dev/null +++ b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile @@ -0,0 +1,288 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +# +# Primus megatron-lm training image with a candidate RCCL built from source. +# +# base : rocm/primus:v26.4 (ROCm 7.2.x, torch 2.10, python3.12) +# + : RCCL built from ${RCCL_REPO}@${RCCL_BRANCH}[/${RCCL_COMMIT}] for +# ${BUILD_GPU_TARGETS} (gfx942=MI300X/MI325X, gfx950=MI350X/MI355X) +# + : candidate librccl installed over EVERY location torch resolves +# + : optional rdma-core ${RDMA_CORE_VERSION} from source (Broadcom Thor2 +# bnxt_re EFAULT fix on some fabrics; empty = keep base rdma-core) +# +# WHY install the candidate librccl into BOTH the system rocm lib dir AND +# torch/lib (not just LD_LIBRARY_PATH=/opt/rccl/lib): +# the primus v26.x base ships NO bundled librccl under torch/lib and +# libtorch_hip.so links the SYSTEM /opt/rocm/lib/librccl.so.1 (verified via ldd +# on v26.3). LD_LIBRARY_PATH alone is not enough, so we overwrite the system +# librccl and drop a copy into torch/lib (RPATH=$ORIGIN) as belt-and-suspenders. +# This is also safe on earlier v26.2/v26.3 bases. +# +# The git-clone source keeps .git, so RCCL's git_version.cmake bakes the real +# commit into the binary as a separate rcclGitHash string ":" +# (e.g. "HEAD:c67fbe4" for a detached checkout) — no tarball/hash-override hack +# required. The final gate confirms one of those baked short hashes is a prefix +# of the recorded RCCL_BUILT_SHA. +# +# IMAGE LAYOUT: RCCL is compiled in a throwaway `rccl-builder` stage and only +# the installed ${RCCL_INSTALL_DIR} tree is COPYed into the final image. This +# keeps the full RCCL source/build tree AND the heavy build toolchain +# (build-essential, cmake, ninja) out of the shipped image. +# +ARG BASE_DOCKER=rocm/primus:v26.4 + +# ---- shared base: apt mirror hygiene (so both stages resolve the same) ------ +FROM ${BASE_DOCKER} AS apt-base +SHELL ["/bin/bash", "-o", "pipefail", "-c"] +USER root +# Force https mirrors + ensure the universe/multiverse components are enabled so +# build deps resolve consistently in both stages. +RUN set -e && \ + sed -i 's|http://archive.ubuntu.com/ubuntu|https://archive.ubuntu.com/ubuntu|g; s|http://security.ubuntu.com/ubuntu|https://security.ubuntu.com/ubuntu|g' /etc/apt/sources.list 2>/dev/null || true && \ + if [ -d /etc/apt/sources.list.d ]; then \ + sed -i 's|http://archive.ubuntu.com/ubuntu|https://archive.ubuntu.com/ubuntu|g; s|http://security.ubuntu.com/ubuntu|https://security.ubuntu.com/ubuntu|g' /etc/apt/sources.list.d/*.list 2>/dev/null || true; \ + sed -i 's|http://archive.ubuntu.com/ubuntu|https://archive.ubuntu.com/ubuntu|g; s|http://security.ubuntu.com/ubuntu|https://security.ubuntu.com/ubuntu|g' /etc/apt/sources.list.d/*.sources 2>/dev/null || true; \ + sed -i -E 's/^Components:.*/Components: main restricted universe multiverse/g' /etc/apt/sources.list.d/*.sources 2>/dev/null || true; \ + fi + +# ---- Stage 1 (throwaway): build + verify candidate RCCL --------------------- +# Everything heavy (full RCCL clone, submodules, build tree, gcc/cmake/ninja) +# lives here and is discarded; only ${RCCL_INSTALL_DIR} is carried forward. +FROM apt-base AS rccl-builder + +ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git +ARG RCCL_BRANCH=develop +ARG RCCL_COMMIT= +ARG RCCL_INSTALL_DIR=/opt/rccl +# gfx942 = MI300X / MI325X ; gfx950 = MI350X / MI355X +ARG BUILD_GPU_TARGETS=gfx942 + +RUN mkdir -p "${RCCL_INSTALL_DIR}" +WORKDIR /opt + +RUN apt-get -o Acquire::ForceIPv4=true \ + -o Acquire::Retries=5 \ + -o Acquire::http::Timeout=30 \ + -o Acquire::https::Timeout=30 \ + update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git \ + ca-certificates \ + cmake \ + ninja-build \ + pkg-config \ + make \ + patch \ + build-essential \ + libatomic1 \ + patchelf && \ + rm -rf /var/lib/apt/lists/* + +# clone the candidate RCCL and record the exact built commit +RUN if [[ -n "${RCCL_COMMIT}" ]]; then \ + git clone "${RCCL_REPO}" /tmp/rccl; \ + else \ + git clone --depth 1 --branch "${RCCL_BRANCH}" "${RCCL_REPO}" /tmp/rccl; \ + fi && \ + cd /tmp/rccl && \ + if [[ -n "${RCCL_BRANCH}" ]]; then git checkout "${RCCL_BRANCH}"; fi && \ + if [[ -n "${RCCL_COMMIT}" ]]; then git checkout "${RCCL_COMMIT}"; fi && \ + if [[ -d projects/rccl ]]; then \ + cd projects/rccl && git submodule update --init --recursive; \ + echo "/tmp/rccl/projects/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + else \ + git submodule update --init --recursive; \ + echo "/tmp/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + fi && \ + git -C "$(cat /tmp/BLD_RCCL_HOME.txt)" rev-parse HEAD > "${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA" && \ + echo "RCCL_BUILT_SHA=$(cat ${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA)" + +RUN set -e && \ + BLD_RCCL_HOME=$(cat /tmp/BLD_RCCL_HOME.txt) && \ + cd "${BLD_RCCL_HOME}" && \ + ./install.sh --amdgpu_targets="${BUILD_GPU_TARGETS}" --prefix="${RCCL_INSTALL_DIR}" + +# ---- build-verification gate: artifact exists + 0 undefined GDAKI/atomic ---- +# Fail first if the built librccl is missing (otherwise `nm` would error, the +# counts would default to 0, and the gate would wrongly pass). The `|| true` on +# the grep counts is only to tolerate "no match" (grep exit 1) under -o pipefail. +RUN set -e; \ + LIB="${RCCL_INSTALL_DIR}/lib/librccl.so"; \ + echo "=== librccl.so ==="; ls -laL "${RCCL_INSTALL_DIR}/lib/" | grep -i rccl || true; \ + [ -e "$LIB" ] || { echo "BUILD GATE FAIL: $LIB is missing"; exit 1; }; \ + nd=$(nm -D --undefined-only "$LIB" | grep -cE "[Gg]daki|GDAKI" || true); \ + na=$(nm -D --undefined-only "$LIB" | grep -E "__atomic_" | grep -vc "@LIBATOMIC" || true); \ + echo "gdaki_undefined=${nd:-0} unversioned_atomic_undefined=${na:-0}"; \ + if [ "${nd:-0}" -ne 0 ] || [ "${na:-0}" -ne 0 ]; then echo "BUILD GATE FAIL: undefined symbols present"; exit 1; fi; \ + echo "GATE PASS (built RCCL @ $(cat ${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA))" + +# ---- final image ------------------------------------------------------------ +FROM apt-base + +ARG BASE_DOCKER +ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git +ARG RCCL_BRANCH=develop +ARG RCCL_INSTALL_DIR=/opt/rccl +# gfx942 = MI300X / MI325X ; gfx950 = MI350X / MI355X +ARG BUILD_GPU_TARGETS=gfx942 +# Optional rdma-core from source (e.g. 63.0 for the Broadcom Thor2 bnxt_re +# EFAULT fix). Empty string keeps the base image's rdma-core untouched. +ARG RDMA_CORE_VERSION= + +ENV WORKSPACE_DIR=/workspace +ENV RCCL_HOME=${RCCL_INSTALL_DIR} +# Prepend the candidate RCCL dir but KEEP the base image's paths (ROCm libs). +# The candidate librccl is also installed over the system path + ldconfig, so +# this is belt-and-suspenders rather than the sole resolution mechanism. +ENV LD_LIBRARY_PATH=${RCCL_INSTALL_DIR}/lib:${LD_LIBRARY_PATH} +ENV NCCL_DEBUG=WARN + +RUN mkdir -p "${WORKSPACE_DIR}" +WORKDIR /opt + +# Bring in ONLY the installed candidate RCCL tree (no source/build artifacts). +COPY --from=rccl-builder ${RCCL_INSTALL_DIR} ${RCCL_INSTALL_DIR} + +# Minimal tools needed to splice the candidate librccl into the image and to run +# the final ELF verification: patchelf (--add-needed), binutils (objdump / nm / +# strings / readelf) and libatomic1 (runtime dep of the candidate librccl). +RUN apt-get -o Acquire::ForceIPv4=true -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + patchelf binutils libatomic1 && \ + rm -rf /var/lib/apt/lists/* + +# ---- Stage 2: install candidate librccl over EVERY location torch resolves -- +# Overwrite the system /opt/rocm librccl AND drop a copy into torch/lib so the +# candidate always wins regardless of how libtorch_hip is linked. Also +# add-needed librocm_smi64 (folded-in rsmifix) so `import torch` resolves +# rsmi_init on bases where the overlay librccl lacks it in NEEDED. +# add_needed() adds librocm_smi64 to NEEDED idempotently (skips if already +# present, so re-running never appends duplicates) and FAILS the build if +# patchelf cannot apply it -- `import torch` depends on this resolving rsmi_init. +RUN set -e && \ + add_needed() { \ + local lib="$1"; \ + if objdump -p "$lib" | grep -qE "NEEDED[[:space:]]+librocm_smi64"; then \ + echo " add_needed: already present in $lib"; return 0; \ + fi; \ + patchelf --add-needed librocm_smi64.so.1 "$lib"; \ + objdump -p "$lib" | grep -qE "NEEDED[[:space:]]+librocm_smi64" || { echo "PATCHELF FAILED on $lib"; exit 1; }; \ + echo " add_needed: applied to $lib"; \ + }; \ + src=$(ls -L "${RCCL_INSTALL_DIR}/lib/librccl.so.1.0" 2>/dev/null || ls -L "${RCCL_INSTALL_DIR}/lib/librccl.so") && \ + echo "candidate librccl: $src" && \ + add_needed "$src" && \ + ROCM_LIB=$(dirname "$(readlink -f /opt/rocm/lib/librccl.so.1)") && \ + echo "system rocm lib dir: $ROCM_LIB" && \ + echo "before:" && ls -la "$ROCM_LIB"/librccl.so* && \ + for f in "$ROCM_LIB"/librccl.so*; do \ + if [ -f "$f" ] && [ ! -L "$f" ]; then echo "overwriting real file: $f"; cp -fL "$src" "$f"; add_needed "$f"; fi; \ + done && \ + ldconfig && \ + echo "after:" && ls -laL "$ROCM_LIB"/librccl.so.1 && \ + TORCH_LIB=$(ls -d /opt/venv/lib/python*/site-packages/torch/lib 2>/dev/null | head -1) && \ + [ -n "$TORCH_LIB" ] && [ -d "$TORCH_LIB" ] && echo "torch lib dir: $TORCH_LIB" && \ + for f in "$TORCH_LIB"/librccl.so*; do \ + if [ -f "$f" ] && [ ! -L "$f" ]; then cp -fL "$src" "$f"; add_needed "$f"; fi; \ + done && \ + cp -fL "$src" "$TORCH_LIB/librccl.so.1.0" && \ + ln -sf librccl.so.1.0 "$TORCH_LIB/librccl.so.1" && \ + ln -sf librccl.so.1 "$TORCH_LIB/librccl.so" && \ + add_needed "$TORCH_LIB/librccl.so.1.0" && \ + if [ ! -e "$TORCH_LIB/librocm_smi64.so" ]; then \ + cp -v -L /opt/rocm/lib/librocm_smi64.so.1 "$TORCH_LIB/librocm_smi64.so"; \ + ln -sfv librocm_smi64.so "$TORCH_LIB/librocm_smi64.so.1"; \ + fi + +# ---- Stage 3 (optional): rdma-core from source ------------------------------ +# Builds only when RDMA_CORE_VERSION is non-empty (e.g. 63.0). Replaces the +# distro rdma-core to pick up the Broadcom Thor2 bnxt_re EFAULT fix. The build +# toolchain it needs is installed here (not in the base image) so the default +# image (RDMA_CORE_VERSION empty) never pays for it. +# +# ORDERING IS LOAD-BEARING: the source rdma-core is installed by REPLACING the +# distro libibverbs/librdmacm packages via `dpkg -r --force-all`. That leaves +# still-installed dependents (libopenmpi3t64, libucx0, rdmacm-utils, ...) with +# dangling deps, so ANY apt command after that point aborts with "Unmet +# dependencies". Therefore every apt operation (install + the docs-tooling +# purge/autoremove) MUST run BEFORE the dpkg removal, while the dpkg state is +# still consistent. Only dpkg + `ninja install` may run afterwards. +RUN if [[ -n "${RDMA_CORE_VERSION}" ]]; then \ + set -e; \ + apt-get -o Acquire::ForceIPv4=true -o Acquire::Retries=5 update; \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git ca-certificates cmake ninja-build build-essential pkg-config make \ + libnl-3-dev libnl-route-3-dev libudev-dev libssl-dev libsystemd-dev \ + python3-docutils pandoc; \ + rm -rf /var/lib/apt/lists/*; \ + git clone --depth 1 --branch "v${RDMA_CORE_VERSION}" https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core; \ + cd /tmp/rdma-core && git log -1 --format='%H %s'; \ + mkdir build && cd build; \ + cmake -GNinja \ + -DCMAKE_INSTALL_PREFIX=/usr \ + -DCMAKE_INSTALL_LIBDIR=lib/x86_64-linux-gnu \ + -DCMAKE_INSTALL_SYSCONFDIR=/etc \ + -DCMAKE_INSTALL_RUNDIR=/run \ + -DCMAKE_BUILD_TYPE=Release \ + -DNO_PYVERBS=1 \ + ..; \ + ninja -j"$(nproc)"; \ + DEBIAN_FRONTEND=noninteractive apt-get purge -y python3-docutils pandoc; \ + DEBIAN_FRONTEND=noninteractive apt-get autoremove -y; \ + rm -rf /var/lib/apt/lists/*; \ + to_remove=""; \ + for p in ibverbs-providers libibverbs1 libibverbs-dev ibverbs-utils \ + librdmacm1 librdmacm-dev rdma-core libibumad3 libibmad5 \ + infiniband-diags; do \ + if dpkg-query -W -f='${Status}' "$p" 2>/dev/null | grep -q "install ok installed"; then \ + to_remove="$to_remove $p"; \ + fi; \ + done; \ + if [[ -n "$to_remove" ]]; then \ + echo "removing distro rdma packages:$to_remove"; \ + dpkg -r --force-all $to_remove; \ + fi; \ + ninja install && ldconfig; \ + cd / && rm -rf /tmp/rdma-core; \ + else \ + echo "RDMA_CORE_VERSION empty — keeping base image rdma-core"; \ + fi + +# ---- final verification: import torch + every librccl == candidate ---------- +RUN set -e && \ + if [[ -n "${RDMA_CORE_VERSION}" ]]; then \ + echo "=== rdma-core ===" && readelf -d /usr/lib/x86_64-linux-gnu/libibverbs.so.1 | grep SONAME && \ + { ls /usr/lib/x86_64-linux-gnu/libibverbs/libbnxt_re-rdmav*.so 2>/dev/null | head || true; }; \ + fi && \ + echo "=== import torch (must succeed) ===" && \ + if [ -x /opt/venv/bin/python3 ]; then PY=/opt/venv/bin/python3; else PY=$(command -v python3); fi && \ + echo "using python: $PY" && \ + "$PY" -c "import torch; print('torch', torch.__version__, 'hip', torch.version.hip)" && \ + echo "=== ldd libtorch_hip rccl resolution ===" && \ + TL=$(ls -d /opt/venv/lib/python*/site-packages/torch/lib | head -1) && \ + ldd "$TL/libtorch_hip.so" | grep -iE "rccl|smi" && \ + BUILT_SHA=$(cat "${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA") && \ + echo "=== ASSERT: every resolvable librccl is the candidate built @ ${BUILT_SHA} ===" && \ + for p in "${RCCL_INSTALL_DIR}/lib/librccl.so.1" /opt/rocm/lib/librccl.so.1 "$TL/librccl.so.1"; do \ + v=$(strings "$p" 2>/dev/null | grep -m1 -oE "RCCL version 2\.[0-9]+\.[0-9]+"); \ + echo "$v" | grep -q "RCCL version 2." || { echo "ASSERT FAIL: $p is not RCCL 2.x ($p -> ${v:-NONE})"; exit 1; }; \ + h=""; \ + for cand in $(strings "$p" 2>/dev/null | grep -oE "[A-Za-z0-9_./-]+:[0-9a-f]{7,40}" | sed 's/.*://' | sort -u); do \ + case "$BUILT_SHA" in "$cand"*) h="$cand"; break;; esac; \ + done; \ + echo " $p -> ${v:-NONE} / baked=${h:-NONE}"; \ + [ -n "$h" ] || { echo "ASSERT FAIL: $p has no baked git commit matching built SHA $BUILT_SHA (upstream 'Unknown' fallback?)"; exit 1; }; \ + done && \ + echo "=== candidate RCCL commit: ${BUILT_SHA} ===" + +# primus_base is derived from BASE_DOCKER so the label cannot drift from the +# actual base when BASE_DOCKER is overridden. +LABEL primus_base="${BASE_DOCKER}" +LABEL rccl_repo="${RCCL_REPO}" +LABEL rccl_branch="${RCCL_BRANCH}" +LABEL build_gpu_targets="${BUILD_GPU_TARGETS}" +LABEL rdma_core_version="${RDMA_CORE_VERSION}" + +WORKDIR ${WORKSPACE_DIR} + +# Record final Python environment for posterity. +RUN pip3 list diff --git a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py new file mode 100644 index 0000000..a514102 --- /dev/null +++ b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py @@ -0,0 +1,149 @@ +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# + +import pandas as pd +import argparse +import csv +import os +import re + +# parse arguments +parser = argparse.ArgumentParser(description='Convert primus pytorch train output format to MAD csv output format') +parser.add_argument("--mode", + type=str, + help="pretrain or finetune") +parser.add_argument("--model", + type=str, + help="model name") +parser.add_argument("--input", + type=str, + help="path to input file") +parser.add_argument("--output", + type=str, + help="path to output file") +parser.add_argument("--precision", + type=str, + help="precision type") +parser.add_argument("--batch_size", + type=str, + help="batch size") +parser.add_argument("--seq_len", + type=str, + help="sequence length") +parser.add_argument("--device", + type=str, + help="device name") +parser.add_argument("--num_gpus", + type=str, + help="number of GPUs") + +# read arguments +args = parser.parse_args() +input_file = args.input +output_file = args.output +print("Input file path: ", input_file) +print("Output file path: ", output_file) + +data = [] + +def find_match(file, search_string): + with open(file, 'r') as file: + content = file.read() + # Updated pattern to match the new log format + # Looks for patterns like "TFLOP/s/GPU): 322.5/320.6" or "tokens/s/GPU): 725.1/720.6" + # and extracts the first value (before the slash) + pattern = fr"{re.escape(search_string)}\):\s*(\d+\.?\d*)/" + matches = re.findall(pattern, content) + if matches: + # Return the last 2 values if they exist (one from each run) + if len(matches) >= 2: + result = [matches[-2], matches[-1]] + print(f"Found {len(matches)} matches for '{search_string}', using last 2: {result}") + else: + result = [matches[-1]] + print(f"Found {len(matches)} match for '{search_string}': {result}") + return result + else: + print(f"Warning: No matches found for '{search_string}' pattern") + return [] + +def find_match_running_avg(file, search_string): + """Return the running-average value (the number AFTER the slash) from the + LAST logged iteration, e.g. '1885.6' from 'tokens/s/GPU): 2355.1/1885.6'. + + Primus logs throughput as 'instant/running_avg'. The instantaneous value + has large run-to-run jitter (MoE routing imbalance, collective stragglers, + kernel warmup), so the trailing running average is a more stable figure. + Returns None when no match is present.""" + with open(file, 'r') as f: + content = f.read() + pattern = fr"{re.escape(search_string)}\):\s*\d+\.?\d*/(\d+\.?\d*)" + matches = re.findall(pattern, content) + if matches: + print(f"Found {len(matches)} running-avg matches for '{search_string}', using last: {matches[-1]}") + return matches[-1] + return None + +if args.model == "Llama-3.1-8B" or args.model == "Llama-3.1-70B" or \ + args.model == "Llama-3.1-405B" or \ + args.model == "Llama-2-7B" or args.model == "Llama-2-70B" or \ + args.model == "Mixtral-8x7B" or args.model == "Mixtral-8x22B-proxy" or \ + args.model == "DeepSeek-V2-lite" or args.model == "DeepSeek-V3-proxy" or \ + args.model == "Llama-3.1-70B-proxy" or args.model == "Llama-3.3-70B" or \ + args.model == "Qwen2.5-7B" or args.model == "Qwen2.5-72B": + tok_per_s_per_gpu_list = find_match(input_file, "tokens/s/GPU") + TFLOPS_per_gpu_list = find_match(input_file, "TFLOP/s/GPU") + + data = [] + # Write separate rows for each run + for i, (tps, tflops) in enumerate(zip(tok_per_s_per_gpu_list, TFLOPS_per_gpu_list)): + run_label = f"run_{i+1}" if len(tok_per_s_per_gpu_list) > 1 else "" + data.extend([ + {'model': args.model, 'performance': tps, 'metric': 'tok_per_s_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': run_label}, + {'model': args.model, 'performance': tflops, 'metric': 'TFLOPS_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': run_label} + ]) + + # Additionally emit the trailing running-average throughput (value after the + # slash in the last logged iteration). Kept as separate, distinctly-named + # rows (suffix '_avg', run='avg') so the instantaneous metrics above are + # preserved and never confused with the steady-state average. + # Applied to all models in this branch; the average is a steadier figure + # given per-iteration throughput jitter (largest for gated_delta_net MoE). + tok_avg = find_match_running_avg(input_file, "tokens/s/GPU") + tflops_avg = find_match_running_avg(input_file, "TFLOP/s/GPU") + if tok_avg is not None: + data.append({'model': args.model, 'performance': tok_avg, 'metric': 'tok_per_s_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) + if tflops_avg is not None: + data.append({'model': args.model, 'performance': tflops_avg, 'metric': 'TFLOPS_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) + +with open(output_file, mode='w', newline='') as file: + print("Preparing to write performance data...") + print("Data: ", data) + writer = csv.DictWriter(file, fieldnames=['model','performance','metric','mode','precision','batch_size','seq_len','device','num_gpus','run']) + writer.writeheader() + writer.writerows(data) + print("Completed writing to output file") + diff --git a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh new file mode 100755 index 0000000..f03df53 --- /dev/null +++ b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh @@ -0,0 +1,610 @@ +#!/bin/bash +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +## Usage: +#./primus_megatron-lm_benchmark_report.sh -m $model_name +## example: +## Pretrain Llama 3.1 70B +#./primus_megatron-lm_benchmark_report.sh -m Llama-3.1-70B -p BF16 +## Pretrain Llama 2 7B +#./primus_megatron-lm_benchmark_report.sh -m Llama-2-7B -p FP8 +## Pretrain DeepSeek V2 Lite +#./primus_megatron-lm_benchmark_report.sh -m DeepSeek-V2-lite +## Pretrain DeepSeek V3 +#./primus_megatron-lm_benchmark_report.sh -m DeepSeek-V3 +## Pretrain Mixtral 8x7B +#./primus_megatron-lm_benchmark_report.sh -m Mixtral-8x7B -l 4 +MODEL_REPO="" +MODE="pretrain" +DATATYPE="BF16" + +usage() { + echo "Usage: $0 -m -p -l " + echo "\nOptions:" + echo " -m Model repository (Llama-2-7B, Llama-2-70B, Llama-3.1-8B, Llama-3.1-70B, Llama-3.1-405B, DeepSeek-V2-lite, DeepSeek-V3-proxy, Mixtral-8x7B, Mixtral-8x22B-proxy)" + echo " -p Precision type (FP8 or BF16)" + exit 1 +} + +# Parse command-line arguments +while getopts "m:p:" opt; do + case "$opt" in + m) MODEL_REPO="$OPTARG" ;; + p) DATATYPE="$OPTARG" ;; + *) usage ;; + esac +done + +echo "=hyper params start=" +echo $MODEL_REPO +echo $DATATYPE +echo "=hyper params end=" + +# Validate inputs +if [[ -z "$MODEL_REPO" ]]; then + echo "Error: Missing required arguments." + usage +fi + +if [[ "$DATATYPE" != "FP8" && "$DATATYPE" != "BF16" ]]; then + echo "Error: Datatype must be either FP8 or BF16." + exit 1 +fi + +# Run benchmark (Placeholder for actual script execution) +echo "Running primus training benchmark with the following parameters:" +echo " Model Repository: $MODEL_REPO" + +# config environment +export MOCK_DATA=1 + +# set performance output paths +TRAIN_LOG="$(pwd)/primus-megatron-$MODEL_REPO-pretrain.csv" +echo "TRAIN LOG: $TRAIN_LOG" + +PERF_LOG="$(pwd)/../perf_primus-megatron-$MODEL_REPO.csv" +echo "PERF LOG: $PERF_LOG" +ls $(pwd) +perf_script="$(pwd)/primus_megatron-lm_benchmark_report.py" + +# Detect device. Newer ROCm stacks may format rocminfo differently, +# so keep rocminfo first and fall back to amd-smi and the injected arch. +detect_device() { + local device="" + local arch="" + + if [[ -x /opt/rocm/bin/rocminfo ]]; then + device=$(/opt/rocm/bin/rocminfo 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') + arch=$(/opt/rocm/bin/rocminfo 2>/dev/null | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + fi + + if [[ -z "$device" && -x /opt/rocm/bin/amd-smi ]]; then + device=$(/opt/rocm/bin/amd-smi 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') + fi + + if [[ -z "$arch" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then + arch=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') + fi + + case "$device" in + MI300X|MI325X|MI350X|MI355X) + ;; + *) + case "$arch" in + gfx942) device="MI300X" ;; + gfx950) device="MI355X" ;; + *) device="" ;; + esac + ;; + esac + + echo "$device" +} + +DEVICE=$(detect_device) +echo "DEVICE found: $DEVICE" +echo "GPU DEVICE name: $DEVICE" +if [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1 + export NVTE_CK_IS_V3_ATOMIC_FP32=1 +fi + +# Set common environment variables +# Keep compatibility with SLURM/madengine distributed env and only fallback to single-node defaults. +NNODES="${NNODES:-1}" +export NNODES +export CPUS_PER_TASK="${CPUS_PER_TASK:-128}" +export HSA_NO_SCRATCH_RECLAIM=1 +export NVTE_CK_USES_BWD_V3=1 +GPUS_PER_NODE="${GPUS_PER_NODE:-${NPROC_PER_NODE:-8}}" +if ! [[ "$NNODES" =~ ^[0-9]+$ ]]; then NNODES=1; fi +if ! [[ "$GPUS_PER_NODE" =~ ^[0-9]+$ ]]; then GPUS_PER_NODE=8; fi +NUM_GPUS=$((NNODES * GPUS_PER_NODE)) +echo "[INFO] Distributed params: NNODES=$NNODES GPUS_PER_NODE=$GPUS_PER_NODE NUM_GPUS=$NUM_GPUS" + +normalize_global_batch_size() { + local mbs="$1" + local gbs="$2" + local dp="$3" + local unit=$((mbs * dp)) + if (( unit <= 0 )); then + echo "$gbs" + return + fi + if (( gbs % unit == 0 )); then + echo "$gbs" + return + fi + local adjusted=$(( ((gbs + unit - 1) / unit) * unit )) + echo "$adjusted" +} + +# Launch a single Primus pretrain run. Every rank must execute an identical +# command so the model shape stays consistent across nodes; a mismatch makes +# multi-node ranks rendezvous with different shapes and deadlocks NCCL. Pass +# the config first, then any shape/proxy/batch overrides -- they reach every +# rank because all ranks run this exact command. +run_primus() { + local config="$1"; shift + bash runner/primus-cli direct \ + --log_file "/tmp/primus_$MODEL_REPO.log" \ + -- train pretrain \ + --config "$config" "$@" 2>&1 | tee "$TRAIN_LOG" +} + +cd /workspace/Primus + +# run models +if [ "$MODEL_REPO" == "Llama-3.1-8B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/llama3.1_8B-$DATATYPE-pretrain.yaml + SEQ_LEN=8192 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + MBS=4 + GBS=512 + run_primus "$EXP" + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + MBS=2 + GBS=128 + run_primus "$EXP" + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Llama-3.1-70B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" # BF16 training only + export EXP=examples/megatron/configs/$DEVICE/llama3.1_70B-$DATATYPE-pretrain.yaml + SEQ_LEN=8192 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + MBS=3 + GBS=24 + original_gbs=$GBS + GBS=$(normalize_global_batch_size "$MBS" "$GBS" "$NUM_GPUS") + if [[ "$GBS" != "$original_gbs" ]]; then + echo "[INFO] Adjusted global batch size for distributed run: ${original_gbs} -> ${GBS} (MBS=${MBS}, NUM_GPUS=${NUM_GPUS})" + fi + run_primus "$EXP" --micro_batch_size $MBS --global_batch_size $GBS + elif [ "$DATATYPE" == "BF16" ]; then + MBS=4 + GBS=32 + original_gbs=$GBS + GBS=$(normalize_global_batch_size "$MBS" "$GBS" "$NUM_GPUS") + if [[ "$GBS" != "$original_gbs" ]]; then + echo "[INFO] Adjusted global batch size for distributed run: ${original_gbs} -> ${GBS} (MBS=${MBS}, NUM_GPUS=${NUM_GPUS})" + fi + run_primus "$EXP" --micro_batch_size $MBS --global_batch_size $GBS + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=3 + GBS=24 + original_gbs=$GBS + GBS=$(normalize_global_batch_size "$MBS" "$GBS" "$NUM_GPUS") + if [[ "$GBS" != "$original_gbs" ]]; then + echo "[INFO] Adjusted global batch size for distributed run: ${original_gbs} -> ${GBS} (MBS=${MBS}, NUM_GPUS=${NUM_GPUS})" + fi + run_primus "$EXP" --micro_batch_size $MBS --global_batch_size $GBS + fi + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Llama-3.1-405B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/llama3.1_405B-$DATATYPE-pretrain.yaml + SEQ_LEN=8192 + MBS=1 + GBS=8 + if [[ "$DATATYPE" == "FP8" ]]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO. Only BF16 is supported." + elif [[ ! -f "$EXP" ]]; then + echo "Error: Config file not found: $EXP" + echo "Hint: add llama3.1_405B-$DATATYPE-pretrain.yaml in Primus configs." + else + run_primus "$EXP" + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Llama-3.1-70B-proxy" ]; then + echo "[INFO] $MODEL_REPO TRAINING" # FP8 training only + export EXP=examples/megatron/configs/$DEVICE/llama3.1_70B-$DATATYPE-pretrain.yaml + SEQ_LEN=8192 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + echo "Error: Llama-3.1-70B-proxy model is not supported on $DEVICE. To train use the full Llama-3.1-70B model." + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + MBS=3 + GBS=24 + run_primus "$EXP" --num_layers 40 --fp8 hybrid --micro_batch_size $MBS --global_batch_size $GBS --no_fp8_weight_transpose_cache true + elif [ "$DATATYPE" == "BF16" ]; then + echo "Error: Datatype BF16 is not supported for $MODEL_REPO on $DEVICE. Only FP8 is supported." + fi + fi + + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Llama-3.3-70B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/llama3.3_70B-$DATATYPE-pretrain.yaml + SEQ_LEN=8192 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=6 + GBS=48 + run_primus "$EXP" + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=2 + GBS=16 + run_primus "$EXP" + fi + fi + + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Llama-2-7B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/llama2_7B-$DATATYPE-pretrain.yaml + SEQ_LEN=4096 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + export MBS=13 + export GBS=416 + run_primus "$EXP" + elif [ "$DATATYPE" == "BF16" ]; then + export MBS=10 + export GBS=640 + run_primus "$EXP" + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + export MBS=4 + export GBS=256 + if [ "$DATATYPE" == "FP8" ]; then + run_primus "$EXP" + elif [ "$DATATYPE" == "BF16" ]; then + run_primus "$EXP" + fi + fi + + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Llama-2-70B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/llama2_70B-$DATATYPE-pretrain.yaml + SEQ_LEN=4096 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=17 + GBS=272 + run_primus "$EXP" + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=7 + GBS=56 + run_primus "$EXP" + fi + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "DeepSeek-V2-lite" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/deepseek_v2_lite-$DATATYPE-pretrain.yaml + SEQ_LEN=4096 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + MBS=12 + GBS=768 + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + run_primus "$EXP" + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=4 + GBS=640 + run_primus "$EXP" + fi + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/deepseek_v3-$DATATYPE-pretrain.yaml + SEQ_LEN=4096 + # DeepEP (Primus-Turbo) is opt-in via PRIMUS_USE_DEEPEP=1 and applies to every + # supported device/precision branch below. DeepEP is incompatible with + # moe_shared_expert_overlap (Primus ROCm validate_args asserts it), so disable + # that overlap when enabling DeepEP. Default (unset/0) keeps alltoall over RCCL. + DEEPEP_ARGS="" + if [[ "${PRIMUS_USE_DEEPEP:-0}" == "1" ]]; then + DEEPEP_ARGS="--use_turbo_deepep true --moe_shared_expert_overlap false" + echo "[INFO] DeepEP enabled via PRIMUS_USE_DEEPEP=1 -> ${DEEPEP_ARGS}" + fi + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=8 + GBS=64 + run_primus "$EXP" --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS $DEEPEP_ARGS + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=3 + GBS=192 + run_primus "$EXP" --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS $DEEPEP_ARGS + fi + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Mixtral-8x7B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/mixtral_8x7B_v0.1-$DATATYPE-pretrain.yaml + SEQ_LEN=4096 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=4 + GBS=256 + run_primus "$EXP" + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=2 + GBS=32 + run_primus "$EXP" + fi + fi + + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Mixtral-8x22B-proxy" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + LAYERS=4 # default proxy model uses 4 layers + echo "[INFO] Proxy model uses $LAYERS layers" + export EXP=examples/megatron/configs/$DEVICE/mixtral_8x22B_v0.1-$DATATYPE-pretrain.yaml + SEQ_LEN=8192 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=2 + GBS=16 + run_primus "$EXP" --num_layers 4 --pipeline_model_parallel_size 1 --micro_batch_size $MBS --global_batch_size $GBS + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=1 + GBS=16 + run_primus "$EXP" --num_layers 4 --pipeline_model_parallel_size 1 --micro_batch_size $MBS --global_batch_size $GBS + fi + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +elif [ "$MODEL_REPO" == "Qwen2.5-7B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/qwen2.5_7B-$DATATYPE-pretrain.yaml + SEQ_LEN=2048 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [[ "$DATATYPE" == "BF16" ]]; then + MBS=16 + GBS=768 + run_primus "$EXP" + elif [[ "$DATATYPE" == "FP8" ]]; then + MBS=20 + GBS=800 + run_primus "$EXP" + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + MBS=10 + GBS=640 + if [ "$DATATYPE" == "FP8" ]; then + run_primus "$EXP" + elif [ "$DATATYPE" == "BF16" ]; then + run_primus "$EXP" + fi + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + fi + +elif [ "$MODEL_REPO" == "Qwen2.5-72B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" + export EXP=examples/megatron/configs/$DEVICE/qwen2.5_72B-$DATATYPE-pretrain.yaml + SEQ_LEN=2048 + if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=16 + GBS=256 + run_primus "$EXP" + fi + elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + if [ "$DATATYPE" == "FP8" ]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." + elif [ "$DATATYPE" == "BF16" ]; then + MBS=4 + GBS=32 + run_primus "$EXP" + fi + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + +else + echo "Error: Unsupported training mode." + exit 1 +fi diff --git a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh new file mode 100755 index 0000000..17c18d3 --- /dev/null +++ b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh @@ -0,0 +1,41 @@ +#!/bin/bash +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +export HF_HOME=/workspace/huggingface + +# Parse named arguments +while [[ "$#" -gt 0 ]]; do + case $1 in + -m) MODEL_NAME="$2"; shift ;; + *) echo "Unknown parameter passed: $1"; usage ;; + esac + shift +done + +echo "[INFO] Primus setup script starting in directory $(pwd)" + +cd /workspace/Primus +git pull \ No newline at end of file diff --git a/scripts/primus_scaleout/megatron-lm/run.sh b/scripts/primus_scaleout/megatron-lm/run.sh new file mode 100755 index 0000000..6d040ca --- /dev/null +++ b/scripts/primus_scaleout/megatron-lm/run.sh @@ -0,0 +1,143 @@ +#!/bin/bash +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# + +export HF_TOKEN=$MAD_SECRETS_HFTOKEN + +# Parse named arguments +while [[ "$#" -gt 0 ]]; do + case $1 in + --model_repo) MODEL_REPO="$2"; shift ;; + *) echo "Unknown parameter passed: $1"; usage ;; + esac + shift +done + +if [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-8b" ]]; then + model="Llama-3.1-8B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b" ]]; then + model="Llama-3.1-70B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-405b" ]]; then + model="Llama-3.1-405B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b-proxy" ]]; then + model="Llama-3.1-70B-proxy" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.3-70b" ]]; then + model="Llama-3.3-70B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-2-7b" ]]; then + model="Llama-2-7B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-2-70b" ]]; then + model="Llama-2-70B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_deepseek-v2-lite-16b" ]]; then + model="DeepSeek-V2-lite" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_deepseek-v2" ]]; then + model="DeepSeek-V2" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_deepseek-v3-proxy" ]]; then + model="DeepSeek-V3-proxy" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_mixtral-8x7b" ]]; then + model="Mixtral-8x7B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_mixtral-8x22b-proxy" ]]; then + model="Mixtral-8x22B-proxy" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_qwen2.5-7b" ]]; then + model="Qwen2.5-7B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_qwen2.5-72b" ]]; then + model="Qwen2.5-72B" +fi + +# Run primus pytorch setup script +echo "Running setup script to download tokenizers" +bash ./primus_megatron-lm_benchmark_setup.sh -m $model + +# Detect device. Newer ROCm stacks may format rocminfo differently, +# so keep rocminfo first and fall back to amd-smi and the injected arch. +detect_device() { + local device="" + local arch="" + + if [[ -x /opt/rocm/bin/rocminfo ]]; then + device=$(/opt/rocm/bin/rocminfo 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') + arch=$(/opt/rocm/bin/rocminfo 2>/dev/null | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + fi + + if [[ -z "$device" && -x /opt/rocm/bin/amd-smi ]]; then + device=$(/opt/rocm/bin/amd-smi 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') + fi + + if [[ -z "$arch" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then + arch=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') + fi + + case "$device" in + MI300X|MI325X|MI350X|MI355X) + ;; + *) + case "$arch" in + gfx942) device="MI300X" ;; + gfx950) device="MI355X" ;; + *) device="" ;; + esac + ;; + esac + + echo "$device" +} + +DEVICE=$(detect_device) +echo "GPU DEVICE name: $DEVICE" + +# Define supported datatypes based on device and model +datatypes=() + +if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then + # MI355X/MI350X support + if [[ "$model" == "Llama-3.1-70B-proxy" ]]; then + echo "Skipping $model - Not supported on $DEVICE" + elif [[ "$model" == "Llama-3.1-8B" || "$model" == "Llama-3.1-70B" || "$model" == "Llama-2-7B" || "$model" == "Qwen2.5-7B" ]]; then + datatypes=("BF16" "FP8") + else + # Most other models only support BF16 on MI355X/MI350X + datatypes=("BF16") + fi +elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then + # MI300X/MI325X support + if [[ "$model" == "Llama-3.1-70B-proxy" ]]; then + datatypes=("FP8") # Only FP8 supported + elif [[ "$model" == "Llama-3.1-8B" || "$model" == "Llama-2-7B" || "$model" == "Qwen2.5-7B" ]]; then + datatypes=("BF16" "FP8") # Both supported + else + # Most large models only support BF16 on MI300X/MI325X + datatypes=("BF16") + fi +else + # Unknown device, try both + datatypes=("BF16" "FP8") +fi + +# datatypes=("FP8") +# Loop through supported combinations +for datatype in "${datatypes[@]}"; do + echo "Running: $model - $datatype" + ./primus_megatron-lm_benchmark_report.sh -m $model -p $datatype +done From cf07399b52372ab895abff489f8f748a2c7eb9b3 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Wed, 1 Jul 2026 16:33:27 +0000 Subject: [PATCH 03/25] Add mad-slurm-multinode skill and associated templates Introduce the mad-slurm-multinode skill, which deploys and runs madengine performance tests on an unprepared SLURM cluster from scratch and launches multi-node runs from manifest templates. Key additions: - `SKILL.md` plus reference docs (cluster-types, deploy-bootstrap, manifests, launch-and-results, gotchas) and helper scripts (`detect_cluster_env.sh`, `preflight.sh`, `validate_manifest.sh`). - Cluster-agnostic `mad.env` templates for the CX7/Mellanox-RoCE, AMD-AINIC/Pollara, and Broadcom-Thor2-RoCE archetypes. - Manifest templates for the supported workloads: `sglang_disagg_deepseek-r1` and Primus Megatron-LM scaleout `primus_llama-3.1-8b`/`-70b` (on `rocm/primus:v26.4`). - Sanitized end-to-end example walkthroughs (no real node names, queues, or tokens) for each archetype. Co-authored-by: Ilia Kosarev --- .claude/skills/mad-slurm-multinode/SKILL.md | 329 ++++++++++++++++++ .../assets/mad.env/mad.env.amd-ainic.template | 68 ++++ .../assets/mad.env/mad.env.cx7-roce.template | 48 +++ .../mad.env/mad.env.thor2-bnxt.template | 50 +++ .../primus_llama-3.1-70b.template.json | 109 ++++++ .../primus_llama-3.1-8b.template.json | 109 ++++++ .../sglang_disagg_deepseek-r1.template.json | 152 ++++++++ .../examples/amd-ainic-walkthrough.md | 93 +++++ .../examples/cx7-roce-walkthrough.md | 84 +++++ .../examples/thor2-bnxt-walkthrough.md | 103 ++++++ .../references/cluster-types.md | 125 +++++++ .../references/deploy-bootstrap.md | 143 ++++++++ .../mad-slurm-multinode/references/gotchas.md | 223 ++++++++++++ .../references/launch-and-results.md | 104 ++++++ .../references/manifests.md | 150 ++++++++ .../scripts/detect_cluster_env.sh | 104 ++++++ .../mad-slurm-multinode/scripts/preflight.sh | 90 +++++ .../scripts/validate_manifest.sh | 147 ++++++++ .gitignore | 3 + 19 files changed, 2234 insertions(+) create mode 100644 .claude/skills/mad-slurm-multinode/SKILL.md create mode 100644 .claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template create mode 100644 .claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template create mode 100644 .claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template create mode 100644 .claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json create mode 100644 .claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json create mode 100644 .claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json create mode 100644 .claude/skills/mad-slurm-multinode/examples/amd-ainic-walkthrough.md create mode 100644 .claude/skills/mad-slurm-multinode/examples/cx7-roce-walkthrough.md create mode 100644 .claude/skills/mad-slurm-multinode/examples/thor2-bnxt-walkthrough.md create mode 100644 .claude/skills/mad-slurm-multinode/references/cluster-types.md create mode 100644 .claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md create mode 100644 .claude/skills/mad-slurm-multinode/references/gotchas.md create mode 100644 .claude/skills/mad-slurm-multinode/references/launch-and-results.md create mode 100644 .claude/skills/mad-slurm-multinode/references/manifests.md create mode 100755 .claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh create mode 100755 .claude/skills/mad-slurm-multinode/scripts/preflight.sh create mode 100644 .claude/skills/mad-slurm-multinode/scripts/validate_manifest.sh diff --git a/.claude/skills/mad-slurm-multinode/SKILL.md b/.claude/skills/mad-slurm-multinode/SKILL.md new file mode 100644 index 0000000..cfb1a82 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/SKILL.md @@ -0,0 +1,329 @@ +--- +name: mad-slurm-multinode +description: Deploy and run madengine performance tests on an unprepared SLURM cluster from scratch. Use when the user wants to perf-test a model via a madengine run engine (any engine that has a template under assets/manifests/), set up a fresh node (clone MAD + madengine, conda/miniforge env, pip install), build a mad.env, pick/adapt a run manifest, and launch a multi-node SLURM run. Covers the cluster archetypes documented in references/cluster-types.md (such as CX7/Mellanox-RoCE, AMD-AINIC/Pollara, and Broadcom-Thor2-RoCE), required interactive inputs (node host, work dir, data dir, MAD_DOCKER_BUILDS, SLURM partition/account/qos/reservation/nodelist), and the cluster-specific communication-backend env vars (RCCL/NCCL over the right transport, e.g. RDMA) that decide whether the right transport is exercised. +compatibility: Requires git, docker, conda/miniforge, a SLURM CLI (sbatch/sinfo), rocm-smi or nvidia-smi, and internet access for cloning and image pulls. +metadata: + author: mkuznet1 (mikhail.kuznetsov@amd.com) + version: "0.1" +--- + +# mad-slurm-multinode + +Bring up madengine on a fresh SLURM node and run a Llama-3.1 perf test +end to end: prerequisites -> clone -> conda env -> install -> mad.env -> +manifest -> `madengine run`. Covers the cluster archetypes documented in +[references/cluster-types.md](references/cluster-types.md). + +> **Skill paths — read first.** When this skill activates you are given the +> absolute path to its directory (the folder containing this `SKILL.md`). All +> `scripts/`, `references/`, and `assets/` paths below are relative to THAT +> directory, not to `$WORKDIR`. The workflow `cd`s into `$WORKDIR`, so set the +> skill dir once and reference files through it: +> +> ```bash +> export SKILL_DIR="" +> ``` +> +> Then use `"$SKILL_DIR/scripts/..."`, `"$SKILL_DIR/assets/..."`, etc. Markdown +> links like `references/foo.md` are also relative to `$SKILL_DIR`. + +## When this fires + +- "Run a perf test with madengine on ``" / "set up madengine on this cluster". +- A Llama-3.1 perf workload that has a template under `assets/manifests/`, on SLURM. +- A node that has nothing installed yet (no conda, no repos, no rundir). + +Out of scope: +- RCCL build-vs-build validation delta campaigns (build image A vs B and + compare the transport) — that is a separate workflow, out of scope here. +- Accumulating many iterations across multiple SLURM allocations (large + multi-allocation campaigns) — out of scope here. +- Single-GPU local (non-SLURM) smoke tests -> just `madengine run --tags ...` directly. + +## Responsibilities — who does what + +The skill stays in its lane and follows the steps rather than re-deriving the +workload. The split is: + +- **mad-slurm-multinode (this skill)** configures one run: brings up the node + (repos, conda env), writes `mad.env` and the manifest (paths, SLURM selectors, + transport vars), launches `madengine run`, and reads the perf CSV. +- **madengine** orchestrates: renders the sbatch script from the manifest, + ensures the docker image per node (load tar / build / pull), runs the + multi-node job, and aggregates results. +- **MAD** holds the Dockerfiles and the per-model `run.sh`. Inside the + container, `run.sh` runs the actual workload **and acquires/prepares the + dataset and model weights** (data prep, model download, etc.). The + skill points it at the right paths and launches; staging datasets or + downloading models by hand is `run.sh`'s job, not a pre-step the skill adds. + +**Do not fix madengine or MAD code.** When a failure is rooted in the madengine +orchestrator or the MAD repo (a Dockerfile, `run.sh`, model code, an +orchestration bug), the skill reports it to the user with the evidence and +stops — it does not patch those repos unless the user explicitly asks. The +skill's own edits stay confined to `rundir/` (`mad.env`, the manifest) and the +user-provided inputs. A known per-workload workaround that only sets env/manifest +values (see [references/gotchas.md](references/gotchas.md)) is applied in the +manifest, not by editing madengine/MAD. + +## Required inputs — ask before exploring + +The first action of a run, before any other tool call (Shell, Read, +transcript/past-session search, or web search), is to list the required inputs +that are missing or not clearly identifiable and resolve them in a single +`AskQuestion` round. A missing or unrecognized input is a question to the user, +not a research task — inferring it from prior sessions, `agent-transcripts`, +repo history, or the web is a known failure mode of this skill and is skipped. +Tool calls begin once the required inputs are in hand. + +These come from the user, and the skill asks for them rather than assuming. +The first four have no safe defaults, so the skill asks; a value already in the +conversation gets reused. + +1. **Compute node hostname or IP** to run on (e.g. the login/jump node you SSH + to) — where the bootstrap, build, and `sbatch` submission happen. +2. **Working directory** `$WORKDIR` — the root that holds the cloned repos + (`MAD`, `madengine`), the `rundir/` (the filled `mad.env`, the + filled manifest, `slurm_output/`, and run logs), and the perf-result CSVs. + Picking it explicitly keeps one run's artifacts together and reusable. +3. **Cluster archetype** — one of the archetypes in + [references/cluster-types.md](references/cluster-types.md). You may run + `scripts/detect_cluster_env.sh` to *propose* it, but confirm with the user. +4. **Data root** (`MAD_DATAHOME` + caches) and **`MAD_DOCKER_BUILDS` dir** — + state what each holds so the user can point them at the right place: + - the **data root** holds datasets, tokenizer, and model weights, plus the + `HF_HOME`/`TORCH_HOME`/pip caches the run reads and writes; + - **`MAD_DOCKER_BUILDS`** is the shared image-tar cache: rank 0 saves the + built image there and every worker loads the same tar, so it lives on + shared FS visible to every node (see Gotchas). +5. **SLURM specifics, on demand**: `partition`, `account`, `qos`, + `reservation`, `nodelist`/`exclude`, node count. These are cluster-private + and are not stored in this skill — requested per run. +6. **Branch per repo** for `MAD` and `madengine`. No branch is + assumed — the skill asks which branch to use for each repo and `git switch`es + to it (see "Repo + branch rule"). + +The HF token file `~/.huggingface/token` provides gated Llama-3.1 access; +`mad.env` reads it into `MAD_SECRETS_HFTOKEN` and madengine forwards it to the +container (see [references/manifests.md](references/manifests.md) "Secrets"), so +confirm it exists before launch. + +**Probe one node before asking about cluster shape.** The cluster-shape values +(GPU arch, HCA list, GID index, management iface) are discoverable on the node +itself, so the skill allocates one node from the user-provided partition and +runs the probe there rather than interrogating the user: + +```bash +srun -p [--reservation ] [--nodelist ] -N1 \ + bash "$SKILL_DIR/scripts/detect_cluster_env.sh" +``` + +Only the cluster-private selectors (partition / account / qos / reservation / +nodelist) come from the user; the probe reports the rest. The same holds for +`scripts/preflight.sh`, which reflects whichever node it runs on — for +compute-node values it runs through `srun` on an allocated node. + +Defaults you MAY assume unless told otherwise (state them when you do): +- conda env name `madenv`, Python 3.12. + +**Repo + branch rule (every run):** the repo and the branch for both +`MAD` and `madengine` are confirmed with the user before cloning or +switching. No branch name is assumed — not a default, not one inferred from a +prior session — the branch is asked and then `git switch`ed to. An existing +checkout is left as-is rather than switched silently: when the directory already +exists, the skill shows its current branch +(`git -C branch --show-current`) and asks before changing it. + +## Workflow + +Track progress with this checklist: + +- [ ] Step 0: preflight (`scripts/preflight.sh`) +- [ ] Step 1: clone repos (idempotent) +- [ ] Step 2: conda env (idempotent) +- [ ] Step 3: `pip install -e ./madengine` +- [ ] Step 4: rundir + mad.env (filled + sourced) +- [ ] Step 5: pick + adapt manifest +- [ ] Step 6: launch + collect results + +Full step detail (commands, idempotency guards) is in +[references/deploy-bootstrap.md](references/deploy-bootstrap.md). The short +version: + +Each step below ends with a **Guard** — the enumerated conditions are the only +branches to consider. If a guard matches, take its action; if none match, the +step passed and the next step follows. This keeps the path deterministic and +avoids re-analyzing a step that already succeeded. + +### Run logs + +Every long command (`madengine run`, docker build, conda/pip install) streams +its stdout+stderr to `/.cursor.logs//-.log` +via `tee`, which keeps a run reproducible and inspectable after the fact +(several bugs here surfaced only in these logs). Kinds map to subdirectories +(`.cursor.logs/build/`, `.cursor.logs/run/`, `.cursor.logs/install/`); a +filename carries a UTC timestamp and a short slug and holds no secrets. Example: + +```bash +mkdir -p "$WORKDIR/rundir/.cursor.logs/run" +madengine run --manifest-file run_manifest_.json --live-output \ + -o perf_.csv \ + 2>&1 | tee "$WORKDIR/rundir/.cursor.logs/run/$(date -u +%Y%m%dT%H%M%SZ)-primus-8b.log" +``` + +The log files are `*.log`, which `.gitignore` already ignores, so the +`.cursor.logs/` contents never land in git. + +### Step 0 — Preflight + +```bash +bash "$SKILL_DIR/scripts/preflight.sh" +``` + +Checks Python >= 3.10, docker (present + daemon reachable), conda/miniforge, +SLURM CLI (`sinfo`/`sbatch`), and a GPU SMI (`rocm-smi`/`nvidia-smi`). Fix +any FAIL before continuing. If docker or SLURM is missing this is not a valid +target node — stop and tell the user. + +**Guard:** +- **If Step 0 reports any FAIL** → stop, report which check failed, do not continue. +- **If docker or SLURM is missing** → this node is not a valid target; tell the user and stop. + +### Steps 1-3 — Bootstrap (only do what's missing) + +```bash +cd "$WORKDIR" +# Repo URLs are confirmed with the user; clone only if missing. +[ -d MAD ] || git clone --recursive +[ -d madengine ] || git clone --recursive +# Branch is ASKED per repo (no default). For an existing checkout, show the +# current branch first and ask before switching: +# git -C MAD branch --show-current +# git -C madengine branch --show-current +( cd MAD && git switch "" && git submodule update --init --recursive ) +( cd madengine && git switch "" && git submodule update --init --recursive ) + +# conda: install miniforge only if `conda` is absent (see deploy-bootstrap.md) +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +After cloning or switching, the chosen manifest's `dockerfile` and the model +`scripts`/`run.sh` resolve under the `MAD` checkout (later +`$MODEL_DIR`). A missing path stops the run with a report (likely a wrong +branch, uninitialized submodules, or a different layout) rather than a silent +search — ask the user for the correct branch. + +**Guard:** +- **If a clone, `git switch`, or submodule init fails** → stop, report, ask the user for the correct repo/branch; do not guess another branch. +- **If `pip install -e ./madengine` fails** → stop and report; a madengine bug is not patched here (see Responsibilities). +- **If the manifest's `dockerfile`/`run.sh` do not resolve under `$MODEL_DIR`** → stop and report (likely wrong branch or uninitialized submodules); ask the user. + +### Step 4 — rundir + mad.env + +```bash +cd "$WORKDIR" && mkdir -p rundir && cd rundir +# copy the mad.env template for the archetype (one file per archetype under +# assets/mad.env/; see references/cluster-types.md). The template stays unedited: +cp "$SKILL_DIR/assets/mad.env/mad.env..template" ./mad.env +# FILL every placeholder, then: +source mad.env +``` + +Every `` placeholder is resolved before `source`. The cluster-specific +values (`MAD_SYSTEM_GPU_ARCHITECTURE`, `NCCL_SOCKET_IFNAME`, +`NCCL_IB_GID_INDEX`, and the per-manifest `NCCL_IB_HCA` list) are confirmed +against the actual node — defaults in the template are archetype-typical, not +guaranteed. +Run `bash "$SKILL_DIR/scripts/detect_cluster_env.sh"` to propose them and use +the matrix in [references/cluster-types.md](references/cluster-types.md) to +interpret. + +**Guard:** +- **If a `` value is unknown** → ask the user; do not infer it from a prior session, the repo, or the web. +- **If `source mad.env` errors, or `MODEL_DIR`/`MAD_DOCKER_BUILDS` come back empty** → stop and report before continuing. + +### Step 5 — Manifest + +Copy the matching template from `assets/manifests/` into `rundir/`, then adapt. +That directory holds one `*.template.json` per workload/size; pick the one named +for the requested workload. + +```bash +cp "$SKILL_DIR/assets/manifests/.template.json" run_manifest_.json +# fill the manifest, then statically validate it (GPU-free; reads $MODEL_DIR): +bash "$SKILL_DIR/scripts/validate_manifest.sh" run_manifest_.json +``` + +`validate_manifest.sh` is a read-only static check: JSON validity, leftover +`` placeholders, `NCCL_IB_HCA` set and equal in both env blocks, +network-interface consistency, `slurm.nodes == distributed.nnodes` (+ nodelist +cardinality), no stray `MAD_SECRETS_HFTOKEN`, AINIC transport-var symmetry, and +(with `mad.env` sourced) that the `dockerfile`/`run.sh` resolve under +`$MODEL_DIR`. It exits non-zero on any FAIL. + +What the run sets per run (cluster-private, kept out of the templates): +`deployment_config.slurm.{partition,account,qos,reservation,nodelist,exclude,nodes}`, +node count consistency (`slurm.nodes` == `distributed.nnodes`), the +`NCCL_IB_HCA` device list, and host paths. Field-by-field guidance: +[references/manifests.md](references/manifests.md). + +The manifest's `dockerfile` and the model `scripts`/`run.sh` resolve under +`$MODEL_DIR` (`[ -f "$MODEL_DIR/" ] && [ -f "$MODEL_DIR/" ]`). +A missing path stops the run with a report rather than a silent search. + +**Guard:** +- **If no template matches the requested engine** → ask the user; do not invent a manifest. +- **If `scripts/validate_manifest.sh` reports any FAIL** → fix it before launching (it covers JSON validity, leftover placeholders, `NCCL_IB_HCA`, iface consistency, `nodes == nnodes`, a stray HF token, and dockerfile/run.sh resolution under `$MODEL_DIR`). + +### Step 6 — Launch + results + +```bash +cd "$WORKDIR/rundir" +source mad.env +# -o writes the aggregated perf to a per-run CSV (instead of the default +# perf.csv), so each workload keeps its own result and parallel runs never +# clobber a shared file: +madengine run --manifest-file run_manifest_.json --live-output \ + -o perf_.csv +``` + +Pre-building or pulling the image is unnecessary. With `local_image: true`, +madengine ensures the image itself on each node: if it is not already present +locally it loads it from the `MAD_DOCKER_BUILDS` tar, otherwise builds it from +the manifest's `dockerfile` (and falls back to `docker pull`), then `docker +save`s it to `MAD_DOCKER_BUILDS` so workers load the same tar. Implications: +the first run is slower (it builds once on rank 0), `MAD_DOCKER_BUILDS` lives on +shared FS, and the manifest's `dockerfile` resolves (it lives in the cloned +`MAD`) or the `docker_image` tag is pullable. Details: +[references/launch-and-results.md](references/launch-and-results.md). + +Perf lands in the `-o` file (`perf_.csv`) and the per-model +`multiple_results` CSV. Multi-node aggregation and how to read the result are in +the same file. + +**Guard:** +- **If `madengine run` fails** → read the captured `.cursor.logs/run/...` log and triage with [references/launch-and-results.md](references/launch-and-results.md). +- **If the cause is config/input the skill owns** (mad.env, manifest, paths, transport vars) → fix it in `rundir/` and re-run. +- **If the cause is madengine or MAD code** → stop and report to the user; do not patch those repos unless asked (see Responsibilities). + +## Gotchas + +Cross-cutting and per-workload pitfalls — mad.env sourcing, `MAD_DOCKER_BUILDS` +on shared FS, HF-token handling, `docker_mounts` direction, `NCCL_IB_HCA` +per-cluster, AINIC transport vars, perf-CSV aggregation, and per-workload notes +(sglang_disagg and primus_megatron training) — are in +[references/gotchas.md](references/gotchas.md), read before a run. + +## Examples + +Sanitized end-to-end walkthroughs live in [examples/](examples/); the run +templates they build on live in `assets/` (`assets/mad.env/`, +`assets/manifests/`). New workloads and clusters are added there, not inlined +here. + + + diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template new file mode 100644 index 0000000..5bcb9d7 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template @@ -0,0 +1,68 @@ +# mad.env — AMD-AINIC / Pollara archetype (AMD AINIC rdma* HCAs, gfx950) +# Copy to /rundir/mad.env, resolve every , then `source mad.env`. +# Do NOT commit a filled copy into this skill — it carries cluster-private paths. +# +# AINIC specifics vs CX7: gfx950, eno0, GID 1, ionic drivers, RCCL_AINIC_ROCE. +# Confirm against the node (see references/cluster-types.md + +# scripts/detect_cluster_env.sh). + +# --- secrets (file-based, never inline a literal hf_... token) --- +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token) + +# --- system --- +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 # AINIC/MI355-class default; confirm on node +export MAD_VERBOSE_CONFIG=true + +# --- MAD source (point at the MAD checkout) --- +export MAD_SETUP_MODEL_DIR=false +export MODEL_DIR=/ # e.g. /MAD/ + +# --- shared dir convenience (set once, reuse below) --- +SHARE_DIR= # shared FS visible to every node, e.g. /shared/ + +# --- data / cache roots --- +export MAD_DATAHOME=$SHARE_DIR/data/models +export HF_HOME=$SHARE_DIR/data/cache/huggingface +export TORCH_HOME=$SHARE_DIR/data/cache/torch +export XDG_CACHE_HOME=$SHARE_DIR/data/cache/xdg +export PIP_CACHE_DIR=$SHARE_DIR/data/cache/pip +export TRITON_CACHE_DIR=$SHARE_DIR/data/cache/triton + +# --- image tar cache: MUST be on shared FS visible to every compute node --- +export MAD_DOCKER_BUILDS=$SHARE_DIR/mad_docker_builds + +# --- HF helper aliases --- +export TRANSFORMERS_CACHE=$HF_HOME +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub + +# --- MAD metadata --- +export MAD_DEPLOYMENT_TYPE=slurm +export BUILD_NUMBER=${BUILD_NUMBER:-0} + +# --- AINIC RoCE / NCCL transport (Pollara, vendor 0x1dd8, exposed as rdma0..7) --- +# NO mlx5 devices. NCCL_IB_HCA (the rdma* list) is set PER-MANIFEST. +export NCCL_IB_DISABLE=0 +export NCCL_SOCKET_IFNAME=eno0 # management iface; confirm with `ip link` +export GLOO_SOCKET_IFNAME=eno0 +export NCCL_IB_GID_INDEX=1 # RoCEv2 GID on AINIC; confirm with show_gids + +# --- AINIC transport selection (also required in BOTH manifest env blocks) --- +# Without these RCCL silently falls back to verbs/sockets and certifies the +# wrong codepath. +export RCCL_AINIC_ROCE=1 +export RDMAV_DRIVERS=ionic +export IBV_DRIVERS=ionic + +# --- AINIC driver mounts (manifest-level, documented here for reference) --- +# The ionic verbs driver lives on the host and must be bind-mounted into every +# container. Add to built_models.additional_docker_run_options: +# -v /usr/lib/x86_64-linux-gnu/libionic.so:/usr/lib/x86_64-linux-gnu/libionic.so:ro +# -v /usr/lib/x86_64-linux-gnu/libionic.so.1:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro +# -v /etc/libibverbs.d:/etc/libibverbs.d:ro +# And to context.docker_mounts: +# "/etc/libibverbs.d": "/etc/libibverbs.d" +# "/usr/lib/x86_64-linux-gnu/libibverbs": "/usr/lib/x86_64-linux-gnu/libibverbs" +# See references/cluster-types.md "AINIC driver mounts" for the full rationale. + +# Workload-specific data paths are NOT set here — they live in the workload +# manifest, which keeps this file cluster/cache-only and reusable across workloads. diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template new file mode 100644 index 0000000..738c4fa --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template @@ -0,0 +1,48 @@ +# mad.env — CX7 / Mellanox-RoCE archetype (Mellanox mlx5 HCAs, gfx942) +# Copy to /rundir/mad.env, resolve every , then `source mad.env`. +# Do NOT commit a filled copy into this skill — it carries cluster-private paths. +# +# Verify the cluster-specific values against the actual node (see +# references/cluster-types.md and scripts/detect_cluster_env.sh): +# MAD_SYSTEM_GPU_ARCHITECTURE, NCCL_SOCKET_IFNAME, NCCL_IB_GID_INDEX + +# --- secrets (file-based, never inline a literal hf_... token) --- +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token) + +# --- system --- +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx942 # CX7/MI300-class default; confirm on node +export MAD_VERBOSE_CONFIG=true + +# --- MAD source (point at the MAD checkout) --- +export MAD_SETUP_MODEL_DIR=false +export MODEL_DIR=/ # e.g. /MAD/ + +# --- data / cache roots: keep large artifacts off /home, on shared FS --- +export MAD_DATAHOME=/models +export HF_HOME=/cache/huggingface +export TORCH_HOME=/cache/torch +export XDG_CACHE_HOME=/cache/xdg +export PIP_CACHE_DIR=/cache/pip +export TRITON_CACHE_DIR=/cache/triton + +# --- image tar cache: MUST be on shared FS visible to every compute node --- +export MAD_DOCKER_BUILDS= + +# --- HF helper aliases --- +export TRANSFORMERS_CACHE=$HF_HOME +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub + +# --- MAD metadata --- +export MAD_DEPLOYMENT_TYPE=slurm +export BUILD_NUMBER=${BUILD_NUMBER:-0} + +# --- RoCE / NCCL transport (CX7: Mellanox mlx5 HCAs, RoCEv2 over rdma*) --- +# Bootstrap/control plane is TCP over the management iface; data plane is RDMA +# via NCCL_IB_HCA which is set PER-MANIFEST (device list differs per node). +export NCCL_IB_DISABLE=0 +export NCCL_SOCKET_IFNAME=eth0 # management iface; confirm with `ip link` +export GLOO_SOCKET_IFNAME=eth0 +export NCCL_IB_GID_INDEX=3 # RoCEv2 GID; confirm with show_gids + +# Workload-specific data paths are NOT set here — they live in the workload +# manifest, which keeps this file cluster/cache-only and reusable across workloads. diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template new file mode 100644 index 0000000..7c99305 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template @@ -0,0 +1,50 @@ +# mad.env — Broadcom Thor2 / RoCE archetype (bnxt_re HCAs, gfx950 / MI355X-class) +# Copy to /rundir/mad.env, resolve every , then `source mad.env`. +# Do NOT commit a filled copy into this skill — it carries cluster-private paths. +# +# Verify the cluster-specific values against the actual node (see +# references/cluster-types.md and scripts/detect_cluster_env.sh): +# MAD_SYSTEM_GPU_ARCHITECTURE, NCCL_SOCKET_IFNAME, NCCL_IB_GID_INDEX + +# --- secrets (file-based, never inline a literal hf_... token) --- +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token) + +# --- system --- +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 # MI355X-class default; confirm on node (gfx942 = MI300X) +export MAD_VERBOSE_CONFIG=true + +# --- MAD source (point at the MAD checkout) --- +export MAD_SETUP_MODEL_DIR=false +export MODEL_DIR=/ # e.g. /MAD/ + +# --- data / cache roots: keep large artifacts off /home, on shared FS --- +export MAD_DATAHOME=/models +export HF_HOME=/cache/huggingface +export TORCH_HOME=/cache/torch +export XDG_CACHE_HOME=/cache/xdg +export PIP_CACHE_DIR=/cache/pip +export TRITON_CACHE_DIR=/cache/triton + +# --- image tar cache: MUST be on shared FS visible to every compute node --- +export MAD_DOCKER_BUILDS= + +# --- HF helper aliases --- +export TRANSFORMERS_CACHE=$HF_HOME +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub + +# --- MAD metadata --- +export MAD_DEPLOYMENT_TYPE=slurm +export BUILD_NUMBER=${BUILD_NUMBER:-0} + +# --- RoCE / NCCL transport (Broadcom Thor2 bnxt_re HCAs, RoCEv2) --- +# Bootstrap/control plane is TCP over the management iface (fenic0 here); data +# plane is RDMA via NCCL_IB_HCA which is set PER-MANIFEST (bnxt_re device list). +export NCCL_IB_DISABLE=0 +export NCCL_SOCKET_IFNAME=fenic0 # management iface; confirm with `ip -br link` +export GLOO_SOCKET_IFNAME=fenic0 +export NCCL_IB_GID_INDEX=3 # RoCEv2 GID; confirm with show_gids +export RDMAV_DRIVERS=bnxt_re # Broadcom Thor2 ibverbs provider +export IBV_DRIVERS=bnxt_re + +# Workload-specific data paths are NOT set here — they live in the workload +# manifest, which keeps this file cluster/cache-only and reusable across workloads. diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json new file mode 100644 index 0000000..f108eb1 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json @@ -0,0 +1,109 @@ +{ + "built_images": { + "rocm-primus-llama31-70b": { + "model": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", + "docker_image": "rocm/primus:v26.4->", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", + "base_docker": "rocm/primus:v26.4", + "build_duration": 0, + "local_image": true, + "registry_image": null, + "registry": null, + "gpu_vendor": "AMD" + } + }, + "built_models": { + "rocm-primus-llama31-70b": { + "name": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd", + "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-70B.csv", + "tags": ["pyt", "pretrain", "llama-3.1-70b", "training"], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-70b", + "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 -v /sys:/sys:ro -v /run/udev:/run/udev:ro" + } + }, + "context": { + "docker_env_vars": { + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_IB_HCA": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "IBV_SHOW_WARNINGS": "1" + }, + "docker_mounts": { + "/dev/infiniband": "/dev/infiniband" + }, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "0,1,2,3,4,5,6,7" + }, + "credentials_required": [], + "summary": { + "successful_builds": [], + "failed_builds": [], + "total_build_time": 0, + "successful_pushes": [], + "failed_pushes": [] + }, + "deployment_config": { + "target": "slurm", + "slurm": { + "partition": "", + "account": "", + "qos": "", + "nodes": 2, + "gpus_per_node": 8, + "time": "12:00:00", + "output_dir": "./slurm_output", + "exclusive": true, + "network_interface": "" + }, + "distributed": { + "launcher": "torchrun", + "backend": "nccl", + "port": 29500, + "nnodes": 2, + "nproc_per_node": 8 + }, + "env_vars": { + "HF_HOME": "${HF_HOME}", + "TORCH_HOME": "${TORCH_HOME}", + "XDG_CACHE_HOME": "${XDG_CACHE_HOME}", + "PIP_CACHE_DIR": "${PIP_CACHE_DIR}", + "MAD_DATAHOME": "${MAD_DATAHOME}", + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_IB_HCA": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "NCCL_TIMEOUT": "900", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "TORCH_NCCL_HIGH_PRIORITY": "1", + "OMP_NUM_THREADS": "8", + "MIOPEN_FIND_MODE": "1", + "MIOPEN_USER_DB_PATH": "${XDG_CACHE_HOME}/miopen" + }, + "debug": false, + "docker_gpus": "0,1,2,3,4,5,6,7" + } +} diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json new file mode 100644 index 0000000..bc60a2d --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json @@ -0,0 +1,109 @@ +{ + "built_images": { + "rocm-primus-llama31-8b": { + "model": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", + "docker_image": "rocm/primus:v26.4->", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", + "base_docker": "rocm/primus:v26.4", + "build_duration": 0, + "local_image": true, + "registry_image": null, + "registry": null, + "gpu_vendor": "AMD" + } + }, + "built_models": { + "rocm-primus-llama31-8b": { + "name": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd", + "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-8B.csv", + "tags": ["pyt", "pretrain", "llama-3.1-8b", "training"], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-8b", + "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 -v /sys:/sys:ro -v /run/udev:/run/udev:ro" + } + }, + "context": { + "docker_env_vars": { + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_IB_HCA": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "IBV_SHOW_WARNINGS": "1" + }, + "docker_mounts": { + "/dev/infiniband": "/dev/infiniband" + }, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "0,1,2,3,4,5,6,7" + }, + "credentials_required": [], + "summary": { + "successful_builds": [], + "failed_builds": [], + "total_build_time": 0, + "successful_pushes": [], + "failed_pushes": [] + }, + "deployment_config": { + "target": "slurm", + "slurm": { + "partition": "", + "account": "", + "qos": "", + "nodes": 2, + "gpus_per_node": 8, + "time": "12:00:00", + "output_dir": "./slurm_output", + "exclusive": true, + "network_interface": "" + }, + "distributed": { + "launcher": "torchrun", + "backend": "nccl", + "port": 29500, + "nnodes": 2, + "nproc_per_node": 8 + }, + "env_vars": { + "HF_HOME": "${HF_HOME}", + "TORCH_HOME": "${TORCH_HOME}", + "XDG_CACHE_HOME": "${XDG_CACHE_HOME}", + "PIP_CACHE_DIR": "${PIP_CACHE_DIR}", + "MAD_DATAHOME": "${MAD_DATAHOME}", + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_IB_HCA": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "NCCL_TIMEOUT": "900", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "TORCH_NCCL_HIGH_PRIORITY": "1", + "OMP_NUM_THREADS": "8", + "MIOPEN_FIND_MODE": "1", + "MIOPEN_USER_DB_PATH": "${XDG_CACHE_HOME}/miopen" + }, + "debug": false, + "docker_gpus": "0,1,2,3,4,5,6,7" + } +} diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json new file mode 100644 index 0000000..c8b9167 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json @@ -0,0 +1,152 @@ +{ + "built_images": { + "sglang-disagg-deepseek-r1": { + "model": "sglang_disagg_deepseek-r1", + "docker_image": "smifix-mori--rixl--mooncake-gfx942 (append -rdma62oci when built from the OCI variant Dockerfile)>", + "dockerfile": "docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile", + "base_docker": "lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x", + "build_duration": 0, + "local_image": true, + "registry_image": null, + "registry": null, + "gpu_vendor": "AMD" + } + }, + "built_models": { + "sglang-disagg-deepseek-r1": { + "name": "sglang_disagg_deepseek-r1", + "url": "", + "dockerfile": "docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile", + "scripts": "scripts/sglang_disagg/run.sh", + "data": "", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_sglang-disagg-DeepSeek-R1.csv", + "tags": ["pyt", "sglang", "disagg", "inference", "deepseek-r1"], + "timeout": 21600, + "args": "--model-name DeepSeek-R1 --model-path deepseek-ai/DeepSeek-R1-0528) and peers wait on MODEL_PATH/.stage_done> --run-mori 1 --dp-mode 1 --xp 2 --yd 2 --gpus-per-node 8 --kv-transfer-backend nixl", + "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 --ulimit nofile=1048576:1048576 -v /sys:/sys:ro -v /sys/class/infiniband:/sys/class/infiniband:ro -v /run/udev:/run/udev:ro -v : -v ./slurm_output/run_logs:/run_logs -e SLURM_JOB_ID" + } + }, + "context": { + "skip_perf_collection": false, + "docker_env_vars": { + "MODEL_NAME": "DeepSeek-R1", + "MODEL_PATH": "", + "RUN_MORI": "1", + "DP_MODE": "1", + "xP": "2", + "yD": "2", + "GPUS_PER_NODE": "8", + "KV_TRANSFER_BACKEND": "nixl", + "MORI_JIT_CACHE_DIR": "/tmp/mori_jit", + "MORI_SHMEM_HEAP_SIZE": "17179869184", + "RDMA_CORE_CACHE": "", + "NCCL_DEBUG": "WARN", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "", + "IB_DEVICES": "", + "MORI_RDMA_DEVICES": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "MORI_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "MORI_IB_GID_INDEX": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "TORCH_NCCL_HIGH_PRIORITY": "1", + "GPU_MAX_HW_QUEUES": "2", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "NCCL_TIMEOUT": "900", + "HSA_ENABLE_SDMA": "0", + "HSA_FORCE_FINE_GRAIN_PCIE": "1", + "OMP_NUM_THREADS": "8", + "MIOPEN_FIND_MODE": "1", + "SMOKE": "0", + "BENCHMARK_ITR": "1", + "BENCHMARK_COMBINATIONS": "1024/1024 8192/1024", + "BENCHMARK_CONCURRENCY_LEVELS": "8 16 32 64", + "BENCHMARK_PRECHECK": "0", + "BENCHMARK_FAIL_FAST": "1" + }, + "docker_mounts": { + "/dev/infiniband": "/dev/infiniband", + "/sys/class/infiniband": "/sys/class/infiniband", + "": "", + "/run_logs": "./slurm_output/run_logs" + }, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "0,1,2,3,4,5,6,7" + }, + "credentials_required": [], + "summary": { + "successful_builds": [], + "failed_builds": [], + "total_build_time": 0, + "successful_pushes": [], + "failed_pushes": [] + }, + "deployment_config": { + "target": "slurm", + "slurm": { + "partition": "", + "account": "", + "qos": "", + "nodes": 4, + "gpus_per_node": 8, + "time": "02:30:00", + "output_dir": "./slurm_output", + "exclusive": true, + "network_interface": "" + }, + "distributed": { + "launcher": "sglang-disagg", + "backend": "nccl", + "port": 29500, + "nnodes": 4, + "nproc_per_node": 8, + "sglang_disagg": { + "prefill_nodes": 2, + "decode_nodes": 2 + } + }, + "env_vars": { + "MODEL_NAME": "DeepSeek-R1", + "MODEL_PATH": "", + "RUN_MORI": "1", + "DP_MODE": "1", + "xP": "2", + "yD": "2", + "GPUS_PER_NODE": "8", + "KV_TRANSFER_BACKEND": "nixl", + "MORI_JIT_CACHE_DIR": "/tmp/mori_jit", + "MORI_SHMEM_HEAP_SIZE": "17179869184", + "RDMA_CORE_CACHE": "", + "NCCL_DEBUG": "WARN", + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "", + "IB_DEVICES": "", + "MORI_RDMA_DEVICES": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "MORI_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "MORI_IB_GID_INDEX": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "NCCL_TIMEOUT": "900", + "OMP_NUM_THREADS": "8", + "BENCHMARK_COMBINATIONS": "1024/1024 8192/1024", + "BENCHMARK_CONCURRENCY_LEVELS": "8 16 32 64", + "BENCHMARK_FAIL_FAST": "1" + }, + "monitor": true, + "debug": false, + "docker_gpus": "0,1,2,3,4,5,6,7", + "gpus_per_node": 8 + } +} diff --git a/.claude/skills/mad-slurm-multinode/examples/amd-ainic-walkthrough.md b/.claude/skills/mad-slurm-multinode/examples/amd-ainic-walkthrough.md new file mode 100644 index 0000000..880b1ba --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/examples/amd-ainic-walkthrough.md @@ -0,0 +1,93 @@ +# Example walkthrough — Primus Llama-3.1-8B on an AMD-AINIC / Pollara cluster + +> Sanitized end-to-end walkthrough: no real node names, reservations, or tokens. +> Replace every `<...>` with your values. + +Scenario: a 4-node Primus 8B run on a reserved AINIC partition. + +## 0. Inputs gathered from the requester + +| Input | Example value | +|--------------------|----------------------------------------| +| Compute/jump node | `` | +| `$WORKDIR` | `/run` | +| Archetype | AMD-AINIC / Pollara | +| Shared root | `` (e.g. /shared/) | +| `MAD_DOCKER_BUILDS`| `/mad_docker_builds` | +| SLURM | partition ``, reservation ``, nodelist `,,,` | +| Nodes | 4 | + +## 1. Preflight + bootstrap + +```bash +export SKILL_DIR="" +bash "$SKILL_DIR/scripts/preflight.sh" # expect gfx950, docker ok, sbatch ok +bash "$SKILL_DIR/scripts/detect_cluster_env.sh" # expect rdma0..7, ionic, GID ~1, eno0 +cd /run +[ -d MAD ] || git clone --recursive +( cd MAD && git switch "" && git submodule update --init --recursive ) +[ -d madengine ] || git clone --recursive +( cd madengine && git switch "" && git submodule update --init --recursive ) +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +## 2. rundir + mad.env (from mad.env.amd-ainic.template) + +Filled values (illustrative): +```bash +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 +SHARE_DIR= +export MODEL_DIR=$SHARE_DIR/MAD/ +export MAD_DOCKER_BUILDS=$SHARE_DIR/mad_docker_builds +export NCCL_SOCKET_IFNAME=eno0 +export NCCL_IB_GID_INDEX=1 +export RCCL_AINIC_ROCE=1 +export RDMAV_DRIVERS=ionic +export IBV_DRIVERS=ionic +``` +```bash +cd /run && mkdir -p rundir && cd rundir +cp "$SKILL_DIR/assets/mad.env/mad.env.amd-ainic.template" ./mad.env +# edit placeholders ... +source mad.env +``` + +## 3. Manifest (from primus_llama-3.1-8b.template.json) + +Edits applied to `run_manifest_primus_8b_ainic.json`: +- `slurm.partition=`, `reservation=`, + `nodelist=,,,`, `nodes=4` (remove `account` if unused) +- `distributed.nnodes=4` +- `NCCL_IB_HCA=rdma0:1,rdma1:1,rdma2:1,rdma3:1,rdma4:1,rdma5:1,rdma6:1,rdma7:1` + in BOTH env blocks +- `NCCL_SOCKET_IFNAME=eno0`, `GLOO_SOCKET_IFNAME=eno0`, `NCCL_IB_GID_INDEX=1` +- `RDMAV_DRIVERS=ionic`, `IBV_DRIVERS=ionic`, `RCCL_AINIC_ROCE=1` in BOTH blocks +- `slurm.network_interface=eno0` +- `docker_image=rocm/primus:v26.2-rccl--` + +For a custom RCCL drop you can also set `context.docker_build_arg` with +`BUILD_GPU_TARGETS=gfx950`, `RCCL_REPO`, `RCCL_COMMIT` — this rebuilds RCCL via +the `rccl_overlay` Dockerfile shipped in MAD. +```bash +python3 -m json.tool run_manifest_primus_8b_ainic.json >/dev/null && echo "valid" +``` + +## 4. Launch + read result + +```bash +source mad.env +madengine run --manifest-file run_manifest_primus_8b_ainic.json --live-output +squeue -u $USER +``` +Confirm the AINIC/ionic transport was selected in the rank-0 RCCL debug log +before trusting the number (otherwise it fell back to verbs/sockets — see +references/cluster-types.md). Aggregated perf in `rundir/perf.csv`. + +## Notes + +- 4 nodes minimum is recommended for AINIC to exercise multi-rail collectives. +- Hold the same node set across compared runs (`nodelist`) for apples-to-apples. +- Deep AINIC transport validation / RCCL build-vs-build deltas are a separate + workflow, out of scope for this perf-run skill. diff --git a/.claude/skills/mad-slurm-multinode/examples/cx7-roce-walkthrough.md b/.claude/skills/mad-slurm-multinode/examples/cx7-roce-walkthrough.md new file mode 100644 index 0000000..03effd2 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/examples/cx7-roce-walkthrough.md @@ -0,0 +1,84 @@ +# Example walkthrough — Primus Llama-3.1-8B on a CX7 / Mellanox-RoCE cluster + +> Sanitized end-to-end walkthrough: no real node names, SLURM queues, accounts, +> or tokens. Replace every `<...>` with your values. + +Scenario: bring up a fresh CX7 login node and run a 2-node Primus 8B perf test. + +## 0. Inputs gathered from the requester + +| Input | Example value | +|--------------------|----------------------------------------| +| Compute/login node | `` | +| `$WORKDIR` | `~/source/run1` | +| Archetype | CX7 / Mellanox-RoCE | +| Data root | `` (shared FS) | +| `MAD_DOCKER_BUILDS`| `/mad_docker_builds` | +| SLURM | partition ``, account ``, qos `` | +| Nodes | 2 | + +## 1. Preflight + bootstrap + +```bash +export SKILL_DIR="" +bash "$SKILL_DIR/scripts/preflight.sh" # expect gfx942, docker ok, sbatch ok +cd ~/source/run1 +[ -d MAD ] || git clone --recursive +( cd MAD && git switch "" && git submodule update --init --recursive ) +[ -d madengine ] || git clone --recursive +( cd madengine && git switch "" && git submodule update --init --recursive ) +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +## 2. rundir + mad.env (from mad.env.cx7-roce.template) + +Filled values (illustrative): +```bash +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx942 +export MODEL_DIR=~/source/run1/MAD/ +export MAD_DATAHOME=/models +export MAD_DOCKER_BUILDS=/mad_docker_builds +export NCCL_SOCKET_IFNAME=eth0 +export NCCL_IB_GID_INDEX=3 +``` +```bash +cd ~/source/run1 && mkdir -p rundir && cd rundir +cp "$SKILL_DIR/assets/mad.env/mad.env.cx7-roce.template" ./mad.env +# edit placeholders ... +source mad.env +[ -d "$MODEL_DIR/scripts" ] && echo "MODEL_DIR ok" +``` + +## 3. Manifest (from primus_llama-3.1-8b.template.json) + +Edits applied to `run_manifest_primus_8b.json`: +- `slurm.partition=`, `account=`, `qos=`, `nodes=2` +- `distributed.nnodes=2` +- `NCCL_IB_HCA=mlx5_0:1,mlx5_1:1` in BOTH `context.docker_env_vars` and + `deployment_config.env_vars` +- `NCCL_SOCKET_IFNAME=eth0`, `GLOO_SOCKET_IFNAME=eth0`, `NCCL_IB_GID_INDEX=3` +- `RDMAV_DRIVERS=mlx5`, `IBV_DRIVERS=mlx5`, deleted the `RCCL_AINIC_ROCE` key +- `docker_image=rocm/primus:v26.2-rccl--` +```bash +python3 -m json.tool run_manifest_primus_8b.json >/dev/null && echo "manifest valid" +``` + +## 4. Launch + read result + +```bash +source mad.env +madengine run --manifest-file run_manifest_primus_8b.json --live-output +squeue -u $USER # watch the 2-node job +``` +Result lands in `rundir/perf.csv` (aggregated) and +`perf_primus-megatron-Llama-3.1-8B.csv`. Report tok/s/gpu and TFLOPS/gpu from +the aggregated CSV (node_0's local CSV can read empty — see +references/launch-and-results.md). + +## Notes + +- 70B: same flow with `primus_llama-3.1-70b.template.json`; bump walltime/nodes. +- Data staging (download + preprocessing) is handled by the MAD `run.sh`; the + skill just fills the workload's data mounts and points the env vars at them. diff --git a/.claude/skills/mad-slurm-multinode/examples/thor2-bnxt-walkthrough.md b/.claude/skills/mad-slurm-multinode/examples/thor2-bnxt-walkthrough.md new file mode 100644 index 0000000..b368243 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/examples/thor2-bnxt-walkthrough.md @@ -0,0 +1,103 @@ +# Example walkthrough — Primus Llama-3.1-70B on a Broadcom-Thor2-RoCE cluster + +> Sanitized end-to-end walkthrough: no real node names, SLURM queues, accounts, +> or tokens. Replace every `<...>` with your values. + +Scenario: bring up a fresh Broadcom-Thor2-RoCE login node and run a 16-node Primus 70B perf +test over RoCEv2 on Broadcom Thor2 (`bnxt_re`) NICs. + +## 0. Inputs gathered from the requester + +| Input | Example value | +|--------------------|----------------------------------------| +| Compute/login node | `` | +| `$WORKDIR` | `` (shared NFS) | +| Archetype | Broadcom-Thor2-RoCE | +| Data root | `` (shared NFS) | +| `MAD_DOCKER_BUILDS`| `/mad_docker_builds` (shared NFS)| +| SLURM | partition `` (no account/qos) | +| Nodes | 16 (8 GPU/node) | + +## 1. Preflight + bootstrap + +```bash +export SKILL_DIR="" +bash "$SKILL_DIR/scripts/preflight.sh" # expect gfx950, docker ok, sbatch ok +cd +[ -d MAD ] || git clone --recursive +( cd MAD && git switch "" && git submodule update --init --recursive ) +[ -d madengine ] || git clone https://github.com/ROCm/madengine --recursive +( cd madengine && git switch develop && git submodule update --init --recursive ) # PR #142 merged +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +Confirm the data-plane HCAs and GID on an allocated node before filling the env: +```bash +srun -p -N1 bash "$SKILL_DIR/scripts/detect_cluster_env.sh" +# expect: gfx950, bnxt_re0..7, mgmt iface fenic0, RoCEv2 GID 3 +``` + +## 2. rundir + mad.env (from mad.env.thor2-bnxt.template) + +Filled values (illustrative): +```bash +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 +export MODEL_DIR=/MAD/ +export MAD_DATAHOME=/models +export MAD_DOCKER_BUILDS=/mad_docker_builds +export NCCL_SOCKET_IFNAME=fenic0 +export NCCL_IB_GID_INDEX=3 +export RDMAV_DRIVERS=bnxt_re +export IBV_DRIVERS=bnxt_re +``` +```bash +cd && mkdir -p rundir && cd rundir +cp "$SKILL_DIR/assets/mad.env/mad.env.thor2-bnxt.template" ./mad.env +# edit placeholders ... +source mad.env +[ -d "$MODEL_DIR/scripts" ] && echo "MODEL_DIR ok" +``` + +## 3. Manifest (from primus_llama-3.1-70b.template.json) + +Edits applied to `run_manifest_primus_70b.json`: +- `slurm.partition=`, removed `account`/`qos` keys (cluster has none), + `nodes=16`, and **added `"skip_gpus_directive": true`** (this cluster's SLURM rejects + the `--gpus-per-node` directive — see references/cluster-types.md) +- `distributed.nnodes=16` +- `NCCL_IB_HCA=bnxt_re0:1,bnxt_re1:1,bnxt_re2:1,bnxt_re3:1,bnxt_re4:1,bnxt_re5:1,bnxt_re6:1,bnxt_re7:1` + in BOTH `context.docker_env_vars` and `deployment_config.env_vars` +- `NCCL_SOCKET_IFNAME=fenic0`, `GLOO_SOCKET_IFNAME=fenic0`, `NCCL_IB_GID_INDEX=3` +- `RDMAV_DRIVERS=bnxt_re`, `IBV_DRIVERS=bnxt_re`, deleted the `RCCL_AINIC_ROCE` key +- added the bnxt_re tuning vars to both env blocks (QPS/inline/merge/adaptive/ + gdr/dmabuf — see references/cluster-types.md "Broadcom-Thor2 specifics") +- added `"/etc/libibverbs.d": "/etc/libibverbs.d"` to `context.docker_mounts` +- `docker_image=rocm/primus:v26.3-rccl--`, `base_docker=rocm/primus:v26.3` +```bash +bash "$SKILL_DIR/scripts/validate_manifest.sh" run_manifest_primus_70b.json +``` + +## 4. Launch + read result + +```bash +source mad.env +madengine run --manifest-file run_manifest_primus_70b.json --live-output \ + -o perf_primus_70b.csv +squeue -u $USER # watch the 16-node job (PR #142 auto-selects healthy nodes) +``` +Result lands in `rundir/perf_primus_70b.csv` (aggregated) and +`perf_primus-megatron-Llama-3.1-70B.csv`. Report tok/s/gpu and TFLOPS/gpu from +the aggregated CSV (a worker node's local CSV can read header-only — see +references/launch-and-results.md). + +## Notes + +- PR #142 (merged in madengine `develop`) runs a GPU health check and + auto-selects healthy nodes; leave `nodelist` unset to use it, or pin nodes by + filling `slurm.nodelist` (cardinality must equal `nodes`). +- With `local_image: true` + `MAD_DOCKER_BUILDS` on shared NFS, rank 0 saves the + image tar once and every worker loads it — no manual `docker load` per node. +- The Primus 70B recipe may sweep BF16+FP8; a BF16 cold-start flake can mark the + run FAILURE and skip perf collection. Re-run for a clean perf CSV. diff --git a/.claude/skills/mad-slurm-multinode/references/cluster-types.md b/.claude/skills/mad-slurm-multinode/references/cluster-types.md new file mode 100644 index 0000000..a819d40 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/cluster-types.md @@ -0,0 +1,125 @@ +# Cluster archetypes and the env vars that differ + +Keywords: CX7 Mellanox mlx5 RoCE, AMD AINIC Pollara ionic rdma, Broadcom +Thor2 bnxt_re RoCE, gfx942 gfx950, NCCL_IB_HCA, NCCL_SOCKET_IFNAME, +NCCL_IB_GID_INDEX, RCCL_AINIC_ROCE, RDMAV_DRIVERS, IBV_DRIVERS, eth0 eno0 fenic0, +show_gids, ibv_devices, skip_gpus_directive + +Three archetypes seen in practice. The differences below decide whether the run +exercises the intended RDMA transport. Values are archetype-*typical* and are +confirmed on the actual node (`scripts/detect_cluster_env.sh`). + +## Difference matrix + +| Variable | CX7 / Mellanox-RoCE | AMD-AINIC / Pollara | Broadcom-Thor2-RoCE | +|------------------------------|--------------------------------------|----------------------------------------|----------------------------------------| +| `MAD_SYSTEM_GPU_ARCHITECTURE`| `gfx942` | `gfx950` | `gfx950` (MI355X; `gfx942` = MI300X) | +| HCA hardware | Mellanox CX7 `mlx5_*` | AMD AINIC `rdma0..7` (vendor 0x1dd8) | Broadcom Thor2 `bnxt_re0..7` | +| `NCCL_IB_HCA` | `mlx5_0:1,mlx5_1:1` (or 8-rail mlx5 list) | `rdma0:1,rdma1:1,...,rdma7:1` | `bnxt_re0:1,bnxt_re1:1,...,bnxt_re7:1` | +| `NCCL_SOCKET_IFNAME` | `eth0` | `eno0` | `fenic0` | +| `GLOO_SOCKET_IFNAME` | `eth0` | `eno0` | `fenic0` | +| `NCCL_IB_GID_INDEX` | `3` | `1` | `3` | +| `RDMAV_DRIVERS` / `IBV_DRIVERS` | `mlx5` | `ionic` | `bnxt_re` | +| `RCCL_AINIC_ROCE` | (unset) | `1` (required) | (unset) | +| Shared FS paths | shared inference FS + off-/home scratch | shared team FS (mounted on all nodes) | shared NFS work/data root on all nodes | +| SLURM selection | `partition`+`account`+`qos`, sometimes `exclude` | usually `partition`+`reservation`+`nodelist` | `partition`, `exclude`/`nodelist`; `skip_gpus_directive` (see below) | +| AINIC driver mounts | (not needed) | add to every manifest's `built_models.additional_docker_run_options` and `context.docker_mounts` — see **AINIC driver mounts** below | (not needed) | + +## How to confirm each value on the node + +These values live on the node, so the probe runs there before the user is +asked. The login/jump node and the compute nodes can differ, so compute-node +values come from running `scripts/detect_cluster_env.sh` through `srun` on an +allocated node (`srun -p [--reservation ] [--nodelist ] +-N1 bash scripts/detect_cluster_env.sh`); only the cluster-private selectors +(partition / account / qos / reservation / nodelist) come from the user. + + +- GPU arch: `rocm-smi --showhw` or `rocminfo | grep gfx` -> the `gfx9xx` target. +- IB/RDMA HCAs: `ibv_devices` (lists `mlx5_*` or `rdma*`). The `NCCL_IB_HCA` + list should enumerate the GPU-attached HCAs (`:1` = port 1). +- Management iface: `ip -br link` / `ip -br addr` -> the routable host iface + used for bootstrap (`eth0` vs `eno0`). This is not the data-plane device. +- RoCEv2 GID index: `show_gids` (or read + `/sys/class/infiniband//ports/1/gid_attrs/types/*`); pick the RoCEv2 + (v2) GID index. CX7 commonly 3, AINIC commonly 1. +- Driver: if `ibv_devices` shows `mlx5_*` -> `mlx5`; if `rdma*` (ionic) -> + `ionic` and you are on AINIC (set `RCCL_AINIC_ROCE=1`). + +## AINIC driver mounts + +The ionic verbs driver lives on the host (`/usr/lib/x86_64-linux-gnu/libionic.so*`, +`/etc/libibverbs.d/`) and must be bind-mounted into every container so that +ibverbs inside the container can find and load the ionic provider. + +**`built_models.additional_docker_run_options`** — add these `-v` flags: +``` +-v /usr/lib/x86_64-linux-gnu/libionic.so:/usr/lib/x86_64-linux-gnu/libionic.so:ro +-v /usr/lib/x86_64-linux-gnu/libionic.so.1:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro +-v /etc/libibverbs.d:/etc/libibverbs.d:ro +``` + +**`context.docker_mounts`** — add: +```json +"/etc/libibverbs.d": "/etc/libibverbs.d", +"/usr/lib/x86_64-linux-gnu/libibverbs": "/usr/lib/x86_64-linux-gnu/libibverbs" +``` + +Without these mounts `libibverbs` inside the container finds no provider, reports +0 IB devices, and RCCL silently falls back to TCP sockets. + +## Broadcom-Thor2 specifics + +Broadcom-Thor2 nodes carry Broadcom Thor2 NICs exposed as `bnxt_re0..7` ibverbs devices +over RoCEv2. Two things differ from the CX7/AINIC archetypes: + +**1. `skip_gpus_directive` in the manifest's `deployment_config.slurm`.** This cluster's +SLURM rejects the generated `--gpus-per-node` directive (the partition does not +advertise GPU GRES the way the default sbatch template assumes), so the job is +refused before launch. Set `"skip_gpus_directive": true` so madengine emits the +sbatch script without that directive and relies on `exclusive`/`nproc_per_node` +instead. Symptom when missing: sbatch is rejected with an invalid `--gpus`/GRES +error and no job id is returned. + +**2. bnxt_re transport tuning (set in BOTH manifest env blocks).** The Broadcom +provider needs a conservative QP/feature profile to run RoCEv2 reliably; the +defaults can hang or fall back. Validated set: + +```json +"NCCL_IB_QPS_PER_CONNECTION": "1", +"NCCL_IB_USE_INLINE": "0", +"NCCL_IB_MERGE_NICS": "0", +"NCCL_IB_SPLIT_DATA_ON_QPS": "0", +"NCCL_IB_ADAPTIVE_ROUTING": "0", +"NCCL_GDR_FLUSH_DISABLE": "1", +"NCCL_DMABUF_ENABLE": "0", +"NCCL_IB_ROCE_VERSION_NUM": "2" +``` + +Also point ibverbs at the host provider inside the container (the image carries +the bnxt_re provider, but the host `/etc/libibverbs.d` is mounted to be safe): +`LIBIBVERBS_DRIVER_PATH=/usr/lib/x86_64-linux-gnu/libibverbs` and add +`"/etc/libibverbs.d": "/etc/libibverbs.d"` to `context.docker_mounts`. Confirm +the device list with `ibv_devices` (expect `bnxt_re*`) and the RoCEv2 GID with +`show_gids` (commonly index 3). + +## MIOpen first-run compile time on gfx950 (AINIC) + +On gfx950 (MI355X) MIOpen compiles GPU kernels on first use, which can take +**20–40 minutes**. This affects any workload that starts MIOpen (inference, training). + +- Set `BARRIER_TIMEOUT_S` high enough (≥ `7200`) for the first run. +- Mount a **persistent** MIOpen cache dir so subsequent runs skip compilation: + - env var: `MIOPEN_USER_DB_PATH` → in-container path (e.g. `/miopen_cache`) + - mount: container `/miopen_cache` → host shared-FS dir +- Once the cache is warm, `BARRIER_TIMEOUT_S` can be reduced to a few minutes. + +## Why this matters + +If `NCCL_IB_HCA` names devices that do not exist on the node, RCCL/NCCL +initializes zero NICs and either aborts (`Failed to initialize any NET +plugin`) or silently falls back to TCP sockets — making the perf number a +measurement of the wrong path. On AINIC, omitting `RCCL_AINIC_ROCE=1` / +`RDMAV_DRIVERS=ionic` has the same effect (fallback to verbs/sockets). These +appear in BOTH `context.docker_env_vars` and `deployment_config.env_vars` of the +manifest (and in `mad.env` host env as a belt-and-suspenders default). diff --git a/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md b/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md new file mode 100644 index 0000000..5cf2e9f --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md @@ -0,0 +1,143 @@ +# Deploy bootstrap (Steps 0-4) + +Keywords: clone MAD madengine, miniforge conda madenv, pip install -e, +rundir, source mad.env, idempotent setup, git switch, python 3.12, docker check + +Detailed, idempotent bring-up of a fresh SLURM node. Run each step only if its +result is missing — re-running a completed step stays safe. Substitute +`$WORKDIR` with the user-provided working directory. `$SKILL_DIR` is the +absolute path to this skill's directory (the folder with `SKILL.md`); export it +once because the steps below `cd` into `$WORKDIR`, so bare `scripts/...` paths +would not resolve. + +## Step 0 — Preflight + +```bash +bash "$SKILL_DIR/scripts/preflight.sh" +``` + +Hard requirements (FAIL -> stop, this is not a valid target node): +- `docker` present and `docker info` succeeds (daemon reachable, user in group). +- A SLURM client: `sinfo` and `sbatch` on PATH. +- `git` (needed to clone/switch the repos in Step 1). `preflight.sh` treats a + missing `git` as a hard FAIL. + +Soft requirements (warn, can be installed in later steps): +- Python >= 3.10 (madengine needs 3.10+; we create a 3.12 conda env anyway). +- conda/miniforge (installed in Step 2 if absent). +- A GPU SMI (`rocm-smi` for AMD, `nvidia-smi` for NVIDIA) for arch detection. +- HF token at `~/.huggingface/token` (required at run time for gated Llama-3.1). + +## Step 1 — Clone repos (idempotent) + +```bash +cd "$WORKDIR" + +# MAD source -> MODEL_DIR. Repo URL is confirmed with the user; clone if missing. +if [ ! -d MAD ]; then + git clone --recursive +fi +# Branch is ASKED (no default). Switch, then sync submodules. +( cd MAD && git switch "" && git submodule update --init --recursive ) + +# madengine. Repo URL is confirmed with the user; clone if missing. +if [ ! -d madengine ]; then + git clone --recursive +fi +# Branch is ASKED (no default). +( cd madengine && git switch "" && git submodule update --init --recursive ) +``` + +Confirm the repo and branch with the user before cloning or switching. No branch +name is assumed (not a default, not one inferred from a prior session); the +branch is asked per repo and then `git switch`ed to. When the dirs already +exist, the skill shows the current branch +(`git -C branch --show-current`) and asks before changing it rather than +switching silently. An existing checkout is left intact, since the user may have +local edits. + +After cloning or switching, the assets the chosen manifest references resolve +under the `MAD` checkout (this becomes `$MODEL_DIR`). Confirm the +manifest's `dockerfile` and the model `scripts`/`run.sh` exist: + +```bash +# example for primus_llama-3.1-8b — adjust to the chosen manifest's fields +( cd MAD \ + && [ -f docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile ] \ + && [ -f scripts/primus_scaleout/megatron-lm/run.sh ] \ + || echo "STOP: manifest dockerfile/run.sh not found under MAD" ) +``` + +A missing path stops the run with a report rather than a silent search. The +usual causes are a wrong branch, uninitialized submodules, or a different repo +layout, so the next step is to ask the user for the correct branch. + +## Step 2 — conda env (idempotent) + +```bash +# an existing conda/miniforge may just be missing from PATH — source it first +if ! command -v conda >/dev/null 2>&1; then + for cp in "$HOME/miniforge3" "$HOME/miniconda3" "$HOME/mambaforge" \ + "/opt/conda" "${CONDA_PREFIX:-}" "${MAMBA_ROOT_PREFIX:-}"; do + if [ -n "$cp" ] && [ -f "$cp/etc/profile.d/conda.sh" ]; then + source "$cp/etc/profile.d/conda.sh"; break + fi + done +fi + +# install miniforge only if conda is still missing +if ! command -v conda >/dev/null 2>&1; then + wget -q https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh + bash Miniforge3-Linux-x86_64.sh -b -p "$HOME/miniforge3" + source "$HOME/miniforge3/etc/profile.d/conda.sh" +fi + +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +``` + +Notes: +- An already-installed conda/miniforge that is simply off PATH gets sourced from + its common install prefix (`$HOME/miniforge3`, `$HOME/miniconda3`, + `$HOME/mambaforge`, `/opt/conda`, `$CONDA_PREFIX`, `$MAMBA_ROOT_PREFIX`) + rather than reinstalled, which avoids duplicate installs. `preflight.sh` + reports the same find. +- `-b` runs the miniforge installer unattended; `-p` sets the prefix. +- If conda was just installed in this shell, `source .../conda.sh` (or open a + new shell) before `conda activate`. + +## Step 3 — Install madengine + +```bash +cd "$WORKDIR" +pip install -e ./madengine +madengine --help >/dev/null && echo "madengine OK" +``` + +Editable install so the user's branch changes are picked up without reinstall. + +## Step 4 — rundir + mad.env + +```bash +cd "$WORKDIR" +mkdir -p rundir && cd rundir + +# copy the archetype template (see cluster-types.md to choose): +# CX7/Mellanox-RoCE -> mad.env.cx7-roce.template +# AMD-AINIC/Pollara -> mad.env.amd-ainic.template +cp "$SKILL_DIR/assets/mad.env/mad.env..template" ./mad.env +# resolve every , confirm node-specific values, then: +source mad.env +``` + +After `source mad.env`, sanity-check the critical exports in the SAME shell: + +```bash +echo "MODEL_DIR=$MODEL_DIR" +echo "MAD_DOCKER_BUILDS=$MAD_DOCKER_BUILDS" +echo "MAD_SYSTEM_GPU_ARCHITECTURE=$MAD_SYSTEM_GPU_ARCHITECTURE" +[ -d "$MODEL_DIR/scripts" ] || echo "WARNING: MODEL_DIR/scripts missing — run.sh will not be found" +``` + +The HF token file (`~/.huggingface/token`) exists and is valid — Llama-3.1 +is gated, so a missing file is created before sourcing. diff --git a/.claude/skills/mad-slurm-multinode/references/gotchas.md b/.claude/skills/mad-slurm-multinode/references/gotchas.md new file mode 100644 index 0000000..46f5a53 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/gotchas.md @@ -0,0 +1,223 @@ +# Gotchas — cross-cutting and per-workload + +Keywords: source mad.env MODEL_DIR run-script, MAD_DOCKER_BUILDS shared storage, +MAD_SECRETS_HFTOKEN HF 401 single-quoted, docker_mounts container_path host_path -v, +RCCL_AINIC_ROCE RDMAV_DRIVERS ionic, +NCCL_IB_HCA mlx5 rdma per-cluster, perf.csv login-node aggregation, +slurm.nodes distributed.nnodes nodelist world size, +heterogeneous nodes NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME network_interface, +routable interface eth0 eth1 IPv6 link-local fe80 gloo connect timeout subnet + +Cross-cutting and per-workload pitfalls observed in real runs. SKILL.md links +here; this file is read before a run. + +## Cross-cutting + +- **`source mad.env` in the same shell before `madengine run`.** `MODEL_DIR` is + exported there; without it the container's `scripts/` directory resolves empty + and the run dies with a missing run-script error. +- **`MAD_DOCKER_BUILDS` lives on shared storage** visible to every compute + node. The login/build node saves the image tar there; workers load it from + the same path. A node-local path makes workers rebuild (or fail). +- **The HF token lives in `mad.env`, not in the manifest** — a stray + `MAD_SECRETS_HFTOKEN` key in the manifest causes HF 401s (madengine renders it + single-quoted, so the container gets the literal `${MAD_SECRETS_HFTOKEN}`). + Full explanation in [references/manifests.md](manifests.md) + ("Secrets (HF token)"). +- **`docker_mounts` format is `{container_path: host_path}`.** madengine + renders each entry as `-v :`. Swapping the two + sides (a tempting mistake when the paths differ) hands Docker a non-existent + host path, which it silently creates as an empty directory, and the workload + fails with a "path does not exist" error. The direction is worth re-checking + whenever a host data directory maps to a container-internal path that differs + from the host path. See + [references/manifests.md](manifests.md) ("docker_mounts direction"). +- **`NCCL_IB_HCA` is per-cluster, not portable.** CX7 uses `mlx5_*`; AINIC + uses `rdma0..7`. Copying an mlx5 list onto an AINIC node (or vice versa) + inits zero NICs. Verify on the node. +- **AINIC transport vars appear in BOTH manifest blocks.** + `RCCL_AINIC_ROCE=1`, `RDMAV_DRIVERS=ionic`, `IBV_DRIVERS=ionic` go in + `context.docker_env_vars` AND `deployment_config.env_vars`. Missing from + either, RCCL falls back to verbs/sockets and the run measures the + wrong transport. +- **Multi-node perf CSV: trust the login-node aggregation.** Throughput is + often printed only on the last global rank, so node-0's local CSV can look + empty even on a healthy run. Read the aggregated `perf.csv` in `rundir`, + not a single node's file. +- **`slurm.nodes` equals `distributed.nnodes`** and matches `--nodelist` + cardinality, or sbatch/torchrun disagree on world size. +- **Node environments can be heterogeneous across a cluster — don't trust a + single detect probe.** The interface that carries the routable control-plane + IP is not guaranteed to have the same name on every node (e.g. one node routes + on `eth0` while another routes on `eth1`), and a given named interface may hold + only a non-routable IPv6 link-local address (`fe80::...`) on some nodes. If you + pin `NCCL_SOCKET_IFNAME`/`GLOO_SOCKET_IFNAME`/`network_interface` by name based + on one probe node, the bootstrap can silently fail on the actually-allocated + nodes. Typical symptom: torchrun rendezvous and `world_size` form fine (that + goes over the hostname's routable IP), but `initializing torch distributed` + hangs and gloo times out connecting peers over `fe80::` link-local addresses. + Mitigation: verify the routable interface on the *allocated* nodes (not just + the probe node), and prefer selecting the interface by its routable subnet + rather than hard-coding an interface name. This is an environment + (cluster-provisioning) inconsistency, not a workload bug — flag it to the + cluster owner if a uniform-environment guarantee is expected. + +## sglang_disagg + +Keywords: launcher sglang-disagg run.sh xP yD RUN_MORI DP_MODE KV_TRANSFER_BACKEND +nixl mooncake, detokenizer hang health check No response from detokenizer, +MoRI overlay #366 inter-node decode freeze, RCCL overlay rsmi_init libtorch_hip +undefined symbol torch broken, rocm720 base librocm_smi64, patchelf add-needed +DT_NEEDED smifix, single full-overlay Dockerfile no base chaining, oci-rdma62 +rdma-core v62 variant, mooncake baked launcher build-layer removed runtime pip, +self-discover node IPs rendezvous IP_SYNC_TIMEOUT SGLANG_NODE_IPS not forwarded, +per-node docker load shared tar, perf CSV rank0 BENCHMARK_FAIL_FAST, circuit +breaker prefill workers Service Unavailable BENCHMARK_POINT_RETRIES transient +sweep retry. + +- **The launcher is `sglang-disagg` with `scripts/sglang_disagg/run.sh` as the + entrypoint** (PR #142 native launcher). `run.sh` reads topology from + `--xp/--yd` (or `xP`/`yD` env) and `--kv-transfer-backend` (`nixl`/`mooncake`); + with `--run-mori 1` it execs `sglang_disagg_mori_io_ep.sh`, otherwise + `sglang_disagg_server.sh`. There is no `slurm_multi` wrapper — one `madengine + run` is launched per node and the roles (proxy / prefill x xP / decode x yD) + are derived from `NODE_RANK` + the ordered IP list. Do not reintroduce a + `*_mn` slurm wrapper; it duplicates topology logic and drifts from `run.sh`. + +- **A newer RCCL build drops the `rocm_smi` `DT_NEEDED` and breaks + `import torch` — re-add it with `patchelf`.** The RCCL stage of the single + full-overlay Dockerfile + (`docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile` — base + SGLang + RCCL + MoRI + NIXL/Mooncake KV-transfer + Mooncake pip in one file, no + base-image chaining) uses a rocm-systems RCCL (e.g. develop `78e8ba0`) on top of + the rocm720 sglang base and prepends the overlay lib dir to `LD_LIBRARY_PATH`, so + the overlay `librccl.so` resolves before the base one. That build links `librccl` + *without* the `DT_NEEDED librocm_smi64.so` the base librccl carried, so + `rocm_smi` is never transitively loaded and every later stage (and the run) dies + on `import torch` with `libtorch_hip.so: undefined symbol: rsmi_init`. Fix + (already baked into the Dockerfile): re-add the dependency on the overlay librccl + during the RCCL stage — + `patchelf --add-needed librocm_smi64.so. "$(readlink -f .../librccl.so)"` — + then sanity-check `python3 -c "import torch"` *with the overlay librccl resolved + first* (the exact case that regressed). With the smifix the full overlay + **base -> RCCL (+smifix) -> MoRI -> NIXL/Mooncake KV-transfer -> Mooncake pip** + runs green (validated 4-node DeepSeek-R1, 56/56 sweep points). If you don't need + a specific RCCL version, skip the RCCL stage and the base RCCL also works; if you + keep it, apply the smifix. The Dockerfile sanity-checks `import torch` after each + stage — a failure at the NIXL stage usually means an earlier stage (RCCL) is the + real culprit, not NIXL. + +- **Keep runtime installs out of the launcher — bake them into the image.** The + launcher (`scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh`) used to + `pip install` py-spy/flask/pyyaml and build rdma-core v62 at job start; that + mutated the container runtime env and silently corrupted the Python libs (the + model flags then stopped parsing and the server came up with defaults). Those pip + deps and the Mooncake KV backend are now baked into the full-overlay Dockerfile; + rdma-core v62 lives only in the `*.oci-rdma62.*` variant (for OCI-CX7 hosts whose + RDMA stack needs a newer libibverbs/librdmacm/libmlx5 than the base ships — no + host `libibverbs` bind-mounts needed). Do not reintroduce a runtime build-layer + in the launcher. + +- **The MoRI overlay (#366, e.g. commit `a14e6992`) is the one that fixes the + mid-decode inter-node freeze.** (Commit hashes here and below — RCCL, MoRI, + NIXL — are illustrative pins valid when written; they can drift after a squash, + so track the PR/branch as the source of truth, not the SHA.) Without it, multi-node decode hangs partway + through generation (observed ~token 31 during warmup): the MoE expert-parallel + all-to-all stalls, the decode worker stops responding, and the proxy reports + `No response from detokenizer` / the benchmark stream aborts + (`ClientPayloadError` / `TransferEncodingError`). It is `RUN_MORI=1` + + `MORI_RDMA_DEVICES`/`MORI_SOCKET_IFNAME`/`MORI_IB_GID_INDEX` that route this + transport — set them alongside the NCCL/IB vars. + +- **`No response from detokenizer` on idle decode ranks can be a false + positive.** In a `yD>1` decode pool, idle decode ranks can trip the + detokenizer health check even when generation is healthy (sglang upstream + #20756). A *recent* sglang base fixes this; pin a base new enough to include it + (v0.5.12.post1+ worked) rather than chasing it at the transport layer. A true + hang (above) and this false positive look similar in the proxy log — confirm by + checking whether the *active* decode rank is still emitting tokens + (`py-spy dump` / `gstack` on the decode PID) before blaming the network. + +- **`run.sh` self-discovers the rank-ordered node IPs in-container — do NOT set + `--ipaddrs` / `IPADDRS`.** madengine's container runner does not forward + `SGLANG_NODE_IPS`, so a manifest that relied on it fell back to + `IPADDRS=localhost` and `DP_MODE=1` failed (roles dialed loopback and never + connected). `run.sh` now derives the rank-ordered IP list from the forwarded + `MASTER_ADDR`/`NODE_RANK`/`NNODES`: the SLURM nodelist when `scontrol`/`getent` + exist, else a rank-0 TCP rendezvous on `MASTER_ADDR`. `NNODES` falls back to `xP+yD` (proxy/router is + co-located on rank 0, no separate node). Leave `IPADDRS` / `--ipaddrs` unset in + the manifest; only override them for a reproduction launched outside madengine. + Raise `IP_SYNC_TIMEOUT` (default 1800s) if image-load skew delays peers. + +- **Let madengine distribute the image — don't hand-roll per-node `docker + load`.** For manifest-driven runs madengine already fans the image out on every + node via `container_runner._ensure_local_image_available()` (inspect → `docker + load` the `MAD_DOCKER_BUILDS/.tar` → build → pull); rank 0 `docker save`s + the tar, a TCP barrier syncs, then workers load it (see `launch-and-results.md` + and ROCm/madengine@d28f2f5). So with `local_image: true` + `MAD_DOCKER_BUILDS` + on shared FS the ~20 GB overlay image is distributed for you — no manual + `docker save`/`docker load` step is required. A manual per-node `docker load` + is only relevant for the standalone holder/`srun` reproduction that launches the + container *outside* madengine's dispatcher. + +- **Perf CSV is produced only on rank 0; make the benchmark fail-fast.** Set + `MAD_COLLECT_METRICS=true`/`MAD_SKIP_PERF_COLLECTION=false` on rank 0 only, and + `BENCHMARK_FAIL_FAST=1` so a zero-result sweep aborts instead of writing an + empty `perf_sglang-disagg-DeepSeek-R1.csv`. The sweep is + `BENCHMARK_COMBINATIONS` (isl/osl) x `BENCHMARK_CONCURRENCY_LEVELS`; read the + rank-0 aggregated CSV, not a non-zero node's local file. + +- **A single transient gateway circuit-open should not nuke the whole sweep — + retry points.** Occasionally a point fails at warmup with + `Service Unavailable ... No available prefill workers (all circuits open or + unhealthy)` even though the servers recover seconds later: the sgl-model-gateway + circuit breaker opened on a transient prefill detok hiccup. Under + `BENCHMARK_FAIL_FAST=1` that one blip otherwise aborts a ~50-minute run. The + sweep retries each point `BENCHMARK_POINT_RETRIES` times (default 2) with a + `BENCHMARK_RETRY_COOLDOWN_SECONDS` (default 45s) pause so the breaker can + re-close; a point that fails *every* attempt still fails the run, so real + regressions are not masked. Distinguish this transient from a true hard detok + freeze (heartbeat frozen for minutes, never recovering) by checking + `last_heartbeat` in `/run_logs//prefill_NODE*.log` / + `decode_NODE*.log` — a hard freeze needs a real fix (e.g. the MoRI overlay), + not a retry. + +## primus_megatron (training) + +Keywords: primus megatron scaleout MoE training, primus_turbo DeepEP token dispatcher, +use_turbo_deepep moe_enable_deepep flex alltoall, moe_shared_expert_overlap assert, +print_rank_last throughput last global rank, multi-node perf collection rank-0. + +- **DeepEP is opt-in and OFF by default (dispatcher = `alltoall` over RCCL).** + The baseline MoE token dispatcher is Megatron `alltoall` (a `torch.distributed` + all-to-all over the EP group, carried by RCCL on the RoCE/IB fabric) — *not* + MoRI. To switch to the Primus-Turbo **DeepEP** dispatcher, set + `PRIMUS_USE_DEEPEP=1` in the manifest env (a primus scaleout manifest can carry + it defaulting to `"0"`), which makes + `scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh` + (invoked by the model `run.sh`) pass `--use_turbo_deepep true` to Primus. Primus + then (`_is_turbo_deepep_enabled`) auto-sets `moe_enable_deepep=True` and + `moe_token_dispatcher_type='flex'` and swaps in + `PrimusTurboDeepEPTokenDispatcher`. DeepEP lives inside `primus_turbo` (module + `primus_turbo.pytorch.deep_ep`), not as a standalone `deep_ep` package. + +- **DeepEP requires `tensor_model_parallel_size == 1` + `enable_primus_turbo`, and + is incompatible with `moe_shared_expert_overlap`.** With DeepEP on, Primus ROCm + `validate_args` asserts + `AssertionError: DeepEP not support moe_shared_expert_overlap, please set + moe_shared_expert_overlap=False` (DeepSeek-V3 enables shared-expert overlap by + default). So enabling DeepEP must also pass `--moe_shared_expert_overlap false` + (done automatically alongside `--use_turbo_deepep true`). Measured 4-node x + 8-GPU DeepSeek-V3 proxy config (MI300X, BF16, seq 4096): DeepEP ~295-299 TFLOP/s/GPU, + ~13.1-13.3k tokens/s/GPU vs alltoall ~270 / ~12.05k (~+9%). + +- **Throughput is printed on the global last rank (`print_rank_last`), not rank 0 + — read the aggregated perf, not node 0's local CSV.** On a multi-node run the + Megatron/Primus training_log line (`TFLOP/s/GPU`, `tokens/s/GPU`) is emitted by + the last global rank, which usually lands on the *last* node, while madengine's + designated metric collector is rank 0 (node 0). madengine's login-node + `collect_results` handles this (it picks the richest per-node `multiple_results` + CSV via `_select_best_multiple_results_csv`, and treats an empty local perf on a + `skip_perf_collection` SLURM node as SUCCESS). Use a madengine that has this + multi-node aggregation (present in ROCm/madengine `develop`); otherwise the run + trains fine but is mis-reported as FAILED with an empty `perf.csv`. diff --git a/.claude/skills/mad-slurm-multinode/references/launch-and-results.md b/.claude/skills/mad-slurm-multinode/references/launch-and-results.md new file mode 100644 index 0000000..df1bb58 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/launch-and-results.md @@ -0,0 +1,104 @@ +# Launch and results + +Keywords: madengine run --manifest-file --live-output, perf.csv perf_super, +multiple_results, slurm_output node_N.out, multi-node aggregation, tok/s/gpu, +TFLOPS, sbatch squeue, MODEL_DIR scripts missing + +## Launch + +From `/rundir`, in a shell where `mad.env` has been sourced: + +```bash +cd "$WORKDIR/rundir" +source mad.env # same shell that runs madengine +madengine run --manifest-file run_manifest_.json --live-output +``` + +Useful flags (`madengine run --help` for the full list): +- `--manifest-file/-m ` — run an existing manifest (skips the separate + build phase; the image is still ensured at run time, see below). +- `--live-output/-l` — stream container/job output in real time. +- `--output/-o ` — performance output CSV (default `perf.csv`). +- `--summary-output/-s ` — JSON run summary. +- `--keep-alive` — leave containers up after the run (debugging). +- `--keep-model-dir` — keep the staged model dir (debugging). +- `--verbose/-v` — verbose logging. + +`madengine` renders the sbatch script from the `deployment_config` and submits +the SLURM job itself. Watch the job with `squeue -u $USER`; per-node +stdout/stderr land under `slurm_output/` (e.g. `node_0.out`, `node_1.out`). + +## Where the image comes from (madengine handles it — no pre-build needed) + +For a `local_image: true` manifest, the execution phase calls +`_ensure_local_image_available()` on each node (`container_runner.py`), so no +pre-build or pre-pull is needed. Per node the logic is: + +1. `docker image inspect ` — already present? Then nothing to do. +2. On the primary node (NODE_RANK 0), if the image is missing: + - if a tar exists at `MAD_DOCKER_BUILDS/.tar` -> `docker load` it; + - else **build it from the manifest `dockerfile`** (`docker build -f --pull ...`), + and if that build fails, **fall back to `docker pull `**; + - if `MAD_DOCKER_BUILDS` is set and the tar is absent -> `docker save` the + image into that tar. +3. A TCP barrier syncs all nodes; then worker nodes (rank > 0) `docker load` + the tar (when `MAD_DOCKER_BUILDS` is set) or build/pull independently. + +Implications for the agent: +- **No manual build/pull step.** Just run `madengine run -m`. +- **First run is slower** when the image doesn't exist yet — rank 0 builds it + once, then everyone reuses the tar. +- **`MAD_DOCKER_BUILDS` lives on shared FS** so rank 0's tar is visible to + workers (otherwise every node rebuilds/pulls). This is the same gotcha as in + `SKILL.md`. +- **The build path needs the manifest's `dockerfile` to resolve** — it lives in + the cloned `MAD` (the dockerfile path is relative to the MAD repo). + If the dockerfile is missing/unresolvable, madengine falls back to + `docker pull `, which then requires that tag to exist in a + reachable registry. +- Keep the manifest's `docker_image` tag meaningful: it is both the build tag + and the pull fallback tag, and it names the tar in `MAD_DOCKER_BUILDS`. + +## Where results land + +- `perf.csv` (or `--output` path) in `rundir` — the aggregated performance row. +- `perf_super*.csv` / `perf_entry*.csv` — intermediate per-entry files. +- The model's `multiple_results` CSV (e.g. + `perf_primus-megatron-Llama-3.1-8B.csv`) — per-model detail. +- `slurm_output/node_*.out` — raw logs (throughput banner, NCCL/RCCL transport + lines, tracebacks). + +## Multi-node aggregation gotcha + +Throughput is frequently emitted only by the LAST global rank (often the last +node), not node 0. So node 0's local perf CSV can be empty even on a fully +healthy run. The login-node aggregated `perf.csv` in `rundir` is the source of +truth over any single node's file. A node-0-only "empty perf" check is a false +negative — the aggregation across nodes picks the non-empty result. + +## Reading the number + +- Primus: report tok/s/gpu and TFLOPS/gpu (from the aggregated CSV / rank-last + `.out` throughput banner). +- sglang_disagg: request throughput / latency from the disagg benchmark CSV. + +## Quick failure triage + +- Run dies immediately with a missing run-script / empty `scripts/` -> + `MODEL_DIR` not exported. Re-`source mad.env` in the launching shell and + confirm `[ -d "$MODEL_DIR/scripts" ]`. +- Workers rebuild the image or fail to find it -> `MAD_DOCKER_BUILDS` is not on + shared storage visible to every node. Point it at shared FS and re-run. +- All ranks exit in the first minutes with `Failed to initialize any NET + plugin` / RCCL falls back to sockets -> wrong `NCCL_IB_HCA` for the node, or + AINIC transport vars missing in one of the two env blocks. See + cluster-types.md. +- Data/gated-model error -> HF token missing/expired (`~/.huggingface/token`). + A rendered `docker run` showing `MAD_SECRETS_HFTOKEN='${MAD_SECRETS_HFTOKEN}'` + means a manifest wrongly redeclared the token; remove that key and re-source + `mad.env` (see manifests.md "Secrets"). +- "path does not exist" / an unexpectedly empty mounted dir -> swapped + `docker_mounts`. Entries are keyed `{container_path: host_path}`, so the host + path is the value (see manifests.md "docker_mounts direction"). +- `MODEL_PATH is missing` -> wrong model host path; the model + resolves at `$MODEL_DIR/$MODEL_NAME` inside the container. diff --git a/.claude/skills/mad-slurm-multinode/references/manifests.md b/.claude/skills/mad-slurm-multinode/references/manifests.md new file mode 100644 index 0000000..5a160ef --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/manifests.md @@ -0,0 +1,150 @@ +# Manifest anatomy and per-run adaptation + +Keywords: run_manifest json, built_images local_image, docker_env_vars, +deployment_config slurm partition account qos reservation nodelist exclude, +distributed nnodes nproc_per_node launcher torchrun, multiple_results, +docker_mounts, ${VAR} expansion, MAD_SECRETS_HFTOKEN + +A run manifest is a build+run description madengine executes with +`madengine run --manifest-file `. The templates in `assets/manifests/` +are sanitized — fill them per run and write the result into `/rundir/`. +A filled, cluster-private manifest stays in the rundir, not back in the skill. + +## Top-level keys + +- `built_images` — the image(s) to run. `docker_image` is the local tag. + `local_image: true` means the image is treated as locally provided: at run + time madengine ensures it per node (load from `MAD_DOCKER_BUILDS` tar, else + build from `dockerfile`, else `docker pull`) — see launch-and-results.md + ("Where the image comes from"). `dockerfile` is the path used for that build. +- `built_models` — model/run metadata: `scripts` (run.sh path, resolved under + `MODEL_DIR`), `args` (`--model_repo ...`), `multiple_results` (the per-model + perf CSV name to look for), `additional_docker_run_options` (privileged, + shm-size, device mounts, ulimits). +- `context` — `docker_env_vars` (env inside the container), + `docker_mounts` (bind mounts, keyed `{container_path: host_path}`; madengine + renders each entry as `-v :`), `docker_build_arg`, + `docker_gpus`. +- `deployment_config` — `slurm` block, `distributed` block, `env_vars` + (env on the SLURM job / host side). + +## Verify assets resolve under MODEL_DIR + +The `built_images.*.dockerfile` and `built_models.*.scripts` (the `run.sh`) +paths are relative to `$MODEL_DIR` — the cloned `MAD` checkout. After +cloning or switching its branch, both resolve there +(`[ -f "$MODEL_DIR/" ] && [ -f "$MODEL_DIR/" ]`). A missing +path stops the run with a report (a wrong branch, uninitialized submodules, or a +different layout) rather than a silent search; the next step is to ask the user +for the correct branch. + +## Secrets (HF token) + +The HF token comes from `mad.env`, not the manifest. `mad.env` exports +`MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token)`, and madengine forwards every +`MAD_SECRETS_*` from the sourced host env into the container's env +(`core/context.py` reads `MAD_SECRETS*` from `os.environ`; host-env values take +priority over manifest entries). A manifest that also sets +`"MAD_SECRETS_HFTOKEN": "${MAD_SECRETS_HFTOKEN}"` triggers an HF 401: madengine +renders that entry single-quoted in `docker run`, so the container receives the +literal string instead of the token. The templates carry no token key, and +keeping it that way is the fix. + +## Fill checklist (every run) + +Cluster-private — request from the user, keep out of the skill: +- `deployment_config.slurm.partition` +- `deployment_config.slurm.account` (omit the key if the cluster has no accounts) +- `deployment_config.slurm.qos` +- `deployment_config.slurm.reservation` (AINIC reserved campaigns) — **NOTE**: madengine + 2.0.0.post53 silently drops this field from the generated sbatch script. After submit, fix + with `scontrol update jobid= Reservation=`. Alternatively, `export + SBATCH_RESERVATION=` in the shell before `madengine run`. +- `deployment_config.slurm.nodelist` and/or `exclude` (specific node names) + +Cluster-shape — confirm on the node (see cluster-types.md): +- `NCCL_IB_HCA` (in BOTH env blocks) — the mlx5/rdma device list +- `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME` / `slurm.network_interface` — the + same iface value in all three places (`context.docker_env_vars`, `slurm`, and + `deployment_config.env_vars`); a mismatch splits control-plane routing. +- `NCCL_IB_GID_INDEX` +- `RDMAV_DRIVERS` / `IBV_DRIVERS` (+ `RCCL_AINIC_ROCE=1` on AINIC, delete on CX7) +- `MAD_SYSTEM_GPU_ARCHITECTURE` / `BUILD_GPU_TARGETS` + +Run-shape — the user's run choice: +- `slurm.nodes` equals `distributed.nnodes` equals the cardinality of + `nodelist` (if used). A mismatch -> sbatch/torchrun disagree on world size. +- `slurm.time` walltime, `gpus_per_node`, `nproc_per_node`. +- `docker_image` tag (the RCCL/build tag you intend to test). + +Paths — from `mad.env` via `${VAR}` expansion where possible +(`${HF_HOME}`, `${MAD_DATAHOME}`, ...), explicit `` otherwise. +madengine expands `${VAR}` from the sourced host env, so keep mad.env and the +manifest consistent. Any workload-specific data paths are explicit +`/...` in the workload manifest itself, not exported +from mad.env (which stays cluster/cache-only). + +## Placeholder convention in templates + +JSON has no comments, so placeholders carry inline guidance: +- `""` — a transport/cluster value; + replace the whole string with the archetype value from the matrix in + [cluster-types.md](cluster-types.md). The templates are cluster-agnostic, so + the per-archetype values live only in that matrix, not inline in the manifest. +- `"RCCL_AINIC_ROCE": ""` — set per archetype + (AINIC: `"1"`); delete the key entirely on an archetype that does not need it. +- `${VAR}` — leave as-is; madengine expands from the sourced env. + +Re-validate after editing with the static checker (GPU-free, includes JSON +validity plus the fill checklist above): +`bash "$SKILL_DIR/scripts/validate_manifest.sh" .json` (source +`mad.env` first so it can also resolve the dockerfile/run.sh under `$MODEL_DIR`). + +## docker_mounts direction + +`docker_mounts` is keyed `{container_path: host_path}` and madengine renders +each entry as `-v :` (`execution/container_runner.py` +emits `-v {value}:{key}`; `orchestration/run_orchestrator.py` reads the dict as +`{container_path: host_path}`). The key is the path inside the container; the +value is the path on the host. Swapping the two sides points Docker at a +non-existent host path, which Docker then creates as an empty directory, and the +workload fails with a missing-path error (for example a `MODEL_PATH is missing` +error). When a host data path and its container mount point +differ, the host path belongs on the value side. + +## Per-workload notes + +- **Primus** (`primus_pyt_megatron_lm_train_llama-3.1-{8b,70b}_scaleout`): uses the + `rccl_overlay` Dockerfile; `launcher: torchrun`, `nproc_per_node: 8`. + `multiple_results` = `perf_primus-megatron-Llama-3.1-{8B,70B}.csv`. 70B is + heavier — give it more walltime / more nodes. +- **sglang_disagg** (`sglang_disagg_deepseek-r1`): disaggregated + prefill/decode serving of DeepSeek-R1 on SGLang. + `launcher: sglang-disagg`, `scripts/sglang_disagg/run.sh` entrypoint, + `nproc_per_node: 8`, typically 4 nodes (e.g. xP=2 prefill, yD=2 decode; the + proxy/load-balancer rides rank 0). Extra knobs: `xP`, `yD`, `RUN_MORI` (MoE + expert-parallel all-to-all via MoRI — keep `1` for multi-node), `DP_MODE`, + `KV_TRANSFER_BACKEND` (`nixl` or `mooncake`), `GPUS_PER_NODE`, the `MORI_*` + transport vars (mirror the `NCCL_IB_*` data-plane values), and the benchmark + sweep (`BENCHMARK_COMBINATIONS`, `BENCHMARK_CONCURRENCY_LEVELS`, + `BENCHMARK_FAIL_FAST`). `run.sh` self-discovers the rank-ordered node IP list + in-container (SLURM nodelist, else a rank-0 TCP rendezvous on the forwarded + `MASTER_ADDR`), so do NOT set `IPADDRS` / `--ipaddrs`: madengine does not + forward `SGLANG_NODE_IPS`, and a stale hand-filled list breaks role addressing + (tune `IP_SYNC_TIMEOUT`, default 1800s, for image-load skew instead). The image + is a single full-overlay build + (`docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile`): base + SGLang + RCCL + MoRI + NIXL/Mooncake KV-transfer + Mooncake pip in one + Dockerfile (no base-image chaining). For OCI-CX7 hosts whose RDMA stack needs a + newer rdma-core than the base ships, build the `*.oci-rdma62.*` variant instead + (bakes rdma-core v62 in, so no host `libibverbs` bind-mounts are needed). The + full overlay re-adds the `rocm_smi` `DT_NEEDED` + (`patchelf --add-needed librocm_smi64.so.`) so a newer RCCL on the rocm720 + base does not break `import torch` (see [gotchas.md](gotchas.md#sglang_disagg)). + Perf lands in `perf_sglang-disagg-DeepSeek-R1.csv`, collected on rank 0 only. + +## Scaling a manifest (e.g. 2-node -> 4-node) + +Change `slurm.nodes`, `distributed.nnodes`, and the `nodelist` cardinality +together. Everything else (env, mounts, image) stays the same. Real examples +exist at 2 and 4 nodes for primus and sglang_disagg. diff --git a/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh b/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh new file mode 100755 index 0000000..9d2d2dd --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh @@ -0,0 +1,104 @@ +#!/usr/bin/env bash +# Propose cluster-specific values for mad.env / manifests by inspecting the node. +# Read-only. The agent confirms these with the user before writing configs; +# they are best-effort guesses, not authoritative. +# +# This reflects the node it runs on. The login/jump node and the compute nodes +# can differ (HCAs, GPU arch, interfaces), so for compute-node values run it on +# an allocated node, e.g.: +# srun -p [--reservation ] [--nodelist ] -N1 \ +# bash detect_cluster_env.sh +set -u + +echo "== mad-slurm-multinode cluster env detection (proposals only) ==" + +# --- GPU architecture --- +arch="" +if command -v rocminfo >/dev/null 2>&1; then + arch=$(rocminfo 2>/dev/null | grep -m1 -oE 'gfx[0-9a-f]+') +elif command -v rocm-smi >/dev/null 2>&1; then + arch=$(rocm-smi --showhw 2>/dev/null | grep -m1 -oE 'gfx[0-9a-f]+') +fi +echo "MAD_SYSTEM_GPU_ARCHITECTURE : ${arch:-}" + +# --- RDMA / IB HCAs --- +hcas="" +if command -v ibv_devices >/dev/null 2>&1; then + hcas=$(ibv_devices 2>/dev/null | awk 'NR>2 {print $1}' | grep -E '^(mlx5_|rdma)' | sort) +fi +if [ -z "$hcas" ] && [ -d /sys/class/infiniband ]; then + hcas=$(ls /sys/class/infiniband 2>/dev/null | grep -E '^(mlx5_|rdma)' | sort) +fi +if [ -n "$hcas" ]; then + list=$(echo "$hcas" | sed 's/$/:1/' | paste -sd, -) + echo "NCCL_IB_HCA (all ports :1) : $list" + if echo "$hcas" | grep -q '^rdma'; then + fam="AINIC (ionic) -> RDMAV_DRIVERS=ionic, IBV_DRIVERS=ionic, RCCL_AINIC_ROCE=1, GID likely 1" + elif echo "$hcas" | grep -q '^mlx5'; then + fam="CX7/Mellanox (mlx5) -> RDMAV_DRIVERS=mlx5, IBV_DRIVERS=mlx5, GID likely 3" + else + fam="unknown vendor -> archetype not auto-classified; confirm with the user (see references/cluster-types.md)" + fi + echo "archetype guess : $fam" + echo " (note: NCCL_IB_HCA usually lists only the GPU-attached HCAs; trim as needed)" +else + echo "NCCL_IB_HCA : " +fi + +# --- HCA <-> netdev / NUMA affinity (management vs data-plane) --- +# Some HCAs back a routed management iface (e.g. mlx5_1->eth0, mlx5_6->eth1) +# rather than a GPU rail. NCCL_IB_HCA should list only the data-plane HCAs, so +# split them by their backing netdev: eth*/eno*/bond* = management (skip), +# rdma*/ib*/none = data-plane (use). Confirm against GPU<->NIC affinity below. +echo "HCA -> netdev / NUMA affinity:" +if command -v ibdev2netdev >/dev/null 2>&1; then + ibdev2netdev 2>/dev/null | sed 's/^/ /' +fi +mgmt_hcas=""; data_hcas="" +if [ -d /sys/class/infiniband ]; then + for d in /sys/class/infiniband/*; do + [ -e "$d" ] || continue + dev=$(basename "$d") + numa=$(cat "$d/device/numa_node" 2>/dev/null) + nd=$(ls "$d/device/net" 2>/dev/null | paste -sd, -) + printf ' %-10s numa=%-3s netdev=%s\n' "$dev" "${numa:-?}" "${nd:-none}" + if echo "$nd" | grep -qE '^(eth|eno|ens|enp|em|bond)'; then + mgmt_hcas="$mgmt_hcas $dev" + else + data_hcas="$data_hcas $dev" + fi + done +fi +if [ -n "$data_hcas" ]; then + dlist=$(echo $data_hcas | tr ' ' '\n' | sed 's/$/:1/' | paste -sd, -) + echo " management HCAs (skip) :${mgmt_hcas:- none}" + echo " data-plane HCAs (use) :$data_hcas" + echo " NCCL_IB_HCA (data-plane :1) : $dlist" + echo " (heuristic: a HCA backing a routed eth*/bond iface is management;" + echo " rdma*/ib* rails attached to the GPUs are data-plane — confirm by affinity)" +fi + +# --- GPU PCIe bus / NUMA (to match data-plane rails to GPUs) --- +if command -v rocm-smi >/dev/null 2>&1; then + echo "GPU PCIe bus (rocm-smi --showbus):" + rocm-smi --showbus 2>/dev/null | grep -iE 'GPU|bus' | sed 's/^/ /' | head -20 +fi + +# --- management interface --- +echo "candidate NCCL_SOCKET_IFNAME :" +if command -v ip >/dev/null 2>&1; then + ip -br addr 2>/dev/null | awk '$2=="UP" && $1!="lo" {print " "$1" "$3}' +else + echo " " +fi +echo " (pick the routable management iface, e.g. eth0 on CX7, eno0 on AINIC)" + +# --- RoCEv2 GID index --- +if command -v show_gids >/dev/null 2>&1; then + echo "RoCEv2 GID candidates (show_gids, look for v2):" + show_gids 2>/dev/null | grep -iE 'v2|RoCE' | head -10 | sed 's/^/ /' +else + echo "NCCL_IB_GID_INDEX : " +fi + +echo "== confirm all of the above with the user before writing mad.env/manifest ==" diff --git a/.claude/skills/mad-slurm-multinode/scripts/preflight.sh b/.claude/skills/mad-slurm-multinode/scripts/preflight.sh new file mode 100755 index 0000000..63a6892 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/scripts/preflight.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash +# Step 0 preflight for mad-slurm-multinode. +# Prints a PASS/WARN/FAIL table of prerequisites. Exits non-zero if any HARD +# requirement (docker daemon, SLURM client) is missing. +set -u + +pass=0; warn=0; fail=0 +row() { printf ' %-6s %-22s %s\n' "$1" "$2" "$3"; } +ok() { row "PASS" "$1" "$2"; pass=$((pass+1)); } +wn() { row "WARN" "$1" "$2"; warn=$((warn+1)); } +bad() { row "FAIL" "$1" "$2"; fail=$((fail+1)); } + +echo "== mad-slurm-multinode preflight ==" + +# Python >= 3.10 (soft: a conda env is created later anyway) +if command -v python3 >/dev/null 2>&1; then + pv=$(python3 -c 'import sys;print("%d.%d"%sys.version_info[:2])' 2>/dev/null) + if python3 -c 'import sys;exit(0 if sys.version_info[:2]>=(3,10) else 1)' 2>/dev/null; then + ok "python3" "$pv (>= 3.10)" + else + wn "python3" "$pv (< 3.10; conda env 3.12 will be created)" + fi +else + wn "python3" "not found (conda env 3.12 will be created)" +fi + +# Docker (HARD): present + daemon reachable +if command -v docker >/dev/null 2>&1; then + if docker info >/dev/null 2>&1; then + ok "docker" "daemon reachable" + else + bad "docker" "present but 'docker info' failed (daemon/perms?)" + fi +else + bad "docker" "not found" +fi + +# SLURM client (HARD) +if command -v sbatch >/dev/null 2>&1 && command -v sinfo >/dev/null 2>&1; then + ok "slurm" "$(sinfo --version 2>/dev/null | head -1)" +else + bad "slurm" "sbatch/sinfo not found" +fi + +# conda / miniforge (soft: installed in Step 2 if missing) +if command -v conda >/dev/null 2>&1; then + ok "conda" "$(conda --version 2>/dev/null)" +else + # conda is not on PATH; an existing install may just need sourcing + cphit="" + for cp in "$HOME/miniforge3" "$HOME/miniconda3" "$HOME/mambaforge" \ + "/opt/conda" "${CONDA_PREFIX:-}" "${MAMBA_ROOT_PREFIX:-}"; do + if [ -n "$cp" ] && [ -f "$cp/etc/profile.d/conda.sh" ]; then cphit="$cp"; break; fi + done + if [ -n "$cphit" ]; then + wn "conda" "found at $cphit, not on PATH — source $cphit/etc/profile.d/conda.sh" + else + wn "conda" "not found (miniforge installed in Step 2)" + fi +fi + +# GPU SMI (soft: used for arch detection) +if command -v rocm-smi >/dev/null 2>&1; then + ok "gpu-smi" "rocm-smi (AMD)" +elif command -v nvidia-smi >/dev/null 2>&1; then + ok "gpu-smi" "nvidia-smi (NVIDIA)" +else + wn "gpu-smi" "neither rocm-smi nor nvidia-smi found" +fi + +# git (HARD: needed to clone/switch repos) +if command -v git >/dev/null 2>&1; then + ok "git" "$(git --version 2>/dev/null)" +else + bad "git" "not found" +fi + +# HF token file (soft: required at run time for gated Llama-3.1) +if [ -s "$HOME/.huggingface/token" ]; then + ok "hf-token" "~/.huggingface/token present" +else + wn "hf-token" "~/.huggingface/token missing (needed before launch)" +fi + +echo "== summary: $pass PASS, $warn WARN, $fail FAIL ==" +if [ "$fail" -gt 0 ]; then + echo "Hard requirements failed. This node is not ready for madengine SLURM runs." + exit 1 +fi +exit 0 diff --git a/.claude/skills/mad-slurm-multinode/scripts/validate_manifest.sh b/.claude/skills/mad-slurm-multinode/scripts/validate_manifest.sh new file mode 100644 index 0000000..f156dd6 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/scripts/validate_manifest.sh @@ -0,0 +1,147 @@ +#!/usr/bin/env bash +# GPU-free static validator for a filled run manifest. +# Usage: bash validate_manifest.sh +# Reads $MODEL_DIR (export it via `source mad.env`) to also resolve the +# dockerfile / run.sh paths. Read-only; prints a PASS/WARN/FAIL table and exits +# non-zero if any FAIL. No GPU, no docker, no network needed. +set -u + +manifest="${1:-}" +if [ -z "$manifest" ]; then + echo "usage: $0 " >&2 + exit 2 +fi +if ! command -v python3 >/dev/null 2>&1; then + echo "python3 required for validate_manifest.sh" >&2 + exit 2 +fi + +exec python3 - "$manifest" <<'PY' +import json, os, re, sys + +path = sys.argv[1] +model_dir = os.environ.get("MODEL_DIR", "") + +passes = warns = fails = 0 +def ok(m): + global passes; passes += 1; print(f" PASS {m}") +def warn(m): + global warns; warns += 1; print(f" WARN {m}") +def fail(m): + global fails; fails += 1; print(f" FAIL {m}") + +print(f"== mad-slurm-multinode manifest validation: {path} ==") + +if not os.path.isfile(path): + print(f" FAIL manifest file not found: {path}") + print("== 0 PASS, 0 WARN, 1 FAIL ==") + sys.exit(1) + +raw = open(path, encoding="utf-8", errors="replace").read() + +# 1) JSON validity +try: + m = json.loads(raw) + ok("valid JSON") +except Exception as e: + print(f" FAIL invalid JSON: {e}") + print("== 0 PASS, 0 WARN, 1 FAIL ==") + sys.exit(1) + +# 2) no unfilled placeholders +ph = re.findall(r']*>', raw) +if ph: + fail(f"{len(ph)} unfilled placeholder(s) remain, e.g. {ph[0]!r}") +else: + ok("no placeholders left") + +ctx = (m.get("context") or {}).get("docker_env_vars") or {} +dep = m.get("deployment_config") or {} +slurm = dep.get("slurm") or {} +dist = dep.get("distributed") or {} +denv = dep.get("env_vars") or {} + +def is_ph(v): + return isinstance(v, str) and (" 1: + fail(f"network interface mismatch across blocks: {present}") +elif vals: + ok(f"network interface consistent ({vals[0]})") + +# 5) node count consistency +nodes, nnodes = slurm.get("nodes"), dist.get("nnodes") +if nodes is not None and nnodes is not None: + if nodes != nnodes: + fail(f"slurm.nodes ({nodes}) != distributed.nnodes ({nnodes})") + else: + ok(f"slurm.nodes == distributed.nnodes ({nodes})") +nl = slurm.get("nodelist") +if isinstance(nl, str) and nl and not is_ph(nl): + if "[" in nl: + warn(f"nodelist uses a bracket range; cardinality not auto-checked ({nl})") + elif nodes is not None: + count = len([x for x in nl.split(",") if x.strip()]) + if count != nodes: + fail(f"nodelist has {count} node(s) but slurm.nodes={nodes}") + else: + ok(f"nodelist cardinality matches slurm.nodes ({count})") + +# 6) stray HF token in the manifest -> HF 401 +if "MAD_SECRETS_HFTOKEN" in ctx or "MAD_SECRETS_HFTOKEN" in denv: + fail("MAD_SECRETS_HFTOKEN is declared in the manifest -> HF 401; " + "remove it (the token comes from mad.env)") +else: + ok("no MAD_SECRETS_HFTOKEN key in manifest") + +# 7) AINIC transport vars must be symmetric across both env blocks +for k in ("RCCL_AINIC_ROCE", "RDMAV_DRIVERS", "IBV_DRIVERS"): + inc = (k in ctx) and not is_ph(ctx.get(k)) + ind = (k in denv) and not is_ph(denv.get(k)) + if inc != ind: + fail(f"{k} set in only one env block (context={inc}, env_vars={ind}); " + "transport vars must be in BOTH") + +# 8) dockerfile / run.sh resolve under MODEL_DIR (if available) +asset_paths = [] +for v in (m.get("built_images") or {}).values(): + if isinstance(v, dict) and v.get("dockerfile"): + asset_paths.append(("dockerfile", v["dockerfile"])) +for v in (m.get("built_models") or {}).values(): + if isinstance(v, dict) and v.get("scripts"): + asset_paths.append(("scripts", v["scripts"])) +if model_dir and os.path.isdir(model_dir): + for kind, p in asset_paths: + full = os.path.join(model_dir, p) + if os.path.exists(full): + ok(f"{kind} resolves under MODEL_DIR: {p}") + else: + fail(f"{kind} not found under MODEL_DIR: {p}") +else: + warn(f"MODEL_DIR not set/usable ('{model_dir}') -> " + "dockerfile/run.sh path checks skipped") + +print(f"== validate_manifest: {passes} PASS, {warns} WARN, {fails} FAIL ==") +sys.exit(1 if fails else 0) +PY diff --git a/.gitignore b/.gitignore index 1285852..b7c59cf 100644 --- a/.gitignore +++ b/.gitignore @@ -72,3 +72,6 @@ venv.bak/ *.log *.out *.html + +# Skill manifest templates are checked in despite *.json above +!.claude/skills/mad-slurm-multinode/assets/manifests/*.template.json From 91fbb75046003f21aa44e3676e5ed711f216e289 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Wed, 1 Jul 2026 16:33:43 +0000 Subject: [PATCH 04/25] Add new model configurations for SGLang and Primus Megatron-LM training Add model entries in `models.json` for the new inference and training workloads. Key additions: - SGLang disaggregated DeepSeek-R1 inference configuration. - Primus Megatron-LM scaleout training configurations for Llama-3.1 8B/70B/405B, each with its Dockerfile, scripts, and perf-result tracking. Co-authored-by: Ilia Kosarev --- models.json | 83 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/models.json b/models.json index af914f4..9f95de4 100644 --- a/models.json +++ b/models.json @@ -1977,6 +1977,29 @@ "args": "--model_repo deepseek-ai/DeepSeek-R1-Distill-Qwen-32B --test_option latency --num_gpu 8 --datatype bfloat16 --dataset random --batch_size 1,8,32 --lat_input_output_len '128:128;128:1024;1024:128;1024:1024'" }, + { + "name": "sglang_disagg_deepseek-r1", + "url": "", + "dockerfile": "docker/sglang_disagg_inference", + "scripts": "scripts/sglang_disagg/run.sh", + "n_gpus": "8", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_sglang-disagg-DeepSeek-R1.csv", + "tags": [ + "pyt", + "sglang", + "inference", + "disaggregated", + "deepseek-r1", + "moe", + "ep-internode", + "rccl-tp", + "mori-a2a" + ], + "timeout": 14400, + "args": "" + }, { "name": "pyt_hy_video", "url": "", @@ -2660,6 +2683,66 @@ "args": "--workload z_image", "multiple_results": "results.csv" }, + { + "name": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay", + "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-8B.csv", + "tags": [ + "pyt", + "pretrain", + "llama3", + "training", + "primus", + "scaleout" + ], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-8b" + }, + { + "name": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay", + "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-70B.csv", + "tags": [ + "pyt", + "pretrain", + "llama3", + "training", + "primus", + "scaleout" + ], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-70b" + }, + { + "name": "primus_pyt_megatron_lm_train_llama-3.1-405b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay", + "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-405B.csv", + "tags": [ + "pyt", + "pretrain", + "llama3", + "training", + "primus", + "scaleout" + ], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-405b" + }, { "name": "pyt_sglang_disagg_mori_io_llama-3.1-8b", "url": "", From 72e7a712f14167de00fa982ddc5f2d253d535f8e Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Fri, 3 Jul 2026 12:34:33 +0000 Subject: [PATCH 05/25] Consolidate primus_scaleout megatron-lm scripts into primus/megatron-lm Fold the scaleout-only capabilities into the main Primus Megatron-LM benchmark scripts and remove the duplicated scripts/primus_scaleout tree (~940 lines of near-copy): - multinode-aware NNODES/GPUS_PER_NODE/NUM_GPUS with global_batch_size renormalization; single-node behavior is unchanged - run_primus() launch helper - amd-smi and MAD_SYSTEM_GPU_ARCHITECTURE device-detection fallback - Llama-3.1-405B branch with a missing-config guard - DeepSeek-V3 DeepEP opt-in via PRIMUS_USE_DEEPEP - trailing running-average throughput rows in the perf report (new run column) Repoint the three *_scaleout models in models.json to scripts/primus/megatron-lm/run.sh, keeping their rccl_overlay Dockerfile and scaleout tag. --- models.json | 6 +- .../primus_megatron-lm_benchmark_report.py | 38 +- .../primus_megatron-lm_benchmark_report.sh | 147 ++++- scripts/primus/megatron-lm/run.sh | 11 + .../primus_megatron-lm_benchmark_report.py | 149 ----- .../primus_megatron-lm_benchmark_report.sh | 610 ------------------ .../primus_megatron-lm_benchmark_setup.sh | 41 -- scripts/primus_scaleout/megatron-lm/run.sh | 143 ---- 8 files changed, 165 insertions(+), 980 deletions(-) delete mode 100644 scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py delete mode 100755 scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh delete mode 100755 scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh delete mode 100755 scripts/primus_scaleout/megatron-lm/run.sh diff --git a/models.json b/models.json index 9f95de4..7c35e1d 100644 --- a/models.json +++ b/models.json @@ -2687,7 +2687,7 @@ "name": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", "url": "", "dockerfile": "docker/primus_megatron_train_rccl_overlay", - "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "scripts": "scripts/primus/megatron-lm/run.sh", "n_gpus": "-1", "owner": "mad.support@amd.com", "training_precision": "", @@ -2707,7 +2707,7 @@ "name": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", "url": "", "dockerfile": "docker/primus_megatron_train_rccl_overlay", - "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "scripts": "scripts/primus/megatron-lm/run.sh", "n_gpus": "-1", "owner": "mad.support@amd.com", "training_precision": "", @@ -2727,7 +2727,7 @@ "name": "primus_pyt_megatron_lm_train_llama-3.1-405b_scaleout", "url": "", "dockerfile": "docker/primus_megatron_train_rccl_overlay", - "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "scripts": "scripts/primus/megatron-lm/run.sh", "n_gpus": "-1", "owner": "mad.support@amd.com", "training_precision": "", diff --git a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py index d1d1f3b..f75bb89 100644 --- a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py +++ b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py @@ -114,7 +114,25 @@ def find_match(file_path, search_string): print(f"Warning: No matches found for '{search_string}' pattern") return None +def find_match_running_avg(file_path, search_string): + """Return the running-average value (the number AFTER the slash) from the + LAST logged iteration, e.g. '1885.6' from 'tokens/s/GPU): 2355.1/1885.6'. + + Primus logs throughput as 'instant/running_avg'. The instantaneous value + has large run-to-run jitter (MoE routing imbalance, collective stragglers, + kernel warmup), so the trailing running average is a steadier figure, + especially for multi-node scaleout runs. Returns None when no match.""" + with open(file_path, 'r') as f: + content = f.read() + pattern = fr"{re.escape(search_string)}\):\s*\d+\.?\d*/(\d+\.?\d*)" + matches = re.findall(pattern, content) + if matches: + print(f"Found {len(matches)} running-avg matches for '{search_string}', using last: {matches[-1]}") + return matches[-1] + return None + if args.model == "Llama-3.1-8B" or args.model == "Llama-3.1-70B" or \ + args.model == "Llama-3.1-405B" or \ args.model == "Llama-2-7B" or args.model == "Llama-2-70B" or \ args.model == "Mixtral-8x7B" or args.model == "Mixtral-8x22B-proxy" or \ args.model == "DeepSeek-V2-lite" or args.model == "DeepSeek-V3-proxy" or \ @@ -135,9 +153,23 @@ def find_match(file_path, search_string): # Write data for metrics that are found (don't require both to be present) # Skip tokens/s/GPU for Qwen-3-32B if tok_per_s_per_gpu is not None and args.model != "Qwen-3-32B": - data.append({'model': args.model, 'performance': tok_per_s_per_gpu, 'metric': 'tok_per_s_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus}) + data.append({'model': args.model, 'performance': tok_per_s_per_gpu, 'metric': 'tok_per_s_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': ''}) if TFLOPS_per_gpu is not None: - data.append({'model': args.model, 'performance': TFLOPS_per_gpu, 'metric': 'TFLOPS_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus}) + data.append({'model': args.model, 'performance': TFLOPS_per_gpu, 'metric': 'TFLOPS_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': ''}) + + # Additionally emit the trailing running-average throughput (value after the + # slash in the last logged iteration). Kept as separate, distinctly-named + # rows (metric suffix '_avg', run='avg') so the instantaneous metrics above + # are preserved and never confused with the steady-state average. This is a + # steadier figure for multi-node scaleout runs where per-iteration jitter is + # largest. Skipped for Qwen-3-32B, matching the tokens/s/GPU handling above. + if args.model != "Qwen-3-32B": + tok_avg = find_match_running_avg(input_file, "tokens/s/GPU") + if tok_avg is not None: + data.append({'model': args.model, 'performance': tok_avg, 'metric': 'tok_per_s_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) + tflops_avg = find_match_running_avg(input_file, "TFLOP/s/GPU") + if tflops_avg is not None: + data.append({'model': args.model, 'performance': tflops_avg, 'metric': 'TFLOPS_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) if not os.path.exists(output_file) or os.stat(output_file).st_size == 0: mode = 'w' # Write if file doesn't exist or is empty @@ -146,7 +178,7 @@ def find_match(file_path, search_string): with open(output_file, mode=mode, newline='') as file: print("Preparing to write performance data...") print("Data: ", data) - writer = csv.DictWriter(file, fieldnames=['model','performance','metric','mode','precision','batch_size','global_batch_size','posttrain_type','seq_len','device','num_gpus']) + writer = csv.DictWriter(file, fieldnames=['model','performance','metric','mode','precision','batch_size','global_batch_size','posttrain_type','seq_len','device','num_gpus','run']) if mode == 'w': writer.writeheader() writer.writerows(data) diff --git a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh index a4604f0..ac84117 100755 --- a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh +++ b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh @@ -140,8 +140,17 @@ export PATH="${ROCM_BIN}:${PATH}" # Run rocminfo and grep for "AMD Instinct" DEVICE=$("${ROCM_BIN}/rocminfo" | grep "AMD Instinct" | head -n1 | awk '{print $5}') echo "DEVICE found: $DEVICE" +# Fall back to amd-smi when the rocminfo output format changes or rocminfo is +# unavailable (some newer ROCm/scaleout stacks). +if [[ -z "$DEVICE" && -x "${ROCM_BIN}/amd-smi" ]]; then + DEVICE=$("${ROCM_BIN}/amd-smi" 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') +fi if [[ -z "$DEVICE" || ("$DEVICE" != "MI300X" && "$DEVICE" != "MI355X") ]]; then ARCH=$("${ROCM_BIN}/rocminfo" | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + # Fall back to the arch injected by madengine when rocminfo cannot report it. + if [[ -z "$ARCH" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then + ARCH=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') + fi case "$ARCH" in "gfx942") DEVICE="MI300X" ;; "gfx950") DEVICE="MI355X" ;; @@ -161,13 +170,66 @@ else CONFIG_DEVICE="$DEVICE" fi -# Set common environment variables -export NNODES=1 -export CPUS_PER_TASK=128 +# Set common environment variables. +# Keep compatibility with SLURM/madengine distributed env and only fall back to +# single-node defaults: single-node runs behave exactly as before, while +# multi-node scaleout runs pick up the injected topology. +NNODES="${NNODES:-1}" +export NNODES +export CPUS_PER_TASK="${CPUS_PER_TASK:-128}" export HSA_NO_SCRATCH_RECLAIM=1 export NVTE_CK_USES_BWD_V3=1 -GPUS_PER_NODE=8 # default to 8 GPUs per node -NUM_GPUS=8 +GPUS_PER_NODE="${GPUS_PER_NODE:-${NPROC_PER_NODE:-8}}" +if ! [[ "$NNODES" =~ ^[0-9]+$ ]]; then NNODES=1; fi +if ! [[ "$GPUS_PER_NODE" =~ ^[0-9]+$ ]]; then GPUS_PER_NODE=8; fi +NUM_GPUS=$((NNODES * GPUS_PER_NODE)) +echo "[INFO] Distributed params: NNODES=$NNODES GPUS_PER_NODE=$GPUS_PER_NODE NUM_GPUS=$NUM_GPUS" + +# Round global batch size up to a multiple of micro_batch_size * data_parallel so +# the global batch stays divisible across an arbitrary number of ranks. +normalize_global_batch_size() { + local mbs="$1" + local gbs="$2" + local dp="$3" + local unit=$((mbs * dp)) + if (( unit <= 0 )); then + echo "$gbs" + return + fi + if (( gbs % unit == 0 )); then + echo "$gbs" + return + fi + echo "$(( ((gbs + unit - 1) / unit) * unit ))" +} + +# Emit the extra CLI flags needed to keep a run valid on multiple nodes: when +# NUM_GPUS > 8 the config global_batch_size (tuned for a single 8-GPU node) is +# renormalized and passed as an override. Prints nothing for single-node runs, +# so behavior there is byte-for-byte identical to before. +scaleout_gbs_override() { + local mbs="$1" + local gbs="$2" + if (( NUM_GPUS > 8 )) && [[ "$mbs" =~ ^[0-9]+$ && "$gbs" =~ ^[0-9]+$ ]]; then + local adjusted + adjusted=$(normalize_global_batch_size "$mbs" "$gbs" "$NUM_GPUS") + if [[ "$adjusted" != "$gbs" ]]; then + echo "[INFO] Adjusted global batch size for distributed run: ${gbs} -> ${adjusted} (MBS=${mbs}, NUM_GPUS=${NUM_GPUS})" >&2 + fi + printf -- '--global_batch_size %s' "$adjusted" + fi +} + +# Launch a single Primus pretrain run. Every rank must execute an identical +# command so the model shape stays consistent across nodes; a mismatch makes +# multi-node ranks rendezvous with different shapes and deadlocks the collectives. +run_primus() { + local config="$1"; shift + bash runner/primus-cli direct \ + --log_file "/tmp/primus_$MODEL_REPO.log" \ + -- train pretrain \ + --config "$config" "$@" 2>&1 | tee "$TRAIN_LOG" +} cd /workspace/Primus @@ -180,26 +242,18 @@ if [ "$MODEL_REPO" == "Llama-3.1-8B" ]; then MBS=$(grep -E '^\s*micro_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [[ "$DATATYPE" == "MXFP4" ]]; then - NVTE_USE_CAST_TRANSPOSE_TRITON=0 bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + NVTE_USE_CAST_TRANSPOSE_TRITON=0 run_primus "$EXP" $GBS_OVERRIDE else - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then if [[ "$DATATYPE" == "MXFP8" || "$DATATYPE" == "MXFP4" ]]; then echo "Error: Datatype $DATATYPE is not supported for $MODEL_REPO on $DEVICE. Only supported on MI355X/MI350X." else - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi fi if [ -f "$TRAIN_LOG" ]; then @@ -221,26 +275,18 @@ elif [ "$MODEL_REPO" == "Llama-3.1-70B" ]; then MBS=$(grep -E '^\s*micro_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [ "$DATATYPE" == "FP8" ]; then - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE elif [ "$DATATYPE" == "BF16" ]; then - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then if [ "$DATATYPE" == "FP8" ]; then echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." elif [ "$DATATYPE" == "BF16" ]; then - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi fi if [ -f "$TRAIN_LOG" ]; then @@ -287,6 +333,36 @@ elif [ "$MODEL_REPO" == "Llama-3.1-70B-proxy" ]; then echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$GBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG fi +elif [ "$MODEL_REPO" == "Llama-3.1-405B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" # BF16 training only + export EXP=examples/megatron/configs/$CONFIG_DEVICE/llama3.1_405B-$DATATYPE-pretrain.yaml + + SEQ_LEN=8192 + if [[ -f "$EXP" ]]; then + MBS=$(grep -E '^\s*micro_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') + GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') + echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + fi + if [[ "$DATATYPE" == "FP8" ]]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO. Only BF16 is supported." + elif [[ ! -f "$EXP" ]]; then + echo "Error: Config file not found: $EXP" + echo "Hint: add llama3.1_405B-$DATATYPE-pretrain.yaml in Primus configs." + else + GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") + run_primus "$EXP" $GBS_OVERRIDE + fi + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --global_batch_size $GBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,global_batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$GBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$GBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + elif [ "$MODEL_REPO" == "Llama-3.3-70B" ]; then echo "[INFO] $MODEL_REPO TRAINING" export EXP=examples/megatron/configs/$CONFIG_DEVICE/llama3.3_70B-$DATATYPE-pretrain.yaml @@ -455,6 +531,15 @@ elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then SEQ_LEN=4096 # Set MBS/GBS through command line for proxy models echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + # DeepEP (Primus-Turbo) is opt-in via PRIMUS_USE_DEEPEP=1 and applies to every + # supported device/precision branch below. DeepEP is incompatible with + # moe_shared_expert_overlap (Primus ROCm validate_args asserts it), so disable + # that overlap when enabling DeepEP. Default (unset/0) keeps alltoall over RCCL. + DEEPEP_ARGS="" + if [[ "${PRIMUS_USE_DEEPEP:-0}" == "1" ]]; then + DEEPEP_ARGS="--use_turbo_deepep true --moe_shared_expert_overlap false" + echo "[INFO] DeepEP enabled via PRIMUS_USE_DEEPEP=1 -> ${DEEPEP_ARGS}" + fi if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [ "$DATATYPE" == "FP8" ]; then echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." @@ -478,7 +563,7 @@ elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then --recompute_layer_ids null \ --overlap_grad_reduce false \ --overlap_param_gather false \ - --gradient_accumulation_fusion false 2>&1 | tee $TRAIN_LOG + --gradient_accumulation_fusion false $DEEPEP_ARGS 2>&1 | tee $TRAIN_LOG fi elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then if [ "$DATATYPE" == "FP8" ]; then @@ -489,7 +574,7 @@ elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then bash runner/primus-cli direct \ --log_file /tmp/primus_$MODEL_REPO.log \ -- train pretrain \ - --config $EXP --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS 2>&1 | tee $TRAIN_LOG + --config $EXP --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS $DEEPEP_ARGS 2>&1 | tee $TRAIN_LOG fi fi if [ -f "$TRAIN_LOG" ]; then diff --git a/scripts/primus/megatron-lm/run.sh b/scripts/primus/megatron-lm/run.sh index 6d8267f..afe4254 100755 --- a/scripts/primus/megatron-lm/run.sh +++ b/scripts/primus/megatron-lm/run.sh @@ -40,6 +40,8 @@ if [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-8b" ]]; then model="Llama-3.1-8B" elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b" ]]; then model="Llama-3.1-70B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-405b" ]]; then + model="Llama-3.1-405B" elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b-proxy" ]]; then model="Llama-3.1-70B-proxy" elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.3-70b" ]]; then @@ -114,8 +116,17 @@ export PATH="${ROCM_BIN}:${PATH}" # Detect device DEVICE=$("${ROCM_BIN}/rocminfo" | grep "AMD Instinct" | head -n1 | awk '{print $5}') +# Fall back to amd-smi when the rocminfo output format changes or rocminfo is +# unavailable (some newer ROCm/scaleout stacks). +if [[ -z "$DEVICE" && -x "${ROCM_BIN}/amd-smi" ]]; then + DEVICE=$("${ROCM_BIN}/amd-smi" 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') +fi if [ -z "$DEVICE" ]; then ARCH=$("${ROCM_BIN}/rocminfo" | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + # Fall back to the arch injected by madengine when rocminfo cannot report it. + if [[ -z "$ARCH" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then + ARCH=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') + fi case "$ARCH" in "gfx942") DEVICE="MI300X" ;; "gfx950") DEVICE="MI355X" ;; diff --git a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py deleted file mode 100644 index a514102..0000000 --- a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.py +++ /dev/null @@ -1,149 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2025 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -import pandas as pd -import argparse -import csv -import os -import re - -# parse arguments -parser = argparse.ArgumentParser(description='Convert primus pytorch train output format to MAD csv output format') -parser.add_argument("--mode", - type=str, - help="pretrain or finetune") -parser.add_argument("--model", - type=str, - help="model name") -parser.add_argument("--input", - type=str, - help="path to input file") -parser.add_argument("--output", - type=str, - help="path to output file") -parser.add_argument("--precision", - type=str, - help="precision type") -parser.add_argument("--batch_size", - type=str, - help="batch size") -parser.add_argument("--seq_len", - type=str, - help="sequence length") -parser.add_argument("--device", - type=str, - help="device name") -parser.add_argument("--num_gpus", - type=str, - help="number of GPUs") - -# read arguments -args = parser.parse_args() -input_file = args.input -output_file = args.output -print("Input file path: ", input_file) -print("Output file path: ", output_file) - -data = [] - -def find_match(file, search_string): - with open(file, 'r') as file: - content = file.read() - # Updated pattern to match the new log format - # Looks for patterns like "TFLOP/s/GPU): 322.5/320.6" or "tokens/s/GPU): 725.1/720.6" - # and extracts the first value (before the slash) - pattern = fr"{re.escape(search_string)}\):\s*(\d+\.?\d*)/" - matches = re.findall(pattern, content) - if matches: - # Return the last 2 values if they exist (one from each run) - if len(matches) >= 2: - result = [matches[-2], matches[-1]] - print(f"Found {len(matches)} matches for '{search_string}', using last 2: {result}") - else: - result = [matches[-1]] - print(f"Found {len(matches)} match for '{search_string}': {result}") - return result - else: - print(f"Warning: No matches found for '{search_string}' pattern") - return [] - -def find_match_running_avg(file, search_string): - """Return the running-average value (the number AFTER the slash) from the - LAST logged iteration, e.g. '1885.6' from 'tokens/s/GPU): 2355.1/1885.6'. - - Primus logs throughput as 'instant/running_avg'. The instantaneous value - has large run-to-run jitter (MoE routing imbalance, collective stragglers, - kernel warmup), so the trailing running average is a more stable figure. - Returns None when no match is present.""" - with open(file, 'r') as f: - content = f.read() - pattern = fr"{re.escape(search_string)}\):\s*\d+\.?\d*/(\d+\.?\d*)" - matches = re.findall(pattern, content) - if matches: - print(f"Found {len(matches)} running-avg matches for '{search_string}', using last: {matches[-1]}") - return matches[-1] - return None - -if args.model == "Llama-3.1-8B" or args.model == "Llama-3.1-70B" or \ - args.model == "Llama-3.1-405B" or \ - args.model == "Llama-2-7B" or args.model == "Llama-2-70B" or \ - args.model == "Mixtral-8x7B" or args.model == "Mixtral-8x22B-proxy" or \ - args.model == "DeepSeek-V2-lite" or args.model == "DeepSeek-V3-proxy" or \ - args.model == "Llama-3.1-70B-proxy" or args.model == "Llama-3.3-70B" or \ - args.model == "Qwen2.5-7B" or args.model == "Qwen2.5-72B": - tok_per_s_per_gpu_list = find_match(input_file, "tokens/s/GPU") - TFLOPS_per_gpu_list = find_match(input_file, "TFLOP/s/GPU") - - data = [] - # Write separate rows for each run - for i, (tps, tflops) in enumerate(zip(tok_per_s_per_gpu_list, TFLOPS_per_gpu_list)): - run_label = f"run_{i+1}" if len(tok_per_s_per_gpu_list) > 1 else "" - data.extend([ - {'model': args.model, 'performance': tps, 'metric': 'tok_per_s_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': run_label}, - {'model': args.model, 'performance': tflops, 'metric': 'TFLOPS_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': run_label} - ]) - - # Additionally emit the trailing running-average throughput (value after the - # slash in the last logged iteration). Kept as separate, distinctly-named - # rows (suffix '_avg', run='avg') so the instantaneous metrics above are - # preserved and never confused with the steady-state average. - # Applied to all models in this branch; the average is a steadier figure - # given per-iteration throughput jitter (largest for gated_delta_net MoE). - tok_avg = find_match_running_avg(input_file, "tokens/s/GPU") - tflops_avg = find_match_running_avg(input_file, "TFLOP/s/GPU") - if tok_avg is not None: - data.append({'model': args.model, 'performance': tok_avg, 'metric': 'tok_per_s_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) - if tflops_avg is not None: - data.append({'model': args.model, 'performance': tflops_avg, 'metric': 'TFLOPS_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) - -with open(output_file, mode='w', newline='') as file: - print("Preparing to write performance data...") - print("Data: ", data) - writer = csv.DictWriter(file, fieldnames=['model','performance','metric','mode','precision','batch_size','seq_len','device','num_gpus','run']) - writer.writeheader() - writer.writerows(data) - print("Completed writing to output file") - diff --git a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh deleted file mode 100755 index f03df53..0000000 --- a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh +++ /dev/null @@ -1,610 +0,0 @@ -#!/bin/bash -############################################################################### -# -# MIT License -# -# Copyright (c) 2025 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# -## Usage: -#./primus_megatron-lm_benchmark_report.sh -m $model_name -## example: -## Pretrain Llama 3.1 70B -#./primus_megatron-lm_benchmark_report.sh -m Llama-3.1-70B -p BF16 -## Pretrain Llama 2 7B -#./primus_megatron-lm_benchmark_report.sh -m Llama-2-7B -p FP8 -## Pretrain DeepSeek V2 Lite -#./primus_megatron-lm_benchmark_report.sh -m DeepSeek-V2-lite -## Pretrain DeepSeek V3 -#./primus_megatron-lm_benchmark_report.sh -m DeepSeek-V3 -## Pretrain Mixtral 8x7B -#./primus_megatron-lm_benchmark_report.sh -m Mixtral-8x7B -l 4 -MODEL_REPO="" -MODE="pretrain" -DATATYPE="BF16" - -usage() { - echo "Usage: $0 -m -p -l " - echo "\nOptions:" - echo " -m Model repository (Llama-2-7B, Llama-2-70B, Llama-3.1-8B, Llama-3.1-70B, Llama-3.1-405B, DeepSeek-V2-lite, DeepSeek-V3-proxy, Mixtral-8x7B, Mixtral-8x22B-proxy)" - echo " -p Precision type (FP8 or BF16)" - exit 1 -} - -# Parse command-line arguments -while getopts "m:p:" opt; do - case "$opt" in - m) MODEL_REPO="$OPTARG" ;; - p) DATATYPE="$OPTARG" ;; - *) usage ;; - esac -done - -echo "=hyper params start=" -echo $MODEL_REPO -echo $DATATYPE -echo "=hyper params end=" - -# Validate inputs -if [[ -z "$MODEL_REPO" ]]; then - echo "Error: Missing required arguments." - usage -fi - -if [[ "$DATATYPE" != "FP8" && "$DATATYPE" != "BF16" ]]; then - echo "Error: Datatype must be either FP8 or BF16." - exit 1 -fi - -# Run benchmark (Placeholder for actual script execution) -echo "Running primus training benchmark with the following parameters:" -echo " Model Repository: $MODEL_REPO" - -# config environment -export MOCK_DATA=1 - -# set performance output paths -TRAIN_LOG="$(pwd)/primus-megatron-$MODEL_REPO-pretrain.csv" -echo "TRAIN LOG: $TRAIN_LOG" - -PERF_LOG="$(pwd)/../perf_primus-megatron-$MODEL_REPO.csv" -echo "PERF LOG: $PERF_LOG" -ls $(pwd) -perf_script="$(pwd)/primus_megatron-lm_benchmark_report.py" - -# Detect device. Newer ROCm stacks may format rocminfo differently, -# so keep rocminfo first and fall back to amd-smi and the injected arch. -detect_device() { - local device="" - local arch="" - - if [[ -x /opt/rocm/bin/rocminfo ]]; then - device=$(/opt/rocm/bin/rocminfo 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') - arch=$(/opt/rocm/bin/rocminfo 2>/dev/null | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') - fi - - if [[ -z "$device" && -x /opt/rocm/bin/amd-smi ]]; then - device=$(/opt/rocm/bin/amd-smi 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') - fi - - if [[ -z "$arch" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then - arch=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') - fi - - case "$device" in - MI300X|MI325X|MI350X|MI355X) - ;; - *) - case "$arch" in - gfx942) device="MI300X" ;; - gfx950) device="MI355X" ;; - *) device="" ;; - esac - ;; - esac - - echo "$device" -} - -DEVICE=$(detect_device) -echo "DEVICE found: $DEVICE" -echo "GPU DEVICE name: $DEVICE" -if [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1 - export NVTE_CK_IS_V3_ATOMIC_FP32=1 -fi - -# Set common environment variables -# Keep compatibility with SLURM/madengine distributed env and only fallback to single-node defaults. -NNODES="${NNODES:-1}" -export NNODES -export CPUS_PER_TASK="${CPUS_PER_TASK:-128}" -export HSA_NO_SCRATCH_RECLAIM=1 -export NVTE_CK_USES_BWD_V3=1 -GPUS_PER_NODE="${GPUS_PER_NODE:-${NPROC_PER_NODE:-8}}" -if ! [[ "$NNODES" =~ ^[0-9]+$ ]]; then NNODES=1; fi -if ! [[ "$GPUS_PER_NODE" =~ ^[0-9]+$ ]]; then GPUS_PER_NODE=8; fi -NUM_GPUS=$((NNODES * GPUS_PER_NODE)) -echo "[INFO] Distributed params: NNODES=$NNODES GPUS_PER_NODE=$GPUS_PER_NODE NUM_GPUS=$NUM_GPUS" - -normalize_global_batch_size() { - local mbs="$1" - local gbs="$2" - local dp="$3" - local unit=$((mbs * dp)) - if (( unit <= 0 )); then - echo "$gbs" - return - fi - if (( gbs % unit == 0 )); then - echo "$gbs" - return - fi - local adjusted=$(( ((gbs + unit - 1) / unit) * unit )) - echo "$adjusted" -} - -# Launch a single Primus pretrain run. Every rank must execute an identical -# command so the model shape stays consistent across nodes; a mismatch makes -# multi-node ranks rendezvous with different shapes and deadlocks NCCL. Pass -# the config first, then any shape/proxy/batch overrides -- they reach every -# rank because all ranks run this exact command. -run_primus() { - local config="$1"; shift - bash runner/primus-cli direct \ - --log_file "/tmp/primus_$MODEL_REPO.log" \ - -- train pretrain \ - --config "$config" "$@" 2>&1 | tee "$TRAIN_LOG" -} - -cd /workspace/Primus - -# run models -if [ "$MODEL_REPO" == "Llama-3.1-8B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/llama3.1_8B-$DATATYPE-pretrain.yaml - SEQ_LEN=8192 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - MBS=4 - GBS=512 - run_primus "$EXP" - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - MBS=2 - GBS=128 - run_primus "$EXP" - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Llama-3.1-70B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" # BF16 training only - export EXP=examples/megatron/configs/$DEVICE/llama3.1_70B-$DATATYPE-pretrain.yaml - SEQ_LEN=8192 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - MBS=3 - GBS=24 - original_gbs=$GBS - GBS=$(normalize_global_batch_size "$MBS" "$GBS" "$NUM_GPUS") - if [[ "$GBS" != "$original_gbs" ]]; then - echo "[INFO] Adjusted global batch size for distributed run: ${original_gbs} -> ${GBS} (MBS=${MBS}, NUM_GPUS=${NUM_GPUS})" - fi - run_primus "$EXP" --micro_batch_size $MBS --global_batch_size $GBS - elif [ "$DATATYPE" == "BF16" ]; then - MBS=4 - GBS=32 - original_gbs=$GBS - GBS=$(normalize_global_batch_size "$MBS" "$GBS" "$NUM_GPUS") - if [[ "$GBS" != "$original_gbs" ]]; then - echo "[INFO] Adjusted global batch size for distributed run: ${original_gbs} -> ${GBS} (MBS=${MBS}, NUM_GPUS=${NUM_GPUS})" - fi - run_primus "$EXP" --micro_batch_size $MBS --global_batch_size $GBS - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=3 - GBS=24 - original_gbs=$GBS - GBS=$(normalize_global_batch_size "$MBS" "$GBS" "$NUM_GPUS") - if [[ "$GBS" != "$original_gbs" ]]; then - echo "[INFO] Adjusted global batch size for distributed run: ${original_gbs} -> ${GBS} (MBS=${MBS}, NUM_GPUS=${NUM_GPUS})" - fi - run_primus "$EXP" --micro_batch_size $MBS --global_batch_size $GBS - fi - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Llama-3.1-405B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/llama3.1_405B-$DATATYPE-pretrain.yaml - SEQ_LEN=8192 - MBS=1 - GBS=8 - if [[ "$DATATYPE" == "FP8" ]]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO. Only BF16 is supported." - elif [[ ! -f "$EXP" ]]; then - echo "Error: Config file not found: $EXP" - echo "Hint: add llama3.1_405B-$DATATYPE-pretrain.yaml in Primus configs." - else - run_primus "$EXP" - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Llama-3.1-70B-proxy" ]; then - echo "[INFO] $MODEL_REPO TRAINING" # FP8 training only - export EXP=examples/megatron/configs/$DEVICE/llama3.1_70B-$DATATYPE-pretrain.yaml - SEQ_LEN=8192 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - echo "Error: Llama-3.1-70B-proxy model is not supported on $DEVICE. To train use the full Llama-3.1-70B model." - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - MBS=3 - GBS=24 - run_primus "$EXP" --num_layers 40 --fp8 hybrid --micro_batch_size $MBS --global_batch_size $GBS --no_fp8_weight_transpose_cache true - elif [ "$DATATYPE" == "BF16" ]; then - echo "Error: Datatype BF16 is not supported for $MODEL_REPO on $DEVICE. Only FP8 is supported." - fi - fi - - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Llama-3.3-70B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/llama3.3_70B-$DATATYPE-pretrain.yaml - SEQ_LEN=8192 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=6 - GBS=48 - run_primus "$EXP" - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=2 - GBS=16 - run_primus "$EXP" - fi - fi - - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Llama-2-7B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/llama2_7B-$DATATYPE-pretrain.yaml - SEQ_LEN=4096 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - export MBS=13 - export GBS=416 - run_primus "$EXP" - elif [ "$DATATYPE" == "BF16" ]; then - export MBS=10 - export GBS=640 - run_primus "$EXP" - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - export MBS=4 - export GBS=256 - if [ "$DATATYPE" == "FP8" ]; then - run_primus "$EXP" - elif [ "$DATATYPE" == "BF16" ]; then - run_primus "$EXP" - fi - fi - - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Llama-2-70B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/llama2_70B-$DATATYPE-pretrain.yaml - SEQ_LEN=4096 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=17 - GBS=272 - run_primus "$EXP" - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=7 - GBS=56 - run_primus "$EXP" - fi - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "DeepSeek-V2-lite" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/deepseek_v2_lite-$DATATYPE-pretrain.yaml - SEQ_LEN=4096 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - MBS=12 - GBS=768 - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - run_primus "$EXP" - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=4 - GBS=640 - run_primus "$EXP" - fi - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/deepseek_v3-$DATATYPE-pretrain.yaml - SEQ_LEN=4096 - # DeepEP (Primus-Turbo) is opt-in via PRIMUS_USE_DEEPEP=1 and applies to every - # supported device/precision branch below. DeepEP is incompatible with - # moe_shared_expert_overlap (Primus ROCm validate_args asserts it), so disable - # that overlap when enabling DeepEP. Default (unset/0) keeps alltoall over RCCL. - DEEPEP_ARGS="" - if [[ "${PRIMUS_USE_DEEPEP:-0}" == "1" ]]; then - DEEPEP_ARGS="--use_turbo_deepep true --moe_shared_expert_overlap false" - echo "[INFO] DeepEP enabled via PRIMUS_USE_DEEPEP=1 -> ${DEEPEP_ARGS}" - fi - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=8 - GBS=64 - run_primus "$EXP" --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS $DEEPEP_ARGS - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=3 - GBS=192 - run_primus "$EXP" --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS $DEEPEP_ARGS - fi - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Mixtral-8x7B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/mixtral_8x7B_v0.1-$DATATYPE-pretrain.yaml - SEQ_LEN=4096 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=4 - GBS=256 - run_primus "$EXP" - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=2 - GBS=32 - run_primus "$EXP" - fi - fi - - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Mixtral-8x22B-proxy" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - LAYERS=4 # default proxy model uses 4 layers - echo "[INFO] Proxy model uses $LAYERS layers" - export EXP=examples/megatron/configs/$DEVICE/mixtral_8x22B_v0.1-$DATATYPE-pretrain.yaml - SEQ_LEN=8192 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=2 - GBS=16 - run_primus "$EXP" --num_layers 4 --pipeline_model_parallel_size 1 --micro_batch_size $MBS --global_batch_size $GBS - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=1 - GBS=16 - run_primus "$EXP" --num_layers 4 --pipeline_model_parallel_size 1 --micro_batch_size $MBS --global_batch_size $GBS - fi - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -elif [ "$MODEL_REPO" == "Qwen2.5-7B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/qwen2.5_7B-$DATATYPE-pretrain.yaml - SEQ_LEN=2048 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [[ "$DATATYPE" == "BF16" ]]; then - MBS=16 - GBS=768 - run_primus "$EXP" - elif [[ "$DATATYPE" == "FP8" ]]; then - MBS=20 - GBS=800 - run_primus "$EXP" - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - MBS=10 - GBS=640 - if [ "$DATATYPE" == "FP8" ]; then - run_primus "$EXP" - elif [ "$DATATYPE" == "BF16" ]; then - run_primus "$EXP" - fi - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - fi - -elif [ "$MODEL_REPO" == "Qwen2.5-72B" ]; then - echo "[INFO] $MODEL_REPO TRAINING" - export EXP=examples/megatron/configs/$DEVICE/qwen2.5_72B-$DATATYPE-pretrain.yaml - SEQ_LEN=2048 - if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=16 - GBS=256 - run_primus "$EXP" - fi - elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - if [ "$DATATYPE" == "FP8" ]; then - echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." - elif [ "$DATATYPE" == "BF16" ]; then - MBS=4 - GBS=32 - run_primus "$EXP" - fi - fi - if [ -f "$TRAIN_LOG" ]; then - echo "[INFO] Benchmarking" - python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS - rm $TRAIN_LOG - else - echo "[INFO] Training log not found - configuration not supported." - echo "model,performance,metric,mode,precision,batch_size,seq_len,device,num_gpus" > $PERF_LOG - echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG - fi - -else - echo "Error: Unsupported training mode." - exit 1 -fi diff --git a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh b/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh deleted file mode 100755 index 17c18d3..0000000 --- a/scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_setup.sh +++ /dev/null @@ -1,41 +0,0 @@ -#!/bin/bash -############################################################################### -# -# MIT License -# -# Copyright (c) 2025 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# -export HF_HOME=/workspace/huggingface - -# Parse named arguments -while [[ "$#" -gt 0 ]]; do - case $1 in - -m) MODEL_NAME="$2"; shift ;; - *) echo "Unknown parameter passed: $1"; usage ;; - esac - shift -done - -echo "[INFO] Primus setup script starting in directory $(pwd)" - -cd /workspace/Primus -git pull \ No newline at end of file diff --git a/scripts/primus_scaleout/megatron-lm/run.sh b/scripts/primus_scaleout/megatron-lm/run.sh deleted file mode 100755 index 6d040ca..0000000 --- a/scripts/primus_scaleout/megatron-lm/run.sh +++ /dev/null @@ -1,143 +0,0 @@ -#!/bin/bash -############################################################################### -# -# MIT License -# -# Copyright (c) 2025 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export HF_TOKEN=$MAD_SECRETS_HFTOKEN - -# Parse named arguments -while [[ "$#" -gt 0 ]]; do - case $1 in - --model_repo) MODEL_REPO="$2"; shift ;; - *) echo "Unknown parameter passed: $1"; usage ;; - esac - shift -done - -if [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-8b" ]]; then - model="Llama-3.1-8B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b" ]]; then - model="Llama-3.1-70B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-405b" ]]; then - model="Llama-3.1-405B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b-proxy" ]]; then - model="Llama-3.1-70B-proxy" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.3-70b" ]]; then - model="Llama-3.3-70B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-2-7b" ]]; then - model="Llama-2-7B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-2-70b" ]]; then - model="Llama-2-70B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_deepseek-v2-lite-16b" ]]; then - model="DeepSeek-V2-lite" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_deepseek-v2" ]]; then - model="DeepSeek-V2" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_deepseek-v3-proxy" ]]; then - model="DeepSeek-V3-proxy" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_mixtral-8x7b" ]]; then - model="Mixtral-8x7B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_mixtral-8x22b-proxy" ]]; then - model="Mixtral-8x22B-proxy" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_qwen2.5-7b" ]]; then - model="Qwen2.5-7B" -elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_qwen2.5-72b" ]]; then - model="Qwen2.5-72B" -fi - -# Run primus pytorch setup script -echo "Running setup script to download tokenizers" -bash ./primus_megatron-lm_benchmark_setup.sh -m $model - -# Detect device. Newer ROCm stacks may format rocminfo differently, -# so keep rocminfo first and fall back to amd-smi and the injected arch. -detect_device() { - local device="" - local arch="" - - if [[ -x /opt/rocm/bin/rocminfo ]]; then - device=$(/opt/rocm/bin/rocminfo 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') - arch=$(/opt/rocm/bin/rocminfo 2>/dev/null | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') - fi - - if [[ -z "$device" && -x /opt/rocm/bin/amd-smi ]]; then - device=$(/opt/rocm/bin/amd-smi 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') - fi - - if [[ -z "$arch" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then - arch=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') - fi - - case "$device" in - MI300X|MI325X|MI350X|MI355X) - ;; - *) - case "$arch" in - gfx942) device="MI300X" ;; - gfx950) device="MI355X" ;; - *) device="" ;; - esac - ;; - esac - - echo "$device" -} - -DEVICE=$(detect_device) -echo "GPU DEVICE name: $DEVICE" - -# Define supported datatypes based on device and model -datatypes=() - -if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then - # MI355X/MI350X support - if [[ "$model" == "Llama-3.1-70B-proxy" ]]; then - echo "Skipping $model - Not supported on $DEVICE" - elif [[ "$model" == "Llama-3.1-8B" || "$model" == "Llama-3.1-70B" || "$model" == "Llama-2-7B" || "$model" == "Qwen2.5-7B" ]]; then - datatypes=("BF16" "FP8") - else - # Most other models only support BF16 on MI355X/MI350X - datatypes=("BF16") - fi -elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then - # MI300X/MI325X support - if [[ "$model" == "Llama-3.1-70B-proxy" ]]; then - datatypes=("FP8") # Only FP8 supported - elif [[ "$model" == "Llama-3.1-8B" || "$model" == "Llama-2-7B" || "$model" == "Qwen2.5-7B" ]]; then - datatypes=("BF16" "FP8") # Both supported - else - # Most large models only support BF16 on MI300X/MI325X - datatypes=("BF16") - fi -else - # Unknown device, try both - datatypes=("BF16" "FP8") -fi - -# datatypes=("FP8") -# Loop through supported combinations -for datatype in "${datatypes[@]}"; do - echo "Running: $model - $datatype" - ./primus_megatron-lm_benchmark_report.sh -m $model -p $datatype -done From f90526aeaf5e39570539d86aa8d0625b57f1c917 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Mon, 6 Jul 2026 11:23:31 +0000 Subject: [PATCH 06/25] sglang_disagg_server: co-locate router with prefill0 (Variant A) Align the Mooncake/NIXL launcher topology with the MoRI EP launcher (sglang_disagg_mori_io_ep.sh): the router/proxy now runs on NODE_RANK=0 alongside the first prefill server instead of requiring a dedicated proxy node. Total nodes = xP + yD (was 1 + xP + yD), so node counting is consistent across both KV-transfer backends. - IP_ARRAY[0..xP-1] are prefill nodes, IP_ARRAY[xP..xP+yD-1] are decode. - Backend sglang servers listen on SERVER_PORT (3000); the router listens on ROUTER_PORT (2322) so prefill0 and the router can share NODE_RANK=0. - Proxy readiness loops and the socket barrier are reindexed accordingly. - Add an explicit out-of-range guard for NODE_RANK. --- scripts/sglang_disagg/sglang_disagg_server.sh | 71 ++++++++++++++----- 1 file changed, 54 insertions(+), 17 deletions(-) diff --git a/scripts/sglang_disagg/sglang_disagg_server.sh b/scripts/sglang_disagg/sglang_disagg_server.sh index 3db53d7..f21df58 100755 --- a/scripts/sglang_disagg/sglang_disagg_server.sh +++ b/scripts/sglang_disagg/sglang_disagg_server.sh @@ -114,15 +114,25 @@ python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ IFS=',' read -ra IP_ARRAY <<< "$IPADDRS" +# Colocated topology (matches sglang_disagg_mori_io_ep.sh): the router/proxy runs +# on NODE_RANK=0 alongside the first prefill server, so no dedicated proxy node is +# needed. Total nodes = xP + yD. +# IP_ARRAY[0..xP-1] -> prefill nodes (NODE_RANK 0..xP-1) +# IP_ARRAY[xP..xP+yD-1] -> decode nodes (NODE_RANK xP..xP+yD-1) +# Backend sglang servers listen on SERVER_PORT; the router listens on ROUTER_PORT. +# They must differ so prefill0 and the router can share NODE_RANK=0. +SERVER_PORT="${SERVER_PORT:-3000}" +ROUTER_PORT="${ROUTER_PORT:-2322}" + PREFILL_ARGS="" DECODE_ARGS="" -for ((i=1; i<=$xP && i<${#IP_ARRAY[@]}; i++)); do - PREFILL_ARGS+=" --prefill http://${IP_ARRAY[$i]}:2322 " +for ((i=0; i<$xP && i<${#IP_ARRAY[@]}; i++)); do + PREFILL_ARGS+=" --prefill http://${IP_ARRAY[$i]}:${SERVER_PORT} " done -for ((i=$xP+1; i<${#IP_ARRAY[@]}; i++)); do - DECODE_ARGS+=" --decode http://${IP_ARRAY[$i]}:2322 " +for ((i=$xP; i<${#IP_ARRAY[@]}; i++)); do + DECODE_ARGS+=" --decode http://${IP_ARRAY[$i]}:${SERVER_PORT} " done # ============================================================================= @@ -139,10 +149,31 @@ if [ "$NODE_RANK" -eq 0 ]; then echo "CLUSTER INFO ====================================" echo "================================================" - echo "${host_name}:${host_ip} is Proxy Node" + echo "${host_name}:${host_ip} is Prefill Node 0 + Proxy/Router (co-located)" echo "${PREFILL_ARGS} are Proxy's Prefill" echo "${DECODE_ARGS} are Proxy's Decode" echo "================================================" + + # Launch the first prefill server on this node, co-located with the router. + PREFILL_CMD="MC_TE_METRIC=true python3 -m sglang.launch_server \ + --model-path $MODEL_PATH \ + --disaggregation-mode prefill \ + --disaggregation-ib-device ${IBDEVICES} \ + --host ${host_ip} \ + --port ${SERVER_PORT} \ + --stream-output \ + --trust-remote-code \ + --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" + + if [[ -n "$PREFILL_MODEL_CONFIG" ]]; then + PREFILL_CMD="$PREFILL_CMD $PREFILL_MODEL_CONFIG" + fi + + eval "$PREFILL_CMD" \ + 2>&1 | tee /run_logs/${SLURM_JOB_ID}/prefill_NODE${NODE_RANK}.log >/dev/null & + + node0_prefill_pid=$! + echo "Proxy server is waiting for prefill and decode nodes to be ready ..." \ | tee /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null sleep 20 @@ -150,7 +181,7 @@ if [ "$NODE_RANK" -eq 0 ]; then SLEEP_SECONDS=10 SEARCH_SIGNAL="The server is fired up and ready to roll!" SECONDS=0 - for ((i=1; i<=$xP && i<${#IP_ARRAY[@]}; i++)); do + for ((i=0; i<$xP && i<${#IP_ARRAY[@]}; i++)); do LOG_FILE=/run_logs/${SLURM_JOB_ID}/prefill_NODE${i}.log #wait until prefill nodes get ready until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do @@ -164,7 +195,7 @@ if [ "$NODE_RANK" -eq 0 ]; then done done - for ((i=$xP+1; i<${#IP_ARRAY[@]}; i++)); do + for ((i=$xP; i<${#IP_ARRAY[@]}; i++)); do LOG_FILE=/run_logs/${SLURM_JOB_ID}/decode_NODE${i}.log #wait until decode nodes get ready until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do @@ -185,7 +216,7 @@ if [ "$NODE_RANK" -eq 0 ]; then ${PREFILL_ARGS} \ ${DECODE_ARGS} \ --host 0.0.0.0 \ - --port 2322 \ + --port ${ROUTER_PORT} \ 2>&1 | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null & proxy_pid=$! @@ -193,7 +224,7 @@ if [ "$NODE_RANK" -eq 0 ]; then echo "Waiting for all prefill and decode servers to be up . . ." python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ --node-ips ${IPADDRS} \ - --node-ports 2322 + --node-ports ${SERVER_PORT} echo "Proxy Server Ready for benchmarking on ${host_name}:${host_ip}" @@ -204,7 +235,10 @@ if [ "$NODE_RANK" -eq 0 ]; then echo "Killing the proxy server" kill $proxy_pid -elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -le "$xP" ]; then + echo "Killing the co-located prefill server" + kill $node0_prefill_pid + +elif [ "$NODE_RANK" -ge 1 ] && [ "$NODE_RANK" -lt "$xP" ]; then echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-'default'})" echo "Using prefill config: $PREFILL_MODEL_CONFIG" echo "Using KV transfer backend: ${KV_TRANSFER_BACKEND}" @@ -214,7 +248,7 @@ elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -le "$xP" ]; then --disaggregation-mode prefill \ --disaggregation-ib-device ${IBDEVICES} \ --host ${host_ip} \ - --port 2322 \ + --port ${SERVER_PORT} \ --stream-output \ --trust-remote-code \ --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" @@ -231,17 +265,17 @@ elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -le "$xP" ]; then echo "Waiting for proxy server to be up..." python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ --node-ips ${MASTER_ADDR} \ - --node-ports 2322 + --node-ports ${ROUTER_PORT} echo "Waiting until proxy server closes..." python $MOONCAKE_COOKBOOK_PATH/socket_wait.py \ --remote-ip ${MASTER_ADDR} \ - --remote-port 2322 + --remote-port ${ROUTER_PORT} echo "Killing the prefill server" kill $prefill_pid -else +elif [ "$NODE_RANK" -ge "$xP" ] && [ "$NODE_RANK" -le "$((xP + yD - 1))" ]; then echo "${host_name}:${host_ip} is Decode Node (Model: ${MODEL_NAME:-'default'})" echo "Using decode config: $DECODE_MODEL_CONFIG" echo "Using KV transfer backend: ${KV_TRANSFER_BACKEND}" @@ -251,7 +285,7 @@ else --disaggregation-mode decode \ --disaggregation-ib-device ${IBDEVICES} \ --host ${host_ip} \ - --port 2322 \ + --port ${SERVER_PORT} \ --stream-output \ --trust-remote-code \ --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" @@ -268,16 +302,19 @@ else echo "Waiting for proxy server to be up..." python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ --node-ips ${MASTER_ADDR} \ - --node-ports 2322 + --node-ports ${ROUTER_PORT} echo "Waiting until proxy server closes..." python $MOONCAKE_COOKBOOK_PATH/socket_wait.py \ --remote-ip ${MASTER_ADDR} \ - --remote-port 2322 + --remote-port ${ROUTER_PORT} echo "Killing the decode server" kill $decode_pid +else + echo "ERROR: NODE_RANK=${NODE_RANK} out of range (expected 0..$((xP + yD - 1))) for xP=${xP} yD=${yD}" >&2 + exit 1 fi # ============================================================================= From bd5dc07b6eb193676497b2291bc44f857a1a33b1 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Mon, 6 Jul 2026 11:24:32 +0000 Subject: [PATCH 07/25] parse_to_csv: tag backend by RUN_MORI first, then KV backend _get_run_metadata checked DP_MODE before RUN_MORI, so a Mooncake/NIXL run (RUN_MORI=0) with DP_MODE=1 was mislabelled "mori_dp". Decide on RUN_MORI first (DP_MODE only distinguishes mori_dp vs mori_io within the MoRI path), and fall back to the actual KV_TRANSFER_BACKEND (mooncake/nixl) otherwise. Defaults for RUN_MORI/DP_MODE aligned with run.sh (both 1). --- scripts/sglang_disagg/parse_to_csv.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/scripts/sglang_disagg/parse_to_csv.py b/scripts/sglang_disagg/parse_to_csv.py index 951d2b5..d8b3074 100644 --- a/scripts/sglang_disagg/parse_to_csv.py +++ b/scripts/sglang_disagg/parse_to_csv.py @@ -163,15 +163,19 @@ def _get_run_metadata(pipeline: str = "sglang"): xP = os.environ.get("xP", "1") yD = os.environ.get("yD", "1") - dp_mode = os.environ.get("DP_MODE", "0") - run_mori = os.environ.get("RUN_MORI", "0") + dp_mode = os.environ.get("DP_MODE", "1") + run_mori = os.environ.get("RUN_MORI", "1") + kv_backend = os.environ.get("KV_TRANSFER_BACKEND", "mooncake").lower() gpus = os.environ.get("GPUS_PER_NODE", "8") - # Determine backend tag - if dp_mode == "1": - backend = "mori_dp" - elif run_mori == "1": - backend = "mori_io" + # Determine backend tag. RUN_MORI is the launcher's own switch and takes + # priority; DP_MODE only distinguishes mori_dp vs mori_io *within* the MoRI + # path. When MoRI is off, tag by the actual KV transfer backend so a + # Mooncake/NIXL run with DP_MODE=1 is not mislabelled as mori_dp. + if run_mori == "1": + backend = "mori_dp" if dp_mode == "1" else "mori_io" + elif kv_backend in ("mooncake", "nixl"): + backend = kv_backend else: backend = "mooncake" From a3c8bb55a70bffa84c98c5afe10fc5798be5bd6e Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Mon, 6 Jul 2026 11:37:58 +0000 Subject: [PATCH 08/25] sglang_disagg: require shared /run_logs; benchmark honors RUN_LOG_DIR R1: /run_logs must be a shared (NFS-backed) mount passed through the docker env because the disaggregated launchers write per-node readiness logs to /run_logs/${SLURM_JOB_ID} and the rank-0 proxy greps them across nodes. Disaggregated P/D always spans >=2 nodes (rank-0 prefill/proxy + at least one decode node), so a node-local /tmp fallback would silently break that cross-node rendezvous. run.sh now requires /run_logs and fails fast with an actionable message when it is missing or not writable (the previous silent /tmp fallback is removed). C3: benchmark_xPyD.sh no longer re-hardcodes /run_logs; it honors the RUN_LOG_DIR exported by run.sh (single source of truth), so its log lands in the same directory and the tee cannot fail the sweep. --- scripts/sglang_disagg/benchmark_xPyD.sh | 5 ++++- scripts/sglang_disagg/run.sh | 28 +++++++++++++++---------- 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/scripts/sglang_disagg/benchmark_xPyD.sh b/scripts/sglang_disagg/benchmark_xPyD.sh index 598b917..c28d18f 100644 --- a/scripts/sglang_disagg/benchmark_xPyD.sh +++ b/scripts/sglang_disagg/benchmark_xPyD.sh @@ -3,7 +3,10 @@ set -uo pipefail timestamp=$(date "+%Y%m%d_%H%M%S") RUN_LOG_JOB_ID="${SLURM_JOB_ID:-0}" -RUN_LOG_DIR="/run_logs/${RUN_LOG_JOB_ID}" +# Honor RUN_LOG_DIR exported by run.sh (single source of truth for the log +# location, including its /run_logs-vs-fallback decision). Only derive the +# default when this script is invoked standalone without run.sh in the env. +RUN_LOG_DIR="${RUN_LOG_DIR:-/run_logs/${RUN_LOG_JOB_ID}}" mkdir -p "$RUN_LOG_DIR" 2>/dev/null || true LOG="${RUN_LOG_DIR}/benchmark_${RUN_LOG_JOB_ID}_${timestamp}_xP${xP}_yD${yD}_${MODEL_NAME}" diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh index fe90926..32f96cb 100755 --- a/scripts/sglang_disagg/run.sh +++ b/scripts/sglang_disagg/run.sh @@ -50,19 +50,25 @@ export NODE_RANK="${NODE_RANK:-${SGLANG_NODE_RANK:-${SLURM_PROCID:-0}}}" # Persist run.sh output to the shared /run_logs mount so failures are # diagnosable from the submission node (madengine discards model stdout). # The dir is created from inside the container (root); make it world-writable -# so an ordinary user outside the container can clean it up afterwards. Fall -# back to /tmp when /run_logs is absent/not writable, otherwise the tee in the -# process substitution can't open its file and we lose the ERROR output we -# need for the post-mortem. +# so an ordinary user outside the container can clean it up afterwards. +# +# /run_logs MUST be a shared (NFS-backed) mount passed through the docker env: +# the disaggregated launchers write per-node readiness logs +# (prefill_NODE*/decode_NODE*/proxy_NODE*) to /run_logs/${SLURM_JOB_ID} and the +# rank-0 proxy greps them across nodes. Disaggregated P/D always spans >=2 nodes +# (rank-0 prefill/proxy + at least one decode node), so a node-local fallback +# would silently break that cross-node rendezvous. Require /run_logs and fail +# fast if it is missing or not writable. RUN_LOG_DIR="/run_logs/${SLURM_JOB_ID:-local}" -if mkdir -p "${RUN_LOG_DIR}" 2>/dev/null && [[ -w "${RUN_LOG_DIR}" ]]; then - chmod 0777 /run_logs 2>/dev/null || true - chmod -R 0777 "${RUN_LOG_DIR}" 2>/dev/null || true -else - RUN_LOG_DIR="/tmp/run_logs/${SLURM_JOB_ID:-local}" - mkdir -p "${RUN_LOG_DIR}" 2>/dev/null || true - chmod -R 0777 "${RUN_LOG_DIR}" 2>/dev/null || true +if ! mkdir -p "${RUN_LOG_DIR}" 2>/dev/null || [[ ! -w "${RUN_LOG_DIR}" ]]; then + echo "FATAL: /run_logs is not writable. It must be a shared (NFS-backed)" >&2 + echo " mount passed through the docker env: the disaggregated launchers" >&2 + echo " write per-node readiness logs to /run_logs/\${SLURM_JOB_ID} and the" >&2 + echo " rank-0 proxy greps them across nodes. Mount it and re-run." >&2 + exit 1 fi +chmod 0777 /run_logs 2>/dev/null || true +chmod -R 0777 "${RUN_LOG_DIR}" 2>/dev/null || true export RUN_LOG_DIR exec > >(tee -a "${RUN_LOG_DIR}/run_sh_rank${NODE_RANK}.log") 2>&1 export xP="${xP:-${SGLANG_DISAGG_PREFILL_NODES:-1}}" From ae1527feec99b8edb377675359f6ebca5d0b3918 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Mon, 6 Jul 2026 13:17:18 +0000 Subject: [PATCH 09/25] sglang_disagg manifest: /run_logs mount via ${SLURM_SUBMIT_DIR} (shared) The DeepSeek-R1 disagg template mounted /run_logs from a bare relative ./slurm_output/run_logs, mirroring the PR-174 review finding that a non-shared /run_logs breaks the cross-node readiness rendezvous (run.sh now requires /run_logs and fails fast otherwise, see commit a3c8bb5). A relative path resolves against whatever CWD the per-node madengine process happens to have, which is not guaranteed to be the same shared rundir on every node. Use ${SLURM_SUBMIT_DIR} instead of a manual placeholder or a literal $WORKDIR. SLURM sets SLURM_SUBMIT_DIR to the directory `sbatch` was invoked from; madengine's slurm deployment runs `sbatch` with no `cwd=` override (deployment/slurm.py), and the skill's Step 6 always launches from $WORKDIR/rundir, so SLURM_SUBMIT_DIR already equals the shared rundir with zero manual filling. It propagates end-to-end (job script -> srun, no --export restriction -> final `docker run`, since additional_docker_run_options is concatenated unquoted and console.sh() runs with env=None i.e. the full inherited environment) so it reliably expands at every node. $WORKDIR itself is only a doc-only convention; nothing in mad.env/madengine actually exports it as a real env var. Drop the redundant /run_logs entry from docker_mounts: it was always excluded at runtime (container_runner._extract_additional_mount_targets skips duplicate container targets already present in additional_docker_run_options), and docker_mounts values are rendered through shlex.quote(), which single-quotes the whole string and would have silently blocked ${SLURM_SUBMIT_DIR} expansion there anyway. Document the mechanism (and that run.sh's /run_logs check only verifies writability, not actual cross-node sharedness) in references/gotchas.md so it isn't rediscovered per-cluster. --- .../sglang_disagg_deepseek-r1.template.json | 5 ++-- .../mad-slurm-multinode/references/gotchas.md | 26 +++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json index c8b9167..b6ab3ca 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json @@ -26,7 +26,7 @@ "tags": ["pyt", "sglang", "disagg", "inference", "deepseek-r1"], "timeout": 21600, "args": "--model-name DeepSeek-R1 --model-path deepseek-ai/DeepSeek-R1-0528) and peers wait on MODEL_PATH/.stage_done> --run-mori 1 --dp-mode 1 --xp 2 --yd 2 --gpus-per-node 8 --kv-transfer-backend nixl", - "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 --ulimit nofile=1048576:1048576 -v /sys:/sys:ro -v /sys/class/infiniband:/sys/class/infiniband:ro -v /run/udev:/run/udev:ro -v : -v ./slurm_output/run_logs:/run_logs -e SLURM_JOB_ID" + "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 --ulimit nofile=1048576:1048576 -v /sys:/sys:ro -v /sys/class/infiniband:/sys/class/infiniband:ro -v /run/udev:/run/udev:ro -v : -v ${SLURM_SUBMIT_DIR}/slurm_output/run_logs:/run_logs -e SLURM_JOB_ID" } }, "context": { @@ -74,8 +74,7 @@ "docker_mounts": { "/dev/infiniband": "/dev/infiniband", "/sys/class/infiniband": "/sys/class/infiniband", - "": "", - "/run_logs": "./slurm_output/run_logs" + "": "" }, "docker_build_arg": {}, "gpu_vendor": "AMD", diff --git a/.claude/skills/mad-slurm-multinode/references/gotchas.md b/.claude/skills/mad-slurm-multinode/references/gotchas.md index 46f5a53..ae3ae1f 100644 --- a/.claude/skills/mad-slurm-multinode/references/gotchas.md +++ b/.claude/skills/mad-slurm-multinode/references/gotchas.md @@ -2,6 +2,8 @@ Keywords: source mad.env MODEL_DIR run-script, MAD_DOCKER_BUILDS shared storage, MAD_SECRETS_HFTOKEN HF 401 single-quoted, docker_mounts container_path host_path -v, +run_logs shared NFS mount, SLURM_SUBMIT_DIR sbatch cwd rundir, WORKDIR not +exported shlex.quote docker_mounts additional_docker_run_options, RCCL_AINIC_ROCE RDMAV_DRIVERS ionic, NCCL_IB_HCA mlx5 rdma per-cluster, perf.csv login-node aggregation, slurm.nodes distributed.nnodes nodelist world size, @@ -32,6 +34,30 @@ here; this file is read before a run. whenever a host data directory maps to a container-internal path that differs from the host path. See [references/manifests.md](manifests.md) ("docker_mounts direction"). +- **Prefer `${SLURM_SUBMIT_DIR}` over `$WORKDIR`/relative paths in + `additional_docker_run_options` host mounts — it is a real, always-set SLURM + var that already equals the shared `rundir`.** `$WORKDIR` is a doc-only + convention (nothing in mad.env/madengine actually exports it), and a + relative host path (e.g. `./slurm_output/run_logs`) resolves against + whatever CWD the per-node `madengine run` process happens to have. SLURM + itself sets `SLURM_SUBMIT_DIR` to the directory `sbatch` was invoked from + (`deployment/slurm.py` runs `sbatch` without a `cwd=` override, and Step 6 + always launches from `$WORKDIR/rundir`), and that value is inherited + end-to-end — through the generated job script, `srun` (no `--export` + restriction), and the final `docker run` (`additional_docker_run_options` + is concatenated unquoted and `console.sh()` runs with `env=None`, i.e. full + inherited environment) — so `${SLURM_SUBMIT_DIR}` reliably expands to + `rundir` at every node with zero manual filling. This only works for + `additional_docker_run_options`, NOT `docker_mounts`: madengine renders + `docker_mounts` values through `shlex.quote()`, which single-quotes the + whole string and blocks `$VAR`/`${VAR}` expansion outright — don't duplicate + a shared-path mount there. The `sglang_disagg_deepseek-r1` template's + `/run_logs` mount (`-v ${SLURM_SUBMIT_DIR}/slurm_output/run_logs:/run_logs`) + is the reference example: `scripts/sglang_disagg/run.sh` requires `/run_logs` + and fails fast (no `/tmp` fallback) if it is missing or not writable, but + only checks writability, not actual cross-node sharedness — a node-local dir + that happens to exist would pass that check and silently break the + cross-node readiness rendezvous (see the `sglang_disagg` section below). - **`NCCL_IB_HCA` is per-cluster, not portable.** CX7 uses `mlx5_*`; AINIC uses `rdma0..7`. Copying an mlx5 list onto an AINIC node (or vice versa) inits zero NICs. Verify on the node. From d0fa66958751cb91bc833b0dfcfc327957130b38 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Mon, 6 Jul 2026 13:58:01 +0000 Subject: [PATCH 10/25] Primus Llama-3.1-8B scaleout: v26.4 RCCL overlay sweep + NCCL_NET_PLUGIN Adapt the fixes from mkuznet1/MAD#1 to the current branch. - docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile: after the targeted /opt/rocm + torch/lib overwrites, sweep every non-symlink librccl.so* on disk and overwrite it with the candidate (+add-needed librocm_smi64), and tighten the final gate to assert EVERY on-disk librccl (not just 3 fixed paths) is RCCL 2.x with a baked git hash matching the built SHA. v26.4 splits ROCm libs into _rocm_sdk_libraries/lib (runtime) and _rocm_sdk_devel/lib (dev); torch maps librccl.so.1 from _rocm_sdk_libraries at runtime, which the targeted overwrites did not cover, so the stock base librccl silently ran and invalidated the overlay validation. - assets/manifests/primus_llama-3.1-8b.template.json: set NCCL_NET_PLUGIN=none in both env blocks. rocm/primus:v26.4 defaults NCCL_NET_PLUGIN=librccl-anp.so, which deadlocks RCCL init when the image carries a bundled RCCL overlay. - references/gotchas.md: document the NCCL_NET_PLUGIN pitfall, and the GBS normalization requirement (adapted to this branch: the GBS fix from the PR's scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh is already generalized here by scaleout_gbs_override in scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh after the primus_scaleout consolidation, so no per-branch code change is ported). Ref: https://github.com/mkuznet1/MAD/pull/1 --- .../primus_llama-3.1-8b.template.json | 2 ++ .../mad-slurm-multinode/references/gotchas.md | 30 ++++++++++++++++++- ...n_train_rccl_overlay.ubuntu.amd.Dockerfile | 20 +++++++++++-- 3 files changed, 48 insertions(+), 4 deletions(-) diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json index bc60a2d..84dec88 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json @@ -41,6 +41,7 @@ "RDMAV_DRIVERS": "", "IBV_DRIVERS": "", "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", "IBV_SHOW_WARNINGS": "1" }, "docker_mounts": { @@ -96,6 +97,7 @@ "RDMAV_DRIVERS": "", "IBV_DRIVERS": "", "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", "NCCL_TIMEOUT": "900", "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", "TORCH_NCCL_HIGH_PRIORITY": "1", diff --git a/.claude/skills/mad-slurm-multinode/references/gotchas.md b/.claude/skills/mad-slurm-multinode/references/gotchas.md index ae3ae1f..c2fa9a4 100644 --- a/.claude/skills/mad-slurm-multinode/references/gotchas.md +++ b/.claude/skills/mad-slurm-multinode/references/gotchas.md @@ -220,7 +220,7 @@ print_rank_last throughput last global rank, multi-node perf collection rank-0. MoRI. To switch to the Primus-Turbo **DeepEP** dispatcher, set `PRIMUS_USE_DEEPEP=1` in the manifest env (a primus scaleout manifest can carry it defaulting to `"0"`), which makes - `scripts/primus_scaleout/megatron-lm/primus_megatron-lm_benchmark_report.sh` + `scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh` (invoked by the model `run.sh`) pass `--use_turbo_deepep true` to Primus. Primus then (`_is_turbo_deepep_enabled`) auto-sets `moe_enable_deepep=True` and `moe_token_dispatcher_type='flex'` and swaps in @@ -247,3 +247,31 @@ print_rank_last throughput last global rank, multi-node perf collection rank-0. `skip_perf_collection` SLURM node as SUCCESS). Use a madengine that has this multi-node aggregation (present in ROCm/madengine `develop`); otherwise the run trains fine but is mis-reported as FAILED with an empty `perf.csv`. + +- **Per-model global batch size must be normalized to the world size, or Megatron + aborts at startup on non-standard node counts.** The config `global_batch_size` + is tuned for a single 8-GPU node; at any node count where + `GBS % (MBS * world_size) != 0` (e.g. a 3-node/24-GPU run: `512 % (4*24) != 0`) + Megatron aborts with `global batch size ... is not divisible by micro batch + size ... times data parallel size`. In our consolidated + `scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh` this is + handled generically for *every* model branch: `scaleout_gbs_override` reads + MBS/GBS from the run's config YAML and, when `NUM_GPUS > 8`, rounds GBS up to + the next multiple of `MBS * NUM_GPUS` (via `normalize_global_batch_size`) and + passes it as an explicit `--global_batch_size` override; single-node runs emit + no override so behavior there is byte-for-byte unchanged. Any new per-model + branch that hardcodes GBS should route it through `scaleout_gbs_override` to + stay node-count-portable. (Upstream fix for the 8B branch: + [mkuznet1/MAD#1](https://github.com/mkuznet1/MAD/pull/1) — our tree already + generalizes it, so no per-branch change is needed here.) + +- **`rocm/primus:v26.4` auto-loads `librccl-anp.so`, which deadlocks a bundled + RCCL overlay — set `NCCL_NET_PLUGIN=none`.** The v26.4 base image ships an + environment default of `NCCL_NET_PLUGIN=librccl-anp.so` (the ANP net plugin). + When the image carries a *bundled* RCCL overlay (the whole point of the + `rccl_overlay` Dockerfile), that plugin is incompatible with the overlay + `librccl` and RCCL init hangs at the first collective — the run never starts and + eventually times out. Set `NCCL_NET_PLUGIN=none` in the manifest env (both + `context.docker_env_vars` and `deployment_config.env_vars`, like the other + transport vars) to disable the plugin and let the bundled `librccl` drive the + IB/RoCE net path directly. The primus template ships this key set to `none`. diff --git a/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile index 7ff95b7..70ccde6 100644 --- a/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile +++ b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile @@ -191,7 +191,21 @@ RUN set -e && \ if [ ! -e "$TORCH_LIB/librocm_smi64.so" ]; then \ cp -v -L /opt/rocm/lib/librocm_smi64.so.1 "$TORCH_LIB/librocm_smi64.so"; \ ln -sfv librocm_smi64.so "$TORCH_LIB/librocm_smi64.so.1"; \ - fi + fi && \ + echo "=== global sweep: overwrite EVERY real librccl on disk with candidate ===" && \ + canon_src=$(readlink -f "$src") && \ + for f in $(find / -name 'librccl.so*' -not -type l 2>/dev/null); do \ + [ "$(readlink -f "$f")" = "$canon_src" ] && continue; \ + echo " overwrite: $f"; cp -fL "$src" "$f"; add_needed "$f"; \ + done && \ + ldconfig +# WHY the global sweep: v26.4 splits ROCm libs into _rocm_sdk_libraries/lib +# (runtime .so) and _rocm_sdk_devel/lib (dev). torch maps librccl.so.1 from +# _rocm_sdk_libraries at RUNTIME (soname wins once loaded), which the targeted +# /opt/rocm + torch/lib overwrites above do NOT cover -- so the candidate is +# built and shipped but the STOCK base librccl is what actually runs. Overwriting +# every non-symlink librccl on disk makes whichever copy the loader maps the +# candidate, independent of ROCm's per-version layout churn. # ---- Stage 3 (optional): rdma-core from source ------------------------------ # Builds only when RDMA_CORE_VERSION is non-empty (e.g. 63.0). Replaces the @@ -261,8 +275,8 @@ RUN set -e && \ TL=$(ls -d /opt/venv/lib/python*/site-packages/torch/lib | head -1) && \ ldd "$TL/libtorch_hip.so" | grep -iE "rccl|smi" && \ BUILT_SHA=$(cat "${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA") && \ - echo "=== ASSERT: every resolvable librccl is the candidate built @ ${BUILT_SHA} ===" && \ - for p in "${RCCL_INSTALL_DIR}/lib/librccl.so.1" /opt/rocm/lib/librccl.so.1 "$TL/librccl.so.1"; do \ + echo "=== ASSERT: EVERY librccl on disk is the candidate built @ ${BUILT_SHA} ===" && \ + for p in $(find / -name 'librccl.so*' -not -type l 2>/dev/null); do \ v=$(strings "$p" 2>/dev/null | grep -m1 -oE "RCCL version 2\.[0-9]+\.[0-9]+"); \ echo "$v" | grep -q "RCCL version 2." || { echo "ASSERT FAIL: $p is not RCCL 2.x ($p -> ${v:-NONE})"; exit 1; }; \ h=""; \ From af1dfbb227d1d5367e8588fd07d548ecdebfdaf5 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Mon, 6 Jul 2026 14:01:01 +0000 Subject: [PATCH 11/25] mad-slurm-multinode skill: repoint scaleout run.sh paths to primus/megatron-lm The primus_scaleout/megatron-lm tree was consolidated into scripts/primus/megatron-lm (commit 72e7a71; models.json repointed the *_scaleout models to scripts/primus/megatron-lm/run.sh), but the skill still referenced the removed path. Fix the stale `scripts/primus_scaleout/megatron-lm/run.sh` references: - primus_llama-3.1-8b.template.json / primus_llama-3.1-70b.template.json: built_models.scripts -> scripts/primus/megatron-lm/run.sh - references/deploy-bootstrap.md: the asset-existence check -> same path Model names (..._scaleout), the scaleout terminology, and the scaleout_gbs_override helper are intentionally left unchanged; only the filesystem paths moved. --- .../assets/manifests/primus_llama-3.1-70b.template.json | 2 +- .../assets/manifests/primus_llama-3.1-8b.template.json | 2 +- .../skills/mad-slurm-multinode/references/deploy-bootstrap.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json index f108eb1..4207fd7 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json @@ -17,7 +17,7 @@ "name": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", "url": "", "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd", - "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "scripts": "scripts/primus/megatron-lm/run.sh", "n_gpus": "-1", "owner": "mad.support@amd.com", "training_precision": "", diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json index 84dec88..180d96d 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json @@ -17,7 +17,7 @@ "name": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", "url": "", "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd", - "scripts": "scripts/primus_scaleout/megatron-lm/run.sh", + "scripts": "scripts/primus/megatron-lm/run.sh", "n_gpus": "-1", "owner": "mad.support@amd.com", "training_precision": "", diff --git a/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md b/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md index 5cf2e9f..25bdf7e 100644 --- a/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md +++ b/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md @@ -64,7 +64,7 @@ manifest's `dockerfile` and the model `scripts`/`run.sh` exist: # example for primus_llama-3.1-8b — adjust to the chosen manifest's fields ( cd MAD \ && [ -f docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile ] \ - && [ -f scripts/primus_scaleout/megatron-lm/run.sh ] \ + && [ -f scripts/primus/megatron-lm/run.sh ] \ || echo "STOP: manifest dockerfile/run.sh not found under MAD" ) ``` From ddcc909a6926098d4639e0a0737caa470051d9a6 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Mon, 6 Jul 2026 16:45:12 +0000 Subject: [PATCH 12/25] RCCL overlay: fix amdclang++ device-compile on v26.4 (symlink clang into _rocm_sdk_devel/bin) Adapted from mkuznet1/MAD PR#1 commit 6475da3f. v26.4 ships ROCm as pip wheels; the amdclang++ wrapper in _rocm_sdk_devel/bin/ looks for its clang++/clang-23 helpers next to itself, but they live only in _rocm_sdk_devel/lib/llvm/bin/. RCCL's device-code compile calls bin/amdclang++ directly (--offload-device-only) and fails: amdclang++: binary '.../_rocm_sdk_devel/bin/clang++' does not exist. Symlink clang/clang++/clang-23 into bin/ so the wrapper resolves. Guarded to be a no-op on v26.3 / non-wheel bases. --- ...n_train_rccl_overlay.ubuntu.amd.Dockerfile | 45 ++++++++++++++----- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile index 70ccde6..113413a 100644 --- a/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile +++ b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile @@ -96,6 +96,19 @@ RUN if [[ -n "${RCCL_COMMIT}" ]]; then \ git -C "$(cat /tmp/BLD_RCCL_HOME.txt)" rev-parse HEAD > "${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA" && \ echo "RCCL_BUILT_SHA=$(cat ${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA)" +# v26.4 compat: ROCm ships as pip wheels and the amdclang++ wrapper in +# _rocm_sdk_devel/bin/ resolves its clang++/clang-23 helpers relative to its own +# directory -- but those binaries live ONLY in _rocm_sdk_devel/lib/llvm/bin/. +# RCCL's device-code compile invokes bin/amdclang++ directly and dies with +# "amdclang++: binary '.../_rocm_sdk_devel/bin/clang++' does not exist". Symlink +# the helpers into bin/ so the wrapper resolves. No-op on v26.3 / non-wheel bases. +# (Fixing CC/CXX is NOT enough: --offload-device-only still calls bin/amdclang++.) +RUN SDK="$(ls -d /opt/venv/lib/python*/site-packages/_rocm_sdk_devel 2>/dev/null || true)" && \ + if [ -n "$SDK" ] && [ ! -e "$SDK/bin/clang++" ] && [ -d "$SDK/lib/llvm/bin" ]; then \ + for b in clang clang++ clang-23; do ln -sf ../lib/llvm/bin/$b "$SDK/bin/$b"; done; \ + echo "[v26.4-fix] symlinked clang/clang++/clang-23 into $SDK/bin"; \ + fi + RUN set -e && \ BLD_RCCL_HOME=$(cat /tmp/BLD_RCCL_HOME.txt) && \ cd "${BLD_RCCL_HOME}" && \ @@ -171,14 +184,18 @@ RUN set -e && \ src=$(ls -L "${RCCL_INSTALL_DIR}/lib/librccl.so.1.0" 2>/dev/null || ls -L "${RCCL_INSTALL_DIR}/lib/librccl.so") && \ echo "candidate librccl: $src" && \ add_needed "$src" && \ - ROCM_LIB=$(dirname "$(readlink -f /opt/rocm/lib/librccl.so.1)") && \ - echo "system rocm lib dir: $ROCM_LIB" && \ - echo "before:" && ls -la "$ROCM_LIB"/librccl.so* && \ - for f in "$ROCM_LIB"/librccl.so*; do \ - if [ -f "$f" ] && [ ! -L "$f" ]; then echo "overwriting real file: $f"; cp -fL "$src" "$f"; add_needed "$f"; fi; \ - done && \ - ldconfig && \ - echo "after:" && ls -laL "$ROCM_LIB"/librccl.so.1 && \ + if [ -e "/opt/rocm/lib/librccl.so.1" ]; then \ + ROCM_LIB=$(dirname "$(readlink -f /opt/rocm/lib/librccl.so.1)"); \ + echo "system rocm lib dir: $ROCM_LIB"; \ + ls -la "$ROCM_LIB"/librccl.so* 2>/dev/null || true; \ + for f in "$ROCM_LIB"/librccl.so*; do \ + [ -f "$f" ] && [ ! -L "$f" ] && { echo "overwriting real file: $f"; cp -fL "$src" "$f"; add_needed "$f"; }; \ + done; \ + ldconfig; \ + echo "after:"; ls -laL "$ROCM_LIB"/librccl.so.1 2>/dev/null || true; \ + else \ + echo "[rccl-overlay] /opt/rocm/lib/librccl.so.1 not present -- targeted /opt/rocm/lib overwrite skipped (global sweep covers it)"; \ + fi && \ TORCH_LIB=$(ls -d /opt/venv/lib/python*/site-packages/torch/lib 2>/dev/null | head -1) && \ [ -n "$TORCH_LIB" ] && [ -d "$TORCH_LIB" ] && echo "torch lib dir: $TORCH_LIB" && \ for f in "$TORCH_LIB"/librccl.so*; do \ @@ -189,8 +206,16 @@ RUN set -e && \ ln -sf librccl.so.1 "$TORCH_LIB/librccl.so" && \ add_needed "$TORCH_LIB/librccl.so.1.0" && \ if [ ! -e "$TORCH_LIB/librocm_smi64.so" ]; then \ - cp -v -L /opt/rocm/lib/librocm_smi64.so.1 "$TORCH_LIB/librocm_smi64.so"; \ - ln -sfv librocm_smi64.so "$TORCH_LIB/librocm_smi64.so.1"; \ + smi_src=$(ls /opt/rocm/lib/librocm_smi64.so.1 2>/dev/null || \ + ls /opt/venv/lib/python*/site-packages/_rocm_sdk_libraries/lib/librocm_smi64.so.1 2>/dev/null || \ + find /opt -name 'librocm_smi64.so.1' -not -type l 2>/dev/null | head -1 || true); \ + if [ -n "$smi_src" ]; then \ + cp -v -L "$smi_src" "$TORCH_LIB/librocm_smi64.so"; \ + ln -sfv librocm_smi64.so "$TORCH_LIB/librocm_smi64.so.1"; \ + echo "librocm_smi64 copied from: $smi_src"; \ + else \ + echo "[rccl-overlay] librocm_smi64.so.1 not found -- skipping copy (add-needed already applied to candidate)"; \ + fi; \ fi && \ echo "=== global sweep: overwrite EVERY real librccl on disk with candidate ===" && \ canon_src=$(readlink -f "$src") && \ From 112b524003a6b64079bd964a40c65803052a921e Mon Sep 17 00:00:00 2001 From: Ilia Kosarev Date: Tue, 7 Jul 2026 16:02:40 +0200 Subject: [PATCH 13/25] Fix SGLang disagg on ionic/RoCE fabrics (AINIC transport env) (#2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix SGLang disagg on ionic/RoCE fabrics: preserve AINIC transport env MOONCAKE_COOKBOOK's set_env_vars.sh assumes a Mellanox/mlx5 fabric and, on ionic-based RoCE clusters (mia1), disables RDMA (NCCL_IB_DISABLE=1), hardcodes mlx5 device names, and mangles NCCL_SOCKET_IFNAME. Save the launcher-provided transport iface before sourcing, then re-assert NCCL_IB_DISABLE=0, IBDEVICES from IB_DEVICES, and NCCL/GLOO socket ifnames afterward. Also derive host_ip from the transport interface instead of the first hostname -I address. Validated on mia1 (DeepSeek-R1-0528, 1P+1D disagg, TP=8, RCCL 78e8ba0 overlay + AINIC ionic RDMA), SLURM job 11211: both servers fired up, router bound on 2322, full benchmark sweep 100% request success. Co-Authored-By: Claude Opus 4 * sglang_disagg_server: source host_ip from IPADDRS[NODE_RANK] Re-deriving the node IP via ip addr/hostname -I risks picking a different NIC on multi-homed nodes, causing a mismatch between what peers registered in the router/barrier (built from IPADDRS) and what --host / socket_barrier --local-ip binds to. Source host_ip from IPADDRS[NODE_RANK] instead — the rank-ordered, post-rendezvous list that run.sh already resolved and passed in. Keep a hostname -I fallback for environments where IPADDRS is unset. Remove the now-redundant IFS/read of IP_ARRAY in the Cluster Topology section; the array is already populated at host_ip derivation time. Suggested-by: mkuznet1 Co-Authored-By: Claude Sonnet 4 --------- Co-authored-by: ilkosare@amd.com Co-authored-by: Claude Opus 4 --- scripts/sglang_disagg/sglang_disagg_server.sh | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/scripts/sglang_disagg/sglang_disagg_server.sh b/scripts/sglang_disagg/sglang_disagg_server.sh index f21df58..b6d6419 100755 --- a/scripts/sglang_disagg/sglang_disagg_server.sh +++ b/scripts/sglang_disagg/sglang_disagg_server.sh @@ -33,11 +33,31 @@ esac pip install py-spy pip install --ignore-installed --force-reinstall flask +# AINIC/ionic RoCE fix: MOONCAKE_COOKBOOK's set_env_vars.sh assumes a Mellanox/mlx5 +# fabric. On ionic-based RoCE clusters (e.g. mia1) it (1) hardcodes mlx5 IB device +# names, (2) forces NCCL_IB_DISABLE=1, and (3) mangles NCCL_SOCKET_IFNAME via an +# `ip route` awk parse. Save the transport iface set by the launcher, then re-assert +# the correct ionic values after sourcing so RDMA over ionic actually engages. +_SAVED_NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-}" + source $MOONCAKE_COOKBOOK_PATH/set_env_vars.sh #trap 'echo "Error occurred. Cleaning up..."; exit 0' ERR -host_ip=$(hostname -I | awk '{print $1}') +# Re-assert RDMA/socket env clobbered by set_env_vars.sh (see note above). +export NCCL_IB_DISABLE=0 +[[ -n "${IB_DEVICES:-}" ]] && export IBDEVICES="${IB_DEVICES}" +export NCCL_SOCKET_IFNAME="${_SAVED_NCCL_SOCKET_IFNAME:-eno0}" +export GLOO_SOCKET_IFNAME="${_SAVED_NCCL_SOCKET_IFNAME%%,*}" + +# This node's IP is already resolved by run.sh and handed to us (rank-ordered, +# post-rendezvous) via IPADDRS. Reuse IPADDRS[NODE_RANK] so the address we bind +# and advertise (--host, socket_barrier --local-ip) matches exactly what peers +# registered for us in the router/barrier; re-deriving it here can pick a +# different NIC on multi-homed nodes and desync the two. +IFS=',' read -ra IP_ARRAY <<< "$IPADDRS" +host_ip="${IP_ARRAY[NODE_RANK]:-$(hostname -I | awk '{print $1}')}" host_name=$(hostname) +unset _SAVED_NCCL_SOCKET_IFNAME # ============================================================================= # Model-Specific Configuration Maps @@ -112,7 +132,7 @@ python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ # Cluster Topology Configuration # ============================================================================= -IFS=',' read -ra IP_ARRAY <<< "$IPADDRS" +# IP_ARRAY was already parsed from IPADDRS above (where host_ip is derived). # Colocated topology (matches sglang_disagg_mori_io_ep.sh): the router/proxy runs # on NODE_RANK=0 alongside the first prefill server, so no dedicated proxy node is From 90eeb94c4c6d617d5ba484823c387abea7e5b85f Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Wed, 8 Jul 2026 10:48:27 +0000 Subject: [PATCH 14/25] Remove sglang_disagg_server.sh, route run.sh fully to mori_io_ep.sh sglang_disagg_mori_io_ep.sh is a functional superset of server.sh (models.yaml-driven config, mori/mooncake/nixl backends via KV_TRANSFER_BACKEND, DP_MODE support), and run_xPyD_models.slurm already always uses it. run.sh was the only remaining consumer of server.sh, so route both RUN_MORI branches to the unified launcher and drop the duplicate. Also port server.sh's KV_TRANSFER_BACKEND allowlist guard (against eval injection) into mori_io_ep.sh, and update the README and mad-slurm-multinode skill docs accordingly. Addresses PR #174 review feedback (basemam #1: stale/duplicate server.sh). --- .../mad-slurm-multinode/references/gotchas.md | 14 +- scripts/sglang_disagg/README.MD | 3 +- scripts/sglang_disagg/run.sh | 9 +- .../sglang_disagg/sglang_disagg_mori_io_ep.sh | 11 + scripts/sglang_disagg/sglang_disagg_server.sh | 348 ------------------ 5 files changed, 24 insertions(+), 361 deletions(-) delete mode 100755 scripts/sglang_disagg/sglang_disagg_server.sh diff --git a/.claude/skills/mad-slurm-multinode/references/gotchas.md b/.claude/skills/mad-slurm-multinode/references/gotchas.md index c2fa9a4..e285e1d 100644 --- a/.claude/skills/mad-slurm-multinode/references/gotchas.md +++ b/.claude/skills/mad-slurm-multinode/references/gotchas.md @@ -103,12 +103,14 @@ sweep retry. - **The launcher is `sglang-disagg` with `scripts/sglang_disagg/run.sh` as the entrypoint** (PR #142 native launcher). `run.sh` reads topology from - `--xp/--yd` (or `xP`/`yD` env) and `--kv-transfer-backend` (`nixl`/`mooncake`); - with `--run-mori 1` it execs `sglang_disagg_mori_io_ep.sh`, otherwise - `sglang_disagg_server.sh`. There is no `slurm_multi` wrapper — one `madengine - run` is launched per node and the roles (proxy / prefill x xP / decode x yD) - are derived from `NODE_RANK` + the ordered IP list. Do not reintroduce a - `*_mn` slurm wrapper; it duplicates topology logic and drifts from `run.sh`. + `--xp/--yd` (or `xP`/`yD` env) and `--kv-transfer-backend` (`mori`/`mooncake`/ + `nixl`), then always execs the unified `sglang_disagg_mori_io_ep.sh` (it + branches internally on `KV_TRANSFER_BACKEND`; the older Mooncake/NIXL-only + `sglang_disagg_server.sh` was retired as a functional subset). There is no + `slurm_multi` wrapper — one `madengine run` is launched per node and the + roles (proxy / prefill x xP / decode x yD) are derived from `NODE_RANK` + the + ordered IP list. Do not reintroduce a `*_mn` slurm wrapper; it duplicates + topology logic and drifts from `run.sh`. - **A newer RCCL build drops the `rocm_smi` `DT_NEEDED` and breaks `import torch` — re-add it with `patchelf`.** The RCCL stage of the single diff --git a/scripts/sglang_disagg/README.MD b/scripts/sglang_disagg/README.MD index 20dbc0d..fa3c6e3 100644 --- a/scripts/sglang_disagg/README.MD +++ b/scripts/sglang_disagg/README.MD @@ -56,8 +56,7 @@ docker build --build-arg MORI_COMMIT= -t sglang_disagg_pd_image -f sglang_d |------|-------------| | `run.sh` | **madengine entrypoint** (models.json `scripts` target). Bridges madengine's `sglang-disagg` launcher env to the launchers below | | `run_xPyD_models.slurm` | Standalone SLURM script to launch docker containers on all nodes via `sbatch` | -| `sglang_disagg_mori_io_ep.sh` | Container entrypoint (MoRI EP) — starts prefill/decode servers, proxy, and benchmark | -| `sglang_disagg_server.sh` | Container entrypoint (Mooncake/NIXL) | +| `sglang_disagg_mori_io_ep.sh` | Unified container entrypoint (MoRI EP / Mooncake / NIXL, selected via `KV_TRANSFER_BACKEND`) — starts prefill/decode servers, proxy, and benchmark | | `models.yaml` | Model-specific CLI flags for all supported models | | `mori_ep_env.sh` | RDMA/NCCL/Gloo environment variables | | `benchmark_xPyD.sh` | Concurrency sweep benchmark using sglang bench_serving | diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh index 32f96cb..cb57c60 100755 --- a/scripts/sglang_disagg/run.sh +++ b/scripts/sglang_disagg/run.sh @@ -216,8 +216,7 @@ echo " RUN_MORI=$RUN_MORI DP_MODE=$DP_MODE KV_TRANSFER_BACKEND=$KV_TRANSFER_ echo " MASTER_ADDR=$MASTER_ADDR IPADDRS=$IPADDRS" echo "==============================================================" -if [[ "$RUN_MORI" == "1" ]]; then - exec bash "$SCRIPT_DIR/sglang_disagg_mori_io_ep.sh" -else - exec bash "$SCRIPT_DIR/sglang_disagg_server.sh" -fi +# sglang_disagg_mori_io_ep.sh is a functional superset of the retired +# sglang_disagg_server.sh (per-model config from models.yaml, mori/mooncake/nixl +# via KV_TRANSFER_BACKEND, DP_MODE support) — route both RUN_MORI values there. +exec bash "$SCRIPT_DIR/sglang_disagg_mori_io_ep.sh" diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index 9b739c2..d986956 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -187,6 +187,17 @@ source "${SCRIPT_DIR}/mori_ep_env.sh" # KV transfer backend: default mori, switchable to mooncake or nixl. # Kept out of models.yaml so model config is backend-agnostic. _TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mori}" + +# _TRANSFER_BACKEND is interpolated into an eval'd launch command below; restrict +# it to known-good values to avoid invalid backends and shell-token injection. +case "$_TRANSFER_BACKEND" in + mori|mooncake|nixl) ;; + *) + echo "ERROR: unsupported KV_TRANSFER_BACKEND='$_TRANSFER_BACKEND' (expected: mori|mooncake|nixl)" >&2 + exit 1 + ;; +esac + PREFILL_MODEL_CONFIG+=" --disaggregation-transfer-backend ${_TRANSFER_BACKEND}" DECODE_MODEL_CONFIG+=" --disaggregation-transfer-backend ${_TRANSFER_BACKEND}" export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG diff --git a/scripts/sglang_disagg/sglang_disagg_server.sh b/scripts/sglang_disagg/sglang_disagg_server.sh deleted file mode 100755 index b6d6419..0000000 --- a/scripts/sglang_disagg/sglang_disagg_server.sh +++ /dev/null @@ -1,348 +0,0 @@ -#!/bin/bash -# SGLang Disaggregated Server Launcher with Model-Specific Configurations -# ============================================================================= - -# ============================================================================= -# Environment Configuration -# ============================================================================= - -MASTER_ADDR="${MASTER_ADDR:-localhost}" -MASTER_PORT="${MASTER_PORT:-23731}" -NODE_RANK="${NODE_RANK:-0}" -MODEL_PATH=$MODEL_PATH -MODEL_NAME="${MODEL_NAME:-}" -xP="${xP:-1}" -yD="${yD:-1}" -IPADDRS="${IPADDRS:-localhost}" -KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mooncake}" - -# KV_TRANSFER_BACKEND is interpolated into an eval'd launch command below; restrict -# it to known-good values to avoid invalid backends and shell-token injection. -case "$KV_TRANSFER_BACKEND" in - mori|mooncake|nixl) ;; - *) - echo "ERROR: unsupported KV_TRANSFER_BACKEND='$KV_TRANSFER_BACKEND' (expected: mori|mooncake|nixl)" >&2 - exit 1 - ;; -esac - -# ============================================================================= -# Dependencies and Environment Setup -# ============================================================================= - -pip install py-spy -pip install --ignore-installed --force-reinstall flask - -# AINIC/ionic RoCE fix: MOONCAKE_COOKBOOK's set_env_vars.sh assumes a Mellanox/mlx5 -# fabric. On ionic-based RoCE clusters (e.g. mia1) it (1) hardcodes mlx5 IB device -# names, (2) forces NCCL_IB_DISABLE=1, and (3) mangles NCCL_SOCKET_IFNAME via an -# `ip route` awk parse. Save the transport iface set by the launcher, then re-assert -# the correct ionic values after sourcing so RDMA over ionic actually engages. -_SAVED_NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-}" - -source $MOONCAKE_COOKBOOK_PATH/set_env_vars.sh -#trap 'echo "Error occurred. Cleaning up..."; exit 0' ERR - -# Re-assert RDMA/socket env clobbered by set_env_vars.sh (see note above). -export NCCL_IB_DISABLE=0 -[[ -n "${IB_DEVICES:-}" ]] && export IBDEVICES="${IB_DEVICES}" -export NCCL_SOCKET_IFNAME="${_SAVED_NCCL_SOCKET_IFNAME:-eno0}" -export GLOO_SOCKET_IFNAME="${_SAVED_NCCL_SOCKET_IFNAME%%,*}" - -# This node's IP is already resolved by run.sh and handed to us (rank-ordered, -# post-rendezvous) via IPADDRS. Reuse IPADDRS[NODE_RANK] so the address we bind -# and advertise (--host, socket_barrier --local-ip) matches exactly what peers -# registered for us in the router/barrier; re-deriving it here can pick a -# different NIC on multi-homed nodes and desync the two. -IFS=',' read -ra IP_ARRAY <<< "$IPADDRS" -host_ip="${IP_ARRAY[NODE_RANK]:-$(hostname -I | awk '{print $1}')}" -host_name=$(hostname) -unset _SAVED_NCCL_SOCKET_IFNAME - -# ============================================================================= -# Model-Specific Configuration Maps -# ============================================================================= - -declare -A MODEL_PREFILL_CONFIGS=( - ["Qwen3-32B"]="--tp-size 8" - ["Mixtral-8x7B-v0.1"]="--tp-size 8" - ["Llama-3.1-8B-Instruct"]="--tp-size 8" - ["Llama-3.1-405B-Instruct-FP8-KV"]="--tp-size 8" - ["amd-Llama-3.3-70B-Instruct-FP8-KV"]="--tp-size 8" - ["DeepSeek-V3"]="--tp-size 8" - ["DeepSeek-R1"]="--tp-size 8" -) - -declare -A MODEL_DECODE_CONFIGS=( - ["Qwen3-32B"]="--tp-size 8" - ["Mixtral-8x7B-v0.1"]="--tp-size 8" - ["Llama-3.1-8B-Instruct"]="--tp-size 8" - ["Llama-3.1-405B-Instruct-FP8-KV"]="--tp-size 8" - ["amd-Llama-3.3-70B-Instruct-FP8-KV"]="--tp-size 8" - ["DeepSeek-V3"]="--tp-size 8" - ["DeepSeek-R1"]="--tp-size 8" -) - -# ============================================================================= -# Configuration Selection Functions -# ============================================================================= - -get_model_config() { - local mode="$1" - local model_name="$2" - - if [[ "$mode" == "prefill" ]]; then - if [[ -n "${MODEL_PREFILL_CONFIGS[$model_name]}" ]]; then - echo "${MODEL_PREFILL_CONFIGS[$model_name]}" - else - echo "--tp-size 4" - fi - elif [[ "$mode" == "decode" ]]; then - if [[ -n "${MODEL_DECODE_CONFIGS[$model_name]}" ]]; then - echo "${MODEL_DECODE_CONFIGS[$model_name]}" - else - echo "--tp-size 4" - fi - fi -} - -if [[ -z "$MODEL_NAME" ]]; then - echo "Warning: MODEL_NAME not set, using default configurations" - PREFILL_MODEL_CONFIG="--tp-size 4" - DECODE_MODEL_CONFIG="--tp-size 4" -else - PREFILL_MODEL_CONFIG=$(get_model_config "prefill" "$MODEL_NAME") - DECODE_MODEL_CONFIG=$(get_model_config "decode" "$MODEL_NAME") - echo "Using model-specific configuration for: $MODEL_NAME" -fi - -# ============================================================================= -# Container Synchronization -# ============================================================================= - -echo "Waiting at the container creation barrier on $host_name" -python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --local-ip ${host_ip} \ - --local-port 4342 \ - --enable-port \ - --node-ips ${IPADDRS} \ - --node-ports 4342 - -# ============================================================================= -# Cluster Topology Configuration -# ============================================================================= - -# IP_ARRAY was already parsed from IPADDRS above (where host_ip is derived). - -# Colocated topology (matches sglang_disagg_mori_io_ep.sh): the router/proxy runs -# on NODE_RANK=0 alongside the first prefill server, so no dedicated proxy node is -# needed. Total nodes = xP + yD. -# IP_ARRAY[0..xP-1] -> prefill nodes (NODE_RANK 0..xP-1) -# IP_ARRAY[xP..xP+yD-1] -> decode nodes (NODE_RANK xP..xP+yD-1) -# Backend sglang servers listen on SERVER_PORT; the router listens on ROUTER_PORT. -# They must differ so prefill0 and the router can share NODE_RANK=0. -SERVER_PORT="${SERVER_PORT:-3000}" -ROUTER_PORT="${ROUTER_PORT:-2322}" - -PREFILL_ARGS="" -DECODE_ARGS="" - -for ((i=0; i<$xP && i<${#IP_ARRAY[@]}; i++)); do - PREFILL_ARGS+=" --prefill http://${IP_ARRAY[$i]}:${SERVER_PORT} " -done - -for ((i=$xP; i<${#IP_ARRAY[@]}; i++)); do - DECODE_ARGS+=" --decode http://${IP_ARRAY[$i]}:${SERVER_PORT} " -done - -# ============================================================================= -# Node Role Assignment and Server Launch -# ============================================================================= - -if [ "$NODE_RANK" -eq 0 ]; then - echo "NODE INFO =======================================" - echo "================================================" - echo "Node List : ${SLURM_JOB_NODELIST}" - echo "Node IPs : ${IPADDRS}" - echo "Model Name : ${MODEL_NAME:-'Not specified'}" - echo "================================================" - - echo "CLUSTER INFO ====================================" - echo "================================================" - echo "${host_name}:${host_ip} is Prefill Node 0 + Proxy/Router (co-located)" - echo "${PREFILL_ARGS} are Proxy's Prefill" - echo "${DECODE_ARGS} are Proxy's Decode" - echo "================================================" - - # Launch the first prefill server on this node, co-located with the router. - PREFILL_CMD="MC_TE_METRIC=true python3 -m sglang.launch_server \ - --model-path $MODEL_PATH \ - --disaggregation-mode prefill \ - --disaggregation-ib-device ${IBDEVICES} \ - --host ${host_ip} \ - --port ${SERVER_PORT} \ - --stream-output \ - --trust-remote-code \ - --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" - - if [[ -n "$PREFILL_MODEL_CONFIG" ]]; then - PREFILL_CMD="$PREFILL_CMD $PREFILL_MODEL_CONFIG" - fi - - eval "$PREFILL_CMD" \ - 2>&1 | tee /run_logs/${SLURM_JOB_ID}/prefill_NODE${NODE_RANK}.log >/dev/null & - - node0_prefill_pid=$! - - echo "Proxy server is waiting for prefill and decode nodes to be ready ..." \ - | tee /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null - sleep 20 - TIMEOUT_SECONDS=4000 - SLEEP_SECONDS=10 - SEARCH_SIGNAL="The server is fired up and ready to roll!" - SECONDS=0 - for ((i=0; i<$xP && i<${#IP_ARRAY[@]}; i++)); do - LOG_FILE=/run_logs/${SLURM_JOB_ID}/prefill_NODE${i}.log - #wait until prefill nodes get ready - until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do - if [ $SECONDS -ge $TIMEOUT_SECONDS ]; then - echo "FATAL: awaited ${SECONDS}s; readiness signal not found for prefill ${i} (${LOG_FILE}). Aborting before launching the router." \ - | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >&2 - exit 1 - fi - sleep $SLEEP_SECONDS - SECONDS=$(( SECONDS + SLEEP_SECONDS)) - done - done - - for ((i=$xP; i<${#IP_ARRAY[@]}; i++)); do - LOG_FILE=/run_logs/${SLURM_JOB_ID}/decode_NODE${i}.log - #wait until decode nodes get ready - until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do - if [ $SECONDS -ge $TIMEOUT_SECONDS ]; then - echo "FATAL: awaited ${SECONDS}s; readiness signal not found for decode ${i} (${LOG_FILE}). Aborting before launching the router." \ - | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >&2 - exit 1 - fi - sleep $SLEEP_SECONDS - SECONDS=$(( SECONDS + SLEEP_SECONDS)) - done - done - - sleep 10 - - python -m sglang_router.launch_router \ - --pd-disaggregation \ - ${PREFILL_ARGS} \ - ${DECODE_ARGS} \ - --host 0.0.0.0 \ - --port ${ROUTER_PORT} \ - 2>&1 | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null & - - proxy_pid=$! - - echo "Waiting for all prefill and decode servers to be up . . ." - python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --node-ips ${IPADDRS} \ - --node-ports ${SERVER_PORT} - - echo "Proxy Server Ready for benchmarking on ${host_name}:${host_ip}" - - sleep 10 - cd /opt/mooncake-cookbook - bash /opt/mooncake-cookbook/benchmark_xPyD.sh - - echo "Killing the proxy server" - kill $proxy_pid - - echo "Killing the co-located prefill server" - kill $node0_prefill_pid - -elif [ "$NODE_RANK" -ge 1 ] && [ "$NODE_RANK" -lt "$xP" ]; then - echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-'default'})" - echo "Using prefill config: $PREFILL_MODEL_CONFIG" - echo "Using KV transfer backend: ${KV_TRANSFER_BACKEND}" - - PREFILL_CMD="MC_TE_METRIC=true python3 -m sglang.launch_server \ - --model-path $MODEL_PATH \ - --disaggregation-mode prefill \ - --disaggregation-ib-device ${IBDEVICES} \ - --host ${host_ip} \ - --port ${SERVER_PORT} \ - --stream-output \ - --trust-remote-code \ - --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" - - if [[ -n "$PREFILL_MODEL_CONFIG" ]]; then - PREFILL_CMD="$PREFILL_CMD $PREFILL_MODEL_CONFIG" - fi - - eval "$PREFILL_CMD" \ - 2>&1 | tee /run_logs/${SLURM_JOB_ID}/prefill_NODE${NODE_RANK}.log >/dev/null & - - prefill_pid=$! - - echo "Waiting for proxy server to be up..." - python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --node-ips ${MASTER_ADDR} \ - --node-ports ${ROUTER_PORT} - - echo "Waiting until proxy server closes..." - python $MOONCAKE_COOKBOOK_PATH/socket_wait.py \ - --remote-ip ${MASTER_ADDR} \ - --remote-port ${ROUTER_PORT} - - echo "Killing the prefill server" - kill $prefill_pid - -elif [ "$NODE_RANK" -ge "$xP" ] && [ "$NODE_RANK" -le "$((xP + yD - 1))" ]; then - echo "${host_name}:${host_ip} is Decode Node (Model: ${MODEL_NAME:-'default'})" - echo "Using decode config: $DECODE_MODEL_CONFIG" - echo "Using KV transfer backend: ${KV_TRANSFER_BACKEND}" - - DECODE_CMD="python3 -m sglang.launch_server \ - --model-path ${MODEL_PATH} \ - --disaggregation-mode decode \ - --disaggregation-ib-device ${IBDEVICES} \ - --host ${host_ip} \ - --port ${SERVER_PORT} \ - --stream-output \ - --trust-remote-code \ - --disaggregation-transfer-backend ${KV_TRANSFER_BACKEND}" - - if [[ -n "$DECODE_MODEL_CONFIG" ]]; then - DECODE_CMD="$DECODE_CMD $DECODE_MODEL_CONFIG" - fi - - eval "$DECODE_CMD" \ - 2>&1 | tee /run_logs/${SLURM_JOB_ID}/decode_NODE${NODE_RANK}.log >/dev/null & - - decode_pid=$! - - echo "Waiting for proxy server to be up..." - python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --node-ips ${MASTER_ADDR} \ - --node-ports ${ROUTER_PORT} - - echo "Waiting until proxy server closes..." - python $MOONCAKE_COOKBOOK_PATH/socket_wait.py \ - --remote-ip ${MASTER_ADDR} \ - --remote-port ${ROUTER_PORT} - - echo "Killing the decode server" - kill $decode_pid - -else - echo "ERROR: NODE_RANK=${NODE_RANK} out of range (expected 0..$((xP + yD - 1))) for xP=${xP} yD=${yD}" >&2 - exit 1 -fi - -# ============================================================================= -# Cleanup -# ============================================================================= - -echo "Killing the etcd server" -kill $etcd_pid - -echo "Script completed successfully" -exit 0 From 54ea1550164a65d718c437aa13d55b48a82afdeb Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Wed, 8 Jul 2026 11:19:07 +0000 Subject: [PATCH 15/25] Merge oci-rdma62 Dockerfile variant into base full-overlay image The two full-overlay Dockerfiles were ~92% identical (270 vs 250 lines, only 23 lines differed): the rdma-core-v62 variant was a straight copy of the base overlay plus one appended stage. Fold that stage into the base Dockerfile behind an empty-by-default ARG (ENABLE_RDMA62), mirroring the existing RDMA_CORE_VERSION-gated optional stage in docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile, and drop the now-redundant variant file. Net: one Dockerfile instead of two near-duplicates, -231 lines. Also drop vendor/cluster-specific naming ("OCI-CX7") from the affected docs and manifest template, describing the requirement generically (hosts needing a newer rdma-core than the base image ships) to match the skill's existing hardware-archetype convention rather than naming a specific cluster. --- .../sglang_disagg_deepseek-r1.template.json | 6 +- .../mad-slurm-multinode/references/gotchas.md | 16 +- .../references/manifests.md | 9 +- ...l_overlay.oci-rdma62.ubuntu.amd.Dockerfile | 270 ------------------ ...ference_full_overlay.ubuntu.amd.Dockerfile | 35 +++ scripts/sglang_disagg/README.MD | 3 +- 6 files changed, 54 insertions(+), 285 deletions(-) delete mode 100644 docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json index b6ab3ca..63a2bb1 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json @@ -2,7 +2,7 @@ "built_images": { "sglang-disagg-deepseek-r1": { "model": "sglang_disagg_deepseek-r1", - "docker_image": "smifix-mori--rixl--mooncake-gfx942 (append -rdma62oci when built from the OCI variant Dockerfile)>", + "docker_image": "smifix-mori--rixl--mooncake-gfx942 (append -rdma62 when built with --build-arg ENABLE_RDMA62=1)>", "dockerfile": "docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile", "base_docker": "lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x", "build_duration": 0, @@ -42,7 +42,7 @@ "KV_TRANSFER_BACKEND": "nixl", "MORI_JIT_CACHE_DIR": "/tmp/mori_jit", "MORI_SHMEM_HEAP_SIZE": "17179869184", - "RDMA_CORE_CACHE": "", + "RDMA_CORE_CACHE": "", "NCCL_DEBUG": "WARN", "NCCL_DEBUG_SUBSYS": "INIT,NET", "NCCL_IB_DISABLE": "0", @@ -124,7 +124,7 @@ "KV_TRANSFER_BACKEND": "nixl", "MORI_JIT_CACHE_DIR": "/tmp/mori_jit", "MORI_SHMEM_HEAP_SIZE": "17179869184", - "RDMA_CORE_CACHE": "", + "RDMA_CORE_CACHE": "", "NCCL_DEBUG": "WARN", "NCCL_IB_DISABLE": "0", "NCCL_IB_HCA": "", diff --git a/.claude/skills/mad-slurm-multinode/references/gotchas.md b/.claude/skills/mad-slurm-multinode/references/gotchas.md index e285e1d..5601676 100644 --- a/.claude/skills/mad-slurm-multinode/references/gotchas.md +++ b/.claude/skills/mad-slurm-multinode/references/gotchas.md @@ -94,8 +94,8 @@ Keywords: launcher sglang-disagg run.sh xP yD RUN_MORI DP_MODE KV_TRANSFER_BACKE nixl mooncake, detokenizer hang health check No response from detokenizer, MoRI overlay #366 inter-node decode freeze, RCCL overlay rsmi_init libtorch_hip undefined symbol torch broken, rocm720 base librocm_smi64, patchelf add-needed -DT_NEEDED smifix, single full-overlay Dockerfile no base chaining, oci-rdma62 -rdma-core v62 variant, mooncake baked launcher build-layer removed runtime pip, +DT_NEEDED smifix, single full-overlay Dockerfile no base chaining, ENABLE_RDMA62 +rdma-core v62 optional stage, mooncake baked launcher build-layer removed runtime pip, self-discover node IPs rendezvous IP_SYNC_TIMEOUT SGLANG_NODE_IPS not forwarded, per-node docker load shared tar, perf CSV rank0 BENCHMARK_FAIL_FAST, circuit breaker prefill workers Service Unavailable BENCHMARK_POINT_RETRIES transient @@ -141,10 +141,14 @@ sweep retry. mutated the container runtime env and silently corrupted the Python libs (the model flags then stopped parsing and the server came up with defaults). Those pip deps and the Mooncake KV backend are now baked into the full-overlay Dockerfile; - rdma-core v62 lives only in the `*.oci-rdma62.*` variant (for OCI-CX7 hosts whose - RDMA stack needs a newer libibverbs/librdmacm/libmlx5 than the base ships — no - host `libibverbs` bind-mounts needed). Do not reintroduce a runtime build-layer - in the launcher. + rdma-core v62 is an optional stage in the SAME Dockerfile, gated by + `--build-arg ENABLE_RDMA62=1` (for hosts whose RDMA stack needs a newer + libibverbs/librdmacm/libmlx5 than the base ships — no host `libibverbs` + bind-mounts needed; unset/default skips the stage entirely, mirroring + `docker/primus_megatron_train_rccl_overlay...Dockerfile`'s `RDMA_CORE_VERSION` + gate). There is no separate rdma-core-variant Dockerfile anymore — it was a + ~92%-identical copy of the base overlay and was merged in to cut duplication. + Do not reintroduce a runtime build-layer in the launcher. - **The MoRI overlay (#366, e.g. commit `a14e6992`) is the one that fixes the mid-decode inter-node freeze.** (Commit hashes here and below — RCCL, MoRI, diff --git a/.claude/skills/mad-slurm-multinode/references/manifests.md b/.claude/skills/mad-slurm-multinode/references/manifests.md index 5a160ef..b164f12 100644 --- a/.claude/skills/mad-slurm-multinode/references/manifests.md +++ b/.claude/skills/mad-slurm-multinode/references/manifests.md @@ -135,10 +135,11 @@ differ, the host path belongs on the value side. is a single full-overlay build (`docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile`): base SGLang + RCCL + MoRI + NIXL/Mooncake KV-transfer + Mooncake pip in one - Dockerfile (no base-image chaining). For OCI-CX7 hosts whose RDMA stack needs a - newer rdma-core than the base ships, build the `*.oci-rdma62.*` variant instead - (bakes rdma-core v62 in, so no host `libibverbs` bind-mounts are needed). The - full overlay re-adds the `rocm_smi` `DT_NEEDED` + Dockerfile (no base-image chaining). For hosts whose RDMA stack needs a + newer rdma-core than the base ships, build the same Dockerfile with + `--build-arg ENABLE_RDMA62=1` (bakes rdma-core v62 in, so no host `libibverbs` + bind-mounts are needed; unset/default skips that stage). The full overlay + re-adds the `rocm_smi` `DT_NEEDED` (`patchelf --add-needed librocm_smi64.so.`) so a newer RCCL on the rocm720 base does not break `import torch` (see [gotchas.md](gotchas.md#sglang_disagg)). Perf lands in `perf_sglang-disagg-DeepSeek-R1.csv`, collected on rank 0 only. diff --git a/docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile b/docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile deleted file mode 100644 index 7d5ae50..0000000 --- a/docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile +++ /dev/null @@ -1,270 +0,0 @@ -# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} -############################################################################### -# SGLang Disaggregated P/D — MERGED full overlay (RCCL + MoRI + KV-transfer) -# -# Single-Dockerfile equivalent of the three chained overlays: -# base sglang -> RCCL overlay (+smifix) -> MoRI overlay -> KV-transfer (RIXL) -# Built in one `docker build` step so madengine's single-dockerfile build path -# produces the whole stack at once (no base_docker chaining between overlays). -# -# Pins (override via --build-arg): -# RCCL : ROCm/rocm-systems develop @ RCCL_COMMIT (default 78e8ba0) + smifix -# MoRI : ROCm/mori @ MORI_COMMIT (default a14e6992, includes #366) -# NIXL : ROCm/RIXL @ RIXL_COMMIT (default f33a5599) + ROCm/ucx @ UCX_COMMIT -# (the AMD NIXL implementation; exposed to SGLang as `nixl` via alias). -# NOTE: this is NOT ai-dynamo/nixl. The rocm KV transport for -# KV_TRANSFER_BACKEND=nixl is ROCm/RIXL; recipe mirrors the proven -# scripts/kvcache_transfer_bench/Dockerfile. -############################################################################### -ARG BASE_DOCKER=lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x -FROM $BASE_DOCKER - -SHELL ["/bin/bash", "-o", "pipefail", "-c"] -USER root - -############################################################################### -# 1) RCCL overlay — rebuild RCCL from source so the RCCL under test wins over -# the base image's librccl. (mirrors sglang_disagg_inference_rccl_overlay) -############################################################################### -ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git -ARG RCCL_BRANCH=develop -ARG RCCL_COMMIT=78e8ba0 -ARG RCCL_INSTALL_DIR=/opt/rccl -ARG BUILD_GPU_TARGETS=gfx942 - -ENV RCCL_HOME=${RCCL_INSTALL_DIR} -# Prepend the overlay RCCL so it wins over the base image's librccl. -ENV LD_LIBRARY_PATH=${RCCL_INSTALL_DIR}/lib:${LD_LIBRARY_PATH} - -RUN mkdir -p "${RCCL_INSTALL_DIR}" -WORKDIR /opt - -RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ - apt-get -o Acquire::Retries=5 update && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - git cmake ninja-build pkg-config make patch && \ - rm -rf /var/lib/apt/lists/* - -RUN if [[ -n "${RCCL_COMMIT}" ]]; then \ - git clone "${RCCL_REPO}" /tmp/rccl; \ - else \ - git clone --depth 1 --branch "${RCCL_BRANCH}" "${RCCL_REPO}" /tmp/rccl; \ - fi && \ - cd /tmp/rccl && \ - if [[ -n "${RCCL_BRANCH}" ]]; then git checkout "${RCCL_BRANCH}" || true; fi && \ - if [[ -n "${RCCL_COMMIT}" ]]; then git checkout "${RCCL_COMMIT}"; fi && \ - if [[ -d projects/rccl ]]; then \ - cd projects/rccl && git submodule update --init --recursive; \ - echo "/tmp/rccl/projects/rccl" > /tmp/BLD_RCCL_HOME.txt; \ - else \ - git submodule update --init --recursive; \ - echo "/tmp/rccl" > /tmp/BLD_RCCL_HOME.txt; \ - fi - -RUN set -e && \ - BLD_RCCL_HOME=$(cat /tmp/BLD_RCCL_HOME.txt) && \ - cd "${BLD_RCCL_HOME}" && \ - ./install.sh --amdgpu_targets="${BUILD_GPU_TARGETS}" --prefix="${RCCL_INSTALL_DIR}" - -RUN rm -rf /tmp/rccl /tmp/BLD_RCCL_HOME.txt - -# Re-add the rocm_smi dependency the rocm-systems RCCL build drops (smifix). -# Newer rocm-systems RCCL builds librccl WITHOUT a DT_NEEDED on librocm_smi64.so; -# torch's libtorch_hip.so relies on librccl to transitively pull in rocm_smi, so -# without this `import torch` dies with "undefined symbol: rsmi_init". -RUN set -e; \ - command -v patchelf >/dev/null 2>&1 || pip install --no-cache-dir patchelf >/dev/null 2>&1 \ - || { apt-get -o Acquire::Retries=5 update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends patchelf && rm -rf /var/lib/apt/lists/*; }; \ - librccl="$(readlink -f ${RCCL_INSTALL_DIR}/lib/librccl.so)"; \ - smi_soname="$(basename "$(ls /opt/rocm*/lib/librocm_smi64.so.[0-9]* 2>/dev/null | grep -E 'librocm_smi64\.so\.[0-9]+$' | head -1)")"; \ - test -n "${smi_soname}" || { echo "ROCM_SMI_SONAME_NOT_FOUND"; exit 1; }; \ - if readelf -d "${librccl}" 2>/dev/null | grep -q "NEEDED.*${smi_soname}"; then \ - echo "RCCL_SMI_NEEDED_ALREADY_PRESENT ${smi_soname}"; \ - else \ - patchelf --add-needed "${smi_soname}" "${librccl}" && echo "RCCL_SMI_NEEDED_ADDED ${smi_soname} -> ${librccl}"; \ - fi - -# Sanity: overlay librccl present AND torch imports with overlay librccl first. -RUN set -e; \ - test -e "${RCCL_INSTALL_DIR}/lib/librccl.so" || { echo "RCCL_OVERLAY_MISSING"; exit 1; }; \ - echo "RCCL_OVERLAY_OK $(ls -l ${RCCL_INSTALL_DIR}/lib/librccl.so*)"; \ - python3 -c "import torch; print('RCCL_OVERLAY_TORCH_OK', torch.__version__)" - -RUN pip list 2>/dev/null | grep -iE "sglang|torch" || true - -############################################################################### -# 2) MoRI overlay — substitutable MoE EP all-to-all (dispatch/combine, IBGDA). -# Default ROCm/mori @ a14e6992 includes #366 (internode decode hang fix). -############################################################################### -ARG MORI_REPO=https://github.com/ROCm/mori.git -ARG MORI_BRANCH=main -ARG MORI_COMMIT=a14e6992ffa95478e83127fe2672afff2840856f -ARG MORI_WHEEL_URL= -ARG MORI_GPU_ARCHS=gfx942 -ARG MORI_VERSION=1.2.0 -ARG MORI_SRC_DIR=/sgl-workspace/mori - -ENV MORI_GPU_ARCHS=${MORI_GPU_ARCHS} \ - MORI_SKIP_PRECOMPILE=1 \ - CMAKE_BUILD_TYPE=Release \ - SETUPTOOLS_SCM_PRETEND_VERSION=${MORI_VERSION} - -WORKDIR /sgl-workspace - -RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ - apt-get -o Acquire::Retries=5 update && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - git cmake ninja-build pkg-config make patch && \ - rm -rf /var/lib/apt/lists/* - -RUN set -e; \ - if [[ -n "${MORI_WHEEL_URL}" ]]; then \ - echo "[mori-overlay] installing prebuilt wheel: ${MORI_WHEEL_URL}"; \ - pip install --no-cache-dir --force-reinstall "${MORI_WHEEL_URL}"; \ - else \ - echo "[mori-overlay] source build ${MORI_REPO}@${MORI_BRANCH}${MORI_COMMIT:+ (${MORI_COMMIT})} archs=${MORI_GPU_ARCHS}"; \ - rm -rf "${MORI_SRC_DIR}"; \ - if [[ -n "${MORI_COMMIT}" ]]; then \ - git clone "${MORI_REPO}" "${MORI_SRC_DIR}"; \ - cd "${MORI_SRC_DIR}" && git checkout "${MORI_COMMIT}"; \ - else \ - git clone --depth 1 --branch "${MORI_BRANCH}" "${MORI_REPO}" "${MORI_SRC_DIR}"; \ - fi; \ - cd "${MORI_SRC_DIR}" && git submodule update --init --recursive || true; \ - pip install --no-cache-dir --force-reinstall .; \ - fi - -# Sanity: MoRI must import and report the overlay version. -RUN python3 -c "import mori, importlib.metadata as m; \ -print('MORI_OVERLAY_OK', getattr(mori,'__version__', m.version('amd_mori')))" \ - || { echo 'MORI_OVERLAY_IMPORT_FAILED'; exit 1; } - -############################################################################### -# 3) KV-transfer overlay — ROCm UCX + ROCm/RIXL (the AMD "nixl" KV transport). -# Mirrors the proven recipe in scripts/kvcache_transfer_bench/Dockerfile. -# SGLang's KV_TRANSFER_BACKEND=nixl imports `nixl`; we alias rixl -> nixl. -############################################################################### -ARG ROCM_PATH=/opt/rocm -ARG KV_WORKSPACE=/sgl-workspace -ARG UCX_REPO=https://github.com/ROCm/ucx.git -ARG UCX_COMMIT=da3fac2a -ARG RIXL_REPO=https://github.com/ROCm/RIXL.git -ARG RIXL_COMMIT=f33a5599 - -ENV UCX_HOME=${KV_WORKSPACE}/ucx -ENV RIXL_HOME=${KV_WORKSPACE}/rixl -ENV PATH=${UCX_HOME}/bin:${PATH} -ENV LD_LIBRARY_PATH=${RIXL_HOME}/lib/x86_64-linux-gnu:${UCX_HOME}/lib:${LD_LIBRARY_PATH} - -WORKDIR ${KV_WORKSPACE} - -RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ - apt-get -o Acquire::Retries=5 update && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - git cmake ninja-build pkg-config make patch autoconf automake libtool \ - cython3 libaio-dev libibverbs-dev librdmacm-dev libpci-dev \ - libgflags-dev libgoogle-glog-dev && \ - rm -rf /var/lib/apt/lists/* - -# meson/ninja/pybind11 build frontends for the RIXL meson build + wheel. -RUN pip install --no-cache-dir -U meson ninja pybind11 pyyaml build wheel - -# --- ROCm UCX --------------------------------------------------------------- -RUN set -e; \ - git clone "${UCX_REPO}" "${UCX_HOME}.src" && cd "${UCX_HOME}.src" && \ - git checkout "${UCX_COMMIT}" && \ - ./autogen.sh && mkdir -p build && cd build && \ - ../configure --prefix="${UCX_HOME}" --enable-shared --disable-static \ - --disable-doxygen-doc --enable-optimizations --enable-devel-headers \ - --with-rocm="${ROCM_PATH}" --with-verbs --with-dm --enable-mt && \ - make -j"$(nproc)" && make install && \ - cd "${KV_WORKSPACE}" && rm -rf "${UCX_HOME}.src" - -# --- ROCm/RIXL (the AMD NIXL implementation) -------------------------------- -RUN set -e; \ - git clone "${RIXL_REPO}" "${RIXL_HOME}" && cd "${RIXL_HOME}" && \ - git checkout "${RIXL_COMMIT}" && (git submodule update --init --recursive || true); \ - meson setup build --prefix="${RIXL_HOME}" -Ducx_path="${UCX_HOME}" -Drocm_path="${ROCM_PATH}" && \ - cd build && ninja && ninja install - -# Install the meson-built RIXL python package into site-packages directly. -# The upstream contrib/build-wheel.sh requires uv + py3.12 + auditwheel, which -# the py3.10 sglang base lacks; `ninja install` already produced the -# cpython-310 bindings under the RIXL prefix, so copy them into site-packages -# and alias `nixl` -> `rixl` for SGLang's KV_TRANSFER_BACKEND=nixl import path. -RUN set -e; \ - echo "=== RIXL python install tree ==="; \ - find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | sed 's/^/PYDIR: /'; \ - pysp="$(find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | head -1)"; \ - test -n "${pysp}" || { echo "RIXL_PY_INSTALL_NOT_FOUND"; find "${RIXL_HOME}" -name '_bindings*.so' -o -name '*.py' | head; exit 1; }; \ - echo "RIXL_PY_SRC=${pysp}"; ls -la "${pysp}"; \ - SP="$(python3 -c 'import site; print(site.getsitepackages()[0])')"; \ - copied=0; \ - for d in "${pysp}"/*/; do \ - name="$(basename "$d")"; \ - case "$name" in *dist-info|__pycache__) continue;; esac; \ - rm -rf "${SP}/${name}"; cp -a "$d" "${SP}/${name}"; echo "INSTALLED_PY_PKG ${name} -> ${SP}/${name}"; copied=1; \ - done; \ - test "$copied" = 1 || { echo "NO_RIXL_PKG_COPIED"; exit 1; }; \ - echo "import rixl, sys; sys.modules['nixl'] = rixl" > "${SP}/nixl_alias.pth"; \ - echo "NIXL_ALIAS_WRITTEN ${SP}/nixl_alias.pth" - -# Sanity: nixl (== rixl) import must succeed; report version + agent symbol. -RUN set -e; \ - python3 -c "import nixl; print('NIXL_IMPORT_OK', getattr(nixl,'__file__','?'))"; \ - python3 -c "import rixl, importlib.metadata as m; print('RIXL_VERSION', m.version('rixl'))" || true; \ - python3 -c "from nixl._api import nixl_agent, nixl_agent_config; print('NIXL_AGENT_OK')" \ - || echo "NIXL_AGENT_API_DIFFERS (verify sglang nixl import path at runtime)" - -############################################################################### -# 4) Runtime python deps + Mooncake KV-transfer backend. -# These were formerly built/pip-installed by the launcher at job start -# (scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh); baked into the image so -# the launcher no longer mutates the runtime environment. -############################################################################### -# Runtime python deps formerly pip-installed by the launcher. -RUN pip install --no-cache-dir py-spy pyyaml pandas \ - && pip install --no-cache-dir --ignore-installed --force-reinstall flask - -# Mooncake KV-transfer backend (KV_TRANSFER_BACKEND=mooncake). Mirrors -# docker/sglang_disagg_inference_kvtransfer_overlay.ubuntu.amd.Dockerfile:62-74. -# Default: pip wheel pin; set MOONCAKE_COMMIT for a source build. -ARG MOONCAKE_VERSION=0.3.6.post1 -ARG MOONCAKE_REPO=https://github.com/kvcache-ai/Mooncake.git -ARG MOONCAKE_COMMIT= -RUN set -e; \ - if [[ -n "${MOONCAKE_COMMIT}" ]]; then \ - echo "[full-overlay] Mooncake source ${MOONCAKE_REPO}@${MOONCAKE_COMMIT}"; \ - git clone "${MOONCAKE_REPO}" /tmp/mooncake && cd /tmp/mooncake && \ - git checkout "${MOONCAKE_COMMIT}" && (git submodule update --init --recursive || true); \ - pip install --no-cache-dir --force-reinstall . && rm -rf /tmp/mooncake; \ - elif [[ -n "${MOONCAKE_VERSION}" ]]; then \ - echo "[full-overlay] Mooncake pip mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ - pip install --no-cache-dir --force-reinstall "mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ - else \ - echo "[full-overlay] Mooncake: keeping base image version"; \ - fi - -# Sanity: mooncake must import alongside nixl/rixl (non-fatal — module path may vary). -RUN python3 -c "import mooncake; print('MOONCAKE_OVERLAY_OK')" \ - || echo "MOONCAKE_IMPORT_DIFFERS (verify sglang mooncake import path at runtime)" - -############################################################################### -# 5) OCI cluster workaround: build + install rdma-core v62 from source. -# The OCI-CX7 host stack needs a newer libibverbs/librdmacm/libmlx5 than the -# base image ships. Formerly built at job start by the launcher; baked here so -# the runtime env is not mutated. This variant is OCI-only; the base overlay -# keeps the image-default rdma-core. -############################################################################### -ARG RDMA_VER=v62.0 -RUN set -e; \ - git clone --branch "${RDMA_VER}" --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ - cd /tmp/rdma-core && mkdir -p build && cd build && \ - cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ - ninja && ninja install && ldconfig && \ - rm -rf /tmp/rdma-core - -# Sanity: rebuilt libibverbs present and python still imports (linker not corrupted). -RUN set -e; \ - ls -l /usr/lib/libibverbs.so* 2>/dev/null || ls -l /usr/lib/*/libibverbs.so* 2>/dev/null || true; \ - python3 -c "import torch; print(\"OCI_RDMA62_TORCH_OK\", torch.__version__)" diff --git a/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile index 5e1f3d4..ad51cd9 100644 --- a/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile +++ b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile @@ -15,6 +15,12 @@ # NOTE: this is NOT ai-dynamo/nixl. The rocm KV transport for # KV_TRANSFER_BACKEND=nixl is ROCm/RIXL; recipe mirrors the proven # scripts/kvcache_transfer_bench/Dockerfile. +# +# Hosts whose RDMA stack needs a newer rdma-core than the base image ships: +# pass --build-arg ENABLE_RDMA62=1 to additionally build + install rdma-core +# v62 from source (newer libibverbs/librdmacm/libmlx5; no host libibverbs +# bind-mounts needed). Default (unset) keeps the image-default rdma-core and +# skips that stage entirely. ############################################################################### ARG BASE_DOCKER=lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x FROM $BASE_DOCKER @@ -248,3 +254,32 @@ RUN set -e; \ # Sanity: mooncake must import alongside nixl/rixl (non-fatal — module path may vary). RUN python3 -c "import mooncake; print('MOONCAKE_OVERLAY_OK')" \ || echo "MOONCAKE_IMPORT_DIFFERS (verify sglang mooncake import path at runtime)" + +############################################################################### +# 5) Optional: build + install rdma-core v62 from source, for hosts whose RDMA +# stack needs a newer libibverbs/librdmacm/libmlx5 than the base image +# ships. Formerly built at job start by the launcher; baked here (gated by +# ENABLE_RDMA62, empty by default) so the runtime env is not mutated and +# the default build never pays for it. Mirrors the RDMA_CORE_VERSION-gated +# stage in docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile. +############################################################################### +ARG ENABLE_RDMA62= +ARG RDMA_VER=v62.0 +RUN if [[ -n "${ENABLE_RDMA62}" ]]; then \ + set -e; \ + git clone --branch "${RDMA_VER}" --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ + cd /tmp/rdma-core && mkdir -p build && cd build && \ + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ + ninja && ninja install && ldconfig && \ + rm -rf /tmp/rdma-core; \ + else \ + echo "ENABLE_RDMA62 unset — keeping base image rdma-core"; \ + fi + +# Sanity: rebuilt libibverbs present (only when ENABLE_RDMA62 was set) and +# python still imports (linker not corrupted). +RUN set -e; \ + if [[ -n "${ENABLE_RDMA62}" ]]; then \ + ls -l /usr/lib/libibverbs.so* 2>/dev/null || ls -l /usr/lib/*/libibverbs.so* 2>/dev/null || true; \ + fi; \ + python3 -c "import torch; print(\"RDMA62_TORCH_OK\" if \"${ENABLE_RDMA62}\" else \"RDMA62_SKIPPED_TORCH_OK\", torch.__version__)" diff --git a/scripts/sglang_disagg/README.MD b/scripts/sglang_disagg/README.MD index fa3c6e3..c16722f 100644 --- a/scripts/sglang_disagg/README.MD +++ b/scripts/sglang_disagg/README.MD @@ -95,8 +95,7 @@ KV-transfer (Mooncake/NIXL) steps on top of `BASE_DOCKER` in a single image | Overlay Dockerfile | Substitutes | Key build args | |--------------------|-------------|----------------| -| `docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile` | RCCL + MoRI + Mooncake/NIXL (merged) | `RCCL_REPO/RCCL_BRANCH/RCCL_COMMIT`, `MORI_REPO/MORI_BRANCH/MORI_COMMIT` or `MORI_WHEEL_URL`, `MOONCAKE_VERSION` / `NIXL_VERSION` | -| `docker/sglang_disagg_inference_full_overlay.oci-rdma62.ubuntu.amd.Dockerfile` | Same, OCI RDMA 6.2 variant | (as above) | +| `docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile` | RCCL + MoRI + Mooncake/NIXL (merged) | `RCCL_REPO/RCCL_BRANCH/RCCL_COMMIT`, `MORI_REPO/MORI_BRANCH/MORI_COMMIT` or `MORI_WHEEL_URL`, `MOONCAKE_VERSION` / `NIXL_VERSION`, `ENABLE_RDMA62` (set to `1` on hosts needing a newer rdma-core than the base image ships, to additionally build rdma-core v62 from source) | > **RCCL note:** a newer rocm-systems RCCL build links `librccl` without > the `librocm_smi64.so` `DT_NEEDED` the base carried. Since the overlay is first From 8b7b8c225be2f3d1d303d303a25fd1e73a6c020b Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 11:05:32 +0000 Subject: [PATCH 16/25] Address PR174 review comments: overlay dockerfile, csv metadata, docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - models.json: point sglang_disagg_deepseek-r1 at the full-overlay Dockerfile instead of the base sglang image, so `madengine run --tags` builds the RCCL/MoRI/NIXL/Mooncake stack this PR is built around (i-kosarev). - sglang_disagg/parse_to_csv.py: stop hardcoding launcher=slurm_multi and gpu_architecture=gfx942 in perf.csv metadata; launcher is now the actual native sglang-disagg path, and gpu_architecture is read from MAD_SYSTEM_GPU_ARCHITECTURE with a gfx942 fallback (Copilot). - sglang_disagg/ip_rendezvous.py: drop the docstring claim that this file is duplicated under scripts/vllm_dissag/ — no such copy exists (Copilot). - mad-slurm-multinode primus 8B/70B manifest templates: fix built_models.dockerfile to include the .Dockerfile suffix, matching built_images.dockerfile in the same file (Copilot). - mad-slurm-multinode detect_cluster_env.sh: recognize Broadcom Thor2 bnxt_re* HCAs in the RDMA probe, not just mlx5_/rdma (Copilot). --- .../assets/manifests/primus_llama-3.1-70b.template.json | 2 +- .../assets/manifests/primus_llama-3.1-8b.template.json | 2 +- .../mad-slurm-multinode/scripts/detect_cluster_env.sh | 6 ++++-- models.json | 2 +- scripts/sglang_disagg/ip_rendezvous.py | 4 ---- scripts/sglang_disagg/parse_to_csv.py | 8 ++++++-- 6 files changed, 13 insertions(+), 11 deletions(-) diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json index 4207fd7..a5a6256 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json @@ -16,7 +16,7 @@ "rocm-primus-llama31-70b": { "name": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", "url": "", - "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", "scripts": "scripts/primus/megatron-lm/run.sh", "n_gpus": "-1", "owner": "mad.support@amd.com", diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json index 180d96d..0d93b7a 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json @@ -16,7 +16,7 @@ "rocm-primus-llama31-8b": { "name": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", "url": "", - "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", "scripts": "scripts/primus/megatron-lm/run.sh", "n_gpus": "-1", "owner": "mad.support@amd.com", diff --git a/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh b/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh index 9d2d2dd..47e5f03 100755 --- a/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh +++ b/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh @@ -24,10 +24,10 @@ echo "MAD_SYSTEM_GPU_ARCHITECTURE : ${arch:-}" # --- RDMA / IB HCAs --- hcas="" if command -v ibv_devices >/dev/null 2>&1; then - hcas=$(ibv_devices 2>/dev/null | awk 'NR>2 {print $1}' | grep -E '^(mlx5_|rdma)' | sort) + hcas=$(ibv_devices 2>/dev/null | awk 'NR>2 {print $1}' | grep -E '^(mlx5_|rdma|bnxt_re)' | sort) fi if [ -z "$hcas" ] && [ -d /sys/class/infiniband ]; then - hcas=$(ls /sys/class/infiniband 2>/dev/null | grep -E '^(mlx5_|rdma)' | sort) + hcas=$(ls /sys/class/infiniband 2>/dev/null | grep -E '^(mlx5_|rdma|bnxt_re)' | sort) fi if [ -n "$hcas" ]; then list=$(echo "$hcas" | sed 's/$/:1/' | paste -sd, -) @@ -36,6 +36,8 @@ if [ -n "$hcas" ]; then fam="AINIC (ionic) -> RDMAV_DRIVERS=ionic, IBV_DRIVERS=ionic, RCCL_AINIC_ROCE=1, GID likely 1" elif echo "$hcas" | grep -q '^mlx5'; then fam="CX7/Mellanox (mlx5) -> RDMAV_DRIVERS=mlx5, IBV_DRIVERS=mlx5, GID likely 3" + elif echo "$hcas" | grep -q '^bnxt_re'; then + fam="Broadcom Thor2 (bnxt_re) -> RDMAV_DRIVERS=bnxt_re, IBV_DRIVERS=bnxt_re, GID: check show_gids" else fam="unknown vendor -> archetype not auto-classified; confirm with the user (see references/cluster-types.md)" fi diff --git a/models.json b/models.json index 7c35e1d..eeb9c12 100644 --- a/models.json +++ b/models.json @@ -1980,7 +1980,7 @@ { "name": "sglang_disagg_deepseek-r1", "url": "", - "dockerfile": "docker/sglang_disagg_inference", + "dockerfile": "docker/sglang_disagg_inference_full_overlay", "scripts": "scripts/sglang_disagg/run.sh", "n_gpus": "8", "owner": "mad.support@amd.com", diff --git a/scripts/sglang_disagg/ip_rendezvous.py b/scripts/sglang_disagg/ip_rendezvous.py index 9091dab..6fcf6eb 100755 --- a/scripts/sglang_disagg/ip_rendezvous.py +++ b/scripts/sglang_disagg/ip_rendezvous.py @@ -13,10 +13,6 @@ on failure an empty line is printed (exit 2), so callers can use `out="$(ip_rendezvous.py ... || true)"` and test for emptiness. -NOTE: this file is intentionally duplicated under scripts/sglang_disagg/ and -scripts/vllm_dissag/ (each model packages its own script dir into the image). -Keep both copies identical. - Usage: ip_rendezvous.py Env: diff --git a/scripts/sglang_disagg/parse_to_csv.py b/scripts/sglang_disagg/parse_to_csv.py index d8b3074..84661ab 100644 --- a/scripts/sglang_disagg/parse_to_csv.py +++ b/scripts/sglang_disagg/parse_to_csv.py @@ -188,8 +188,12 @@ def _get_run_metadata(pipeline: str = "sglang"): "gpus_per_node": gpus, "docker_image": os.environ.get("DOCKER_IMAGE_NAME", ""), "machine_name": os.environ.get("SLURM_JOB_NODELIST", ""), - "launcher": "slurm_multi", - "gpu_architecture": "gfx942", + # This script only runs behind madengine's native sglang-disagg launcher + # (scripts/sglang_disagg/run.sh), so the launcher name is fixed. The GPU + # architecture varies by cluster; madengine exports it, with a gfx942 + # fallback for standalone/manual invocations. + "launcher": "sglang-disagg", + "gpu_architecture": os.environ.get("MAD_SYSTEM_GPU_ARCHITECTURE", "gfx942"), } From fc2ddfbeb25203d25b474be80c9d7bbc864f462d Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 11:30:04 +0000 Subject: [PATCH 17/25] Address second round of PR174 review comments (Copilot) - run.sh: non-rank0 weights-staging wait now reuses _weights_complete (sentinel, or config.json + >=1 shard) for both the poll loop and the post-timeout check, instead of a config.json-only check that could mistake a stale/partial download for a finished one. - run.sh: chmod the shared top-level /run_logs with the sticky bit (1777, like /tmp) instead of plain 0777, so one user can't delete another user's per-job log subdir on the same mount. - models.json: sglang_disagg_deepseek-r1 n_gpus -> "-1" ("all available"), matching every other sglang_disagg_* entry in this file instead of hardcoding 8. --- models.json | 2 +- scripts/sglang_disagg/run.sh | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/models.json b/models.json index eeb9c12..4ce979e 100644 --- a/models.json +++ b/models.json @@ -1982,7 +1982,7 @@ "url": "", "dockerfile": "docker/sglang_disagg_inference_full_overlay", "scripts": "scripts/sglang_disagg/run.sh", - "n_gpus": "8", + "n_gpus": "-1", "owner": "mad.support@amd.com", "training_precision": "", "multiple_results": "perf_sglang-disagg-DeepSeek-R1.csv", diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh index cb57c60..fb32ecb 100755 --- a/scripts/sglang_disagg/run.sh +++ b/scripts/sglang_disagg/run.sh @@ -67,7 +67,7 @@ if ! mkdir -p "${RUN_LOG_DIR}" 2>/dev/null || [[ ! -w "${RUN_LOG_DIR}" ]]; then echo " rank-0 proxy greps them across nodes. Mount it and re-run." >&2 exit 1 fi -chmod 0777 /run_logs 2>/dev/null || true +chmod 1777 /run_logs 2>/dev/null || true chmod -R 0777 "${RUN_LOG_DIR}" 2>/dev/null || true export RUN_LOG_DIR exec > >(tee -a "${RUN_LOG_DIR}/run_sh_rank${NODE_RANK}.log") 2>&1 @@ -200,10 +200,10 @@ if ! _weights_complete; then else echo "[stage_weights] NODE_RANK=${NODE_RANK} waiting for weights at ${MODEL_PATH}" for _i in $(seq 1 720); do - [[ -f "${_SENTINEL}" && -f "${MODEL_PATH}/config.json" ]] && break + _weights_complete && break sleep 10 done - [[ -f "${MODEL_PATH}/config.json" ]] || { echo "[stage_weights] timed out waiting for weights" >&2; exit 1; } + _weights_complete || { echo "[stage_weights] timed out waiting for weights" >&2; exit 1; } fi fi From 7895cbc405c15d3a7b40fc5ea94f2284e30ac2ff Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 11:49:08 +0000 Subject: [PATCH 18/25] Address third round of PR174 review comments (Copilot) - ip_rendezvous.py: rank-0 rendezvous server now only accepts peer reports for rank 1..nnodes-1 (rank 0 is the server's own, already-seeded entry), so a buggy or spoofed rank-0 message from a peer can no longer overwrite the server's own host_ip. - primus_megatron-lm_benchmark_report.sh: renamed the `dp` parameter/param name to `world_size` in normalize_global_batch_size and updated the surrounding comments; the GBS renormalization uses total world size (NUM_GPUS), not the actual Megatron data-parallel degree, which can differ when TP/PP/CP/EP > 1 -- naming now matches the implementation. --- .../primus_megatron-lm_benchmark_report.sh | 15 +++++++++------ scripts/sglang_disagg/ip_rendezvous.py | 5 ++++- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh index ac84117..9d7e0db 100755 --- a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh +++ b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh @@ -185,13 +185,15 @@ if ! [[ "$GPUS_PER_NODE" =~ ^[0-9]+$ ]]; then GPUS_PER_NODE=8; fi NUM_GPUS=$((NNODES * GPUS_PER_NODE)) echo "[INFO] Distributed params: NNODES=$NNODES GPUS_PER_NODE=$GPUS_PER_NODE NUM_GPUS=$NUM_GPUS" -# Round global batch size up to a multiple of micro_batch_size * data_parallel so -# the global batch stays divisible across an arbitrary number of ranks. +# Round global batch size up to a multiple of micro_batch_size * world_size so +# the global batch stays divisible across an arbitrary number of ranks. Note: +# this uses total world size (NUM_GPUS), not the actual Megatron data-parallel +# degree, which can be smaller than world size when TP/PP/CP/EP > 1. normalize_global_batch_size() { local mbs="$1" local gbs="$2" - local dp="$3" - local unit=$((mbs * dp)) + local world_size="$3" + local unit=$((mbs * world_size)) if (( unit <= 0 )); then echo "$gbs" return @@ -205,8 +207,9 @@ normalize_global_batch_size() { # Emit the extra CLI flags needed to keep a run valid on multiple nodes: when # NUM_GPUS > 8 the config global_batch_size (tuned for a single 8-GPU node) is -# renormalized and passed as an override. Prints nothing for single-node runs, -# so behavior there is byte-for-byte identical to before. +# renormalized against total world size and passed as an override. Prints +# nothing for single-node runs, so behavior there is byte-for-byte identical +# to before. scaleout_gbs_override() { local mbs="$1" local gbs="$2" diff --git a/scripts/sglang_disagg/ip_rendezvous.py b/scripts/sglang_disagg/ip_rendezvous.py index 6fcf6eb..a124ad7 100755 --- a/scripts/sglang_disagg/ip_rendezvous.py +++ b/scripts/sglang_disagg/ip_rendezvous.py @@ -68,7 +68,10 @@ def serve(nnodes, host_ip, port, token, deadline): continue wrank = int(msg.get("rank", -1)) wip = str(msg.get("ip", "")).strip() - if 0 <= wrank < nnodes and wip: + # rank 0 is the server itself (by_rank[0] is already seeded with our + # own host_ip above); only accept reports from peers, rank 1..nnodes-1, + # so a buggy/spoofed rank-0 message can't clobber our own entry. + if 1 <= wrank < nnodes and wip: old = conns.get(wrank) if old is not None: try: From 15dfa6912fdb4992cf4aa5b35817fd2eb5420d57 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 12:07:20 +0000 Subject: [PATCH 19/25] Fail fast when PyYAML is missing in sglang_disagg_mori_io_ep.sh This PR moved the PyYAML install for models.yaml parsing from a runtime `pip install pyyaml` into the full-overlay Dockerfile's build layer, but 16 pre-existing pyt_sglang_disagg_* model entries in models.json still use the base (non-overlay) docker/sglang_disagg_inference.ubuntu.amd.Dockerfile via run_xPyD_models.slurm, which always routes through this same script and does not install PyYAML at build time either. Add an explicit `import yaml` check before the eval-based YAML parse, so a missing PyYAML fails with a clear message instead of leaving MODEL_* flags silently unset (Copilot). --- scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index d986956..340e28a 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -87,6 +87,14 @@ if [[ ! -f "$MODELS_YAML" ]]; then exit 1 fi +if ! python3 -c "import yaml" >/dev/null 2>&1; then + echo "ERROR: PyYAML is not installed in this image, but is required to parse" >&2 + echo " ${MODELS_YAML}. Use an image built from" >&2 + echo " docker/sglang_disagg_inference_full_overlay*.Dockerfile (which" >&2 + echo " installs pyyaml at build time), or 'pip install pyyaml'." >&2 + exit 1 +fi + export MODELS_YAML MODEL_NAME PARALLEL_MODE xP GPUS_PER_NODE eval "$(python3 - <<'PY' import os From 166239510aa95ae10192b3714c213fec219e6163 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 12:27:12 +0000 Subject: [PATCH 20/25] Address fourth round of PR174 review comments (Copilot) - primus_megatron-lm_benchmark_report.sh: when scaleout_gbs_override adjusts global_batch_size for NUM_GPUS>8, the GBS variable is now reassigned via the new effective_global_batch_size() helper before the perf CSV is written, in all three scaleout branches (8B/70B/405B). Previously the perf CSV kept reporting the original, pre-override GBS even when the actual run used the renormalized value. - mad.env.{cx7-roce,amd-ainic,thor2-bnxt}.template: warn on stderr when ~/.huggingface/token is missing or empty before exporting MAD_SECRETS_HFTOKEN, so a misconfigured token is caught immediately instead of failing obscurely later during HF downloads. Uses a warning (not exit), since these files are meant to be sourced and an exit would kill the caller's shell. --- .../assets/mad.env/mad.env.amd-ainic.template | 9 +++++- .../assets/mad.env/mad.env.cx7-roce.template | 9 +++++- .../mad.env/mad.env.thor2-bnxt.template | 9 +++++- .../primus_megatron-lm_benchmark_report.sh | 31 ++++++++++++++----- 4 files changed, 48 insertions(+), 10 deletions(-) diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template index 5bcb9d7..49c2f97 100644 --- a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template @@ -7,7 +7,14 @@ # scripts/detect_cluster_env.sh). # --- secrets (file-based, never inline a literal hf_... token) --- -export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token) +# This file is meant to be `source`d, so we warn (not exit) on a missing/empty +# token: an exit here would kill the caller's shell instead of just this file. +if [[ ! -s ~/.huggingface/token ]]; then + echo "WARNING: ~/.huggingface/token is missing or empty; MAD_SECRETS_HFTOKEN" >&2 + echo " will be empty and HF downloads will fail later. Run" >&2 + echo " 'huggingface-cli login' or place a token at that path." >&2 +fi +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token 2>/dev/null) # --- system --- export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 # AINIC/MI355-class default; confirm on node diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template index 738c4fa..8866266 100644 --- a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template @@ -7,7 +7,14 @@ # MAD_SYSTEM_GPU_ARCHITECTURE, NCCL_SOCKET_IFNAME, NCCL_IB_GID_INDEX # --- secrets (file-based, never inline a literal hf_... token) --- -export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token) +# This file is meant to be `source`d, so we warn (not exit) on a missing/empty +# token: an exit here would kill the caller's shell instead of just this file. +if [[ ! -s ~/.huggingface/token ]]; then + echo "WARNING: ~/.huggingface/token is missing or empty; MAD_SECRETS_HFTOKEN" >&2 + echo " will be empty and HF downloads will fail later. Run" >&2 + echo " 'huggingface-cli login' or place a token at that path." >&2 +fi +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token 2>/dev/null) # --- system --- export MAD_SYSTEM_GPU_ARCHITECTURE=gfx942 # CX7/MI300-class default; confirm on node diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template index 7c99305..92053d3 100644 --- a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template @@ -7,7 +7,14 @@ # MAD_SYSTEM_GPU_ARCHITECTURE, NCCL_SOCKET_IFNAME, NCCL_IB_GID_INDEX # --- secrets (file-based, never inline a literal hf_... token) --- -export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token) +# This file is meant to be `source`d, so we warn (not exit) on a missing/empty +# token: an exit here would kill the caller's shell instead of just this file. +if [[ ! -s ~/.huggingface/token ]]; then + echo "WARNING: ~/.huggingface/token is missing or empty; MAD_SECRETS_HFTOKEN" >&2 + echo " will be empty and HF downloads will fail later. Run" >&2 + echo " 'huggingface-cli login' or place a token at that path." >&2 +fi +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token 2>/dev/null) # --- system --- export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 # MI355X-class default; confirm on node (gfx942 = MI300X) diff --git a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh index 9d7e0db..91ddd71 100755 --- a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh +++ b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh @@ -205,20 +205,34 @@ normalize_global_batch_size() { echo "$(( ((gbs + unit - 1) / unit) * unit ))" } +# Pure helper: the possibly-adjusted global batch size for NUM_GPUS>8 (world +# size scaleout), or the original gbs unchanged for single-node runs / values +# that don't parse as plain integers. Callers must reassign their own GBS +# variable to this so the value used for the actual run (and reported to the +# perf CSV) always agree -- see scaleout_gbs_override below. +effective_global_batch_size() { + local mbs="$1" + local gbs="$2" + if (( NUM_GPUS > 8 )) && [[ "$mbs" =~ ^[0-9]+$ && "$gbs" =~ ^[0-9]+$ ]]; then + normalize_global_batch_size "$mbs" "$gbs" "$NUM_GPUS" + else + echo "$gbs" + fi +} + # Emit the extra CLI flags needed to keep a run valid on multiple nodes: when # NUM_GPUS > 8 the config global_batch_size (tuned for a single 8-GPU node) is # renormalized against total world size and passed as an override. Prints # nothing for single-node runs, so behavior there is byte-for-byte identical -# to before. +# to before. Callers should also run `GBS=$(effective_global_batch_size ...)` +# so the reported GBS in the perf CSV matches the value actually used. scaleout_gbs_override() { local mbs="$1" local gbs="$2" - if (( NUM_GPUS > 8 )) && [[ "$mbs" =~ ^[0-9]+$ && "$gbs" =~ ^[0-9]+$ ]]; then - local adjusted - adjusted=$(normalize_global_batch_size "$mbs" "$gbs" "$NUM_GPUS") - if [[ "$adjusted" != "$gbs" ]]; then - echo "[INFO] Adjusted global batch size for distributed run: ${gbs} -> ${adjusted} (MBS=${mbs}, NUM_GPUS=${NUM_GPUS})" >&2 - fi + local adjusted + adjusted=$(effective_global_batch_size "$mbs" "$gbs") + if [[ "$adjusted" != "$gbs" ]]; then + echo "[INFO] Adjusted global batch size for distributed run: ${gbs} -> ${adjusted} (MBS=${mbs}, NUM_GPUS=${NUM_GPUS})" >&2 printf -- '--global_batch_size %s' "$adjusted" fi } @@ -246,6 +260,7 @@ if [ "$MODEL_REPO" == "Llama-3.1-8B" ]; then GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") + GBS=$(effective_global_batch_size "$MBS" "$GBS") if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [[ "$DATATYPE" == "MXFP4" ]]; then NVTE_USE_CAST_TRANSPOSE_TRITON=0 run_primus "$EXP" $GBS_OVERRIDE @@ -279,6 +294,7 @@ elif [ "$MODEL_REPO" == "Llama-3.1-70B" ]; then GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") + GBS=$(effective_global_batch_size "$MBS" "$GBS") if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [ "$DATATYPE" == "FP8" ]; then run_primus "$EXP" $GBS_OVERRIDE @@ -353,6 +369,7 @@ elif [ "$MODEL_REPO" == "Llama-3.1-405B" ]; then echo "Hint: add llama3.1_405B-$DATATYPE-pretrain.yaml in Primus configs." else GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") + GBS=$(effective_global_batch_size "$MBS" "$GBS") run_primus "$EXP" $GBS_OVERRIDE fi if [ -f "$TRAIN_LOG" ]; then From d41959656483edbd33b54e7ef9f6f110cf3c96fb Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 12:40:06 +0000 Subject: [PATCH 21/25] Restore runtime PyYAML install for legacy base-image sglang_disagg models The previous fail-fast fix (15dfa69) only turned a missing PyYAML into a clearer error, but didn't actually restore the ability for the 16 pre-existing pyt_sglang_disagg_* models.json entries (which use the base, non-overlay docker/sglang_disagg_inference.ubuntu.amd.Dockerfile via run_xPyD_models.slurm) to run at all, since that image still lacks PyYAML. Restore a runtime `pip install pyyaml` as a fallback when `import yaml` fails -- matching the behavior on develop before this PR moved the install into the full-overlay Dockerfile's build layer -- and only fail fast if the runtime install itself doesn't succeed (e.g. no network). --- .../sglang_disagg/sglang_disagg_mori_io_ep.sh | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index 340e28a..268b795 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -77,8 +77,6 @@ BARRIER_PORT="${BARRIER_PORT:-4342}" # ============================================================================= # === Model-Specific Configuration from YAML === -# Parse models.yaml into MODEL_* flags. (The runtime build layer was moved into -# the image: docker/sglang_disagg_inference_full_overlay*.Dockerfile.) GPUS_PER_NODE="${GPUS_PER_NODE:-8}" MODELS_YAML="${MODELS_YAML:-${SCRIPT_DIR}/models.yaml}" @@ -88,10 +86,16 @@ if [[ ! -f "$MODELS_YAML" ]]; then fi if ! python3 -c "import yaml" >/dev/null 2>&1; then - echo "ERROR: PyYAML is not installed in this image, but is required to parse" >&2 - echo " ${MODELS_YAML}. Use an image built from" >&2 - echo " docker/sglang_disagg_inference_full_overlay*.Dockerfile (which" >&2 - echo " installs pyyaml at build time), or 'pip install pyyaml'." >&2 + echo "[stage_deps] PyYAML not found; installing at runtime (expected on the" >&2 + echo " base, non-overlay image) ..." >&2 + pip install --no-cache-dir --quiet pyyaml || true +fi +if ! python3 -c "import yaml" >/dev/null 2>&1; then + echo "ERROR: PyYAML is not installed and could not be installed at runtime," >&2 + echo " but is required to parse ${MODELS_YAML}. Use an image built" >&2 + echo " from docker/sglang_disagg_inference_full_overlay*.Dockerfile" >&2 + echo " (which installs pyyaml at build time), or ensure network" >&2 + echo " access for 'pip install pyyaml'." >&2 exit 1 fi From e6224beac6d53853879932341b6b2780cf565f2b Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 13:14:05 +0000 Subject: [PATCH 22/25] Address fifth round of PR174 review comments (Copilot) - run.sh: NNODES no longer falls back to WORLD_SIZE. Elsewhere in this repo WORLD_SIZE means total ranks (GPUS_PER_NODE*NNODES), not node count, so a stray WORLD_SIZE would make ip_rendezvous.py wait for far too many peers and hang. Fall back to SLURM_JOB_NUM_NODES, then xP+yD, both of which are node counts. - sglang_disagg_mori_io_ep.sh: quote the interpolated model_name/config_path in the "model not found" error path of the YAML-to-shell snippet. That string is eval'd by the shell, so unescaped values could break the launcher or inject shell tokens; it now reuses the same shlex.quote helper (moved above its first use) as the rest of the emitted exports. --- scripts/sglang_disagg/run.sh | 2 +- scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh | 15 ++++++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh index fb32ecb..9d54a42 100755 --- a/scripts/sglang_disagg/run.sh +++ b/scripts/sglang_disagg/run.sh @@ -73,7 +73,7 @@ export RUN_LOG_DIR exec > >(tee -a "${RUN_LOG_DIR}/run_sh_rank${NODE_RANK}.log") 2>&1 export xP="${xP:-${SGLANG_DISAGG_PREFILL_NODES:-1}}" export yD="${yD:-${SGLANG_DISAGG_DECODE_NODES:-1}}" -export NNODES="${NNODES:-${SGLANG_DISAGG_TOTAL_NODES:-${WORLD_SIZE:-$((xP + yD))}}}" +export NNODES="${NNODES:-${SGLANG_DISAGG_TOTAL_NODES:-${SLURM_JOB_NUM_NODES:-$((xP + yD))}}}" export MASTER_PORT="${MASTER_PORT:-23731}" export IO_EP_TP_SIZE="${IO_EP_TP_SIZE:-${SGLANG_TP_SIZE:-8}}" # madengine forwards MASTER_ADDR (rank-0 host) into the container; use it as the diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index 268b795..acf8220 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -113,11 +113,20 @@ mode = os.environ["PARALLEL_MODE"] xP = int(os.environ.get("xP", "1")) gpus_per_node = int(os.environ.get("GPUS_PER_NODE", "8")) + +def q(v): + return shlex.quote(str(v if v is not None else "")) + + with open(config_path, "r", encoding="utf-8") as f: models = yaml.safe_load(f) or {} if model_name not in models: - print(f'echo "ERROR: Model {model_name} not found in {config_path}"; exit 1') + # Quote the interpolated values: this string is eval'd by the shell, so an + # unescaped model_name / config_path could break the launcher or inject + # shell tokens. + msg = f"ERROR: Model {model_name} not found in {config_path}" + print(f"echo {q(msg)}; exit 1") sys.exit(0) cfg = models[model_name] or {} @@ -125,10 +134,6 @@ prefill = cfg.get("prefill", {}) or {} decode = cfg.get("decode", {}) or {} -def q(v): - return shlex.quote(str(v if v is not None else "")) - - prefill_flags = prefill.get(mode, "") or "" exports = { From 5569d64e75f651e98d7df0f56b3885e186d752e5 Mon Sep 17 00:00:00 2001 From: mkuznet1 Date: Thu, 9 Jul 2026 15:26:53 +0200 Subject: [PATCH 23/25] Remove stale scripts/vllm_dissag/ reference in run.sh comment The comment above _tcp_discover_ipaddrs claims ip_rendezvous.py is "shared (duplicated, kept identical) with scripts/vllm_dissag/", but no such copy exists there. Same stale reference already removed from the ip_rendezvous.py docstring; drop it here too so it doesn't send maintainers looking for a non-existent file (Copilot). Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- scripts/sglang_disagg/run.sh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh index 9d54a42..084ea33 100755 --- a/scripts/sglang_disagg/run.sh +++ b/scripts/sglang_disagg/run.sh @@ -103,8 +103,7 @@ IP_SYNC_PORT="${IP_SYNC_PORT:-$((20000 + (${SLURM_JOB_ID:-0} % 40000)))}" # rank-0 acts as a tiny TCP rendezvous server: peers connect to MASTER_ADDR and # report their host IP; rank-0 aggregates the rank-ordered list and broadcasts. -# The implementation lives in the sibling ip_rendezvous.py so it can be unit -# tested and is shared (duplicated, kept identical) with scripts/vllm_dissag/. +# The implementation lives in the sibling ip_rendezvous.py so it can be unit tested. _tcp_discover_ipaddrs() { python3 "$SCRIPT_DIR/ip_rendezvous.py" \ "${NODE_RANK}" "${NNODES}" "${HOST_IP}" "${MASTER_ADDR}" "${IP_SYNC_PORT}" "${SLURM_JOB_ID:-0}" From e534875ae732f86a2cf7d3ed9b61dccff6876e60 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 14:21:31 +0000 Subject: [PATCH 24/25] Polish PR174: DP_MODE default, 405B log hygiene, 70B NCCL plugin - scripts/sglang_disagg/run.sh: default DP_MODE to 0 instead of 1, matching README.MD and sglang_disagg_mori_io_ep.sh's own default. With the old default, running this entrypoint for any non-DeepSeek model without an explicit DP_MODE=0 would fail fast against the launcher's DP_MODE=1 allowlist (DeepSeek-V3/R1 only). The registered DeepSeek-R1 model passes --dp-mode 1 explicitly, so its behavior is unchanged. - scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh: - remove a stale TRAIN_LOG at the start of the Llama-3.1-405B branch. That branch skips run_primus when the config is missing or the datatype is unsupported, so without clearing the log a re-run from the same directory could benchmark/parse a stale training log from a previous run. - default MBS/GBS to NA in the 405B branch before writing the perf CSV, so a missing config or unsupported dtype writes explicit NA markers instead of placeholder rows with empty batch_size/global_batch_size fields. - .claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json: add NCCL_NET_PLUGIN=none to both context.docker_env_vars and deployment_config.env_vars, matching the 8B template and references/gotchas.md. On rocm/primus:v26.4 RCCL-overlay images the base librccl-anp.so plugin must be disabled or RCCL can hang on the first collective, so the 70B multi-node perf run never starts. --- .../assets/manifests/primus_llama-3.1-70b.template.json | 2 ++ .../primus/megatron-lm/primus_megatron-lm_benchmark_report.sh | 3 +++ scripts/sglang_disagg/run.sh | 2 +- 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json index a5a6256..4c1aeac 100644 --- a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json @@ -41,6 +41,7 @@ "RDMAV_DRIVERS": "", "IBV_DRIVERS": "", "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", "IBV_SHOW_WARNINGS": "1" }, "docker_mounts": { @@ -96,6 +97,7 @@ "RDMAV_DRIVERS": "", "IBV_DRIVERS": "", "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", "NCCL_TIMEOUT": "900", "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", "TORCH_NCCL_HIGH_PRIORITY": "1", diff --git a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh index 91ddd71..fd6e891 100755 --- a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh +++ b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh @@ -355,6 +355,7 @@ elif [ "$MODEL_REPO" == "Llama-3.1-70B-proxy" ]; then elif [ "$MODEL_REPO" == "Llama-3.1-405B" ]; then echo "[INFO] $MODEL_REPO TRAINING" # BF16 training only export EXP=examples/megatron/configs/$CONFIG_DEVICE/llama3.1_405B-$DATATYPE-pretrain.yaml + rm -f "$TRAIN_LOG" SEQ_LEN=8192 if [[ -f "$EXP" ]]; then @@ -372,6 +373,8 @@ elif [ "$MODEL_REPO" == "Llama-3.1-405B" ]; then GBS=$(effective_global_batch_size "$MBS" "$GBS") run_primus "$EXP" $GBS_OVERRIDE fi + MBS="${MBS:-NA}" + GBS="${GBS:-NA}" if [ -f "$TRAIN_LOG" ]; then echo "[INFO] Benchmarking" python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --global_batch_size $GBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh index 084ea33..0b2b3d1 100755 --- a/scripts/sglang_disagg/run.sh +++ b/scripts/sglang_disagg/run.sh @@ -139,7 +139,7 @@ export MASTER_ADDR="$(echo "$IPADDRS" | cut -d',' -f1)" export MODEL_NAME="${MODEL_NAME:-}" export MODEL_PATH="${MODEL_PATH:-}" export RUN_MORI="${RUN_MORI:-1}" -export DP_MODE="${DP_MODE:-1}" +export DP_MODE="${DP_MODE:-0}" # Default the KV transfer backend from RUN_MORI so the non-MoRI path does not # silently inherit an invalid `mori` backend (mooncake/NIXL use mooncake). if [[ "$RUN_MORI" == "1" ]]; then From 71c894afff8aca4d1fe86efcda6b94d06d018cfa Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 9 Jul 2026 14:37:03 +0000 Subject: [PATCH 25/25] Restore DP_MODE=0 default in parse_to_csv run metadata _get_run_metadata() defaulted DP_MODE to "1", which this PR had changed from the pre-PR default of "0". With RUN_MORI defaulting to "1", a run parsed without DP_MODE exported (standalone/offline parsing) was tagged mori_dp instead of mori_io, disagreeing with run.sh's DP_MODE=0 default. Restore the "0" default so the metadata backend tag matches the launcher. --- scripts/sglang_disagg/parse_to_csv.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/sglang_disagg/parse_to_csv.py b/scripts/sglang_disagg/parse_to_csv.py index 84661ab..4fe659d 100644 --- a/scripts/sglang_disagg/parse_to_csv.py +++ b/scripts/sglang_disagg/parse_to_csv.py @@ -163,7 +163,7 @@ def _get_run_metadata(pipeline: str = "sglang"): xP = os.environ.get("xP", "1") yD = os.environ.get("yD", "1") - dp_mode = os.environ.get("DP_MODE", "1") + dp_mode = os.environ.get("DP_MODE", "0") run_mori = os.environ.get("RUN_MORI", "1") kv_backend = os.environ.get("KV_TRANSFER_BACKEND", "mooncake").lower() gpus = os.environ.get("GPUS_PER_NODE", "8")