Training Benchmark

.github/workflows/docker/Dockerfile_v25.10_ainic_temp

@@ -0,0 +1,105 @@
+# Base image
+FROM docker.io/rocm/megatron-lm:v25.10
+
+# Specify the commit of Primus-Turbo when building: docker build --build-arg PRIMUS_TURBO_COMMIT=xxx .)
+ARG PRIMUS_TURBO_COMMIT
+ARG AINIC_BUNDLE_PATH
+
+# Install basic dependencies
+RUN apt-get update
+
+# Clone and install the Primus-Turbo
+WORKDIR /opt
+RUN mkdir -p /opt && cd /opt && \
+    git clone https://github.com/AMD-AGI/Primus-Turbo.git && \
+    cd Primus-Turbo && \
+    git checkout ${PRIMUS_TURBO_COMMIT} && \
+    git submodule update --init --recursive && \
+    pip3 install -r requirements.txt && \
+    GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation .
+
+RUN apt-get install --reinstall binutils -y && apt-get install numactl -y
+
+WORKDIR /opt
+ENV WORKDIR=/opt
+ENV ROCM_PATH=/opt/rocm
+
+RUN apt-get update && \
+    apt-get install jq dpkg-dev kmod xz-utils \
+    libfmt-dev libboost-all-dev \
+    libibverbs-dev ibverbs-utils infiniband-diags -y
+
+# =============================== Build AINIC Driver ===============================
+# WARNING: Please ensure the following environment variables are correctly set:
+# WARNING:     1. PATH: /usr/sbin must be included.
+# WARNING:     2. LD_LIBRARY_PATH: /usr/lib must be included.
+# WARNING: If these paths are missing, tools and libraries may not function correctly.
+# INFO: Installation completed successfully
+
+COPY ${AINIC_BUNDLE_PATH}/ainic_bundle_1.117.1-a-42.tar.gz /opt/
+RUN cd ${WORKDIR} && \
+    echo "Building ainic bundle... current directory: ${WORKDIR}" && \
+    tar zxf ainic_bundle_1.117.1-a-42.tar.gz && \
+    cd ainic_bundle_1.117.1-a-42 && \
+    tar zxf host_sw_pkg.tar.gz && \
+    cd host_sw_pkg && \
+    ./install.sh --domain=user -y 2>&1 | tee log_install.txt && \
+    cd /opt
+
+# =============================== Test AINIC Driver ===============================
+# ibv_devices
+# rdma link
+# ethtool -i enp9s0
+# ibv_devinfo -vv | grep GID
+
+# =============================== Build UCX ===============================
+RUN cd ${WORKDIR} && wget https://github.com/openucx/ucx/releases/download/v1.18.0/ucx-1.18.0.tar.gz && \
+    mkdir -p ucx-1.18.0 && \
+    tar -zxf ucx-1.18.0.tar.gz -C ucx-1.18.0 --strip-components=1 && \
+    cd ucx-1.18.0 && mkdir build && cd build && \
+    ../configure --prefix=${WORKDIR}/ucx-1.18.0/install --with-rocm=${ROCM_PATH} 2>&1 | tee log_ucx_configure.txt && \
+    make -j 16 2>&1 | tee log_ucx_build.txt && \
+    make install && \
+    cd ${WORKDIR}
+
+ENV UCX_INSTALL_DIR=${WORKDIR}/ucx-1.18.0/install
+
+# =============================== Build MPI ===============================
+RUN cd ${WORKDIR} && \
+    wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.6.tar.gz && \
+    mkdir -p ompi-4.1.6 && \
+    tar -zxf openmpi-4.1.6.tar.gz -C ompi-4.1.6 --strip-components=1 && \
+    cd ompi-4.1.6 && mkdir build && cd build && \
+    ../configure --prefix=${WORKDIR}/ompi-4.1.6/install --with-ucx=${UCX_INSTALL_DIR} \
+        --disable-oshmem --disable-mpi-fortran 2>&1 | tee log_mpi_configure.txt && \
+    make -j 16 2>&1 | tee log_mpi_build.txt && \
+    make install && \
+    cd ${WORKDIR}
+
+ENV MPI_PATH=${WORKDIR}/ompi-4.1.6/install
+
+# =============================== Build RCCL ===============================
+RUN cd ${WORKDIR} && \
+    git clone https://github.com/ROCm/rccl.git && \
+    cd rccl && git checkout drop/2025-08 && \
+    ./install.sh -l --prefix build/ --disable-mscclpp \
+        --disable-msccl-kernel --amdgpu_targets="gfx950" 2>&1 | tee log_rccl_install.txt && \
+    cd ${WORKDIR}
+
+ENV RCCL_HOME=${WORKDIR}/rccl
+
+# =============================== Build AMD ANP ===============================
+
+RUN cd ${WORKDIR} && git clone https://github.com/rocm/amd-anp.git && \
+cd amd-anp && git checkout tags/v1.1.0-5 && \
+sed -i '5a CFLAGS += --offload-arch=gfx950' ./Makefile && head -10 ./Makefile && \
+make -j 16 RCCL_BUILD=${RCCL_HOME}/build/release \
+           MPI_INCLUDE=${MPI_PATH}/include/ \
+           MPI_LIB_PATH=${MPI_PATH}/lib/ \
+           ROCM_PATH=${ROCM_PATH} 2>&1 | tee log_amd_anp_build.txt
+
+# Set the default working directory
+WORKDIR /opt
+
+# check the installed Primus-Turbo package
+RUN python3 -m pip show primus-turbo || true

