Merge pull request GoogleCloudPlatform#4725 from SwarnaBharathiMantena/swarnabm/update_nemofw_ex

Update machine-leaning/a3-ultragpu-8g/nemo-framework to fix segmentation fault error

Author: SwarnaBharathiMantena

examples/machine-learning/a3-ultragpu-8g/nemo-framework/Dockerfile

@@ -1,10 +1,10 @@
-# Copyright 2024 "Google LLC"
+# Copyright 2025 "Google LLC"
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+# http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
@@ -13,25 +13,38 @@
 # limitations under the License.
 
 ARG NEMOFW_VERSION=24.12
+ARG NCCL_GIB_VERSION=v1.1.0
+
+# Stage 1: Use the nccl-plugin-gib container as an installer
+FROM us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib:${NCCL_GIB_VERSION} AS nccl-gib
+RUN /scripts/container_entry.sh install --install-nccl
+
+# Stage 2: Build the final NeMo container
 FROM nvcr.io/nvidia/nemo:$NEMOFW_VERSION
 
-ENV NCCL_DEBUG=INFO,WARN
-ENV NCCL_DEBUG_SUBSYS=NET,COLL,GPU,INIT,ENV
-ENV NCCL_SOCKET_IFNAME=enp0s19,enp192s20
-ENV GLOO_SOCKET_IFNAME=enp0s19,enp192s20
+# Copy the NCCL gIB plugin and other binaries from the installer stage.
+# This replaces the need to manually install and mount /var/lib/gib from the host.
+COPY --from=nccl-gib /var/lib/gib /usr/local/gib
+
+# Set the following to match what is in /usr/local/gib/scripts/set_nccl_env.sh
+ENV NCCL_NET=gIB
 ENV NCCL_CROSS_NIC=0
 ENV NCCL_NET_GDR_LEVEL=PIX
-#ENV PMIX_MCA_gds=^ds12
 ENV NCCL_P2P_NET_CHUNKSIZE=131072
 ENV NCCL_NVLS_CHUNKSIZE=524288
-ENV NCCL_IB_GID_INDEX=3
 ENV NCCL_IB_ADAPTIVE_ROUTING=1
 ENV NCCL_IB_QPS_PER_CONNECTION=4
 ENV NCCL_IB_TC=52
 ENV NCCL_IB_FIFO_TC=84
-ENV NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE=/usr/local/gib/configs/guest_config.txtpb
-ENV NCCL_TUNER_CONFIG_PATH=/usr/local/gib/configs/tuner_config.txtpb
-ENV NCCL_NET=gIB
+# Use A3 Ultra tuner config
+ENV NCCL_TUNER_CONFIG_PATH=/usr/local/gib/configs/tuner_config_a3u.txtpb
+
+
+# Add NCCL DEBUG at WARN level (useful for debugging)
+ENV NCCL_DEBUG=WARN
+ENV NCCL_DEBUG_SUBSYS=NET,COLL,GPU,INIT,ENV
+ENV GLOO_SOCKET_IFNAME=enp0s19,enp192s20
 
+# Set up ldconfig to load nccl-gib
 RUN echo "/usr/local/gib/lib64" >> /etc/ld.so.conf.d/gib.conf && ldconfig
 ENV LD_LIBRARY_PATH=/usr/local/gib/lib64:$LD_LIBRARY_PATH

examples/machine-learning/a3-ultragpu-8g/nemo-framework/README.md

@@ -43,7 +43,7 @@ README
 
    ```shell
    cd launcher_scripts
-   mkdir data
+   mkdir -p data
 
    MAX_STEPS=10
    NUM_NODES=8
@@ -53,8 +53,8 @@ README
        stages=[training] \
        training=gpt3/5b \
        env_vars.TRANSFORMERS_OFFLINE=0 \
-       container=../nemo-24.12.sqsh \
-       container_mounts=[${HOME}/.cache,/usr/local/gib] \
+       container=../nemo-24.12-v1.1.0.sqsh \
+       container_mounts=[${HOME}/.cache] \
        cluster.srun_args=["--container-writable"] \
        training.model.data.data_impl=mock \
        training.model.data.data_prefix=[] \

examples/machine-learning/a3-ultragpu-8g/nemo-framework/setup_nemo.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2024 "Google LLC"
+# Copyright 2025 "Google LLC"
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,12 +19,19 @@
 #SBATCH --exclusive
 
 : "${NEMOFW_VERSION:=24.12}"
+: "${NCCL_GIB_VERSION:=v1.1.0}"
 
-srun docker build --build-arg="NEMOFW_VERSION=${NEMOFW_VERSION}" -t nemo-"${NEMOFW_VERSION}" .
-srun rm -f nemo-"${NEMOFW_VERSION}".sqsh
-srun enroot import dockerd://nemo-"${NEMOFW_VERSION}"
+# This ensures that the docker process on the compute node can access us-docker.pkg.dev
+srun gcloud auth configure-docker us-docker.pkg.dev --quiet
+
+srun docker build \
+	--build-arg="NEMOFW_VERSION=${NEMOFW_VERSION}" \
+	--build-arg="NCCL_GIB_VERSION=${NCCL_GIB_VERSION}" \
+	-t nemo-"${NEMOFW_VERSION}"-"${NCCL_GIB_VERSION}" .
+srun rm -f nemo-"${NEMOFW_VERSION}"-"${NCCL_GIB_VERSION}".sqsh
+srun enroot import dockerd://nemo-"${NEMOFW_VERSION}"-"${NCCL_GIB_VERSION}"
 
 srun \
 	--container-mounts="${PWD}":/workspace/mount_dir,/var/tmp:/var/tmp \
-	--container-image=./nemo-"${NEMOFW_VERSION}".sqsh \
+	--container-image=./nemo-"${NEMOFW_VERSION}"-"${NCCL_GIB_VERSION}".sqsh \
 	bash -c "cp -r /opt/NeMo-Framework-Launcher/requirements.txt /opt/NeMo-Framework-Launcher/launcher_scripts /opt/NeMo-Framework-Launcher/auto_configurator /workspace/mount_dir/"