diff --git a/Dockerfile b/Dockerfile index 31fc4929..4a62c886 100644 --- a/Dockerfile +++ b/Dockerfile @@ -64,7 +64,7 @@ ENV BASE_NAME=cpu ############################################################################### # AMD BASE IMAGE -FROM gdiamos/rocm-base-mi300:v0.9926 AS amd +FROM rocm/primus:v26.3 AS amd ENV BASE_NAME=amd @@ -74,8 +74,16 @@ ENV HIP_FORCE_DEV_KERNARG=1 ARG INSTALL_ROOT=/app/cray WORKDIR ${INSTALL_ROOT} +ENV PATH="/app/venv/bin:$PATH" ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/app/venv/lib/python3.12/site-packages/torch/lib:/usr/local/rdma-lib +RUN pip install uv && \ + uv pip install ninja + +# Open MPI from ROCm base image (used by train_job_entrypoint.sh mpirun launcher) +ENV PATH=$PATH:/opt/ompi-rocm/bin +ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/opt/ompi-rocm/lib + ############################################################################### # FRONTEND BUILD STAGE @@ -267,22 +275,18 @@ WORKDIR ${INSTALL_ROOT} # MAIN IMAGE FROM vllm AS infra -# Build GPU-aware MPI -COPY ./infra/cray_infra/training/gpu_aware_mpi ${INSTALL_ROOT}/infra/cray_infra/training/gpu_aware_mpi -RUN python3 ${INSTALL_ROOT}/infra/cray_infra/training/gpu_aware_mpi/setup.py bdist_wheel --dist-dir=dist && \ - pip install dist/*.whl +ARG INSTALL_ROOT=/app/cray RUN apt-get update -y \ && apt-get install -y build-essential \ less curl wget net-tools vim iputils-ping strace gdb python3-dbg python3-dev \ + dmidecode \ + slurm-wlm libslurm-dev \ && rm -rf /var/lib/apt/lists/* # Setup python path -ENV PYTHONPATH="${PYTHONPATH:-}:${INSTALL_ROOT}/infra" -ENV PYTHONPATH="${PYTHONPATH:-}:${INSTALL_ROOT}/sdk" -ENV PYTHONPATH="${PYTHONPATH:-}:${INSTALL_ROOT}/ml" -ENV PYTHONPATH="${PYTHONPATH:-}:${INSTALL_ROOT}/test" -ENV PYTHONPATH="${PYTHONPATH:-}:${INSTALL_ROOT}/vllm" +ENV INSTALL_ROOT=${INSTALL_ROOT} +ENV PYTHONPATH="${INSTALL_ROOT}/infra:${INSTALL_ROOT}/sdk:${INSTALL_ROOT}/ml:${INSTALL_ROOT}/test:${INSTALL_ROOT}/vllm:${INSTALL_ROOT}" # Megatron dependencies (GPU only) # note this has to happen after vllm because it overrides some packages installed by vllm @@ -290,7 +294,8 @@ COPY ./infra/requirements-megatron.txt ${INSTALL_ROOT}/requirements-megatron.txt COPY ./infra/requirements-megatron-cpu.txt ${INSTALL_ROOT}/requirements-megatron-cpu.txt COPY ./requirements.txt ${INSTALL_ROOT}/requirements.txt -RUN if [ "$VLLM_TARGET_DEVICE" != "cpu" ]; then \ +RUN uv pip install --no-deps --no-compile --no-cache-dir torchao==0.17.0 \ + && if [ "$VLLM_TARGET_DEVICE" != "cpu" ]; then \ uv pip install --no-deps --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-megatron.txt; \ fi && \ if [ "$VLLM_TARGET_DEVICE" != "cuda" ]; then \ @@ -307,12 +312,14 @@ COPY ./test ${INSTALL_ROOT}/test COPY ./ml ${INSTALL_ROOT}/ml COPY ./scripts ${INSTALL_ROOT}/scripts +RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/test/requirements-pytest.txt + WORKDIR ${INSTALL_ROOT} # Build SLURM plugin RUN /app/cray/infra/slurm_src/compile.sh -ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:${PYTHONPATH:-}:/usr/local/lib/slurm +ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:${INSTALL_ROOT}/infra:${INSTALL_ROOT}/sdk:${INSTALL_ROOT}/ml:${INSTALL_ROOT}/test:${INSTALL_ROOT}/vllm:/usr/local/lib/slurm ENV SLURM_CONF=${INSTALL_ROOT}/nfs/slurm.conf ENV VLLM_CPU_MOE_PREPACK=0 diff --git a/Dockerfile_rocm b/Dockerfile_rocm new file mode 100644 index 00000000..fbc131db --- /dev/null +++ b/Dockerfile_rocm @@ -0,0 +1,446 @@ +# Build with: ./scalarlm build-image amd + +ARG COMMON_WORKDIR=/app/cray +ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942" +ARG NIC_BACKEND=all +ARG AINIC_VERSION=1.117.5-a-77 +ARG UBUNTU_CODENAME=noble + +############################################################################### +FROM rocm/primus:v26.2 AS amd + +RUN apt-get update && apt-get install -y bc jq libibverbs-dev rdma-core ibverbs-utils && \ + rm -rf /var/lib/apt/lists/* + +RUN apt-get update -y && apt-get install -y python3-venv slurm-wlm libslurm-dev && \ + rm -rf /var/lib/apt/lists/* + +ENV BASE_NAME=amd +ENV HIP_FORCE_DEV_KERNARG=1 + +ARG INSTALL_ROOT=/app/cray +WORKDIR ${INSTALL_ROOT} + +ENV PATH="/app/venv/bin:/opt/ompi-rocm/bin:${PATH}" +ENV LD_LIBRARY_PATH="/app/venv/lib/python3.12/site-packages/torch/lib:/usr/local/rdma-lib:/opt/ompi-rocm/lib" + +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} + +############################################################################### +# vLLM build base (upstream Dockerfile.rocm `base` stage) +FROM amd AS vllm_base + +ARG COMMON_WORKDIR + +RUN apt-get update -q -y && apt-get install -y --no-install-recommends \ + sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ + apt-transport-https ca-certificates wget curl git \ + libnuma-dev libomp-dev libgrpc-dev libgrpc++-dev libprotobuf-dev \ + protobuf-compiler-grpc libcpprest-dev libaio-dev \ + autoconf libtool pkg-config \ + && rm -rf /var/lib/apt/lists/* + +RUN python -m pip install --upgrade pip + +RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ + && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ + && rm -f /tmp/uv-install.sh \ + && uv --version + +ENV UV_HTTP_TIMEOUT=500 +ENV UV_INDEX_STRATEGY="unsafe-best-match" +ENV UV_LINK_MODE=copy + +RUN uv pip install ninja meson auditwheel patchelf tomlkit + +WORKDIR ${COMMON_WORKDIR} + +############################################################################### +# MoRI runtime + NIC backends (upstream Dockerfile.rocm mori_base) +# AINIC 1.117.5-a-77 on noble; Broadcom bnxt via NIC_BACKEND=all +FROM vllm_base AS mori_base + +ARG NIC_BACKEND=all +ARG AINIC_VERSION=1.117.5-a-77 +ARG UBUNTU_CODENAME=noble + +RUN /bin/bash -lc 'set -euo pipefail; \ + \ + install_ainic() { \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ + rm -rf /var/lib/apt/lists/*; \ + mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + install_bnxt() { \ + install -m 0755 -d /etc/apt/keyrings; \ + curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ + -o /etc/apt/keyrings/broadcom-nic.asc; \ + chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ + > /etc/apt/sources.list.d/broadcom-nic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + bnxt-rocelib=235.2.86.0 \ + ; \ + cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/ || true; \ + ldconfig; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + echo "[MORI] Install MoRI proxy deps"; \ + pip install --quiet --ignore-installed blinker && \ + pip install --quiet quart msgpack aiohttp pyzmq; \ + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ + echo "[MORI] AINIC_VERSION=${AINIC_VERSION}"; \ + \ + case "${NIC_BACKEND}" in \ + none) ;; \ + all) install_ainic; install_bnxt ;; \ + ainic) install_ainic ;; \ + bnxt) install_bnxt ;; \ + *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ + esac' + +ENV MORI_NIC_BACKEND=${NIC_BACKEND} +ENV AINIC_VERSION=${AINIC_VERSION} + +############################################################################### +# Fetch vLLM source (ScalarLM local/remote copy) +FROM mori_base AS fetch_vllm + +ARG COMMON_WORKDIR +ARG INSTALL_ROOT=/app/cray +ARG VLLM_SOURCE=remote +ARG VLLM_BRANCH=main +ARG VLLM_REPO=https://github.com/supermassive-intelligence/vllm-fork.git + +COPY scripts/build-copy-vllm.sh ${INSTALL_ROOT}/build-copy-vllm.sh + +RUN --mount=type=bind,source=./vllm,target=/workspace/vllm,rw \ + bash ${INSTALL_ROOT}/build-copy-vllm.sh ${VLLM_SOURCE} ${COMMON_WORKDIR}/vllm \ + /workspace/vllm ${VLLM_REPO} ${VLLM_BRANCH} + +############################################################################### +# Build vLLM wheel (upstream Dockerfile.rocm build_vllm). +# Rust frontend (build_rust.sh) is optional: ScalarLM vllm-fork may not ship it. +FROM fetch_vllm AS build_vllm + +ARG COMMON_WORKDIR + +COPY scripts/install_protoc.sh /tmp/install_protoc.sh + +RUN cd ${COMMON_WORKDIR}/vllm \ + && if [ -f build_rust.sh ] && [ -f rust-toolchain.toml ] && [ -f rust/Cargo.toml ]; then \ + echo "Building vllm-rs Rust frontend" \ + && apt-get update -q -y \ + && apt-get install -y --no-install-recommends ca-certificates curl unzip \ + && rm -rf /var/lib/apt/lists/* \ + && chmod +x /tmp/install_protoc.sh && /tmp/install_protoc.sh && rm -f /tmp/install_protoc.sh \ + && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain none \ + ; else \ + echo "Skipping vllm-rs prebuild (no build_rust.sh in fetched vLLM source)"; \ + fi + +ENV PATH="/root/.cargo/bin:${PATH}" +ENV CARGO_BUILD_JOBS=4 +ENV CARGO_NET_RETRY=10 +ENV RUSTUP_MAX_RETRIES=10 + +RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ + cd ${COMMON_WORKDIR}/vllm \ + && if [ -f build_rust.sh ] && [ -f rust-toolchain.toml ]; then \ + VLLM_RS_TARGET_PATH=${COMMON_WORKDIR}/vllm/vllm/vllm-rs bash build_rust.sh \ + && test -x ${COMMON_WORKDIR}/vllm/vllm/vllm-rs; \ + fi + +RUN --mount=type=cache,target=/root/.cache/pip \ + cd ${COMMON_WORKDIR}/vllm \ + && if [ -f use_existing_torch.py ]; then python use_existing_torch.py; fi \ + && python -m pip install -r requirements/rocm.txt \ + && python setup.py clean --all \ + && python setup.py bdist_wheel --dist-dir=dist + +FROM scratch AS export_vllm +ARG COMMON_WORKDIR +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 + +############################################################################### +# RIXL / UCX for ROCm RDMA (upstream Dockerfile.rocm build_rixl) +FROM vllm_base AS build_rixl + +ARG RIXL_BRANCH="39be1de8" +ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" +ARG UCX_BRANCH="bfb51733" +ARG UCX_REPO="https://github.com/openucx/ucx.git" +ENV ROCM_PATH=/opt/rocm +ENV UCX_HOME=/usr/local/ucx +ENV RIXL_HOME=/usr/local/rixl + +RUN apt-get -y update && apt-get -y install --no-install-recommends \ + autoconf libtool pkg-config \ + libgrpc-dev libgrpc++-dev libprotobuf-dev protobuf-compiler-grpc \ + libcpprest-dev libaio-dev \ + librdmacm1 librdmacm-dev libibverbs1 libibverbs-dev \ + ibverbs-utils rdmacm-utils ibverbs-providers \ + && rm -rf /var/lib/apt/lists/* + +RUN cd /usr/local/src && \ + git clone ${UCX_REPO} && \ + cd ucx && \ + git checkout ${UCX_BRANCH} && \ + ./autogen.sh && \ + mkdir build && cd build && \ + ../configure \ + --prefix=/usr/local/ucx \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-devel-headers \ + --with-rocm=${ROCM_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt && \ + make -j && \ + make install + +ENV PATH=/usr/local/ucx/bin:$PATH +ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} + +# Ubuntu absl (from gRPC) is too old for RIXL; purge so meson uses the abseil-cpp subproject. +RUN apt-get -y update \ + && apt-get purge -y 'libabsl*' || true \ + && apt-get autoremove -y --purge \ + && rm -rf /var/lib/apt/lists/* + +RUN git clone ${RIXL_REPO} /opt/rixl && \ + cd /opt/rixl && \ + git checkout ${RIXL_BRANCH} && \ + meson setup build --prefix=${RIXL_HOME} \ + -Ducx_path=${UCX_HOME} \ + -Drocm_path=${ROCM_PATH} && \ + cd build && \ + ninja && \ + ninja install + +RUN cd /opt/rixl && \ + sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ + contrib/build-wheel.sh && \ + mkdir -p /app/install && \ + _ucx_install_dir=${UCX_HOME} \ + ./contrib/build-wheel.sh \ + --output-dir /app/install \ + --rocm-dir ${ROCM_PATH} \ + --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ + --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins + +############################################################################### +# Install vLLM + RIXL runtime (upstream `final`, with MoRI NIC backends) +FROM mori_base AS vllm_ready + +ARG COMMON_WORKDIR +ARG NIC_BACKEND=all +ARG AINIC_VERSION=1.117.5-a-77 +ENV ROCM_PATH=/opt/rocm +ENV UCX_HOME=/usr/local/ucx +ENV RIXL_HOME=/usr/local/rixl +ENV PATH=/usr/local/ucx/bin:${PATH} +ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${RIXL_HOME}/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH} + +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install -r requirements/rocm.txt \ + && pip uninstall -y vllm || true \ + && uv pip install *.whl + +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + uv pip install /rixl_install/*.whl + +RUN apt-get update -q -y && apt-get install -y --no-install-recommends \ + librdmacm1 libibverbs1 ibverbs-providers ibverbs-utils \ + ffmpeg libsm6 libxext6 libgl1 libdnnl-dev \ + && rm -rf /var/lib/apt/lists/* + +# vLLM dependencies +COPY ./infra/requirements-vllm.txt ${INSTALL_ROOT}/requirements-vllm.txt +RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-vllm.txt + +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm + +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 +ENV TOKENIZERS_PARALLELISM=false +ENV SAFETENSORS_FAST_GPU=1 +ENV VLLM_TARGET_DEVICE=rocm +ENV MORI_NIC_BACKEND=${NIC_BACKEND} +ENV AINIC_VERSION=${AINIC_VERSION} + +############################################################################### +# FRONTEND BUILD STAGE +FROM vllm_ready AS ui_base + +RUN apt-get update -y && \ + apt-get install -y git curl libgomp1 libcurl4 dnsutils nano + +RUN curl -fsSL https://deb.nodesource.com/setup_24.x | bash - && \ + apt-get install -y nodejs && \ + node --version && \ + npm --version + +ARG INSTALL_ROOT=/app +WORKDIR /app + +ARG UI_SOURCE=remote +ARG UI_BRANCH=main +ARG UI_REPO=https://github.com/supermassive-intelligence/chat-ui-fork.git + +COPY scripts/build-copy-chat-ui.sh ${INSTALL_ROOT}/build-copy-chat-ui.sh + +RUN --mount=type=bind,source=./chat-ui,target=/workspace/chat-ui,rw \ + bash ${INSTALL_ROOT}/build-copy-chat-ui.sh ${UI_SOURCE} ${INSTALL_ROOT}/chat-ui \ + /workspace/chat-ui ${UI_REPO} ${UI_BRANCH} + +RUN npm install -g dotenv-cli + +USER root + +RUN mkdir -p /app/ui && \ + touch /app/ui/.env.local && \ + cp ${INSTALL_ROOT}/chat-ui/.env /app/ui/.env && \ + cp ${INSTALL_ROOT}/chat-ui/entrypoint.sh /app/ui/entrypoint.sh && \ + cp ${INSTALL_ROOT}/chat-ui/package.json /app/ui/package.json && \ + cp ${INSTALL_ROOT}/chat-ui/package-lock.json /app/ui/package-lock.json && \ + chmod +x /app/ui/entrypoint.sh + +FROM node:24.2.0 AS ui_builder + +WORKDIR /app +ARG INSTALL_ROOT=/temp + +USER root +RUN apt-get update -y && apt-get install -y git + +ARG UI_SOURCE=remote +ARG UI_BRANCH=main +ARG UI_REPO=https://github.com/supermassive-intelligence/chat-ui-fork.git + +COPY scripts/build-copy-chat-ui.sh ${INSTALL_ROOT}/build-copy-chat-ui.sh + +RUN --mount=type=bind,source=./chat-ui,target=/workspace/chat-ui,rw \ + bash ${INSTALL_ROOT}/build-copy-chat-ui.sh ${UI_SOURCE} ${INSTALL_ROOT}/chat-ui \ + /workspace/chat-ui ${UI_REPO} ${UI_BRANCH} + +RUN cp ${INSTALL_ROOT}/chat-ui/package-lock.json ${INSTALL_ROOT}/chat-ui/package.json ./ + +ARG APP_BASE=/chat +ARG PUBLIC_APP_COLOR= +ENV BODY_SIZE_LIMIT=15728640 + +RUN --mount=type=cache,target=/app/.npm \ + npm set cache /app/.npm && \ + npm ci + +RUN cp -R ${INSTALL_ROOT}/chat-ui/. /app/ && \ + npm install -D @sveltejs/adapter-static + +RUN git config --global --add safe.directory /app && \ + npm run build + +FROM mongo:7 AS mongo + +FROM ui_base AS local_db + +COPY --from=mongo /usr/bin/mongo* /usr/bin/ + +ENV MONGODB_URL=mongodb://localhost:27017 +USER root +RUN mkdir -p /data/db + +FROM local_db AS ui_final + +ENV INCLUDE_DB=true + +ARG APP_BASE=/chat +ARG PUBLIC_APP_COLOR= +ARG PUBLIC_COMMIT_SHA= +ENV PUBLIC_COMMIT_SHA=${PUBLIC_COMMIT_SHA} +ENV BODY_SIZE_LIMIT=15728640 + +COPY --from=ui_builder /app/build /app/build +COPY --from=ui_builder /app/node_modules /app/node_modules +COPY frontend/entrypoint.sh /app/ui/entrypoint.sh +COPY frontend/.env.local /app/ui/.env.local + +############################################################################### +# MAIN IMAGE +FROM ui_final AS infra + +COPY --from=ui_final --chown=1000 /app /app/ui +COPY --from=ui_final /usr/bin/mongo* /usr/bin/ + +ENV MONGODB_URL=mongodb://localhost:27017 +ENV INCLUDE_DB=true + +ARG INSTALL_ROOT=/app/cray +ARG VLLM_TARGET_DEVICE=rocm +ENV VLLM_TARGET_DEVICE=${VLLM_TARGET_DEVICE} + +RUN apt-get update -y \ + && apt-get install -y build-essential \ + less curl wget net-tools vim iputils-ping strace gdb python3-dbg python3-dev \ + dmidecode \ + slurm-wlm libslurm-dev \ + numactl \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR ${INSTALL_ROOT} + +ENV INSTALL_ROOT=${INSTALL_ROOT} +ENV PYTHONPATH="${INSTALL_ROOT}/infra:${INSTALL_ROOT}/sdk:${INSTALL_ROOT}/ml:${INSTALL_ROOT}/test:${INSTALL_ROOT}/vllm:${INSTALL_ROOT}" + +COPY ./infra/requirements-megatron.txt ${INSTALL_ROOT}/requirements-megatron.txt +COPY ./infra/requirements-megatron-cpu.txt ${INSTALL_ROOT}/requirements-megatron-cpu.txt +COPY ./requirements.txt ${INSTALL_ROOT}/requirements.txt + +RUN uv pip install --no-deps --no-compile --no-cache-dir torchao==0.17.0 \ + && if [ "$VLLM_TARGET_DEVICE" != "cpu" ]; then \ + uv pip install --no-deps --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-megatron.txt; \ + fi && \ + if [ "$VLLM_TARGET_DEVICE" != "cuda" ]; then \ + uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-megatron-cpu.txt; \ + fi && \ + uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements.txt + +RUN mkdir -p ${INSTALL_ROOT}/jobs ${INSTALL_ROOT}/nfs + +COPY ./infra ${INSTALL_ROOT}/infra +COPY ./sdk ${INSTALL_ROOT}/sdk +COPY ./test ${INSTALL_ROOT}/test +COPY ./ml ${INSTALL_ROOT}/ml +COPY ./scripts ${INSTALL_ROOT}/scripts + +RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/test/requirements-pytest.txt + +WORKDIR ${INSTALL_ROOT} + +RUN /app/cray/infra/slurm_src/compile.sh + +ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:${INSTALL_ROOT}/infra:${INSTALL_ROOT}/sdk:${INSTALL_ROOT}/ml:${INSTALL_ROOT}/test:${INSTALL_ROOT}/vllm:/usr/local/lib/slurm +ENV SLURM_CONF=${INSTALL_ROOT}/nfs/slurm.conf +ENV VLLM_CPU_MOE_PREPACK=0 diff --git a/cmd/bashly.yml b/cmd/bashly.yml index 2d50703b..040c27a1 100644 --- a/cmd/bashly.yml +++ b/cmd/bashly.yml @@ -30,7 +30,7 @@ commands: help: Run tests in the container args: - name: test-path - default: "test/infra/*" + default: "test/infra" help: Relative path to the directory or file with test cases flags: - long: --coverage-path diff --git a/cmd/llm_logs_command.sh b/cmd/llm_logs_command.sh index a503ac77..9ee53ad7 100644 --- a/cmd/llm_logs_command.sh +++ b/cmd/llm_logs_command.sh @@ -9,7 +9,7 @@ if [ -z "$model" ]; then model="latest" fi -./cray build-image +./scalarlm build-image declare -a log_command_parts log_command_parts=( diff --git a/cmd/llm_ls_command.sh b/cmd/llm_ls_command.sh index a9371201..5f10fb98 100644 --- a/cmd/llm_ls_command.sh +++ b/cmd/llm_ls_command.sh @@ -1,6 +1,6 @@ inspect_args -./cray build-image +./scalarlm build-image declare -a ls_command_parts ls_command_parts=( diff --git a/cmd/llm_plot_command.sh b/cmd/llm_plot_command.sh index 99dcc442..da30210e 100644 --- a/cmd/llm_plot_command.sh +++ b/cmd/llm_plot_command.sh @@ -6,7 +6,7 @@ if [ -z "$model" ]; then model="latest" fi -./cray build-image +./scalarlm build-image declare -a plot_command_parts plot_command_parts=( diff --git a/cmd/llm_squeue_command.sh b/cmd/llm_squeue_command.sh index 3e157ebd..ce0191d6 100644 --- a/cmd/llm_squeue_command.sh +++ b/cmd/llm_squeue_command.sh @@ -1,6 +1,6 @@ inspect_args -./cray build-image +./scalarlm build-image declare -a squeue_command_parts squeue_command_parts=( diff --git a/cmd/test_command.sh b/cmd/test_command.sh index f7f62441..bb290e4f 100644 --- a/cmd/test_command.sh +++ b/cmd/test_command.sh @@ -9,7 +9,7 @@ if [ -z "$tag" ]; then tag="cray:latest" fi -./cray build-image +./scalarlm build-image declare -a start_slurm_command_parts @@ -28,7 +28,12 @@ if [ "yes" == "$verbose" ]; then pytest_command_parts+=("-rP") fi -pytest_command_parts+=($test_path) +for path in $test_path; do + case "$(basename "$path")" in + __pycache__) continue ;; + esac + pytest_command_parts+=("$path") +done pytest_command="${pytest_command_parts[*]}" TTY=-t diff --git a/deployment/helm/gemma-3-4b-it/Chart.yaml b/deployment/helm/gemma-3-4b-it/Chart.yaml new file mode 100644 index 00000000..31eb25ae --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/Chart.yaml @@ -0,0 +1,5 @@ +apiVersion: v2 +appVersion: 1.0.0 +description: A Helm chart for the ScalarLM service +name: scalarlm +version: 1.0.0 diff --git a/deployment/helm/gemma-3-4b-it/templates/_helpers.tpl b/deployment/helm/gemma-3-4b-it/templates/_helpers.tpl new file mode 100644 index 00000000..837c91b3 --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/_helpers.tpl @@ -0,0 +1,102 @@ +gotemplate +{{- define "scalarlm.fullname" -}} +{{- printf "%s" .Chart.Name | trunc 63 | trimSuffix "-" -}} +{{- end -}} + +{{- define "scalarlm.vllmname" -}} +{{- printf "%s-vllm" .Chart.Name | trunc 63 | trimSuffix "-" -}} +{{- end -}} + +{{- define "scalarlm.megatronname" -}} +{{- printf "%s-megatron" .Chart.Name | trunc 63 | trimSuffix "-" -}} +{{- end -}} + +{{- define "scalarlm.cloudflaredFullName" -}} +{{- printf "%s-cloudflared" .Chart.Name | trunc 63 | trimSuffix "-" -}} +{{- end -}} + +{{- define "scalarlm.labels" -}} +app.kubernetes.io/name: {{ include "scalarlm.fullname" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end -}} + +{{- define "scalarlm.vllmlabels" -}} +app.kubernetes.io/name: {{ include "scalarlm.vllmname" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end -}} + +{{- define "scalarlm.megatronlabels" -}} +app.kubernetes.io/name: {{ include "scalarlm.megatronname" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end -}} + +{{- define "scalarlm.cloudflaredLabels" -}} +app.kubernetes.io/name: {{ include "scalarlm.cloudflaredFullName" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end -}} + +{{- define "scalarlm.scheduling" -}} +{{- $na := .Values.nodeAffinity -}} +{{- if and $na $na.hostnames (not (empty $na.hostnames)) }} +affinity: + nodeAffinity: + {{- if $na.preferred }} + preferredDuringSchedulingIgnoredDuringExecution: + - weight: {{ $na.weight | default 100 }} + preference: + matchExpressions: + - key: kubernetes.io/hostname + operator: In + values: + {{- range $na.hostnames }} + - {{ . | quote }} + {{- end }} + {{- else }} + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: kubernetes.io/hostname + operator: In + values: + {{- range $na.hostnames }} + - {{ . | quote }} + {{- end }} + {{- end }} +{{- else if and $na $na.labels (not (empty $na.labels)) }} +affinity: + nodeAffinity: + {{- if $na.preferred }} + preferredDuringSchedulingIgnoredDuringExecution: + - weight: {{ $na.weight | default 100 }} + preference: + matchExpressions: + {{- range $key, $val := $na.labels }} + - key: {{ $key | quote }} + operator: In + values: + - {{ $val | quote }} + {{- end }} + {{- else }} + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + {{- range $key, $val := $na.labels }} + - key: {{ $key | quote }} + operator: In + values: + - {{ $val | quote }} + {{- end }} + {{- end }} +{{- else if not (empty .Values.affinity) }} +affinity: + {{- toYaml .Values.affinity | nindent 2 }} +{{- end }} +{{- with .Values.nodeSelector }} +nodeSelector: + {{- toYaml . | nindent 2 }} +{{- end }} +{{- with .Values.tolerations }} +tolerations: + {{- toYaml . | nindent 2 }} +{{- end }} +{{- end -}} diff --git a/deployment/helm/gemma-3-4b-it/templates/api_configmap.yaml b/deployment/helm/gemma-3-4b-it/templates/api_configmap.yaml new file mode 100644 index 00000000..db3973ea --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/api_configmap.yaml @@ -0,0 +1,15 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ .Release.Name }}-api-configmap +data: + cray-config.yaml: | + model: {{ .Values.model }} + dtype: {{ .Values.dtype }} + max_model_length: {{ .Values.max_model_length }} + gpu_memory_utilization: {{ .Values.gpu_memory_utilization }} + vllm_api_url: "http://{{ include "scalarlm.vllmname" . }}:{{ .Values.service.vllm_port }}" + server_list: api + max_train_time: {{ .Values.max_train_time }} + + diff --git a/deployment/helm/gemma-3-4b-it/templates/api_deployment.yaml b/deployment/helm/gemma-3-4b-it/templates/api_deployment.yaml new file mode 100644 index 00000000..422d328e --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/api_deployment.yaml @@ -0,0 +1,93 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "scalarlm.fullname" . }} +spec: + replicas: 1 + selector: + matchLabels: + {{- include "scalarlm.labels" . | nindent 6 }} + template: + metadata: + labels: + {{- include "scalarlm.labels" . | nindent 8 }} + spec: + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + hostname: {{ include "scalarlm.fullname" . }} + dnsConfig: + searches: + - {{ include "scalarlm.megatronname" . }}-headless.{{ .Release.Namespace }}.svc.cluster.local + - "{{ .Release.Namespace }}.svc.cluster.local" + - "svc.cluster.local" + - "cluster.local" + options: + - name: ndots + value: "1" + {{- include "scalarlm.scheduling" . | nindent 6 }} + containers: + - name: {{ .Chart.Name }} + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: ["/app/cray/scripts/start_one_server.sh"] + env: + - name: NCCL_IB_GID_INDEX + value: "1" + - name: NCCL_IB_HCA + value: "rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7" + - name: NCCL_SOCKET_IFNAME + value: "eth0" + - name: GLOO_SOCKET_IFNAME + value: "eth0" + - name: NCCL_SOCKET_FAMILY + value: "AF_INET" + - name: NCCL_NET_GDR_LEVEL + value: "SYS" + - name: ROCR_VISIBLE_DEVICES + value: "0" + - name: HIP_VISIBLE_DEVICES + value: "0" + - name: CRAY_TRAIN_DEBUG + value: "0" + - name: HSA_NO_SCRATCH_RECLAIM + value: "1" + ports: + - name: http + containerPort: {{ .Values.service.api_port }} + hostPort: {{ .Values.service.api_port }} + protocol: TCP + volumeMounts: + {{- range .Values.volumes }} + - name: {{ .name }} + mountPath: {{ .path }} + {{- end }} + - name: scalarlm-config + mountPath: /app/cray/cray-config.yaml + subPath: cray-config.yaml + - name: scalarlm-jobs + mountPath: /app/cray/jobs + - name: scalarlm-cache + mountPath: /root/.cache/huggingface + - name: scalarlm-slurm-config + mountPath: /app/cray/nfs + volumes: + - name: scalarlm-jobs + persistentVolumeClaim: + claimName: scalarlm-jobs + - name: scalarlm-cache + persistentVolumeClaim: + claimName: scalarlm-cache + - name: scalarlm-slurm-config + persistentVolumeClaim: + claimName: scalarlm-slurm-config + - name: scalarlm-config + configMap: + name: {{ .Release.Name }}-api-configmap + {{- range .Values.volumes }} + - name: {{ .name }} + hostPath: + path: {{ .hostPath }} + {{- end }} + diff --git a/deployment/helm/gemma-3-4b-it/templates/api_service.yaml b/deployment/helm/gemma-3-4b-it/templates/api_service.yaml new file mode 100644 index 00000000..cf91ea8f --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/api_service.yaml @@ -0,0 +1,19 @@ +apiVersion: v1 +kind: Service +metadata: + name: {{ .Release.Name }} +spec: + type: {{ .Values.service.type }} + ports: + - port: {{ .Values.service.api_port }} + targetPort: 8000 + protocol: TCP + name: http + - port: 6817 + targetPort: 6817 + protocol: TCP + name: slurmctld + externalIPs: + - {{ .Values.service.externalIP }} + selector: + {{- include "scalarlm.labels" . | nindent 4 }} diff --git a/deployment/helm/gemma-3-4b-it/templates/cache_pvc.yaml b/deployment/helm/gemma-3-4b-it/templates/cache_pvc.yaml new file mode 100644 index 00000000..16a1704e --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/cache_pvc.yaml @@ -0,0 +1,13 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: {{ .Release.Name }}-cache + annotations: + helm.sh/resource-policy: keep +spec: + accessModes: + - ReadWriteMany + resources: + requests: + storage: {{ .Values.cache_pvc.size }} + storageClassName: {{ .Values.cache_pvc.storageClass }} diff --git a/deployment/helm/gemma-3-4b-it/templates/cloudflare_deployment.yaml b/deployment/helm/gemma-3-4b-it/templates/cloudflare_deployment.yaml new file mode 100644 index 00000000..81d44d6b --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/cloudflare_deployment.yaml @@ -0,0 +1,24 @@ +# templates/deployment.yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "scalarlm.cloudflaredFullName" . }} +spec: + replicas: 1 + selector: + matchLabels: + {{- include "scalarlm.cloudflaredLabels" . | nindent 6 }} + template: + metadata: + labels: + {{- include "scalarlm.cloudflaredLabels" . | nindent 8 }} + spec: + containers: + - name: {{ .Chart.Name }} + image: "cloudflare/cloudflared:2024.8.3" + imagePullPolicy: {{ .Values.image.pullPolicy }} + args: + - tunnel + - run + - --token + - {{ .Values.cloudflared.tunnelToken }} diff --git a/deployment/helm/gemma-3-4b-it/templates/jobs_pvc.yaml b/deployment/helm/gemma-3-4b-it/templates/jobs_pvc.yaml new file mode 100644 index 00000000..d881f8fe --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/jobs_pvc.yaml @@ -0,0 +1,13 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: {{ .Release.Name }}-jobs + annotations: + helm.sh/resource-policy: keep +spec: + accessModes: + - ReadWriteMany + resources: + requests: + storage: {{ .Values.jobs_pvc.size }} + storageClassName: {{ .Values.jobs_pvc.storageClass }} diff --git a/deployment/helm/gemma-3-4b-it/templates/megatron_configmap.yaml b/deployment/helm/gemma-3-4b-it/templates/megatron_configmap.yaml new file mode 100644 index 00000000..5c02a537 --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/megatron_configmap.yaml @@ -0,0 +1,9 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ .Release.Name }}-megatron-configmap +data: + cray-config.yaml: | + model: {{ .Values.model }} + server_list: megatron + max_train_time: {{ .Values.max_train_time }} diff --git a/deployment/helm/gemma-3-4b-it/templates/megatron_deployment.yaml b/deployment/helm/gemma-3-4b-it/templates/megatron_deployment.yaml new file mode 100644 index 00000000..cb74544b --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/megatron_deployment.yaml @@ -0,0 +1,215 @@ +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: {{ include "scalarlm.megatronname" . }} +spec: + serviceName: {{ include "scalarlm.megatronname" . }}-headless + replicas: {{ .Values.training_gpus }} + selector: + matchLabels: + {{- include "scalarlm.megatronlabels" . | nindent 6 }} + template: + metadata: + annotations: + k8s.v1.cni.cncf.io/networks: kube-amd-network/amd-host-device-sbr-nad + labels: + {{- include "scalarlm.megatronlabels" . | nindent 8 }} + spec: + hostNetwork: false + #hostPID: true + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + dnsConfig: + searches: + - {{ include "scalarlm.megatronname" . }}-headless.{{ .Release.Namespace }}.svc.cluster.local + - "{{ .Release.Namespace }}.svc.cluster.local" + - "svc.cluster.local" + - "cluster.local" + options: + - name: ndots + value: "1" + {{- include "scalarlm.scheduling" . | nindent 6 }} + containers: + - name: {{ .Chart.Name }} + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + #securityContext: + #privileged: true + #allowPrivilegeEscalation: true + command: ["/bin/bash", "-c"] + args: + - | + ulimit -l unlimited + exec /app/cray/scripts/start_one_server.sh + env: + - name: POD_NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + - name: MEGATRON_HEADLESS + value: {{ include "scalarlm.megatronname" . }}-headless + - name: NCCL_IB_GID_INDEX + value: "1" + - name: NCCL_IB_HCA + value: "rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7" + - name: NCCL_SOCKET_IFNAME + value: "eth0" + - name: GLOO_SOCKET_IFNAME + value: "eth0" + - name: NCCL_SOCKET_FAMILY + value: "AF_INET" + - name: NCCL_NET_GDR_LEVEL + value: "SYS" + - name: ROCR_VISIBLE_DEVICES + value: "0" + - name: HIP_VISIBLE_DEVICES + value: "0" + - name: CRAY_TRAIN_DEBUG + value: "0" + - name: HSA_NO_SCRATCH_RECLAIM + value: "1" + volumeMounts: + {{- range .Values.volumes }} + - name: {{ .name }} + mountPath: {{ .path }} + {{- end }} + - name: scalarlm-config + mountPath: /app/cray/cray-config.yaml + subPath: cray-config.yaml + - mountPath: /sys/fs/cgroup + name: cgroup + readOnly: true + - name: scalarlm-jobs + mountPath: /app/cray/jobs + - name: scalarlm-cache + mountPath: /root/.cache/huggingface + - name: scalarlm-slurm-config + mountPath: /app/cray/nfs + - mountPath: /dev/shm + name: dshm + #- name: infiniband-dev + #mountPath: /dev/infiniband + #- name: rdma-cm-dev + #mountPath: /dev/infiniband/rdma_cm + #- name: libibverbs-lib + #mountPath: /usr/lib/x86_64-linux-gnu/libibverbs.so.1 + #readOnly: true + #- name: librdmacm-lib + #mountPath: /usr/lib/x86_64-linux-gnu/librdmacm.so.1 + #readOnly: true + #- name: ibverbs-providers + #mountPath: /usr/lib/x86_64-linux-gnu/libibverbs + #readOnly: true + #- name: local-rdma-libs + #mountPath: /usr/local/rdma-lib + #readOnly: true + #- name: infiniband-class + #mountPath: /sys/class/infiniband + #readOnly: true + #- name: host-sys-class-net + #mountPath: /host/sys/class/net + #readOnly: true + #- name: host-proc-net + #mountPath: /host/proc/net + #readOnly: true + #- name: libionic + #mountPath: /usr/lib/x86_64-linux-gnu/libionic.so.1.1.54.0-184 + #readOnly: true + - name: prestart-scripts + mountPath: /scripts + resources: + requests: + amd.com/gpu: 1 + amd.com/nic: 1 + #amd.com/rdma: 1 + #rdma/rdma_shared_device_1: 1 + limits: + amd.com/gpu: 1 + amd.com/nic: 1 + #amd.com/rdma: 1 + #rdma/rdma_shared_device_1: 1 + securityContext: + capabilities: + add: ["IPC_LOCK", "NET_ADMIN", "NET_RAW", "SYS_ADMIN", "IPC_LOCK", "SYS_RESOURCE", "SYS_RAWIO"] + allowPrivilegeEscalation: true + #privileged: true + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 128Gi + - name: scalarlm-jobs + persistentVolumeClaim: + claimName: scalarlm-jobs + - name: prestart-scripts + configMap: + name: add-routes + - name: scalarlm-cache + persistentVolumeClaim: + claimName: scalarlm-cache + - name: scalarlm-slurm-config + persistentVolumeClaim: + claimName: scalarlm-slurm-config + - name: scalarlm-config + configMap: + name: {{ .Release.Name }}-megatron-configmap + - name: infiniband-dev + hostPath: + path: /dev/infiniband + type: Directory + - name: rdma-cm-dev + hostPath: + path: /dev/infiniband/rdma_cm + type: CharDevice + - name: libibverbs-lib + hostPath: + path: /usr/lib/x86_64-linux-gnu/libibverbs.so.1 + type: File + - name: librdmacm-lib + hostPath: + path: /usr/lib/x86_64-linux-gnu/librdmacm.so.1 + type: File + - name: ibverbs-providers + hostPath: + path: /usr/lib/x86_64-linux-gnu/libibverbs + type: Directory + - name: local-rdma-libs + hostPath: + path: /usr/local/lib + type: Directory + - name: infiniband-class + hostPath: + path: /sys/class/infiniband + type: Directory + - name: host-sys-class-net + hostPath: + path: /sys/class/net + type: Directory + - name: host-proc-net + hostPath: + path: /proc/net + type: Directory + - name: rdma-net-devices + hostPath: + path: /sys/class/net + type: Directory + - name: host-netns + hostPath: + path: /var/run/netns + type: DirectoryOrCreate + - name: libionic + hostPath: + path: /usr/lib/x86_64-linux-gnu/libionic.so.1.1.54.0-184 + type: File + - name: cgroup + hostPath: + path: /sys/fs/cgroup + type: Directory + {{- range .Values.volumes }} + - name: {{ .name }} + hostPath: + path: {{ .hostPath }} + {{- end }} + diff --git a/deployment/helm/gemma-3-4b-it/templates/megratron_service.yaml b/deployment/helm/gemma-3-4b-it/templates/megratron_service.yaml new file mode 100644 index 00000000..ec05ac8f --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/megratron_service.yaml @@ -0,0 +1,15 @@ +apiVersion: v1 +kind: Service +metadata: + name: {{ include "scalarlm.megatronname" . }}-headless + labels: + {{- include "scalarlm.megatronlabels" . | nindent 4 }} +spec: + clusterIP: None # This makes it headless + selector: + {{- include "scalarlm.megatronlabels" . | nindent 4 }} + ports: + - port: 6818 + targetPort: 6818 + protocol: TCP + name: slurmd diff --git a/deployment/helm/gemma-3-4b-it/templates/slurm_config_pvc.yaml b/deployment/helm/gemma-3-4b-it/templates/slurm_config_pvc.yaml new file mode 100644 index 00000000..e7cdbde0 --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/slurm_config_pvc.yaml @@ -0,0 +1,13 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: {{ .Release.Name }}-slurm-config + annotations: + helm.sh/resource-policy: keep +spec: + accessModes: + - ReadWriteMany + resources: + requests: + storage: {{ .Values.slurm_config_pvc.size }} + storageClassName: {{ .Values.slurm_config_pvc.storageClass }} diff --git a/deployment/helm/gemma-3-4b-it/templates/vllm_configmap.yaml b/deployment/helm/gemma-3-4b-it/templates/vllm_configmap.yaml new file mode 100644 index 00000000..26f6d8bb --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/vllm_configmap.yaml @@ -0,0 +1,14 @@ +# templates/configmap.yaml +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ .Release.Name }}-vllm-configmap +data: + cray-config.yaml: | + model: {{ .Values.model }} + max_model_length: {{ .Values.max_model_length }} + gpu_memory_utilization: {{ .Values.gpu_memory_utilization }} + api_url: "http://{{ .Release.Name }}:{{ .Values.service.api_port }}" + server_list: vllm + + diff --git a/deployment/helm/gemma-3-4b-it/templates/vllm_deployment.yaml b/deployment/helm/gemma-3-4b-it/templates/vllm_deployment.yaml new file mode 100644 index 00000000..9257405d --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/vllm_deployment.yaml @@ -0,0 +1,88 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "scalarlm.vllmname" . }} +spec: + replicas: {{ .Values.inference_gpus }} + selector: + matchLabels: + {{- include "scalarlm.vllmlabels" . | nindent 6 }} + template: + metadata: + annotations: + k8s.v1.cni.cncf.io/networks: kube-amd-network/amd-host-device-sbr-nad + labels: + {{- include "scalarlm.vllmlabels" . | nindent 8 }} + spec: + hostNetwork: false + hostIPC: true + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- include "scalarlm.scheduling" . | nindent 6 }} + containers: + - name: {{ .Chart.Name }} + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: ["/bin/bash", "-c"] + args: + - | + chmod +x /scripts/add-routes.sh + bash /scripts/add-routes.sh + exec /app/cray/scripts/start_one_server.sh + env: + - name: UCX_TLS + value: "rc_verbs,ud_verbs,self, rocm_copy, rocm_ipc" + - name: UCX_RC_ENABLE + value: "y" + - name: UCX_LOG_LEVEL + value: "info" + - name: OMPI_MCA_pml + value: "ucx" + ports: + - name: http + containerPort: 8001 + volumeMounts: + {{- range .Values.volumes }} + - name: {{ .name }} + mountPath: {{ .path }} + {{- end }} + - name: scalarlm-config + mountPath: /app/cray/cray-config.yaml + subPath: cray-config.yaml + - name: scalarlm-jobs + mountPath: /app/cray/jobs + - name: scalarlm-cache + mountPath: /root/.cache/huggingface + - name: prestart-scripts + mountPath: /scripts + resources: + requests: + amd.com/gpu: 1 + amd.com/nic: 1 + limits: + amd.com/gpu: 1 + amd.com/nic: 1 + securityContext: + capabilities: + add: ["NET_ADMIN", "NET_RAW", "SYS_ADMIN", "IPC_LOCK", "SYS_RESOURCE", "SYS_RAWIO"] + volumes: + - name: scalarlm-jobs + persistentVolumeClaim: + claimName: scalarlm-jobs + - name: prestart-scripts + configMap: + name: add-routes + - name: scalarlm-cache + persistentVolumeClaim: + claimName: scalarlm-cache + - name: scalarlm-config + configMap: + name: {{ .Release.Name }}-vllm-configmap + {{- range .Values.volumes }} + - name: {{ .name }} + hostPath: + path: {{ .hostPath }} + {{- end }} + diff --git a/deployment/helm/gemma-3-4b-it/templates/vllm_service.yaml b/deployment/helm/gemma-3-4b-it/templates/vllm_service.yaml new file mode 100644 index 00000000..2fa9f0e4 --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/templates/vllm_service.yaml @@ -0,0 +1,15 @@ +apiVersion: v1 +kind: Service +metadata: + name: {{ include "scalarlm.vllmname" . }} +spec: + type: {{ .Values.service.type }} + ports: + - port: {{ .Values.service.vllm_port }} + targetPort: 8001 + protocol: TCP + name: http + externalIPs: + - {{ .Values.service.externalIP }} + selector: + {{- include "scalarlm.vllmlabels" . | nindent 4 }} diff --git a/deployment/helm/gemma-3-4b-it/values.yaml b/deployment/helm/gemma-3-4b-it/values.yaml new file mode 100644 index 00000000..2243ee18 --- /dev/null +++ b/deployment/helm/gemma-3-4b-it/values.yaml @@ -0,0 +1,60 @@ +image: + repository: avesuni/scalarlm-rccl + tag: v0.0.83 + pullPolicy: Always + +nad: + create: false + +networkConfig: + create: false + +# Optional pod scheduling (empty = no constraint). +# Pin to nodes: set nodeAffinity.hostnames and/or affinity. +nodeAffinity: + hostnames: [] + labels: {} + # preferred: false + # weight: 100 +nodeSelector: {} +affinity: {} +tolerations: [] + +service: + type: ClusterIP + api_port: 8000 + vllm_port: 8001 + externalIP: 64.139.222.102 + +jobs_pvc: + storageClass: longhorn + size: 100Gi + +cache_pvc: + storageClass: longhorn + size: 400Gi + +slurm_config_pvc: + storageClass: longhorn + size: 10Gi + +cloudflared: + tunnelToken: eyJhIjoiNDYyNTI5ZDY5NTEyZjg5MTg3OTJjY2M4ZmZkODlhY2QiLCJ0IjoiZTFhZThmZGUtM2ExYi00ZjhlLThhMDgtMDA5Y2YyNWE2MDZkIiwicyI6IlpUaGpabVkzWmpNdE56RXlOaTAwTkRSaExUazBaRGd0TkRrNVlqSmxZems1TkRjNCJ9 + +model: google/gemma-3-4b-it +max_model_length: 4096 +gpu_memory_utilization: 0.95 +dtype: bfloat16 + +training_gpus: 16 +inference_gpus: 1 + +max_train_time: 86400 + +volumes: + - name: modules + hostPath: /lib/modules + path: /lib/modules + - name: sources + hostPath: /usr/src + path: /usr/src diff --git a/deployment/helm/gemma_270m_tw/scalarlm/templates/vllm_configmap.yaml b/deployment/helm/gemma_270m_tw/scalarlm/templates/vllm_configmap.yaml index 26f6d8bb..8c41f2b0 100644 --- a/deployment/helm/gemma_270m_tw/scalarlm/templates/vllm_configmap.yaml +++ b/deployment/helm/gemma_270m_tw/scalarlm/templates/vllm_configmap.yaml @@ -8,7 +8,7 @@ data: model: {{ .Values.model }} max_model_length: {{ .Values.max_model_length }} gpu_memory_utilization: {{ .Values.gpu_memory_utilization }} - api_url: "http://{{ .Release.Name }}:{{ .Values.service.api_port }}" + api_url: "http://{{ include "scalarlm.fullname" . }}:{{ .Values.service.api_port }}" server_list: vllm diff --git a/infra/cray_infra/api/fastapi/routers/megatron_router.py b/infra/cray_infra/api/fastapi/routers/megatron_router.py index e19c5dc4..6efb45c5 100644 --- a/infra/cray_infra/api/fastapi/routers/megatron_router.py +++ b/infra/cray_infra/api/fastapi/routers/megatron_router.py @@ -49,6 +49,9 @@ async def train(request: Request): detail="Invalid request body", ) + if job_config.get("retry") or job_config.get("force_resubmit"): + logger.info("Training resubmit requested (retry=true)") + job_status = await launch_training_job(job_config) return TrainResponse(job_status=job_status, job_config=job_config, deployed=False) diff --git a/infra/cray_infra/api/fastapi/routers/openai_v1_router.py b/infra/cray_infra/api/fastapi/routers/openai_v1_router.py index 9fb2b6a3..3dcedc17 100644 --- a/infra/cray_infra/api/fastapi/routers/openai_v1_router.py +++ b/infra/cray_infra/api/fastapi/routers/openai_v1_router.py @@ -12,11 +12,13 @@ ) from cray_infra.api.fastapi.aiohttp.get_global_session import get_global_session +from cray_infra.one_server.vllm_app_registry import get_vllm_app from cray_infra.util.get_config import get_config from fastapi import APIRouter, Request -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse +import httpx import logging logger = logging.getLogger(__name__) @@ -24,25 +26,130 @@ openai_v1_router = APIRouter() +async def _proxy_vllm_inprocess(path: str, params: dict, error_label: str): + app = get_vllm_app() + if app is None: + return None + + if params.get("stream"): + + async def generator(): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://vllm", timeout=120.0 + ) as client: + async with client.stream("POST", path, json=params) as resp: + if resp.status_code != 200: + error_text = await resp.aread() + logger.error( + f"vLLM {error_label} error ({resp.status_code}): {error_text!r}" + ) + yield ( + f'data: {{"error": "Failed to create {error_label}: ' + f'{error_text.decode()}"}}\n\n' + ) + return + async for chunk in resp.aiter_bytes(): + yield chunk + + return StreamingResponse( + content=generator(), media_type="text/event-stream" + ) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://vllm", timeout=120.0 + ) as client: + response = await client.post(path, json=params) + content_type = response.headers.get("Content-Type", "application/json") + if response.status_code != 200: + logger.error( + f"vLLM {error_label} error ({response.status_code}): {response.content!r}" + ) + return Response( + content=response.content, + status_code=response.status_code, + media_type=content_type, + ) + + +async def _proxy_vllm_http(session, url: str, params: dict, error_label: str): + if params.get("stream"): + + async def generator(): + async with session.post(url, json=params) as resp: + if resp.status != 200: + error_text = await resp.text() + logger.error(f"vLLM {error_label} error ({resp.status}): {error_text}") + yield ( + f'data: {{"error": "Failed to create {error_label}: ' + f'{error_text}"}}\n\n' + ) + return + async for chunk in resp.content.iter_any(): + yield chunk + + return StreamingResponse(content=generator(), media_type="text/event-stream") + + async with session.post(url, json=params) as resp: + body = await resp.read() + content_type = resp.headers.get("Content-Type", "application/json") + if resp.status != 200: + logger.error(f"vLLM {error_label} error ({resp.status}): {body!r}") + return Response(content=body, status_code=resp.status, media_type=content_type) + + +async def _proxy_vllm_post(path: str, url: str, params: dict, error_label: str): + inprocess_response = await _proxy_vllm_inprocess(path, params, error_label) + if inprocess_response is not None: + return inprocess_response + + session = get_global_session() + return await _proxy_vllm_http(session, url, params, error_label) + + +async def _proxy_vllm_get(path: str, url: str, error_label: str): + app = get_vllm_app() + if app is not None: + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://vllm", timeout=120.0 + ) as client: + response = await client.get(path) + if response.status_code == 200: + return response.json() + return JSONResponse( + content={"error": f"Failed to fetch {error_label}: {response.status_code}"}, + status_code=response.status_code, + ) + + session = get_global_session() + async with session.get(url) as resp: + if resp.status == 200: + return await resp.json() + return JSONResponse( + content={"error": f"Failed to fetch {error_label}: {resp.status}"}, + status_code=resp.status, + ) + + @openai_v1_router.get("/models") async def list_models(): """List available models - proxy to vLLM server.""" - session = get_global_session() config = get_config() - async with session.get(config["vllm_api_url"] + "/v1/models") as resp: - if resp.status == 200: - return await resp.json() - else: - return JSONResponse( - content={"error": f"Failed to fetch models: {resp.status}"}, - status_code=resp.status - ) + try: + return await _proxy_vllm_get( + "/v1/models", + config["vllm_api_url"] + "/v1/models", + "models", + ) + except Exception as e: + return JSONResponse(content={"error": str(e)}, status_code=500) @openai_v1_router.post("/completions") async def create_completions(request: CompletionRequest, raw_request: Request): """Create completions - proxy to vLLM server.""" - session = get_global_session() config = get_config() logger.info(f"Received completions request: {request.model_dump(exclude_none=True)}") @@ -69,22 +176,13 @@ async def create_completions(request: CompletionRequest, raw_request: Request): if value is not None and key in allowed_keys } - async def generator(): - async with session.post( - config["vllm_api_url"] + "/v1/completions", - json=params, - ) as resp: - if resp.status != 200: - error_text = await resp.text() - logger.error(f"vLLM completions error ({resp.status}): {error_text}") - yield f'data: {{"error": "Failed to create completion: {error_text}"}}\n\n' - return - - async for chunk in resp.content.iter_any(): - yield chunk - try: - return StreamingResponse(content=generator(), media_type="text/event-stream") + return await _proxy_vllm_post( + "/v1/completions", + config["vllm_api_url"] + "/v1/completions", + params, + "completion", + ) except Exception as e: return JSONResponse(content={"error": str(e)}, status_code=500) @@ -92,7 +190,6 @@ async def generator(): @openai_v1_router.post("/chat/completions") async def create_chat_completions(request: ChatCompletionRequest, raw_request: Request): """Create chat completions - proxy to vLLM server.""" - session = get_global_session() config = get_config() logger.info(f"Received chat completions request: {request.model_dump(exclude_none=True)}") @@ -119,22 +216,12 @@ async def create_chat_completions(request: ChatCompletionRequest, raw_request: R if value is not None and key in allowed_keys } - async def generator(): - async with session.post( - config["vllm_api_url"] + "/v1/chat/completions", - json=params, - ) as resp: - if resp.status != 200: - error_text = await resp.text() - logger.error(f"vLLM chat completions error ({resp.status}): {error_text}") - yield f'data: {{"error": "Failed to create chat completion: {error_text}"}}\n\n' - return - - async for chunk in resp.content.iter_any(): - yield chunk - try: - return StreamingResponse(content=generator(), media_type="text/event-stream") + return await _proxy_vllm_post( + "/v1/chat/completions", + config["vllm_api_url"] + "/v1/chat/completions", + params, + "chat completion", + ) except Exception as e: return JSONResponse(content={"error": str(e)}, status_code=500) -# test diff --git a/infra/cray_infra/api/fastapi/tasks/add_megatron_tasks.py b/infra/cray_infra/api/fastapi/tasks/add_megatron_tasks.py index 355bb53c..1fbdf430 100644 --- a/infra/cray_infra/api/fastapi/tasks/add_megatron_tasks.py +++ b/infra/cray_infra/api/fastapi/tasks/add_megatron_tasks.py @@ -1,6 +1,7 @@ from cray_infra.util.get_config import get_config from cray_infra.training.restart_megatron_jobs import restart_megatron_jobs +from cray_infra.training.sync_training_job_status import sync_training_job_status from cray_infra.training.register_megatron_models import register_megatron_models from cray_infra.training.register_megatron_workers import register_megatron_workers from cray_infra.generate.clear_acked_requests_from_queue import clear_acked_requests_from_queue @@ -27,14 +28,14 @@ async def add_megatron_tasks(app): @repeat_every(seconds=megatron_refresh_period) async def run_megatron_tasks(): try: + await sync_training_job_status() await register_megatron_models() await restart_megatron_jobs() await register_megatron_workers() await clear_acked_requests_from_queue() await setup_frontend() - except Exception as e: + except Exception: print_exception() - raise e await run_megatron_tasks() diff --git a/infra/cray_infra/one_server/create_megatron.py b/infra/cray_infra/one_server/create_megatron.py index e1f17867..79ae2388 100644 --- a/infra/cray_infra/one_server/create_megatron.py +++ b/infra/cray_infra/one_server/create_megatron.py @@ -37,9 +37,8 @@ async def add_megatron_tasks(app): async def run_megatron_tasks(): try: await register_megatron_workers() - except Exception as e: + except Exception: print_exception() - raise e await run_megatron_tasks() diff --git a/infra/cray_infra/one_server/create_vllm.py b/infra/cray_infra/one_server/create_vllm.py index e0820911..04d0cf21 100644 --- a/infra/cray_infra/one_server/create_vllm.py +++ b/infra/cray_infra/one_server/create_vllm.py @@ -2,23 +2,26 @@ import os import torch +from cray_infra.one_server.vllm_app_registry import register_vllm_app from cray_infra.util.get_config import get_config from cray_infra.huggingface.get_hf_token import get_hf_token -from vllm.entrypoints.openai.api_server import build_app, decorate_logs, \ - init_app_state, setup_server, \ - load_log_config, build_async_engine_client - -from vllm.tool_parsers import ToolParserManager +from vllm.entrypoints.openai.api_server import ( + build_app, + build_async_engine_client, + decorate_logs, + init_app_state, + setup_server, +) +from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config from vllm.entrypoints.launcher import serve_http - from vllm.entrypoints.openai.cli_args import make_arg_parser +from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.entrypoints.utils import (log_non_default_args) import vllm.envs as envs -import uvicorn import logging logger = logging.getLogger(__name__) @@ -96,6 +99,7 @@ async def create_vllm(server_status, port): args = parser.parse_args(args=args) + args.host = "0.0.0.0" args.port = port args.model = config["model"] @@ -127,20 +131,22 @@ async def run_server_worker(server_status, listen_address, server_index = client_config.get("client_index", 0) if client_config else 0 - # Load logging config for uvicorn if specified - log_config = load_log_config(args.log_config_file) + log_config = get_uvicorn_log_config(args) if log_config is not None: - uvicorn_kwargs['log_config'] = log_config + uvicorn_kwargs["log_config"] = log_config async with build_async_engine_client( args, client_config=client_config, ) as engine_client: + supported_tasks = await engine_client.get_supported_tasks() + model_config = engine_client.model_config - app = build_app(args) + app = build_app(args, supported_tasks, model_config) server_status.set_app(app) - await init_app_state(engine_client, app.state, args) + await init_app_state(engine_client, app.state, args, supported_tasks) + register_vllm_app(app) logger.info("Starting vLLM API server %d on %s", server_index, listen_address) @@ -159,6 +165,7 @@ async def run_server_worker(server_status, listen_address, ssl_certfile=args.ssl_certfile, ssl_ca_certs=args.ssl_ca_certs, ssl_cert_reqs=args.ssl_cert_reqs, + ssl_ciphers=args.ssl_ciphers, h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, h11_max_header_count=args.h11_max_header_count, **uvicorn_kwargs, diff --git a/infra/cray_infra/one_server/start_cray_server.py b/infra/cray_infra/one_server/start_cray_server.py index 01f29073..8c48212d 100644 --- a/infra/cray_infra/one_server/start_cray_server.py +++ b/infra/cray_infra/one_server/start_cray_server.py @@ -3,6 +3,8 @@ from cray_infra.one_server.create_megatron import create_megatron from cray_infra.one_server.create_generate_worker import create_generate_worker +from cray_infra.one_server.vllm_app_registry import clear_vllm_app + import asyncio import logging @@ -74,6 +76,8 @@ async def shutdown(self): logger.debug(f"Server {server} is cancelled") await server.shutdown() + clear_vllm_app() + def set_app(self, app): self.app = app logger.debug(f"VLLM app set: {app}") diff --git a/infra/cray_infra/one_server/vllm_app_registry.py b/infra/cray_infra/one_server/vllm_app_registry.py new file mode 100644 index 00000000..d594a1b2 --- /dev/null +++ b/infra/cray_infra/one_server/vllm_app_registry.py @@ -0,0 +1,35 @@ +"""Registry for the in-process vLLM FastAPI app used by one_server mode.""" + +import asyncio + +from fastapi import FastAPI + +_vllm_app: FastAPI | None = None +_vllm_ready = asyncio.Event() + + +def register_vllm_app(app: FastAPI) -> None: + global _vllm_app + _vllm_app = app + _vllm_ready.set() + + +def get_vllm_app() -> FastAPI | None: + return _vllm_app + + +async def wait_for_vllm_app(timeout: float | None = 120) -> FastAPI: + try: + await asyncio.wait_for(_vllm_ready.wait(), timeout=timeout) + except asyncio.TimeoutError as exc: + raise TimeoutError("vLLM app did not register within timeout") from exc + + if _vllm_app is None: + raise RuntimeError("vLLM app registered event set but app is missing") + return _vllm_app + + +def clear_vllm_app() -> None: + global _vllm_app + _vllm_app = None + _vllm_ready.clear() diff --git a/infra/cray_infra/one_server/wait_for_vllm.py b/infra/cray_infra/one_server/wait_for_vllm.py index 109d1bc7..cd850aba 100644 --- a/infra/cray_infra/one_server/wait_for_vllm.py +++ b/infra/cray_infra/one_server/wait_for_vllm.py @@ -1,28 +1,64 @@ +from cray_infra.one_server.vllm_app_registry import get_vllm_app, wait_for_vllm_app from cray_infra.util.get_config import get_config import asyncio +import time + import aiohttp +import httpx import logging logger = logging.getLogger(__name__) +DEFAULT_TIMEOUT_SECONDS = 120 + + +async def wait_for_vllm(timeout: float = DEFAULT_TIMEOUT_SECONDS): + deadline = time.monotonic() + timeout -async def wait_for_vllm(): - for _ in range(30): - health_status = await get_vllm_health() - if health_status == 200: + while time.monotonic() < deadline: + app = get_vllm_app() + if app is not None: + if await _inprocess_health(app) == 200: + return + elif await get_vllm_health() == 200: return await asyncio.sleep(1) + raise TimeoutError( + f"vLLM did not become healthy within {timeout} seconds" + ) + + +async def _inprocess_health(app) -> int: + try: + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://vllm" + ) as client: + response = await client.get("/health") + return response.status_code + except Exception as exc: + logger.error(f"Error getting in-process vLLM health: {exc}") + return 500 + async def get_vllm_health(): + app = get_vllm_app() + if app is not None: + return await _inprocess_health(app) + config = get_config() try: async with aiohttp.ClientSession() as session: async with session.get(config["vllm_api_url"] + "/health") as response: return response.status - except Exception as e: - logger.error(f"Error getting health: {e}") + except Exception as exc: + logger.error(f"Error getting vLLM health: {exc}") return 500 + + +async def ensure_vllm_app(timeout: float = DEFAULT_TIMEOUT_SECONDS): + await wait_for_vllm_app(timeout=timeout) diff --git a/infra/cray_infra/slurm/discovery/discover_clusters.py b/infra/cray_infra/slurm/discovery/discover_clusters.py index e09d2ad9..a61a8126 100644 --- a/infra/cray_infra/slurm/discovery/discover_clusters.py +++ b/infra/cray_infra/slurm/discovery/discover_clusters.py @@ -9,6 +9,7 @@ import socket import os import json +import shutil import logging @@ -29,17 +30,27 @@ def main(): discover_clusters() -def discover_clusters(): +def discover_clusters(quiet: bool = False): clean_old_node_info() + if not should_register_node(): + if not quiet: + logger.info( + "Skipping cluster registration on %s (not a megatron training node)", + get_hostname(), + ) + return + node_info = get_node_info() - save_node_info(node_info) + if not save_node_info(node_info, quiet=quiet): + if quiet: + return - cluster_info = get_cluster_info(node_info) + cluster_info = get_cluster_info(node_info, quiet=quiet) - save_cluster_info(cluster_info) + save_cluster_info(cluster_info, quiet=quiet) def setup_logging(): @@ -86,22 +97,44 @@ def get_node_info(): def get_machine_id(): - machine_id = None - try: - return get_board_serial() - except Exception as e: - logger.error(f"Error reading machine ID: {e}") - return machine_id + serial = get_board_serial() + if serial: + return serial + return get_hostname() def get_board_serial() -> str | None: + if shutil.which("dmidecode") is None: + return None + + if not os.access("/dev/mem", os.R_OK): + return None + result = subprocess.run( - ["dmidecode", "-s", "baseboard-serial-number"], capture_output=True, text=True + ["dmidecode", "-s", "baseboard-serial-number"], + capture_output=True, + text=True, ) + if result.returncode != 0: + return None serial = result.stdout.strip() return serial if serial else None +def should_register_node() -> bool: + hostname = get_hostname() + if "megatron" in hostname: + return True + + config = get_config() + server_list = config.get("server_list", "") + return "megatron" in server_list or server_list == "all" + + +def is_training_node(hostname: str) -> bool: + return "megatron" in hostname + + def get_hostname(): return socket.gethostname() @@ -114,25 +147,37 @@ def get_gpu_count(): gpu_count = 0 if torch.cuda.is_available(): gpu_count = torch.cuda.device_count() + if gpu_count == 0: + gpu_count = len(get_gpu_indexes()) return gpu_count -def save_node_info(node_info): +def save_node_info(node_info, quiet: bool = False) -> bool: node_config_path = os.path.join( shared_node_config_directory, f"{node_info['hostname']}.json" ) os.makedirs(shared_node_config_directory, exist_ok=True) + if os.path.exists(node_config_path): + with open(node_config_path, "r") as f: + existing = json.load(f) + if existing == node_info: + return False + with open(node_config_path, "w") as f: json.dump(node_info, f, indent=4) + if not quiet: + logger.info("Updated node registration for %s", node_info["hostname"]) + return True -def get_cluster_info(node_info): + +def get_cluster_info(node_info, quiet: bool = False): all_nodes = load_all_nodes() - controller_info = elect_controller(all_nodes) + controller_info = elect_controller(all_nodes, quiet=quiet) return { "controller_info": controller_info, @@ -143,16 +188,20 @@ def get_cluster_info(node_info): def load_all_nodes(): all_nodes = [] + if not os.path.exists(shared_node_config_directory): + return all_nodes + for filename in os.listdir(shared_node_config_directory): if filename.endswith(".json"): file_path = os.path.join(shared_node_config_directory, filename) with open(file_path, "r") as f: node_info = json.load(f) - all_nodes.append(node_info) + if is_training_node(node_info["hostname"]): + all_nodes.append(node_info) return all_nodes -def elect_controller(all_nodes): +def elect_controller(all_nodes, quiet: bool = False): """ Elects the controller node based on the lowest GPU count. If multiple nodes have the same CPU count, alphabetical order of hostname is used. @@ -166,9 +215,12 @@ def elect_controller(all_nodes): # The first node in the sorted list is the controller controller_node = all_nodes[0] - logger.info( - f"Controller node elected: {controller_node['hostname']} with {controller_node['gpu_count']} GPUs" - ) + if not quiet: + logger.info( + "Controller node elected: %s with %s GPUs", + controller_node["hostname"], + controller_node["gpu_count"], + ) return controller_node @@ -178,12 +230,13 @@ def is_controller(node_info, controller_info): return is_controller -def save_cluster_info(cluster_info): +def save_cluster_info(cluster_info, quiet: bool = False): old_cluster_info = load_cluster_info_file() if old_cluster_info: if old_cluster_info == cluster_info: - logger.info("Cluster info is unchanged, skipping write.") + if not quiet: + logger.info("Cluster info is unchanged, skipping write.") return write_slurm_config(cluster_info) @@ -192,6 +245,9 @@ def save_cluster_info(cluster_info): write_cluster_info_file(cluster_info) reload_slurm_configs() + if not quiet: + logger.info("Cluster info updated and Slurm configs reloaded.") + def load_cluster_info_file(): if not os.path.exists(cluster_info_path): @@ -348,8 +404,7 @@ def get_gpu_indexes(): try: index_str = file[len(card_name) :] if index_str.isdigit(): - print(file[len(card_name) :]) - index_as_int = int(file[len(card_name) :]) + index_as_int = int(index_str) indexes.append(index_as_int) except Exception as e: continue diff --git a/infra/cray_infra/training/distributed.py b/infra/cray_infra/training/distributed.py new file mode 100644 index 00000000..487f9218 --- /dev/null +++ b/infra/cray_infra/training/distributed.py @@ -0,0 +1,146 @@ +import os +import sys +import time + +import torch +import torch.distributed as dist + +from cray_infra.training.train_debug import is_train_debug_enabled + + +def _trace(msg: str) -> None: + if not is_train_debug_enabled(): + return + rank = os.environ.get("RANK", os.environ.get("SLURM_PROCID", "?")) + line = f"[rank={rank}] dist [{time.monotonic():.3f}]: {msg}\n" + sys.stderr.write(line) + sys.stderr.flush() + + +def _cuda_device() -> torch.device: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + return torch.device("cuda", local_rank) + + +def cuda_device() -> torch.device: + return _cuda_device() + + +def _to_cuda(tensor: torch.Tensor) -> torch.Tensor: + if tensor.is_cuda: + return tensor + return tensor.to(_cuda_device()) + + +def init(): + _trace( + f"init() entered pid={os.getpid()} " + f"RANK={os.environ.get('RANK')} LOCAL_RANK={os.environ.get('LOCAL_RANK')} " + f"WORLD_SIZE={os.environ.get('WORLD_SIZE')} " + f"MASTER_ADDR={os.environ.get('MASTER_ADDR')} MASTER_PORT={os.environ.get('MASTER_PORT')}" + ) + if dist.is_initialized(): + _trace("init() skipped, already initialized") + return + if not torch.cuda.is_available(): + raise RuntimeError("CUDA/ROCm is required for distributed training") + + _trace("calling init_process_group(backend=nccl)") + dist.init_process_group(backend="nccl") + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + _trace( + f"set_device({local_rank}) after init " + f"cuda_current_device={torch.cuda.current_device()}" + ) + _trace( + f"init_process_group complete rank={dist.get_rank()} " + f"world_size={dist.get_world_size()} backend={dist.get_backend()}" + ) + + +def _ensure_initialized(): + if not dist.is_initialized(): + init() + + +def get_rank(): + _ensure_initialized() + return dist.get_rank() + + +def get_size(): + _ensure_initialized() + return dist.get_world_size() + + +def barrier(): + if not dist.is_initialized(): + return + dist.barrier() + + +def _ensure_contiguous(tensor): + return tensor if tensor.is_contiguous() else tensor.contiguous() + + +def allgather(sendbuf, recvbuf): + _ensure_initialized() + sendbuf = _ensure_contiguous(_to_cuda(sendbuf)) + recvbuf = _ensure_contiguous(_to_cuda(recvbuf)) + dist.all_gather_into_tensor(recvbuf, sendbuf) + return recvbuf + + +def reduce_scatter(sendbuf, recvbuf): + _ensure_initialized() + sendbuf = _ensure_contiguous(_to_cuda(sendbuf)) + recvbuf = _ensure_contiguous(_to_cuda(recvbuf)) + dist.reduce_scatter_tensor(recvbuf, sendbuf, op=dist.ReduceOp.SUM) + return recvbuf + + +def allreduce(tensor, op=dist.ReduceOp.SUM): + _ensure_initialized() + tensor = _ensure_contiguous(_to_cuda(tensor)) + dist.all_reduce(tensor, op=op) + return tensor + + +def alltoall(sendbuf, recvbuf=None): + _ensure_initialized() + sendbuf = _ensure_contiguous(_to_cuda(sendbuf)) + world_size = get_size() + if sendbuf.numel() % world_size != 0: + raise ValueError("alltoall send buffer numel must be divisible by world size") + + input_list = [_ensure_contiguous(chunk) for chunk in sendbuf.chunk(world_size, dim=0)] + + if recvbuf is None: + recvbuf = torch.empty_like(sendbuf) + recvbuf = _ensure_contiguous(_to_cuda(recvbuf)) + output_list = list(recvbuf.chunk(world_size, dim=0)) + + dist.all_to_all(output_list, input_list) + return recvbuf + + +def send(tensor, dest: int): + _ensure_initialized() + tensor = _ensure_contiguous(_to_cuda(tensor)) + dist.send(tensor, dst=dest) + return tensor + + +def recv(tensor, source: int): + _ensure_initialized() + tensor = _ensure_contiguous(_to_cuda(tensor)) + dist.recv(tensor, src=source) + return tensor + + +def finalize(): + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() diff --git a/infra/cray_infra/training/gpu_aware_mpi/gpu_aware_mpi.cpp b/infra/cray_infra/training/gpu_aware_mpi/gpu_aware_mpi.cpp deleted file mode 100644 index 92e44a81..00000000 --- a/infra/cray_infra/training/gpu_aware_mpi/gpu_aware_mpi.cpp +++ /dev/null @@ -1,259 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include - -static bool mpi_initialized = false; - -void ensure_mpi_initialized() { - if (!mpi_initialized) { - MPI_Init(nullptr, nullptr); - mpi_initialized = true; - } -} - -inline void sync_cuda_if_needed(const torch::Tensor& tensor) { - if (tensor.is_cuda()) { - cudaError_t err = cudaDeviceSynchronize(); - if (err != cudaSuccess) { - throw std::runtime_error(std::string("cudaDeviceSynchronize failed: ") + - cudaGetErrorString(err)); - } - } -} - -inline std::tuple get_typesize(torch::ScalarType dtype) { - switch (dtype) { - case torch::kFloat32: return {MPI_FLOAT, sizeof(float)}; - case torch::kFloat64: return {MPI_DOUBLE, sizeof(double)}; - case torch::kFloat16: return {MPI_SHORT, sizeof(int16_t)}; // treat as raw bytes - case torch::kBFloat16: return {MPI_SHORT, sizeof(int16_t)}; - case torch::kInt32: return {MPI_INT, sizeof(int32_t)}; - case torch::kInt64: return {MPI_LONG_LONG, sizeof(int64_t)}; - case torch::kUInt8: return {MPI_UNSIGNED_CHAR, sizeof(uint8_t)}; - case torch::kInt8: return {MPI_CHAR, sizeof(int8_t)}; - default: - throw std::runtime_error("Unsupported torch::ScalarType for MPI communication"); - } -} - -inline MPI_Datatype get_mpi_datatype(const torch::Tensor& tensor) { - return std::get<0>(get_typesize(tensor.scalar_type())); -} - -void barrier() { - ensure_mpi_initialized(); - MPI_Barrier(MPI_COMM_WORLD); -} - -int get_rank() { - ensure_mpi_initialized(); - int rank; - MPI_Comm_rank(MPI_COMM_WORLD, &rank); - return rank; -} - -int get_size() { - ensure_mpi_initialized(); - int size; - MPI_Comm_size(MPI_COMM_WORLD, &size); - return size; -} - -void finalize_mpi() { - if (mpi_initialized) { - MPI_Finalize(); - mpi_initialized = false; - } -} - -void mpi_send(torch::Tensor& tensor, int dest) { - ensure_mpi_initialized(); - - if (!tensor.is_contiguous()) { - tensor = tensor.contiguous(); - } - - // Flush any pending CUDA work on this tensor's buffer before MPI touches it - sync_cuda_if_needed(tensor); - - MPI_Datatype datatype = get_mpi_datatype(tensor); - int err = MPI_Send(tensor.data_ptr(), tensor.numel(), datatype, dest, 0, MPI_COMM_WORLD); - if (err != MPI_SUCCESS) { - throw std::runtime_error("MPI_Send failed with error code: " + std::to_string(err)); - } -} - -void mpi_recv(torch::Tensor& tensor, int source) { - ensure_mpi_initialized(); - - if (!tensor.is_contiguous()) { - tensor = tensor.contiguous(); - } - - MPI_Datatype datatype = get_mpi_datatype(tensor); - MPI_Status status; - - int err = MPI_Recv(tensor.data_ptr(), tensor.numel(), datatype, source, 0, MPI_COMM_WORLD, &status); - if (err != MPI_SUCCESS) { - throw std::runtime_error("MPI_Recv failed with error code: " + std::to_string(err)); - } - - // Ensure the received data is visible to subsequent CUDA kernels - sync_cuda_if_needed(tensor); - - int recv_count; - MPI_Get_count(&status, datatype, &recv_count); - if (recv_count != tensor.numel()) { - throw std::runtime_error("MPI_Recv: expected " + std::to_string(tensor.numel()) + - " elements but got " + std::to_string(recv_count)); - } - - if (status.MPI_SOURCE != source) { - throw std::runtime_error("MPI_Recv: expected source " + std::to_string(source) + - " but got " + std::to_string(status.MPI_SOURCE)); - } -} - -void mpi_allreduce(torch::Tensor& tensor) { - ensure_mpi_initialized(); - - if (!tensor.is_contiguous()) { - tensor = tensor.contiguous(); - } - - sync_cuda_if_needed(tensor); - - MPI_Datatype datatype = get_mpi_datatype(tensor); - int err = MPI_Allreduce(MPI_IN_PLACE, tensor.data_ptr(), tensor.numel(), - datatype, MPI_SUM, MPI_COMM_WORLD); - if (err != MPI_SUCCESS) { - throw std::runtime_error("MPI_Allreduce failed with error code: " + std::to_string(err)); - } - - sync_cuda_if_needed(tensor); -} - -void mpi_allgather(torch::Tensor& sendbuf, torch::Tensor& recvbuf) { - ensure_mpi_initialized(); - - if (!sendbuf.is_contiguous()) sendbuf = sendbuf.contiguous(); - if (!recvbuf.is_contiguous()) recvbuf = recvbuf.contiguous(); - - int rank = get_rank(); - int world_size = get_size(); - int count = sendbuf.numel(); - - sync_cuda_if_needed(sendbuf); - - // Copy my slice into recvbuf locally - recvbuf.slice(0, rank * count, (rank + 1) * count).copy_(sendbuf); - - auto [datatype, typesize] = get_typesize(sendbuf.scalar_type()); - - // Post all non-blocking sends - std::vector send_reqs(world_size - 1); - int req_idx = 0; - for (int i = 0; i < world_size; ++i) { - if (i == rank) continue; - int err = MPI_Isend(sendbuf.data_ptr(), count, datatype, i, 0, - MPI_COMM_WORLD, &send_reqs[req_idx++]); - if (err != MPI_SUCCESS) - throw std::runtime_error("MPI_Isend failed: " + std::to_string(err)); - } - - // Post and wait on each receive in order - for (int i = 0; i < world_size; ++i) { - if (i == rank) continue; - void* recv_ptr = static_cast(recvbuf.data_ptr()) + i * count * typesize; - MPI_Request recv_req; - MPI_Status recv_status; - - int err = MPI_Irecv(recv_ptr, count, datatype, i, 0, MPI_COMM_WORLD, &recv_req); - if (err != MPI_SUCCESS) - throw std::runtime_error("MPI_Irecv failed: " + std::to_string(err)); - - err = MPI_Wait(&recv_req, &recv_status); - if (err != MPI_SUCCESS) - throw std::runtime_error("MPI_Wait (recv) failed: " + std::to_string(err)); - - int recv_count; - MPI_Get_count(&recv_status, datatype, &recv_count); - if (recv_count != count) - throw std::runtime_error("mpi_allgather: rank " + std::to_string(i) + - " sent " + std::to_string(recv_count) + - " elements, expected " + std::to_string(count)); - } - - // Wait for all sends - std::vector send_statuses(world_size - 1); - int err = MPI_Waitall(world_size - 1, send_reqs.data(), send_statuses.data()); - if (err != MPI_SUCCESS) - throw std::runtime_error("MPI_Waitall (sends) failed: " + std::to_string(err)); - - sync_cuda_if_needed(recvbuf); -} - -void mpi_reduce_scatter(torch::Tensor& sendbuf, torch::Tensor& recvbuf) { - ensure_mpi_initialized(); - - if (!sendbuf.is_contiguous()) sendbuf = sendbuf.contiguous(); - if (!recvbuf.is_contiguous()) recvbuf = recvbuf.contiguous(); - - int world_size; - MPI_Comm_size(MPI_COMM_WORLD, &world_size); - - sync_cuda_if_needed(sendbuf); - - int count = recvbuf.numel(); - std::vector recvcounts(world_size, count); - auto [mpi_dtype, typesize] = get_typesize(sendbuf.scalar_type()); - - int err = MPI_Reduce_scatter(sendbuf.data_ptr(), recvbuf.data_ptr(), - recvcounts.data(), mpi_dtype, MPI_SUM, MPI_COMM_WORLD); - if (err != MPI_SUCCESS) - throw std::runtime_error("MPI_Reduce_scatter failed: " + std::to_string(err)); - - sync_cuda_if_needed(recvbuf); -} - -void mpi_alltoall(torch::Tensor& sendbuf, torch::Tensor& recvbuf) { - ensure_mpi_initialized(); - - if (!sendbuf.is_contiguous()) sendbuf = sendbuf.contiguous(); - if (!recvbuf.is_contiguous()) recvbuf = recvbuf.contiguous(); - - if (recvbuf.numel() != sendbuf.numel()) - throw std::runtime_error("mpi_alltoall: sendbuf and recvbuf must have the same numel"); - - int world_size; - MPI_Comm_size(MPI_COMM_WORLD, &world_size); - - sync_cuda_if_needed(sendbuf); - - int count = sendbuf.numel() / world_size; - auto [mpi_dtype, typesize] = get_typesize(sendbuf.scalar_type()); - - int err = MPI_Alltoall(sendbuf.data_ptr(), count, mpi_dtype, - recvbuf.data_ptr(), count, mpi_dtype, MPI_COMM_WORLD); - if (err != MPI_SUCCESS) - throw std::runtime_error("MPI_Alltoall failed: " + std::to_string(err)); - - sync_cuda_if_needed(recvbuf); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("send", &mpi_send, "MPI Send (GPU-aware)"); - m.def("recv", &mpi_recv, "MPI Recv (GPU-aware)"); - m.def("allreduce", &mpi_allreduce, "MPI Allreduce (GPU-aware)"); - m.def("allgather", &mpi_allgather, "MPI Allgather (GPU-aware)"); - m.def("reduce_scatter",&mpi_reduce_scatter,"MPI Reduce_scatter (GPU-aware)"); - m.def("alltoall", &mpi_alltoall, "MPI Alltoall (GPU-aware)"); - m.def("barrier", &barrier, "MPI Barrier"); - m.def("get_rank", &get_rank, "Get MPI rank"); - m.def("get_size", &get_size, "Get MPI world size"); - m.def("finalize_mpi", &finalize_mpi, "Finalize MPI"); -} diff --git a/infra/cray_infra/training/gpu_aware_mpi/setup.py b/infra/cray_infra/training/gpu_aware_mpi/setup.py deleted file mode 100644 index 9170e1d7..00000000 --- a/infra/cray_infra/training/gpu_aware_mpi/setup.py +++ /dev/null @@ -1,68 +0,0 @@ -import os -from setuptools import setup -from torch.utils import cpp_extension - -def detect_gpu_platform(): - """Determine GPU platform through environment variables and path checks""" - if os.path.exists('/opt/rocm'): - return 'rocm' - if os.path.exists('/usr/local/cuda'): - return 'cuda' - return 'cpu' - -platform = detect_gpu_platform() - -include_dirs = [] -library_dirs = [] -compile_defines = [] -libraries = [] -extra_link_args = [] - -if platform == 'rocm': - os.environ['CXX'] = '/opt/ompi-rocm/bin/mpicxx' - compile_defines.append(('USE_ROCM', '1')) - include_dirs.extend([ - '/opt/rocm/include', - '/opt/ompi-rocm/include' - ]) - library_dirs.extend([ - '/opt/rocm/lib', - '/opt/ompi-rocm/lib' - ]) - extra_link_args.append('-lmpi') -elif platform == 'cuda': - compile_defines.append(('USE_CUDA', '1')) - include_dirs.extend([ - '/usr/local/cuda/include', - '/opt/hpcx/ompi/include' - ]) - library_dirs.extend([ - '/usr/local/cuda/lib64', - '/opt/hpcx/ompi/lib' - ]) - libraries.append('cudart') -elif platform == 'cpu': - os.environ['CXX'] = '/usr/bin/mpicxx' - include_dirs.extend([ - '/usr/lib/aarch64-linux-gnu/openmpi/include', - ]) - library_dirs.extend([ - '/usr/lib/aarch64-linux-gnu/openmpi/lib', - ]) - libraries.append('mpi') - -setup( - name="gpu_aware_mpi", - ext_modules=[cpp_extension.CppExtension( - 'gpu_aware_mpi', - sources=['infra/cray_infra/training/gpu_aware_mpi/gpu_aware_mpi.cpp'], - include_dirs=include_dirs, - library_dirs=library_dirs, - libraries=libraries, - extra_compile_args={ - 'cxx': [f'-D{name}={value}' for name, value in compile_defines] - }, - extra_link_args=extra_link_args, - )], - cmdclass={'build_ext': cpp_extension.BuildExtension} -) diff --git a/infra/cray_infra/training/launch_training_job.py b/infra/cray_infra/training/launch_training_job.py index f6efa17a..8f9a8f0b 100644 --- a/infra/cray_infra/training/launch_training_job.py +++ b/infra/cray_infra/training/launch_training_job.py @@ -1,4 +1,10 @@ from cray_infra.util.get_config import get_config +from cray_infra.training.slurm_jobs import ( + cancel_slurm_jobs_for_train_args, + get_job_name, + is_slurm_job_active_for_train_args, + job_should_block_resubmit, +) import json import os @@ -18,8 +24,23 @@ async def launch_training_job(train_args: Dict): await wait_for_slurm() - if job_already_exists(train_args): - logging.info(f"Job already exists: {train_args['job_directory']}") + force_retry = bool( + train_args.get("retry") or train_args.get("force_resubmit") + ) + + if force_retry: + cancel_slurm_jobs_for_train_args(train_args) + _clear_status_for_retry(train_args) + + if job_already_exists(train_args, force_retry=force_retry): + logging.info(f"Job already active: {train_args['job_directory']}") + return get_existing_job_info(train_args) + + if is_slurm_job_active_for_train_args(train_args): + logging.info( + "SLURM job already queued/running for %s", + get_job_name(train_args), + ) return get_existing_job_info(train_args) make_training_directory(train_args) @@ -43,12 +64,32 @@ async def wait_for_slurm(): time.sleep(1) -def job_already_exists(train_args: Dict): - config = get_config() +def _load_job_status(train_args: Dict): + status_path = os.path.join(train_args["job_directory"], "status.json") + if not os.path.exists(status_path): + return None + try: + with open(status_path, "r") as f: + return json.load(f) + except json.JSONDecodeError: + return None - train_args_path = os.path.join(train_args["job_directory"], "status.json") - return os.path.exists(train_args_path) +def _clear_status_for_retry(train_args: Dict) -> None: + status_path = os.path.join(train_args["job_directory"], "status.json") + if os.path.exists(status_path): + os.remove(status_path) + + +def job_already_exists(train_args: Dict, force_retry: bool = False): + if force_retry: + return False + + status = _load_job_status(train_args) + if status is None: + return False + + return job_should_block_resubmit(status, get_job_name(train_args)) def make_training_directory(train_args: Dict): @@ -70,6 +111,13 @@ def get_training_job_directory(train_args: Dict): def start_slurm_job(train_args): + if is_slurm_job_active_for_train_args(train_args): + logger.info( + "Skipping sbatch; SLURM job already active for %s", + get_job_name(train_args), + ) + return + run_command = create_slurm_run_command(train_args) run_sbatch(run_command, train_args) @@ -90,6 +138,8 @@ def create_slurm_run_command(train_args): run_command += [f"--nodes={node_count}"] logger.info(f"node_count: {node_count}") + run_command += ["--exclusive"] + cpu_per_task = get_cpu_per_task(train_args) run_command += [f"--cpus-per-task={cpu_per_task}"] logger.info(f"cpu_per_task: {cpu_per_task}") @@ -266,23 +316,38 @@ def run_sbatch(run_command, train_args): env=clean_environs, ) + stdout_text = result.stdout.decode("utf-8") + stderr_text = result.stderr.decode("utf-8") + if result.returncode != 0: - result_output = result.stdout.decode("utf-8") + result.stderr.decode("utf-8") + result_output = stdout_text + stderr_text write_job_status("FAILED", train_args, {"output": result_output}) - - job_id = get_job_id_from_sbatch_output(result.stdout.decode("utf-8")) + return + + job_id = get_job_id_from_sbatch_output(stdout_text) + if job_id is None: + write_job_status( + "FAILED", + train_args, + {"output": stdout_text + stderr_text, "error": "Could not parse sbatch job id"}, + ) + return write_job_status("QUEUED", train_args, {"job_id": job_id}) def get_job_id_from_sbatch_output(sbatch_output): logger.info(f"sbatch_output: {sbatch_output}") - return re.search(r"Submitted batch job (\d+)", sbatch_output).group(1) + match = re.search(r"Submitted batch job (\d+)", sbatch_output) + if not match: + return None + return match.group(1) def write_job_status(status, train_args, extra_info): job_status = { "status": status, + "last_updated": time.time(), **extra_info, } diff --git a/infra/cray_infra/training/metrics.py b/infra/cray_infra/training/metrics.py index 8f178c49..102a0614 100644 --- a/infra/cray_infra/training/metrics.py +++ b/infra/cray_infra/training/metrics.py @@ -1,5 +1,5 @@ import torch -from gpu_aware_mpi import get_rank +from cray_infra.training.distributed import get_rank import logging diff --git a/infra/cray_infra/training/register_megatron_workers.py b/infra/cray_infra/training/register_megatron_workers.py index d5f4f62d..8e93cba1 100644 --- a/infra/cray_infra/training/register_megatron_workers.py +++ b/infra/cray_infra/training/register_megatron_workers.py @@ -1,4 +1,4 @@ from cray_infra.slurm.discovery.discover_clusters import discover_clusters async def register_megatron_workers(): - discover_clusters() + discover_clusters(quiet=True) diff --git a/infra/cray_infra/training/restart_megatron_jobs.py b/infra/cray_infra/training/restart_megatron_jobs.py index b46ee1ff..9bde1516 100644 --- a/infra/cray_infra/training/restart_megatron_jobs.py +++ b/infra/cray_infra/training/restart_megatron_jobs.py @@ -1,5 +1,9 @@ from cray_infra.training.training_job_status import TrainingJobStatus from cray_infra.training.launch_training_job import start_slurm_job +from cray_infra.training.slurm_jobs import ( + get_active_slurm_jobs, + is_slurm_job_active, +) from cray_infra.api.fastapi.aiohttp.get_global_session import get_global_session @@ -10,7 +14,7 @@ import os import json import yaml -import subprocess +import time import logging @@ -18,100 +22,98 @@ async def restart_megatron_jobs(): - logger.info("Restarting Megatron jobs") + logger.info("Checking for Megatron jobs to restart") - # Get all the jobs that are running - all_jobs = get_running_jobs() + active_slurm = get_active_slurm_jobs() + logger.info("Active SLURM job names: %s", list(active_slurm.keys())) - # Get slurm jobs that are running - slurm_job_names = await get_slurm_jobs() - - logger.info(f"Slurm jobs running: {slurm_job_names}") - - # Filter out the jobs that are already running - jobs = filter_running_jobs(all_jobs, slurm_job_names) - - # Restart all the jobs - async for job in jobs: + async for job in get_restartable_jobs(): + job_name = os.path.basename(job) + if job_name in active_slurm: + logger.info("Skipping restart for %s: already in squeue", job_name) + continue await restart_job(job) - # Get slurm jobs that are running again - slurm_job_names = await get_slurm_jobs() - - # If any are still running, keep the server alive - if slurm_job_names: - logger.info("Jobs are still running, keeping the server alive") + active_slurm = get_active_slurm_jobs() + if active_slurm: + logger.info("Jobs still running in SLURM, keeping the server alive") await keep_alive() -async def get_running_jobs(): +async def get_restartable_jobs(): config = get_config() + max_restarts = config.get("max_job_restart_count", 3) + cooldown = config.get("job_restart_cooldown_seconds", 300) - # training jobs are in the training job directory - # they are subdirectories which should have a file called status.json - - # get all the directories in the training job directory - # that have a status.json file if not os.path.exists(config["training_job_directory"]): return for path in os.listdir(config["training_job_directory"]): root = os.path.join(config["training_job_directory"], path) - if os.path.exists(os.path.join(root, "status.json")): - try: - with open(os.path.join(root, "status.json")) as f: - status = json.load(f) - if ( - status["status"] == TrainingJobStatus.TRAINING - or status["status"] == TrainingJobStatus.QUEUED - ): - yield root - except Exception as e: - logger.error(f"Error reading status.json for job {root}: {e}") - # print exception info - traceback.print_exc() - - -async def get_slurm_jobs(): - squeue_output = subprocess.check_output( - ["squeue", '--format="%.18i %.9P %.128j %.8u %.8T %.10M %.9l %.6D %R"'] - ) - - # Here is an example of the output of squeue - # JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON) - # 1234 gpu jobname username R 0:10 1 node1 - # 5678 gpu jobname username R 0:10 1 node2 - # 91011 gpu jobname username R 0:10 1 node3 - # 121314 gpu jobname username R 0:10 1 node4 - - # We want to get the job name from the third column - # We also want to remove the header - jobs = squeue_output.decode().split("\n")[1:] - - job_names = [] - - for job in jobs: - if job: - job_fields = job.strip().split() - if len(job_fields) > 3: - job_names.append(job_fields[3]) - - return job_names - - -async def filter_running_jobs(all_jobs, slurm_job_names): - async for job in all_jobs: - if os.path.basename(job) not in slurm_job_names: - yield job + status_path = os.path.join(root, "status.json") + if not os.path.exists(status_path): + continue + try: + with open(status_path) as f: + status = json.load(f) + except Exception as e: + logger.error("Error reading status.json for job %s: %s", root, e) + traceback.print_exc() + continue + + if status.get("status") != TrainingJobStatus.QUEUED: + continue + + restart_count = status.get("restart_count", 0) + if restart_count >= max_restarts: + logger.warning( + "Job %s exceeded max restart count (%s), marking FAILED", + path, + max_restarts, + ) + _mark_job_failed( + root, + status, + f"Exceeded max restart count ({max_restarts})", + ) + continue + + last_restart = status.get("last_restart_at", 0) + if time.time() - last_restart < cooldown: + continue + + yield root + + +def _mark_job_failed(job_directory: str, status: dict, reason: str) -> None: + status["status"] = TrainingJobStatus.FAILED + status["error"] = reason + status["last_updated"] = time.time() + status_path = os.path.join(job_directory, "status.json") + with open(status_path, "w") as f: + json.dump(status, f) async def restart_job(job): - logger.info(f"Restarting job: {job}") + logger.info("Restarting job: %s", job) + + status_path = os.path.join(job, "status.json") + with open(status_path) as f: + status = json.load(f) + + status["restart_count"] = status.get("restart_count", 0) + 1 + status["last_restart_at"] = time.time() + with open(status_path, "w") as f: + json.dump(status, f) - # Get the job config with open(os.path.join(job, "config.yaml")) as f: config = yaml.safe_load(f) + job_name = os.path.basename(job) + if is_slurm_job_active(job_name): + logger.info("SLURM job %s became active before restart, skipping sbatch", job_name) + return + start_slurm_job(config) diff --git a/infra/cray_infra/training/slurm_jobs.py b/infra/cray_infra/training/slurm_jobs.py new file mode 100644 index 00000000..22c58c2b --- /dev/null +++ b/infra/cray_infra/training/slurm_jobs.py @@ -0,0 +1,278 @@ +from cray_infra.training.training_job_status import TrainingJobStatus +from cray_infra.util.get_config import get_config + +import json +import os +import subprocess +import time +import logging + +from typing import Dict, Optional + +logger = logging.getLogger(__name__) + +# SLURM states that mean the job is still in the scheduler queue or running. +_ACTIVE_SQUEUE_STATES = frozenset( + { + "PENDING", + "CONFIGURING", + "RUNNING", + "COMPLETING", + "SUSPENDED", + "PREEMPTED", + "REQUEUED", + "RESIZING", + "REVOKED", + "SIGNALING", + "SPECIAL_EXIT", + "STAGE_OUT", + } +) + +# Terminal sacct states mapped to training job status. +_SACCT_TERMINAL_STATUS = { + "COMPLETED": TrainingJobStatus.COMPLETED, + "FAILED": TrainingJobStatus.FAILED, + "CANCELLED": TrainingJobStatus.FAILED, + "CANCELED": TrainingJobStatus.FAILED, + "TIMEOUT": TrainingJobStatus.FAILED, + "NODE_FAIL": TrainingJobStatus.FAILED, + "OUT_OF_MEMORY": TrainingJobStatus.FAILED, + "PREEMPTED": TrainingJobStatus.FAILED, + "BOOT_FAIL": TrainingJobStatus.FAILED, + "DEADLINE": TrainingJobStatus.FAILED, +} + + +def get_job_name(train_args: Dict) -> str: + return os.path.basename(train_args["job_directory"]) + + +def get_active_slurm_jobs() -> Dict[str, Dict[str, str]]: + """Return {job_name: {job_id, state}} for non-terminal squeue entries.""" + try: + output = subprocess.check_output( + ["squeue", "-h", "-o", "%i %j %T"], + stderr=subprocess.STDOUT, + ) + except subprocess.CalledProcessError as e: + logger.error( + "squeue failed (exit %s): %s", + e.returncode, + e.output.decode("utf-8", errors="replace") if e.output else "", + ) + return {} + + jobs: Dict[str, Dict[str, str]] = {} + for line in output.decode("utf-8", errors="replace").splitlines(): + line = line.strip() + if not line: + continue + parts = line.split(None, 2) + if len(parts) < 2: + continue + job_id, job_name = parts[0], parts[1] + state = parts[2] if len(parts) > 2 else "UNKNOWN" + if state in _ACTIVE_SQUEUE_STATES or state == "UNKNOWN": + jobs[job_name] = {"job_id": job_id, "state": state} + return jobs + + +def is_slurm_job_active(job_name: str) -> bool: + return job_name in get_active_slurm_jobs() + + +def is_slurm_job_active_for_train_args(train_args: Dict) -> bool: + return is_slurm_job_active(get_job_name(train_args)) + + +def cancel_slurm_jobs_for_train_args(train_args: Dict) -> None: + job_name = get_job_name(train_args) + status_path = os.path.join(train_args["job_directory"], "status.json") + + if os.path.exists(status_path): + try: + with open(status_path, "r") as f: + status = json.load(f) + job_id = status.get("job_id") + if job_id: + _scancel_job_id(str(job_id)) + except (json.JSONDecodeError, OSError) as e: + logger.warning("Could not read status for cancel: %s", e) + + try: + subprocess.run( + ["scancel", "--name", job_name], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + except OSError as e: + logger.error("scancel --name=%s failed: %s", job_name, e) + + +def _scancel_job_id(job_id: str) -> None: + try: + subprocess.run( + ["scancel", job_id], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + except OSError as e: + logger.error("scancel %s failed: %s", job_id, e) + + +def get_sacct_state(job_id: str) -> Optional[str]: + try: + output = subprocess.check_output( + [ + "sacct", + "-j", + str(job_id), + "-n", + "-X", + "--parsable2", + "-o", + "State", + ], + stderr=subprocess.PIPE, + ) + except (subprocess.CalledProcessError, OSError): + return None + + for line in output.decode("utf-8", errors="replace").splitlines(): + state = line.strip().split("|")[0].strip() + if state: + # sacct may return State+ suffix e.g. FAILED+OUT_OF_MEMORY + return state.split("+")[0] + return None + + +def _latest_slurm_log_excerpt(job_directory: str, job_id: str, max_chars: int = 4000) -> Optional[str]: + if not job_id: + return None + log_path = os.path.join(job_directory, f"slurm-{job_id}.out") + if not os.path.isfile(log_path): + for name in os.listdir(job_directory): + if name.startswith("slurm-") and name.endswith(".out"): + log_path = os.path.join(job_directory, name) + break + else: + return None + try: + with open(log_path, "rb") as f: + f.seek(0, os.SEEK_END) + size = f.tell() + read_size = min(size, max_chars) + f.seek(-read_size, os.SEEK_END) + return f.read().decode("utf-8", errors="replace") + except OSError: + return None + + +def _load_status(job_directory: str) -> Optional[dict]: + status_path = os.path.join(job_directory, "status.json") + if not os.path.exists(status_path): + return None + try: + with open(status_path, "r") as f: + return json.load(f) + except (json.JSONDecodeError, OSError) as e: + logger.error("Error reading %s: %s", status_path, e) + return None + + +def _save_status(job_directory: str, status: dict) -> None: + status_path = os.path.join(job_directory, "status.json") + status["last_updated"] = time.time() + with open(status_path, "w") as f: + json.dump(status, f) + + +def is_recent_training_activity(status: dict) -> bool: + if status.get("status") != TrainingJobStatus.TRAINING: + return False + config = get_config() + heartbeat_limit = config.get("training_heartbeat_seconds", 600) + last = status.get("last_updated") or status.get("start_time", 0) + return time.time() - last < heartbeat_limit + + +def job_should_block_resubmit(status: Optional[dict], job_name: str) -> bool: + if is_slurm_job_active(job_name): + return True + if status and is_recent_training_activity(status): + return True + return False + + +def sync_training_job_statuses() -> None: + """Update status.json from squeue/sacct for jobs with a SLURM job id.""" + config = get_config() + training_dir = config["training_job_directory"] + if not os.path.isdir(training_dir): + return + + active_slurm = get_active_slurm_jobs() + + for entry in os.listdir(training_dir): + job_directory = os.path.join(training_dir, entry) + if not os.path.isdir(job_directory): + continue + + status = _load_status(job_directory) + if not status: + continue + + current = status.get("status") + if current in ( + TrainingJobStatus.COMPLETED, + TrainingJobStatus.FAILED, + "CANCELLED", + ): + continue + + job_name = entry + job_id = status.get("job_id") + + if job_name in active_slurm: + slurm_info = active_slurm[job_name] + status["job_id"] = slurm_info["job_id"] + status["slurm_state"] = slurm_info["state"] + if current == TrainingJobStatus.QUEUED and slurm_info["state"] == "RUNNING": + # Still QUEUED until training harness sets TRAINING. + pass + _save_status(job_directory, status) + continue + + if not job_id: + continue + + sacct_state = get_sacct_state(str(job_id)) + if not sacct_state: + continue + + terminal = _SACCT_TERMINAL_STATUS.get(sacct_state) + if not terminal: + continue + + if current == TrainingJobStatus.TRAINING and terminal == TrainingJobStatus.COMPLETED: + status["status"] = TrainingJobStatus.COMPLETED + else: + status["status"] = terminal + + status["slurm_state"] = sacct_state + if terminal == TrainingJobStatus.FAILED: + excerpt = _latest_slurm_log_excerpt(job_directory, str(job_id)) + if excerpt: + status["slurm_log_tail"] = excerpt + + _save_status(job_directory, status) + logger.info( + "Synced job %s (slurm %s) -> %s (sacct %s)", + job_name, + job_id, + status["status"], + sacct_state, + ) diff --git a/infra/cray_infra/training/sync_training_job_status.py b/infra/cray_infra/training/sync_training_job_status.py new file mode 100644 index 00000000..2c42b5fa --- /dev/null +++ b/infra/cray_infra/training/sync_training_job_status.py @@ -0,0 +1,12 @@ +from cray_infra.training.slurm_jobs import sync_training_job_statuses + +import logging + +logger = logging.getLogger(__name__) + + +async def sync_training_job_status(): + try: + sync_training_job_statuses() + except Exception: + logger.exception("Failed to sync training job statuses from SLURM") diff --git a/infra/cray_infra/training/train_debug.py b/infra/cray_infra/training/train_debug.py new file mode 100644 index 00000000..ea1be904 --- /dev/null +++ b/infra/cray_infra/training/train_debug.py @@ -0,0 +1,13 @@ +import os + + +def _env_flag(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).lower() in ("1", "true", "yes") + + +def is_train_debug_enabled() -> bool: + return _env_flag("CRAY_TRAIN_DEBUG") + + +def is_fault_handler_enabled() -> bool: + return _env_flag("CRAY_FAULT_HANDLER") or is_train_debug_enabled() diff --git a/infra/cray_infra/training/training_job_context.py b/infra/cray_infra/training/training_job_context.py index 755d3a17..bfc0c54e 100644 --- a/infra/cray_infra/training/training_job_context.py +++ b/infra/cray_infra/training/training_job_context.py @@ -1,4 +1,4 @@ -from gpu_aware_mpi import finalize_mpi +from cray_infra.training.distributed import finalize from cray_infra.training.training_harness import TrainingHarness from cray_infra.training.training_job_status import TrainingJobStatus @@ -17,6 +17,6 @@ def training_job_context(): except Exception as e: harness.update_status(TrainingJobStatus.FAILED, metadata={"error": str(e)}) finally: - finalize_mpi() + finalize() diff --git a/infra/cray_infra/training/training_logs_generator.py b/infra/cray_infra/training/training_logs_generator.py index c9546a29..4e5d735e 100644 --- a/infra/cray_infra/training/training_logs_generator.py +++ b/infra/cray_infra/training/training_logs_generator.py @@ -23,15 +23,14 @@ def training_logs_generator(model_name: str, starting_line_number: int): job_directory = get_job_directory_for_hash(model_name) - # Find the log file inside the job directory, it will be named "slurm-.out, but we don't know the job_id yet log_files = [] for file in os.listdir(job_directory): if file.startswith("slurm-") and file.endswith(".out"): - log_file = os.path.join(job_directory, file) - log_files.append(log_file) + log_files.append(os.path.join(job_directory, file)) + elif file.startswith("rank-") and file.endswith(".log"): + log_files.append(os.path.join(job_directory, file)) - # sort the log files by name log_files.sort() logger.info(f"Found log files: {log_files}") diff --git a/infra/cray_infra/util/default_config.py b/infra/cray_infra/util/default_config.py index 945240eb..19dc08e2 100644 --- a/infra/cray_infra/util/default_config.py +++ b/infra/cray_infra/util/default_config.py @@ -33,6 +33,10 @@ class Config(BaseModel): megatron_refresh_period: int = 30 # seconds + max_job_restart_count: int = 3 + job_restart_cooldown_seconds: int = 300 + training_heartbeat_seconds: int = 600 + vllm_api_url: str = "http://localhost:8001" # vLLM Engine Configuration diff --git a/infra/cray_infra/util/default_job_config.py b/infra/cray_infra/util/default_job_config.py index 7540a73d..d8c52e2a 100644 --- a/infra/cray_infra/util/default_job_config.py +++ b/infra/cray_infra/util/default_job_config.py @@ -23,7 +23,7 @@ class JobConfig(BaseModel): learning_rate: float = 3e-3 batch_size: int = 1 gradient_clip_value: float = 1.0 - gradient_accumulation_steps: int = 4 + gradient_accumulation_steps: int = 1 max_token_block_size: int = 16777216 # 16 mega tokens diff --git a/infra/requirements-megatron-cpu.txt b/infra/requirements-megatron-cpu.txt index b115fab4..79b5481f 100644 --- a/infra/requirements-megatron-cpu.txt +++ b/infra/requirements-megatron-cpu.txt @@ -1,3 +1,3 @@ sentence-transformers -torchao==0.15.0 +torchao==0.17.0 peft diff --git a/infra/requirements-megatron.txt b/infra/requirements-megatron.txt index 9473f481..3c6ef804 100644 --- a/infra/requirements-megatron.txt +++ b/infra/requirements-megatron.txt @@ -1,3 +1,3 @@ -torchao==0.14.1 +torchao==0.17.0 accelerate peft diff --git a/infra/slurm_src/compile.sh b/infra/slurm_src/compile.sh index 6a9091c7..30390c9e 100755 --- a/infra/slurm_src/compile.sh +++ b/infra/slurm_src/compile.sh @@ -29,7 +29,7 @@ else fi # Disable the plugin on the AMD target -if [ $BASE_NAME == "amd" ] || [ $BASE_NAME="nvidia" ]; then +if [ "$BASE_NAME" == "amd" ] || [ "$BASE_NAME" == "nvidia" ]; then sed -i -e 's/CgroupPlugin=cgroup\/docker/CgroupPlugin=cgroup\/v1/g' /app/cray/infra/slurm_configs/cgroup.conf else # Copy the shared object file to the /usr/lib directory diff --git a/ml/__init__.py b/ml/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ml/cray_megatron/collectives/data_parallelism.py b/ml/cray_megatron/collectives/data_parallelism.py index 02a197c9..1b877ba6 100644 --- a/ml/cray_megatron/collectives/data_parallelism.py +++ b/ml/cray_megatron/collectives/data_parallelism.py @@ -1,4 +1,4 @@ -from gpu_aware_mpi import get_rank, get_size +from cray_infra.training.distributed import get_rank, get_size def get_data_parallel_rank(): return get_rank() diff --git a/ml/cray_megatron/collectives/main_rank_only.py b/ml/cray_megatron/collectives/main_rank_only.py index 07ba299e..225966b3 100644 --- a/ml/cray_megatron/collectives/main_rank_only.py +++ b/ml/cray_megatron/collectives/main_rank_only.py @@ -1,5 +1,29 @@ +import sys +import time from functools import wraps -from gpu_aware_mpi import get_rank, barrier + +from cray_infra.training.distributed import get_rank, barrier +from cray_infra.training.train_debug import is_train_debug_enabled + + +def _trace_main_rank_only(msg: str) -> None: + if not is_train_debug_enabled(): + return + if _dist_ready() and get_rank() != 0: + return + rank = get_rank() if _dist_ready() else "?" + line = f"[rank={rank}] main_rank_only [{time.monotonic():.3f}]: {msg}\n" + sys.stderr.write(line) + sys.stderr.flush() + + +def _dist_ready() -> bool: + try: + import torch.distributed as dist + + return dist.is_initialized() + except Exception: + return False def is_main_rank(): return get_rank() == 0 @@ -16,9 +40,13 @@ def wrapper(*args, **kwargs): _in_main_rank_only = True try: + _trace_main_rank_only(f"{func.__name__}: pre-barrier-1") barrier() + _trace_main_rank_only(f"{func.__name__}: post-barrier-1, is_main_rank={is_main_rank()}") result = func(*args, **kwargs) if is_main_rank() else None + _trace_main_rank_only(f"{func.__name__}: pre-barrier-2") barrier() + _trace_main_rank_only(f"{func.__name__}: post-barrier-2") return result finally: _in_main_rank_only = False diff --git a/ml/cray_megatron/main.py b/ml/cray_megatron/main.py index 17ab1164..2210e3c0 100644 --- a/ml/cray_megatron/main.py +++ b/ml/cray_megatron/main.py @@ -1,68 +1,121 @@ +import faulthandler +import logging +import os +import signal +import sys +import time +import traceback + from cray_infra.training.training_job_status import TrainingJobStatus from cray_infra.huggingface.get_hf_token import get_hf_token - from cray_megatron.megatron.training_harness import TrainingHarness +from cray_infra.training.distributed import init, finalize +from cray_infra.training.train_debug import ( + is_fault_handler_enabled, + is_train_debug_enabled, +) -from cray_megatron.collectives.main_rank_only import main_rank_only +logger = logging.getLogger(__name__) + +if is_fault_handler_enabled(): + faulthandler.enable() + faulthandler.dump_traceback_later(timeout=600, repeat=True) + + +def _boot(msg: str) -> None: + if not is_train_debug_enabled(): + return + rank = os.environ.get("RANK", os.environ.get("SLURM_PROCID", "?")) + line = f"[rank={rank}] boot pid={os.getpid()} t={time.monotonic():.3f}: {msg}\n" + sys.stderr.write(line) + sys.stderr.flush() + + +if is_train_debug_enabled(): + _boot("main.py entered") -import traceback -import sys -import os -from gpu_aware_mpi import finalize_mpi -@main_rank_only def print_exception(): exc_type, exc_value, exc_traceback = sys.exc_info() traceback.print_exception(exc_type, exc_value, exc_traceback) -try: - from cray_megatron.megatron.megatron_trainer import MegatronTrainer -except Exception as e: - print_exception() +def main(): + _boot("calling init()") + init() + _boot("init() complete") -import signal -import logging + import torch -logger = logging.getLogger(__name__) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + _boot( + f"set_device({local_rank}) after init " + f"cuda_current_device={torch.cuda.current_device()}" + ) -def main(): + _boot("reconfigure stdout") + sys.stdout.reconfigure(line_buffering=True) + _boot("create TrainingHarness") harness = TrainingHarness() + _boot("get_hf_token") os.environ["HUGGING_FACE_HUB_TOKEN"] = get_hf_token() + _boot("get_hf_token complete") try: + _boot("setup_logging") setup_logging() + _boot("setup_signal_handler") setup_signal_handler(harness) + _boot("import megatron_trainer") + from cray_megatron.megatron.megatron_trainer import MegatronTrainer + + _boot("megatron_trainer imported") + + _boot("create MegatronTrainer") trainer = MegatronTrainer(training_harness=harness) + _boot("MegatronTrainer created") + + _boot("calling trainer.train()") trainer.train() + _boot("trainer.train() complete") except Exception as e: + _boot(f"training failed: {e}") print_exception() harness.update_status( status=TrainingJobStatus.FAILED, metadata={"error": str(e)} ) raise e - finalize_mpi() + _boot("calling finalize()") + finalize() + _boot("finalize() complete") + def setup_logging(): - logging.basicConfig(level=logging.DEBUG) + logging.basicConfig(level=logging.INFO) logging.getLogger("filelock").setLevel(logging.WARNING) + fsdp_level = logging.INFO if is_train_debug_enabled() else logging.WARNING logging.getLogger("cray_megatron.megatron.distribution.fsdp").setLevel( - logging.INFO + fsdp_level ) -def setup_signal_handler(harness): - def signal_handler(sig, frame): - logger.warning("Received signal: ", sig) - harness.update_status(status=TrainingJobStatus.QUEUED) - sys.exit(0) +def setup_signal_handler(harness): + def terminate_handler(sig, frame): + logger.warning("Received termination signal %s", sig) + harness.update_status( + status=TrainingJobStatus.FAILED, + metadata={"error": f"Terminated by signal {sig}"}, + ) + sys.exit(1) - signal.signal(signal.SIGCONT, signal_handler) + signal.signal(signal.SIGTERM, terminate_handler) + signal.signal(signal.SIGINT, terminate_handler) main() diff --git a/ml/cray_megatron/megatron/distribution/apply_distribution_strategy.py b/ml/cray_megatron/megatron/distribution/apply_distribution_strategy.py index 9a856f13..6c1652df 100644 --- a/ml/cray_megatron/megatron/distribution/apply_distribution_strategy.py +++ b/ml/cray_megatron/megatron/distribution/apply_distribution_strategy.py @@ -3,7 +3,7 @@ from cray_megatron.megatron.distribution.no_distribution import NoDistribution from cray_infra.util.get_job_config import get_job_config -from gpu_aware_mpi import get_size, get_rank +from cray_infra.training.distributed import get_size, get_rank import torch @@ -31,6 +31,9 @@ def load_distribution_strategy(): if distribution_strategy == "ddp": logger.info("Using DDP distribution strategy.") strategy["strategy"] = DDP + elif distribution_strategy == "fsdp": + logger.info("Using SimpleFSDP distribution strategy.") + strategy["strategy"] = SimpleFSDP else: logger.warning( f"Unknown distribution strategy '{distribution_strategy}' " diff --git a/ml/cray_megatron/megatron/distribution/ddp.py b/ml/cray_megatron/megatron/distribution/ddp.py index 0fc8b0c0..372f922e 100644 --- a/ml/cray_megatron/megatron/distribution/ddp.py +++ b/ml/cray_megatron/megatron/distribution/ddp.py @@ -1,4 +1,4 @@ -from gpu_aware_mpi import allreduce, get_size +from cray_infra.training.distributed import allreduce, get_size import torch.nn as nn diff --git a/ml/cray_megatron/megatron/distribution/fsdp.py b/ml/cray_megatron/megatron/distribution/fsdp.py index 3400b815..9a9ab896 100644 --- a/ml/cray_megatron/megatron/distribution/fsdp.py +++ b/ml/cray_megatron/megatron/distribution/fsdp.py @@ -1,7 +1,13 @@ import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint -from gpu_aware_mpi import get_size, get_rank, allgather, reduce_scatter +from cray_infra.training.distributed import ( + get_size, + get_rank, + allgather, + reduce_scatter, + cuda_device, +) from collections import defaultdict from cray_infra.training.metrics import get_model_memory_footprint @@ -340,27 +346,32 @@ def shard_tensor(tensor): padded_numel = ((original_numel + world_size - 1) // world_size) * world_size padding = padded_numel - original_numel + device = cuda_device() + # Pad tensor if needed if padding > 0: tensor_padded = torch.cat( [ - tensor.view(-1), - torch.zeros(padding, device=tensor.device, dtype=tensor.dtype), + tensor.view(-1).to(device), + torch.zeros(padding, dtype=tensor.dtype).to(device), ] ) else: - tensor_padded = tensor.view(-1) + tensor_padded = tensor.view(-1).to(device) # Split into equal shards shard_size = padded_numel // world_size start = rank * shard_size shard = tensor_padded[start : start + shard_size].clone() - # Gather metadata from all ranks + # Gather metadata from all ranks on GPU for NCCL. local_metadata = torch.tensor( - [original_numel, *original_shape, shard_size, padding], dtype=torch.long - ) - all_metadata = torch.zeros((world_size, local_metadata.numel()), dtype=torch.long) + [original_numel, *original_shape, shard_size, padding], + dtype=torch.long, + ).to(device) + all_metadata = torch.zeros( + (world_size, local_metadata.numel()), dtype=torch.long + ).to(device) allgather(local_metadata, all_metadata.view(-1)) # Reshape all_metadata back to 2D after allgather @@ -413,11 +424,13 @@ def collectives_all_gather(shard, metadata_dict): world_size = get_size() rank = get_rank() - # Prepare buffers + device = cuda_device() + + # Prepare buffers on GPU for NCCL collectives. orig_dtype = shard.dtype - shard = shard.to(torch.float32) + shard = shard.to(torch.float32).to(device) gathered = torch.zeros( - shard.numel() * world_size, device=shard.device, dtype=torch.float32 + shard.numel() * world_size, dtype=torch.float32, device=device ) # Collective operation in float32 @@ -452,19 +465,21 @@ def collectives_reduce_scatter(tensor, metadata_dict): original_numel, _, shard_size, padding = metadata_dict[rank] + device = cuda_device() + # Pad tensor if needed - tensor_padded = tensor.reshape(-1) + tensor_padded = tensor.reshape(-1).to(device) if padding > 0: tensor_padded = torch.concatenate( [ tensor_padded, - torch.zeros(padding, device=tensor.device, dtype=tensor_padded.dtype), + torch.zeros(padding, dtype=tensor_padded.dtype).to(device), ] ) - # Convert to float32 for the collective + # Convert to float32 for the collective on GPU tensor_padded = tensor_padded.to(torch.float32) - local_shard = torch.zeros(shard_size, device=tensor.device, dtype=torch.float32) + local_shard = torch.zeros(shard_size, dtype=torch.float32, device=device) # Collective operation in float32 reduce_scatter(tensor_padded, local_shard) diff --git a/ml/cray_megatron/megatron/megatron_trainer.py b/ml/cray_megatron/megatron/megatron_trainer.py index 9c9ca26b..4baf9a2a 100644 --- a/ml/cray_megatron/megatron/megatron_trainer.py +++ b/ml/cray_megatron/megatron/megatron_trainer.py @@ -4,13 +4,26 @@ from cray_megatron.megatron.training_loop import TrainingLoop, get_max_steps from cray_megatron.megatron.training_harness import TrainingHarness -import sys - import logging +import os +import sys +import time logger = logging.getLogger(__name__) +from cray_infra.training.train_debug import is_train_debug_enabled + + +def _trace_trainer(msg: str) -> None: + if not is_train_debug_enabled(): + return + rank = os.environ.get("RANK", os.environ.get("SLURM_PROCID", "?")) + line = f"[rank={rank}] trainer [{time.monotonic():.3f}]: {msg}\n" + sys.stderr.write(line) + sys.stderr.flush() + + class MegatronTrainer: def __init__(self, training_harness: TrainingHarness): self.training_harness = training_harness @@ -19,12 +32,23 @@ def train(self): self.train_loop() def train_loop(self): + _trace_trainer("train_loop enter") + _trace_trainer("train_loop calling update_status(TRAINING)") self.training_harness.update_status( status=TrainingJobStatus.TRAINING, metadata={"max_steps": get_max_steps()} ) + _trace_trainer("train_loop update_status(TRAINING) complete") + _trace_trainer("train_loop calling print_logo") print_logo() + _trace_trainer("train_loop print_logo complete") - TrainingLoop(self.training_harness).train() + _trace_trainer("train_loop creating TrainingLoop") + training_loop = TrainingLoop(self.training_harness) + _trace_trainer("train_loop calling TrainingLoop.train()") + training_loop.train() + _trace_trainer("train_loop TrainingLoop.train() complete") + _trace_trainer("train_loop calling update_status(COMPLETED)") self.training_harness.update_status(status=TrainingJobStatus.COMPLETED) + _trace_trainer("train_loop complete") diff --git a/ml/cray_megatron/megatron/training_harness.py b/ml/cray_megatron/megatron/training_harness.py index 4440c630..51af870f 100644 --- a/ml/cray_megatron/megatron/training_harness.py +++ b/ml/cray_megatron/megatron/training_harness.py @@ -3,6 +3,9 @@ from cray_megatron.collectives.main_rank_only import main_rank_only +import sys +import time + import torch import os @@ -13,16 +16,32 @@ logger = logging.getLogger(__name__) +from cray_infra.training.train_debug import is_train_debug_enabled + + +def _trace_harness(msg: str) -> None: + if not is_train_debug_enabled(): + return + rank = os.environ.get("RANK", os.environ.get("SLURM_PROCID", "?")) + line = f"[rank={rank}] harness [{time.monotonic():.3f}]: {msg}\n" + sys.stderr.write(line) + sys.stderr.flush() + + class TrainingHarness: def update_status(self, status, metadata={}): + _trace_harness(f"update_status enter status={status} metadata={metadata}") current_status = get_status() current_status["status"] = status + current_status["last_updated"] = time.time() for key, value in metadata.items(): current_status[key] = value + _trace_harness("update_status calling save_status") save_status(current_status) + _trace_harness("update_status save_status complete") def checkpoint(self, checkpoint_state, checkpoint_name): job_config = get_job_config() @@ -53,6 +72,7 @@ def get_training_job_directory(): @main_rank_only def save_status(job_status): + _trace_harness(f"save_status enter status={job_status.get('status', '?')}") try: contents = json.dumps(job_status) except Exception as e: diff --git a/ml/cray_megatron/megatron/training_loop.py b/ml/cray_megatron/megatron/training_loop.py index faaac5de..e58590d3 100644 --- a/ml/cray_megatron/megatron/training_loop.py +++ b/ml/cray_megatron/megatron/training_loop.py @@ -20,11 +20,23 @@ import time import logging -from gpu_aware_mpi import allreduce, get_size +import os +import sys +from cray_infra.training.distributed import allreduce, get_rank, get_size +from cray_infra.training.train_debug import is_train_debug_enabled logger = logging.getLogger(__name__) +def _trace_loop(msg: str) -> None: + if not is_train_debug_enabled(): + return + rank = os.environ.get("RANK", os.environ.get("SLURM_PROCID", "?")) + line = f"[rank={rank}] training_loop [{time.monotonic():.3f}]: {msg}\n" + sys.stderr.write(line) + sys.stderr.flush() + + class TrainingLoop: def __init__(self, training_harness: TrainingHarness): self.training_harness = training_harness @@ -34,13 +46,20 @@ def __init__(self, training_harness: TrainingHarness): self.training_state = TrainingState() def train(self): + _trace_loop("train() enter") + _trace_loop("train() calling get_model_manager") self.model_manager = get_model_manager() - + _trace_loop("train() calling load_model") self.training_state.model_info = self.model_manager.load_model() + _trace_loop("train() load_model complete") + _trace_loop("train() calling training_loop()") self.training_loop() + _trace_loop("train() training_loop complete") + _trace_loop("train() calling checkpoint()") self.checkpoint() + _trace_loop("train() complete") def training_loop(self): self.on_train_begin() @@ -192,8 +211,8 @@ def training_step_accumulate(self, batch, accum_step, gradient_accumulation_step if not is_nan: scaled_loss.backward() - # Log info for each micro-batch - self.print_microbatch_info(accum_step, avg_loss, start_time) + if gradient_accumulation_steps > 1: + self.print_microbatch_info(accum_step, avg_loss, start_time) return avg_loss @@ -312,10 +331,8 @@ def print_training_step_info(self, loss, step_time): f"- step time {step_time:.3f} seconds" ) - @main_rank_only def print_microbatch_info(self, accum_step, loss, start_time): - # only log if there is more than one microbatch - if get_gradient_accumulation_steps() <= 1: + if get_rank() != 0: return logger.debug( @@ -390,7 +407,7 @@ def get_max_steps(): def get_gradient_accumulation_steps(): job_config = get_job_config() - return job_config.get("gradient_accumulation_steps", 4) + return job_config.get("gradient_accumulation_steps", 1) def get_optimizer(model): diff --git a/ml/cray_megatron/models/load_model.py b/ml/cray_megatron/models/load_model.py index c32be1f6..58bcae3e 100644 --- a/ml/cray_megatron/models/load_model.py +++ b/ml/cray_megatron/models/load_model.py @@ -4,9 +4,9 @@ ) from cray_megatron.collectives.main_rank_only import is_main_rank -from gpu_aware_mpi import get_size, get_rank, allgather +from cray_infra.training.distributed import get_size, get_rank, allgather -from ml.adapters.add_adapters_to_model import add_adapters_to_model +from adapters.add_adapters_to_model import add_adapters_to_model from cray_infra.util.get_job_config import get_job_config from cray_infra.util.get_config import get_config diff --git a/requirements.txt b/requirements.txt index 4994efe0..6f9e237c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,8 +6,8 @@ aiofiles persist-queue matplotlib streaming-form-data==1.15.0 -mpi4py==4.0.3 openai>=1.40.0 # Ensure modern openai package (ensure types module present) +httpx>=0.27.0 openai-harmony humanize atomicwrites diff --git a/scripts/cray b/scripts/cray index 13fd8031..c8875bf4 100755 --- a/scripts/cray +++ b/scripts/cray @@ -189,7 +189,7 @@ cray_test_usage() { # :argument.usage printf " %s\n" "TEST-PATH" printf " Relative path to the directory or file with test cases\n" - printf " %s\n" "Default: test/infra/*" + printf " %s\n" "Default: test/infra" echo fi @@ -690,7 +690,12 @@ cray_test_command() { pytest_command_parts+=("-rP") fi - pytest_command_parts+=($test_path) + for path in $test_path; do + case "$(basename "$path")" in + __pycache__) continue ;; + esac + pytest_command_parts+=("$path") + done pytest_command="${pytest_command_parts[*]}" TTY=-t @@ -1455,7 +1460,7 @@ cray_test_parse_requirements() { done # :command.default_assignments - [[ -n ${args['test-path']:-} ]] || args['test-path']="test/infra/*" + [[ -n ${args['test-path']:-} ]] || args['test-path']="test/infra" [[ -n ${args['--coverage-path']:-} ]] || args['--coverage-path']="/tmp/cray/coverage" [[ -n ${args['--verbose']:-} ]] || args['--verbose']="no" [[ -n ${args['--workers']:-} ]] || args['--workers']="auto" diff --git a/scripts/install_protoc.sh b/scripts/install_protoc.sh new file mode 100644 index 00000000..05876d64 --- /dev/null +++ b/scripts/install_protoc.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Pinned protoc install (from vllm-project/vllm tools/install_protoc.sh). +# Distro protobuf-compiler versions are often too old for vLLM's Rust frontend. + +if [[ $(id -u) -ne 0 ]]; then + echo "Must be run as root" >&2 + exit 1 +fi + +VERSION="${PROTOC_VERSION:-34.2}" + +ARCH="$(uname -m)" +case "${ARCH}" in + aarch64|arm64) URL_ARCH="aarch_64" ;; + x86_64|amd64) URL_ARCH="x86_64" ;; + *) echo "Unsupported arch for protoc binary: ${ARCH}" >&2; exit 1 ;; +esac + +URL="https://github.com/protocolbuffers/protobuf/releases/download/v${VERSION}/protoc-${VERSION}-linux-${URL_ARCH}.zip" +TMPDIR="$(mktemp -d)" +trap 'rm -rf "${TMPDIR}"' EXIT + +echo "Downloading: ${URL}" +curl -fsSL -o "${TMPDIR}/protoc.zip" "${URL}" +unzip -q -o "${TMPDIR}/protoc.zip" -d /usr/local +echo "Installed $(protoc --version)" diff --git a/scripts/start_slurm.sh b/scripts/start_slurm.sh index f2979a13..f40a10cc 100755 --- a/scripts/start_slurm.sh +++ b/scripts/start_slurm.sh @@ -12,8 +12,28 @@ set -Eeuoxa pipefail # Get the directory of this script LOCAL_DIRECTORY="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" +SERVER_LIST=$(python -c "from cray_infra.util.get_config import get_config; print(get_config().get('server_list', ''))" 2>/dev/null || echo "") + +if [[ "$SERVER_LIST" != *"megatron"* ]] && [[ "$SERVER_LIST" != "all" ]]; then + echo "Skipping Slurm on server_list=${SERVER_LIST}" + exit 0 +fi + # Run the slurm discovery service python $LOCAL_DIRECTORY/../infra/cray_infra/slurm/discovery/discover_clusters.py -slurmctld +SLURM_CONF=${SLURM_CONF:-/app/cray/nfs/slurm.conf} +CONTROLLER="" +if [ -f "$SLURM_CONF" ]; then + CONTROLLER=$(grep '^SlurmctldHost=' "$SLURM_CONF" | cut -d= -f2- | tr -d '[:space:]') +fi + +HOSTNAME=$(hostname) +if [ -n "$CONTROLLER" ] && { [ "$HOSTNAME" = "$CONTROLLER" ] || [ "${HOSTNAME%%.*}" = "$CONTROLLER" ]; }; then + echo "Starting slurmctld on controller node ${HOSTNAME}" + slurmctld +else + echo "Skipping slurmctld on ${HOSTNAME} (controller is ${CONTROLLER:-unknown})" +fi + slurmd diff --git a/scripts/train_job_entrypoint.sh b/scripts/train_job_entrypoint.sh index 4161db46..5bea1913 100755 --- a/scripts/train_job_entrypoint.sh +++ b/scripts/train_job_entrypoint.sh @@ -2,19 +2,77 @@ # Safely execute this bash script # e exit on first failure -# x all executed commands are printed to the terminal # u unset variables are errors -# a export all variables to the environment # E any trap on ERR is inherited by shell functions # -o pipefail | produces a failure code if any stage fails -set -Eeuoxa pipefail +set -Eeuo pipefail +if [[ "${CRAY_TRAIN_DEBUG:-0}" == "1" ]]; then + set -x +fi export CRAY_TRAINING_JOB_CONFIG_PATH=REPLACE_CONFIG_PATH # Get the directory of this script LOCAL_DIRECTORY="$( cd "$( dirname "${CRAY_TRAINING_JOB_CONFIG_PATH}" )" >/dev/null 2>&1 && pwd )" +export LOCAL_DIRECTORY -# Put the current ml directory in the python path so that the modules can be imported +# Job ml overrides container; infra comes from Dockerfile PYTHONPATH. export PYTHONPATH=$LOCAL_DIRECTORY/ml:$PYTHONPATH -mpirun --allow-run-as-root python $LOCAL_DIRECTORY/ml/cray_megatron/main.py $* +export PYTHONUNBUFFERED=1 + +NODEFILE=$(mktemp) +scontrol show hostnames "$SLURM_JOB_NODELIST" > "$NODEFILE" +NUM_NODES=$(wc -l < "$NODEFILE") + +export WORLD_SIZE=$SLURM_NTASKS +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_PORT=29500 + +# Ensure essential env vars are set/exported in the api pod. Suggested example for single process per node: +# export NCCL_IB_GID_INDEX=1 +# export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 +# export HSA_NO_SCRATCH_RECLAIM=1 +# export NCCL_SOCKET_IFNAME=eth0 +# export GLOO_SOCKET_IFNAME=eth0 +# export NCCL_SOCKET_FAMILY=AF_INET +# export NCCL_NET_GDR_LEVEL=SYS +# export ROCR_VISIBLE_DEVICES=0 +# export HIP_VISIBLE_DEVICES=0 + +echo "[launcher t=$(date +%s)] LOCAL_DIRECTORY=${LOCAL_DIRECTORY}" >&2 +echo "[launcher t=$(date +%s)] SLURM_JOB_NODELIST=${SLURM_JOB_NODELIST:-?} SLURM_NTASKS=${SLURM_NTASKS:-?} SLURM_NTASKS_PER_NODE=${SLURM_NTASKS_PER_NODE:-?}" >&2 +echo "[launcher t=$(date +%s)] MASTER_ADDR=${MASTER_ADDR} MASTER_PORT=${MASTER_PORT} WORLD_SIZE=${WORLD_SIZE}" >&2 +echo "[launcher t=$(date +%s)] hostfile:" >&2 +cat "$NODEFILE" >&2 + +ulimit -l unlimited +mpirun --allow-run-as-root \ + --hostfile "$NODEFILE" \ + -np "$NUM_NODES" \ + -npernode $SLURM_NTASKS_PER_NODE \ + --bind-to none \ + bash -c ' +export RANK=$OMPI_COMM_WORLD_RANK +export WORLD_SIZE=$OMPI_COMM_WORLD_SIZE +export LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK +export LOCAL_WORLD_SIZE=$OMPI_COMM_WORLD_LOCAL_SIZE + +echo "Launching torchrun rank=$RANK local_rank=$LOCAL_RANK world=$WORLD_SIZE local_world_size $LOCAL_WORLD_SIZE on $(hostname) master=$MASTER_ADDR:$MASTER_PORT" >&2 + +ulimit -l unlimited + +torchrun \ + --nnodes="$WORLD_SIZE" \ + --nproc-per-node=$LOCAL_WORLD_SIZE \ + --node_rank="$RANK" \ + --local-rank="$LOCAL_RANK" \ + --master_addr="$MASTER_ADDR" \ + --master_port="$MASTER_PORT" \ + "$LOCAL_DIRECTORY/ml/cray_megatron/main.py" "$@" +' _ "$@" + +MPIRUN_EXIT=$? +rm -f "$NODEFILE" +echo "[launcher t=$(date +%s)] mpirun finished exit=${MPIRUN_EXIT}" >&2 +exit "$MPIRUN_EXIT" \ No newline at end of file diff --git a/test/benchmark/pytorch/mpi_p2p.py b/test/benchmark/pytorch/mpi_p2p.py index 7ba66d2b..1686fb16 100644 --- a/test/benchmark/pytorch/mpi_p2p.py +++ b/test/benchmark/pytorch/mpi_p2p.py @@ -3,7 +3,7 @@ import time import os -from gpu_aware_mpi import send, recv, barrier, get_rank, get_size, finalize_mpi +from cray_infra.training.distributed import send, recv, barrier, get_rank, get_size, finalize from tqdm import tqdm diff --git a/test/collectives/benchmark-send-recv.py b/test/collectives/benchmark-send-recv.py index 7358dd9f..34443d6e 100644 --- a/test/collectives/benchmark-send-recv.py +++ b/test/collectives/benchmark-send-recv.py @@ -20,7 +20,7 @@ def main(): def get_code(): return """ -from gpu_aware_mpi import get_size, get_rank, send, recv, barrier +from cray_infra.training.distributed import get_size, get_rank, send, recv, barrier import torch import time import statistics diff --git a/test/collectives/test-all-gather.py b/test/collectives/test-all-gather.py index 1c1b9fc0..d54f053c 100644 --- a/test/collectives/test-all-gather.py +++ b/test/collectives/test-all-gather.py @@ -19,7 +19,7 @@ def main(): def get_code(): return """ -from gpu_aware_mpi import get_size, get_rank, allgather, barrier +from cray_infra.training.distributed import get_size, get_rank, allgather, barrier import torch from cray_infra.training.training_job_context import training_job_context diff --git a/test/collectives/test-all-reduce.py b/test/collectives/test-all-reduce.py index 973c5489..04f069c0 100644 --- a/test/collectives/test-all-reduce.py +++ b/test/collectives/test-all-reduce.py @@ -19,7 +19,7 @@ def main(): def get_code(): return """ -from gpu_aware_mpi import get_size, get_rank, allreduce, barrier +from cray_infra.training.distributed import get_size, get_rank, allreduce, barrier import torch from cray_infra.training.training_job_context import training_job_context diff --git a/test/collectives/test-all-to-all.py b/test/collectives/test-all-to-all.py index ed523df9..5a92c2b6 100644 --- a/test/collectives/test-all-to-all.py +++ b/test/collectives/test-all-to-all.py @@ -19,7 +19,7 @@ def main(): def get_code(): return """ -from gpu_aware_mpi import get_size, get_rank, alltoall, barrier +from cray_infra.training.distributed import get_size, get_rank, alltoall, barrier import torch from cray_infra.training.training_job_context import training_job_context diff --git a/test/collectives/test-hello.py b/test/collectives/test-hello.py index 78583b71..151bb037 100644 --- a/test/collectives/test-hello.py +++ b/test/collectives/test-hello.py @@ -19,7 +19,7 @@ def main(): def get_code(): return """ -from gpu_aware_mpi import get_size, get_rank +from cray_infra.training.distributed import get_size, get_rank from cray_infra.training.training_job_context import training_job_context with training_job_context(): diff --git a/test/collectives/test-reduce-scatter.py b/test/collectives/test-reduce-scatter.py index e67c43bd..63502a77 100644 --- a/test/collectives/test-reduce-scatter.py +++ b/test/collectives/test-reduce-scatter.py @@ -19,7 +19,7 @@ def main(): def get_code(): return """ -from gpu_aware_mpi import get_size, get_rank, reduce_scatter, barrier +from cray_infra.training.distributed import get_size, get_rank, reduce_scatter, barrier import torch from cray_infra.training.training_job_context import training_job_context diff --git a/test/collectives/test-send-recv.py b/test/collectives/test-send-recv.py index 32006a5b..66f1a96f 100644 --- a/test/collectives/test-send-recv.py +++ b/test/collectives/test-send-recv.py @@ -19,7 +19,7 @@ def main(): def get_code(): return """ -from gpu_aware_mpi import get_size, get_rank, send, recv, barrier +from cray_infra.training.distributed import get_size, get_rank, send, recv, barrier import torch from cray_infra.training.training_job_context import training_job_context diff --git a/test/infra/distribution_strategy/benchmark_mpi_collectives.py b/test/infra/distribution_strategy/benchmark_mpi_collectives.py index 89e8f711..7f9c5150 100644 --- a/test/infra/distribution_strategy/benchmark_mpi_collectives.py +++ b/test/infra/distribution_strategy/benchmark_mpi_collectives.py @@ -1,89 +1,41 @@ -import argparse -import time -import torch -from gpu_aware_mpi import allgather, allreduce, reduce_scatter, barrier, get_rank, get_size, finalize_mpi - -def create_buffer(arch, size, rank): - if arch == 'cuda': - return torch.ones(size, dtype=torch.float32, device='cuda') - elif arch == 'rocm': - return torch.ones(size, dtype=torch.float32, device='cuda:' + str(rank)) # ROCm uses the same device string - else: - return torch.ones(size, dtype=torch.float32, device='cpu') - -def benchmark_collective(collective_fn, send_size, recv_size, expected_value, all_reduce_flag=False, num_iters=100, warmup=10): - size = get_size() - rank = get_rank() - sendbuf = create_buffer(args.arch, send_size, rank).contiguous() - # Create send/recv buffers - recvbuf = torch.empty(recv_size, dtype=torch.float32, device=sendbuf.device).contiguous() +"""CLI entry point for distributed collective benchmarks.""" - # Warmup iterations - for _ in range(warmup): - collective_fn(sendbuf, recvbuf) - - # Timing the collective operation - barrier() - t0 = time.time() - for _ in range(num_iters): - if all_reduce_flag: - sendbuf = create_buffer(args.arch, send_size, rank).contiguous() - collective_fn(sendbuf, recvbuf) - barrier() - dt = time.time() - t0 +import argparse - # Verify correctness - if all_reduce_flag: - assert torch.allclose(sendbuf, torch.full_like(sendbuf, expected_value), atol=1e-6), "Verification failed" - else: - assert torch.allclose(recvbuf, torch.full_like(recvbuf, expected_value), atol=1e-6), "Verification failed" +from cray_infra.training.distributed import get_rank - # Calculate bandwidth (we use float32 for ReduceScatter and AllReduce internally) - datatype_bytes = 4 - total_data = send_size * datatype_bytes * 2 * size # 4 bytes per float32, 2 for send/recv - bandwidth = (total_data / dt) / 1e9 # GB/s - return bandwidth +from distributed_benchmarks import ( + run_collectives_benchmark, + setup_distributed, + teardown_distributed, +) -if __name__ == "__main__": +def main(): parser = argparse.ArgumentParser() - parser.add_argument('--arch', choices=['cuda', 'rocm', 'cpu'], required=True) - parser.add_argument('--dtype', choices=['float32', 'bfloat16'], default='float32', help="Data type for buffers") + parser.add_argument("--arch", choices=["cuda", "rocm", "cpu"], required=True) + parser.add_argument( + "--dtype", + choices=["float32", "bfloat16"], + default="float32", + help="Reserved for future dtype support", + ) args = parser.parse_args() - args.dtype = torch.float32 if args.dtype == 'float32' else torch.bfloat16 + setup_distributed() + try: + results = run_collectives_benchmark(args.arch) + if get_rank() == 0: + print("\nBenchmark Results (GB/s):") + for name, bw in results.items(): + print(f"{name}: {bw:.2e}") + finally: + teardown_distributed() - data_size = 4194304 - rank = get_rank() - size = get_size() - collectives = { - 'AllGather': (lambda sbuf, rbuf: allgather(sbuf, rbuf), data_size, data_size * size, 1.0, False), - 'ReduceScatter': (lambda sbuf, rbuf: reduce_scatter(sbuf, rbuf), data_size, data_size // size, size * 1.0, False), - 'AllReduce': (lambda sbuf, rbuf: allreduce(sbuf), data_size, 1, size * 1.0, True), - } - - results = {} - - for name, info in collectives.items(): - bw = benchmark_collective(info[0], info[1], info[2], info[3], info[4]) - if rank == 0: - results[name] = bw - - if rank == 0: - print("\nBenchmark Results (GB/s):") - for name, bw in results.items(): - bw_scientific = '{:.2e}'.format(bw) - print(f"{name}: {bw_scientific}") - - finalize_mpi() - - -# For CUDA GPUs -# mpirun --allow-run-as-root --oversubscribe -np 4 python test/infra/distribution_strategy/benchmark_mpi_collectives.py --arch cuda - -# For ROCm GPUs -# mpirun --allow-run-as-root -np 4 python test/infra/distribution_strategy/benchmark_mpi_collectives.py --arch rocm +if __name__ == "__main__": + main() -# For CPU -# mpirun --allow-run-as-root --oversubscribe -np 2 python test/infra/distribution_strategy/benchmark_mpi_collectives.py --arch cpu +# torchrun --nnodes=1 --nproc-per-node=4 test/infra/distribution_strategy/benchmark_mpi_collectives.py --arch cuda +# torchrun --nnodes=1 --nproc-per-node=4 test/infra/distribution_strategy/benchmark_mpi_collectives.py --arch rocm +# torchrun --nnodes=1 --nproc-per-node=2 test/infra/distribution_strategy/benchmark_mpi_collectives.py --arch cpu diff --git a/test/infra/distribution_strategy/benchmark_mpi_sendrecv.py b/test/infra/distribution_strategy/benchmark_mpi_sendrecv.py index 6131e1ee..125f2010 100644 --- a/test/infra/distribution_strategy/benchmark_mpi_sendrecv.py +++ b/test/infra/distribution_strategy/benchmark_mpi_sendrecv.py @@ -1,68 +1,39 @@ -import argparse -import time -import torch -from gpu_aware_mpi import send, recv, barrier, get_rank, get_size, finalize_mpi - -def create_buffer(arch_type, size, rank): - if arch_type == 'cuda': - return torch.ones(size, dtype=torch.float32, device='cuda') - elif arch_type == 'rocm': - return torch.ones(size, dtype=torch.float32, device='cuda:' + str(rank)) # ROCm uses the same device string - else: - return torch.ones(size, dtype=torch.float32, device='cpu') - -def send_recv(sendbuf, i): - sender = i % 2 - rank = get_rank() - barrier() - if rank == sender: - send(sendbuf, (rank + 1) % 2) - else: - recv(sendbuf, (rank + 1) % 2) - barrier() - -def benchmark_send_recv(data_size, num_iters=100, warmup=10): - rank = get_rank() - sendbuf = create_buffer(args.arch_type, data_size, rank).contiguous() - - if rank == 1: - torch.zero_(sendbuf) +"""CLI entry point for distributed send/recv benchmarks.""" - for i in range(0, warmup * 2, 2): - send_recv(sendbuf, i) - - # Timing the collective operation - barrier() - t0 = time.time() - for i in range(0, num_iters * 2, 2): - send_recv(sendbuf, i) - barrier() - dt = time.time() - t0 +import argparse - # Verify correctness - assert torch.allclose(sendbuf, torch.full_like(sendbuf, 1), atol=1e-6), "Verification failed" +from cray_infra.training.distributed import get_rank, get_size - total_data = data_size * 4 * num_iters - bandwidth = (total_data / dt) / 1e9 # GB/s - return bandwidth +from distributed_benchmarks import ( + run_sendrecv_benchmark, + setup_distributed, + teardown_distributed, +) -if __name__ == "__main__": +def main(): parser = argparse.ArgumentParser() - parser.add_argument('--arch_type', choices=['cuda', 'rocm', 'cpu'], required=True) + parser.add_argument("--arch", choices=["cuda", "rocm", "cpu"]) + parser.add_argument( + "--arch_type", + choices=["cuda", "rocm", "cpu"], + help="Deprecated alias for --arch", + ) args = parser.parse_args() + arch = args.arch or args.arch_type + if arch is None: + parser.error("one of --arch or --arch_type is required") - data_size = 262144 * 64 - rank = get_rank() - size = get_size() - - assert size == 2, "This test only works with two ranks." - - print(f"Rank {rank} of {size} running with data size {data_size}") + setup_distributed() + try: + assert get_size() == 2, "This benchmark only works with two ranks." + print(f"Rank {get_rank()} of {get_size()} running with arch={arch}") + bandwidth = run_sendrecv_benchmark(arch) + if get_rank() == 0: + print(f"Send/Recv (GB/s): {bandwidth:.6f}") + finally: + teardown_distributed() - bandwidth = benchmark_send_recv(data_size) - if rank == 0: - print(f"Send/Recv (GB/s): {bandwidth:.6}") - - finalize_mpi() \ No newline at end of file +if __name__ == "__main__": + main() diff --git a/test/infra/distribution_strategy/conftest.py b/test/infra/distribution_strategy/conftest.py new file mode 100644 index 00000000..78a0e406 --- /dev/null +++ b/test/infra/distribution_strategy/conftest.py @@ -0,0 +1,94 @@ +"""Pytest hooks for distributed tests launched via torchrun.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +CRAY_ROOT = Path(__file__).resolve().parents[3] + +_TORCHRUN_CHILD_ENV = "_TORCHRUN_CHILD" +_ARCH_ENV = "DISTRIBUTED_TEST_ARCH" + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "torchrun(nproc, arch): run test under torch.distributed.run", + ) + + +def pytest_collection_modifyitems(items): + for item in items: + if item.get_closest_marker("torchrun") is not None: + item.add_marker(pytest.mark.xdist_group("distributed_torchrun")) + + +def _is_torchrun_child() -> bool: + return os.environ.get(_TORCHRUN_CHILD_ENV) == "1" + + +@pytest.fixture(scope="session", autouse=True) +def distributed_process_group(): + if not _is_torchrun_child(): + yield + return + + from distributed_benchmarks import setup_distributed, teardown_distributed + + setup_distributed() + yield + teardown_distributed() + + +@pytest.fixture +def distributed_arch(request): + arch = os.environ.get(_ARCH_ENV) + if arch is None: + marker = request.node.get_closest_marker("torchrun") + if marker is not None: + arch = marker.kwargs.get("arch") + if arch is None: + pytest.fail("DISTRIBUTED_TEST_ARCH is not set for distributed test worker") + return arch + + +@pytest.hookimpl(tryfirst=True) +def pytest_pyfunc_call(pyfuncitem): + marker = pyfuncitem.get_closest_marker("torchrun") + if marker is None: + return None + if _is_torchrun_child(): + return None + + nproc = marker.kwargs.get("nproc", 2) + arch = marker.kwargs.get("arch", "cpu") + env = { + **os.environ, + _TORCHRUN_CHILD_ENV: "1", + _ARCH_ENV: arch, + } + + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + "--nnodes=1", + f"--nproc-per-node={nproc}", + "-m", + "pytest", + pyfuncitem.nodeid, + "-xvs", + "-p", + "no:xdist", + "-p", + "no:forked", + ] + result = subprocess.run(cmd, env=env, cwd=CRAY_ROOT) + if result.returncode != 0: + pytest.fail( + f"torchrun pytest failed (exit {result.returncode}): {' '.join(cmd)}" + ) + return True diff --git a/test/infra/distribution_strategy/distributed_benchmarks.py b/test/infra/distribution_strategy/distributed_benchmarks.py new file mode 100644 index 00000000..6771865b --- /dev/null +++ b/test/infra/distribution_strategy/distributed_benchmarks.py @@ -0,0 +1,164 @@ +"""Shared distributed collective benchmarks using cray_infra.training.distributed.""" + +import time + +import torch + +from cray_infra.training.distributed import ( + allgather, + allreduce, + barrier, + finalize, + get_rank, + get_size, + init, + recv, + reduce_scatter, + send, +) + +DATA_SIZE_COLLECTIVES = 4_194_304 +DATA_SIZE_SENDRECV = 262_144 * 64 +COLLECTIVE_WARMUP = 10 +COLLECTIVE_ITERS = 100 +SENDRECV_WARMUP = 10 +SENDRECV_ITERS = 100 + + +def create_buffer(arch: str, size: int, rank: int) -> torch.Tensor: + if arch == "cuda": + return torch.ones(size, dtype=torch.float32, device="cuda") + if arch == "rocm": + return torch.ones(size, dtype=torch.float32, device=f"cuda:{rank}") + return torch.ones(size, dtype=torch.float32, device="cpu") + + +def benchmark_collective( + collective_fn, + arch: str, + send_size: int, + recv_size: int, + expected_value: float, + all_reduce_flag: bool = False, + num_iters: int = COLLECTIVE_ITERS, + warmup: int = COLLECTIVE_WARMUP, +) -> float: + size = get_size() + rank = get_rank() + sendbuf = create_buffer(arch, send_size, rank).contiguous() + recvbuf = torch.empty(recv_size, dtype=torch.float32, device=sendbuf.device).contiguous() + + for _ in range(warmup): + collective_fn(sendbuf, recvbuf) + + barrier() + t0 = time.time() + for _ in range(num_iters): + if all_reduce_flag: + sendbuf = create_buffer(arch, send_size, rank).contiguous() + collective_fn(sendbuf, recvbuf) + barrier() + dt = time.time() - t0 + + if all_reduce_flag: + assert torch.allclose( + sendbuf, torch.full_like(sendbuf, expected_value), atol=1e-6 + ), "Verification failed" + else: + assert torch.allclose( + recvbuf, torch.full_like(recvbuf, expected_value), atol=1e-6 + ), "Verification failed" + + datatype_bytes = 4 + total_data = send_size * datatype_bytes * 2 * size + return (total_data / dt) / 1e9 + + +def run_collectives_benchmark(arch: str) -> dict[str, float]: + data_size = DATA_SIZE_COLLECTIVES + size = get_size() + + collectives = { + "AllGather": ( + lambda sbuf, rbuf: allgather(sbuf, rbuf), + data_size, + data_size * size, + 1.0, + False, + ), + "ReduceScatter": ( + lambda sbuf, rbuf: reduce_scatter(sbuf, rbuf), + data_size, + data_size // size, + size * 1.0, + False, + ), + "AllReduce": ( + lambda sbuf, rbuf: allreduce(sbuf), + data_size, + 1, + size * 1.0, + True, + ), + } + + results = {} + for name, info in collectives.items(): + bw = benchmark_collective(info[0], arch, info[1], info[2], info[3], info[4]) + if get_rank() == 0: + results[name] = bw + return results + + +def _send_recv(sendbuf: torch.Tensor, iteration: int) -> None: + sender = iteration % 2 + rank = get_rank() + barrier() + if rank == sender: + send(sendbuf, (rank + 1) % 2) + else: + recv(sendbuf, (rank + 1) % 2) + barrier() + + +def benchmark_send_recv( + arch: str, + data_size: int = DATA_SIZE_SENDRECV, + num_iters: int = SENDRECV_ITERS, + warmup: int = SENDRECV_WARMUP, +) -> float: + rank = get_rank() + sendbuf = create_buffer(arch, data_size, rank).contiguous() + + if rank == 1: + torch.zero_(sendbuf) + + for i in range(0, warmup * 2, 2): + _send_recv(sendbuf, i) + + barrier() + t0 = time.time() + for i in range(0, num_iters * 2, 2): + _send_recv(sendbuf, i) + barrier() + dt = time.time() - t0 + + assert torch.allclose(sendbuf, torch.full_like(sendbuf, 1), atol=1e-6), ( + "Verification failed" + ) + + total_data = data_size * 4 * num_iters + return (total_data / dt) / 1e9 + + +def run_sendrecv_benchmark(arch: str) -> float: + assert get_size() == 2, "Send/recv benchmark requires exactly two ranks" + return benchmark_send_recv(arch) + + +def setup_distributed() -> None: + init() + + +def teardown_distributed() -> None: + finalize() diff --git a/test/infra/distribution_strategy/test_distributed_collectives.py b/test/infra/distribution_strategy/test_distributed_collectives.py new file mode 100644 index 00000000..9ec7f8bd --- /dev/null +++ b/test/infra/distribution_strategy/test_distributed_collectives.py @@ -0,0 +1,44 @@ +"""Pytest coverage for cray_infra.training.distributed collectives. + +Runs under torchrun automatically via the ``@pytest.mark.torchrun`` marker, or +directly, for example:: + + torchrun --nnodes=1 --nproc-per-node=2 \\ + -m pytest test/infra/distribution_strategy/test_distributed_collectives.py -xvs +""" + +import pytest +import torch + +from cray_infra.training.distributed import get_rank + +from distributed_benchmarks import run_collectives_benchmark + + +@pytest.mark.torchrun(nproc=2, arch="cpu") +def test_collectives_cpu(distributed_arch): + results = run_collectives_benchmark(distributed_arch) + if get_rank() == 0: + assert results, "rank 0 should collect benchmark results" + for name, bandwidth in results.items(): + assert bandwidth > 0, f"{name} bandwidth should be positive" + + +@pytest.mark.torchrun(nproc=4, arch="cuda") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_collectives_cuda(distributed_arch): + results = run_collectives_benchmark(distributed_arch) + if get_rank() == 0: + assert results + for bandwidth in results.values(): + assert bandwidth > 0 + + +@pytest.mark.torchrun(nproc=4, arch="rocm") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="ROCm/CUDA device is required") +def test_collectives_rocm(distributed_arch): + results = run_collectives_benchmark(distributed_arch) + if get_rank() == 0: + assert results + for bandwidth in results.values(): + assert bandwidth > 0 diff --git a/test/infra/distribution_strategy/test_distributed_sendrecv.py b/test/infra/distribution_strategy/test_distributed_sendrecv.py new file mode 100644 index 00000000..0dff06f6 --- /dev/null +++ b/test/infra/distribution_strategy/test_distributed_sendrecv.py @@ -0,0 +1,38 @@ +"""Pytest coverage for cray_infra.training.distributed point-to-point ops. + +Runs under torchrun automatically via the ``@pytest.mark.torchrun`` marker, or +directly, for example:: + + torchrun --nnodes=1 --nproc-per-node=2 \\ + -m pytest test/infra/distribution_strategy/test_distributed_sendrecv.py -xvs +""" + +import pytest +import torch + +from cray_infra.training.distributed import get_rank + +from distributed_benchmarks import run_sendrecv_benchmark + + +@pytest.mark.torchrun(nproc=2, arch="cpu") +def test_sendrecv_cpu(distributed_arch): + bandwidth = run_sendrecv_benchmark(distributed_arch) + if get_rank() == 0: + assert bandwidth > 0 + + +@pytest.mark.torchrun(nproc=2, arch="cuda") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_sendrecv_cuda(distributed_arch): + bandwidth = run_sendrecv_benchmark(distributed_arch) + if get_rank() == 0: + assert bandwidth > 0 + + +@pytest.mark.torchrun(nproc=2, arch="rocm") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="ROCm/CUDA device is required") +def test_sendrecv_rocm(distributed_arch): + bandwidth = run_sendrecv_benchmark(distributed_arch) + if get_rank() == 0: + assert bandwidth > 0 diff --git a/test/infra/openai_client.py b/test/infra/openai_client.py index ec68960b..5160a25e 100644 --- a/test/infra/openai_client.py +++ b/test/infra/openai_client.py @@ -30,7 +30,7 @@ async def test_openai_client(self): config = get_config() client = AsyncOpenAI( - base_url=config["api_url"] + "/v1/openai", + base_url=config["api_url"] + "/v1", api_key="token-abc123", )