.github/workflows/docker/Dockerfile_v25.11_ainic

@@ -0,0 +1,114 @@
+# Base image
+FROM docker.io/rocm/primus:v25.11
+
+# Specify the commit of Primus-Turbo when building: docker build --build-arg PRIMUS_TURBO_COMMIT=xxx .)
+ARG PRIMUS_TURBO_COMMIT
+ARG AINIC_BUNDLE_PATH
+
+# Install basic dependencies
+RUN apt-get update
+
+# Clone and install the Primus-Turbo
+WORKDIR /opt
+RUN mkdir -p /opt && cd /opt && \
+    git clone https://github.com/AMD-AGI/Primus-Turbo.git && \
+    cd Primus-Turbo && \
+    git checkout ${PRIMUS_TURBO_COMMIT} && \
+    git submodule update --init --recursive && \
+    pip3 install -r requirements.txt && \
+    GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation .
+
+RUN apt-get install --reinstall binutils -y && apt-get install numactl -y
+
+WORKDIR /opt
+ENV WORKDIR=/opt
+ENV ROCM_PATH=/opt/rocm
+
+RUN apt-get update && \
+    apt-get install jq dpkg-dev kmod xz-utils \
+    libfmt-dev libboost-all-dev \
+    libibverbs-dev ibverbs-utils infiniband-diags -y
+
+# =============================== Build AINIC Driver ===============================
+# WARNING: Please ensure the following environment variables are correctly set:
+# WARNING:     1. PATH: /usr/sbin must be included.
+# WARNING:     2. LD_LIBRARY_PATH: /usr/lib must be included.
+# WARNING: If these paths are missing, tools and libraries may not function correctly.
+# INFO: Installation completed successfully
+
+COPY ${AINIC_BUNDLE_PATH}/ainic_bundle_1.117.5-a-38.tar.gz /opt/
+RUN cd ${WORKDIR} && \
+    echo "Building ainic bundle... current directory: ${WORKDIR}" && \
+    tar zxf ainic_bundle_1.117.5-a-38.tar.gz && \
+    cd ainic_bundle_1.117.5-a-38 && \
+    tar zxf host_sw_pkg.tar.gz && \
+    cd host_sw_pkg && \
+    ./install.sh --domain=user -y 2>&1 | tee log_install.txt && \
+    cd /opt
+
+# =============================== Test AINIC Driver ===============================
+# ibv_devices
+# rdma link
+# ethtool -i enp9s0
+# ibv_devinfo -vv | grep GID
+
+# =============================== Build UCX ===============================
+ENV USE_UCX_VERSION="1.15.0"
+RUN cd ${WORKDIR} && wget https://github.com/openucx/ucx/releases/download/v${USE_UCX_VERSION}/ucx-${USE_UCX_VERSION}.tar.gz && \
+    mkdir -p ucx-${USE_UCX_VERSION} && \
+    tar -zxf ucx-${USE_UCX_VERSION}.tar.gz -C ucx-${USE_UCX_VERSION} --strip-components=1 && \
+    cd ucx-${USE_UCX_VERSION} && mkdir build && cd build && \
+    ../configure --prefix=${WORKDIR}/ucx-${USE_UCX_VERSION}/install --with-rocm=${ROCM_PATH} 2>&1 | tee log_ucx_configure.txt && \
+    make -j 16 2>&1 | tee log_ucx_build.txt && \
+    make install && \
+    cd ${WORKDIR}
+
+ENV UCX_INSTALL_DIR=${WORKDIR}/ucx-${USE_UCX_VERSION}/install
+
+# =============================== Build MPI ===============================
+RUN cd ${WORKDIR} && \
+    wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.6.tar.gz && \
+    mkdir -p ompi-4.1.6 && \
+    tar -zxf openmpi-4.1.6.tar.gz -C ompi-4.1.6 --strip-components=1 && \
+    cd ompi-4.1.6 && mkdir build && cd build && \
+    ../configure --prefix=${WORKDIR}/ompi-4.1.6/install --with-ucx=${UCX_INSTALL_DIR} \
+        --disable-oshmem --disable-mpi-fortran 2>&1 | tee log_mpi_configure.txt && \
+    make -j 16 2>&1 | tee log_mpi_build.txt && \
+    make install && \
+    cd ${WORKDIR}
+
+ENV MPI_PATH=${WORKDIR}/ompi-4.1.6/install
+
+# =============================== Build RCCL ===============================
+    # cd rccl && git checkout rocm-7.1.0 && \
+RUN cd ${WORKDIR} && \
+    git clone https://github.com/ROCm/rccl.git && \
+    cd rccl && git checkout release/rocm-rel-7.1 && \
+    ./install.sh -l --prefix build/ --disable-msccl-kernel \
+        --amdgpu_targets="gfx950" 2>&1 | tee log_rccl_install.txt && \
+    cd ${WORKDIR}
+
+ENV RCCL_HOME=${WORKDIR}/rccl
+
+# =============================== Build AMD ANP ===============================
+
+# sed -i '5a CFLAGS += --offload-arch=gfx950' ./Makefile && head -10 ./Makefile && \
+
+WORKDIR /opt
+# COPY ${AINIC_BUNDLE_PATH}/amd-anp-v1.3.0.patch /opt/
+# git apply /opt/amd-anp-v1.3.0.patch && \
+RUN cd ${WORKDIR} && git clone https://github.com/rocm/amd-anp.git && \
+cd amd-anp && git checkout tags/v1.3.0 && \
+sed -i '5a CFLAGS += --offload-arch=gfx950' ./Makefile && head -10 ./Makefile && \
+make -j 16 RCCL_HOME=${RCCL_HOME} \
+           MPI_INCLUDE=${MPI_PATH}/include/ \
+           MPI_LIB_PATH=${MPI_PATH}/lib/ \
+           ROCM_PATH=${ROCM_PATH} && \
+make RCCL_HOME=$RCCL_HOME ROCM_PATH=${ROCM_PATH} install \
+    2>&1 | tee log_amd_anp_build.txt
+
+# Set the default working directory
+WORKDIR /opt
+
+# check the installed Primus-Turbo package
+RUN python3 -m pip show primus-turbo || true

