Update Dockerfile and README.md instructions for a3mega nemo framework (GoogleCloudPlatform#5164)

Author: mufaqam-gcl

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

@@ -20,10 +20,8 @@ ENV NCCL_FASTRAK_IFNAME=enp6s0,enp7s0,enp13s0,enp14s0,enp134s0,enp135s0,enp141s0
 ENV NCCL_SOCKET_IFNAME=enp0s12
 ENV GLOO_SOCKET_IFNAME=enp0s12
 ENV NCCL_CROSS_NIC=0
-ENV NCCL_ALGO=Ring,Tree
-ENV NCCL_PROTO=Simple
+ENV NCCL_PROTO=Simple,LL128
 ENV NCCL_MIN_NCHANNELS=4
-ENV NCCL_DYNAMIC_CHUNK_SIZE=524288
 ENV NCCL_P2P_NET_CHUNKSIZE=524288
 ENV NCCL_P2P_PCI_CHUNKSIZE=524288
 ENV NCCL_P2P_NVL_CHUNKSIZE=1048576
@@ -36,11 +34,12 @@ ENV NCCL_NET_GDR_LEVEL=PIX
 ENV NCCL_FASTRAK_ENABLE_HOTPATH_LOGGING=0
 ENV NCCL_FASTRAK_USE_LLCM=1
 ENV NCCL_FASTRAK_LLCM_DEVICE_DIRECTORY=/dev/aperture_devices
-ENV NCCL_TUNER_PLUGIN=libnccl-tuner.io
+ENV NCCL_TUNER_PLUGIN=libnccl-tuner.so
 ENV NCCL_TUNER_CONFIG_PATH=/var/lib/tcpxo/lib64/a3plus_tuner_config.textproto
 ENV NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE=/var/lib/tcpxo/lib64/a3plus_guest_config.textproto
 ENV NCCL_FASTRAK_PLUGIN_ACCEPT_TIMEOUT_MS=600000
-ENV NCCL_NVLS_ENABLE=0
+ENV TORCH_NCCL_AVOID_RECORD_STREAMS=0
+ENV NCCL_NVLSTREE_MAX_CHUNKSIZE=131072
 
 RUN echo "/var/lib/tcpxo/lib64" >> /etc/ld.so.conf.d/tcpxo.conf && ldconfig
 ENV LD_LIBRARY_PATH=/var/lib/tcpxo/lib64:$LD_LIBRARY_PATH

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

@@ -39,11 +39,15 @@ README
    ```
 
    This will run an example of training a 5B parameter GPT3 model for 10 steps
-   using mock data as the input.
+   using mock data as the input. Note: Please update `NUM_NODES` to
+   actual number of nodes in the deployment.
 
    ```shell
    cd launcher_scripts
-   mkdir data
+   mkdir -p data/bpe
+
+   wget https://huggingface.co/gpt2/resolve/main/vocab.json -O ./data/bpe/vocab.json
+   wget https://huggingface.co/gpt2/resolve/main/merges.txt -O ./data/bpe/merges.txt
 
    MAX_STEPS=10
    NUM_NODES=8
@@ -52,9 +56,8 @@ README
        launcher_scripts_path=${PWD} \
        stages=[training] \
        training=gpt3/5b \
-       env_vars.TRANSFORMERS_OFFLINE=0 \
        container=../nemofw+tcpxo-24.12.sqsh \
-       container_mounts=[${HOME}/.cache,/var/lib/tcpxo/lib64] \
+       container_mounts=[${HOME}/.cache,/var/lib/tcpxo/lib64,/sys/bus/pci] \
        cluster.srun_args=["--container-writable"] \
        training.model.data.data_impl=mock \
        training.model.data.data_prefix=[] \