Skip to content

Switch collectives from MPI/UCX to NCCL/RCCL (torchdist) #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AVSuni wants to merge 10 commits into
base: main
Choose a base branch
from
Open

Switch collectives from MPI/UCX to NCCL/RCCL (torchdist) #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
017f480
move from mpi/ucx collectives to nccl (pytorch distributed)
AVSuni May 28, 2026
396d5a1
fix build
AVSuni May 28, 2026
d2e8f8e
Fix test errors
AVSuni May 29, 2026
0961953
update torchao
AVSuni May 29, 2026
2f774f4
update training entrypoint script
AVSuni May 29, 2026
8bd199d
add NCCL Helm deployment manifests and harden code for slurm job mana…
AVSuni Jun 3, 2026
336ac26
reduce noise
AVSuni Jun 9, 2026
2fd0dd4
Refactor train job entrypoint, cleanup, pass multi-task per node args
aivanni Jun 11, 2026
fc9336f
adjust env vars on api deployment
aivanni Jun 11, 2026
98ba239
Do not hardcode env vars in train job entrypoint
aivanni Jun 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 19 additions & 12 deletions Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -267,30 +275,27 @@ 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
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 \
Expand All @@ -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

Expand Down
Loading