examples/megatron/configs/MI355X/llama3.1_70B-Zero1-BF16-pretrain.yaml

@@ -0,0 +1,74 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:llama3_70B-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: llama3.1_70B.yaml
+    overrides:
+      # log
+      wandb_project: "Primus_DeepSeek_Pretrain"
+      stderr_sink_level: DEBUG
+
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      train_iters: 50
+      micro_batch_size: 1
+      global_batch_size: 32
+
+      seq_length: 8192
+      max_position_embeddings: 8192
+
+      lr: 1.0e-5
+      min_lr: 0.0
+      lr_warmup_iters: 2
+      lr_decay_iters: null
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: true
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-6
+
+      # parallel
+      tensor_model_parallel_size: 4
+      pipeline_model_parallel_size: 8
+      sequence_parallel: 1
+
+      # data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # ckpt
+      finetune: false
+      auto_continue_train: false
+      load: null
+      no_load_optim: null
+      no_load_rng: null
+      save: null
+      save_interval: 20000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch_dist
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      gradient_accumulation_fusion: true
+
+      # recompute
+      recompute_granularity: full # full, selective
+      recompute_method: block # uniform, block
+      recompute_num_layers: 0 # int
+
+      # Cross entropy flags
+      # cross_entropy_fusion_impl: "te"
+      # cross_entropy_loss_fusion: true

examples/run_local_pretrain.sh

@@ -143,6 +143,20 @@ if [[ "${CLEAN_DOCKER_CONTAINER:-0}" == "1" ]]; then
     fi
 fi
 
+if [[ -n "${DOCKER_LOGIN_KEY}" ]]; then
+    if [[ -z "${DOCKER_LOGIN_USER}" ]]; then
+        echo "DOCKER_LOGIN_KEY is set but DOCKER_LOGIN_USER is empty. Exiting."
+        exit 1
+    fi
+    echo "Logging in to Docker registry as ${DOCKER_LOGIN_USER}..."
+    docker login -u "${DOCKER_LOGIN_USER}" -p "${DOCKER_LOGIN_KEY}"
+    if [[ $? -ne 0 ]]; then
+        echo "Docker login failed! Exiting."
+        exit 1
+    fi
+fi
+
+
 if [[ "${SKIP_TRAIN:-0}" == "1" ]]; then
     echo "Node-${NODE_RANK}: Skipping training container launch."
     exit 0

examples/run_pretrain.sh

@@ -158,12 +158,12 @@ if [ "$USING_AINIC" == "1" ]; then
     LOG_INFO_RANK0 "ANP_HOME_DIR: $ANP_HOME_DIR"
     LOG_INFO_RANK0 "MPI_HOME_DIR: $MPI_HOME_DIR"
 
-    # unset NCCL_IB_GID_INDEX
+    export NCCL_IB_TC=${NCCL_IB_TC:-104}
+    export NCCL_IB_FIFO_TC=${NCCL_IB_FIFO_TC:-184} # 192
+
     export NCCL_IB_GID_INDEX=1
-    # export NCCL_IB_ROCE_VERSION_NUM=2
+    export NCCL_IB_ROCE_VERSION_NUM=2
     export NCCL_MAX_P2P_CHANNELS=56
-    export NCCL_IB_TC=104
-    export NCCL_IB_FIFO_TC=192
     export NET_OPTIONAL_RECV_COMPLETION=1
     export NCCL_IB_USE_INLINE=1
     export RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING=0
@@ -178,7 +178,14 @@ if [ "$USING_AINIC" == "1" ]; then
 
     # v25.09
     export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/lib/x86_64-linux-gnu/libibverbs:${RCCL_HOME_DIR}/build/release:${ANP_HOME_DIR}/build:${MPI_HOME_DIR}/install/lib:$LD_LIBRARY_PATH
-    export LD_PRELOAD=${ANP_HOME_DIR}/build/librccl-net.so:${RCCL_HOME_DIR}/build/release/librccl.so.1.0
+    # if [ -f "${ANP_HOME_DIR}/build/librccl-net.so" ]; then
+    #     export LD_PRELOAD="${ANP_HOME_DIR}/build/librccl-net.so:${RCCL_HOME_DIR}/build/release/librccl.so.1.0"
+    # elif [ -f "${ANP_HOME_DIR}/build/librccl-anp.so" ]; then
+    #     export LD_PRELOAD="${ANP_HOME_DIR}/build/librccl-anp.so:${RCCL_HOME_DIR}/build/release/librccl.so.1.0"
+    # else
+    #     echo "ERROR: Neither librccl-net.so nor librccl-anp.so was found in ${ANP_HOME_DIR}/build"
+    #     exit 1
+    # fi
 else
     export NCCL_IB_GID_INDEX=3
 fi
@@ -477,7 +484,7 @@ fi
 setup_pythonpath() {
     local site_packages
     site_packages=$(python -c "import sysconfig; print(sysconfig.get_paths()['purelib'])")
-    export PYTHONPATH="${site_packages}:${PRIMUS_PATH}:$:${PYTHONPATH}"
+    export PYTHONPATH="${PRIMUS_PATH}:${site_packages}:${PYTHONPATH}"
 }
 
 setup_pythonpath

primus/modules/trainer/megatron/pre_trainer.py

@@ -31,7 +31,7 @@ def get_batch_func(data_iterator, vp_stage=None):
     if not is_first_or_last_pipeline_stage(vp_stage):
         return None, None, None, None, None
 
-    assert data_iterator is not None, f"data_iterator is None vp_stage: {vp_stage}"
+    # assert data_iterator is not None, f"data_iterator is None vp_stage: {vp_stage}"
     # get batches based on the TP rank you are on
     batch = get_batch_on_this_tp_rank(data_iterator)