included llama3.1 8b small llm training scripts

small_llm_pretraining/nemo/Dockerfile.h200

@@ -0,0 +1,134 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:25.01-py3
+FROM ${FROM_IMAGE_NAME}
+
+# Document build setup
+ARG FROM_IMAGE_NAME
+ENV CUSTOM_FROM_IMAGE_NAME ${FROM_IMAGE_NAME}
+
+# Custom libraries version
+WORKDIR /workspace/
+
+ARG GIT_COMMIT_ID
+ENV GIT_COMMIT_ID=$GIT_COMMIT_ID
+
+RUN git config --global user.name "a" && \
+    git config --global user.email "a"
+
+WORKDIR /workspace/
+
+RUN pip install numcodecs==0.13.1
+
+## 1. Apex
+ARG APEX_REVISION=SKIP
+ENV CUSTOM_APEX_REVISION ${APEX_REVISION}
+ARG APEX_MAX_JOBS=4
+
+RUN if [ "${APEX_REVISION}" != SKIP ]; then \
+    git clone https://github.com/NVIDIA/apex && \
+    cd apex && \
+    echo APEX_REVISION=${APEX_REVISION} && \
+    git checkout ${APEX_REVISION} && \
+    echo APEX_COMMIT_HASH=$(git rev-parse HEAD) && \
+    MAX_JOBS=${APEX_MAX_JOBS} NVCC_APPEND_FLAGS="--threads 8" pip install -v --no-build-isolation --no-cache-dir --disable-pip-version-check --config-settings "--build-option=--cpp_ext --cuda_ext --bnp --xentropy --deprecated_fused_adam --deprecated_fused_lamb --fast_multihead_attn --distributed_lamb --fast_layer_norm --transducer --distributed_adam --fmha --fast_bottleneck --nccl_p2p --peer_memory --permutation_search --focal_loss --fused_conv_bias_relu --index_mul_2d --cudnn_gbn --group_norm" . \
+    ; fi
+
+
+
+## 2. Transformer Engine
+ARG TE_REVISION=SKIP
+ENV CUSTOM_TE_REVISION ${TE_REVISION}
+
+RUN if [ "${TE_REVISION}" != SKIP ]; then \
+    pip uninstall -y transformer-engine && \
+    git clone https://github.com/NVIDIA/TransformerEngine.git transformerengine && \
+    cd transformerengine && \
+    git checkout ${TE_REVISION} && \
+    echo TE_COMMIT_HASH=$(git rev-parse HEAD) && \
+    echo $(git rev-parse HEAD) > /TE_COMMIT_HASH.env && \
+    git submodule init && git submodule update && \
+    NVTE_CUDA_ARCHS="90;100" NVTE_UB_WITH_MPI=1 NVTE_FRAMEWORK=pytorch MPI_HOME=/usr/local/mpi pip install --force-reinstall --no-deps . \
+    ; fi
+
+
+## 3. NeMo
+ARG NEMO_REVISION=v2.1.0
+ENV CUSTOM_NEMO_REVISION ${NEMO_REVISION}
+
+# Clone and checkout NeMo at specified version
+RUN git clone https://github.com/NVIDIA/NeMo.git && \
+    cd NeMo && \
+    git checkout ${NEMO_REVISION} && \
+    echo NEMO_COMMIT_HASH=$(git rev-parse HEAD) && \
+    echo $(git rev-parse HEAD) > /NEMO_COMMIT_HASH.env && \
+    pip uninstall -y nemo-toolkit sacrebleu && \
+    # Only keep edits that are necessary (remove AMD-specific workarounds if upstream doesn't need them)
+    sed -i "/mamba-ssm/d" requirements/requirements_nlp.txt && \
+    sed -i 's/tensorstore<0.1.46/tensorstore/g' requirements/requirements_nlp.txt && \
+    sed -i 's/protobuf==3.20.3/protobuf/g' requirements/requirements.txt && \
+    pip install "cython<3.0.0" && \
+    pip install -e ".[llm]" && \
+    pip install -e ".[nlp]"
+
+
+## 3.1 NeMo-Run
+ARG NEMORUN_REVISION=v0.4.0
+ENV CUSTOM_NEMORUN_REVISION ${NEMORUN_REVISION}
+
+RUN git clone https://github.com/NVIDIA/NeMo-Run.git && \
+    cd NeMo-Run && \
+    git checkout ${NEMORUN_REVISION} && \
+    echo NEMORUN_COMMIT_HASH=$(git rev-parse HEAD) && \
+    pip install -e .
+
+
+# Python deps
+# Important this should be done after NeMo, otherwise the pinned transformers==4.40.2 version will be overwritten
+COPY requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt
+
+
+# 4. Megatron-core
+ARG MCORE_REVISION=core_r0.11.0
+ARG MCORE_REPO=https://github.com/NVIDIA/Megatron-LM.git
+ENV CUSTOM_MCORE_REVISION ${MCORE_REVISION}
+
+RUN if [ "${MCORE_REVISION}" != SKIP ]; then \
+    pip uninstall -y megatron-core && \
+    git clone ${MCORE_REPO} Megatron-LM && \
+    cd Megatron-LM && \
+    git checkout ${MCORE_REVISION} && \
+    echo MCORE_COMMIT_HASH=$(git rev-parse HEAD) && \
+    echo $(git rev-parse HEAD) > /MCORE_COMMIT_HASH.env && \
+    pip install . && \
+    cd megatron/core/datasets && \
+    make \
+    ; fi
+
+ENV PYTHONPATH "${PYTHONPATH}:/workspace/Megatron-LM"
+
+
+WORKDIR /workspace/code
+
+# Copy the current state of the code inside the image
+COPY . .

small_llm_pretraining/nemo/Dockerfile.mi325

@@ -0,0 +1,93 @@
+FROM rocm/pytorch:rocm6.4_ubuntu22.04_py3.10_pytorch_release_2.6.0
+
+WORKDIR /workspace
+
+RUN pip install pybind11
+RUN pip install ninja
+RUN pip install packaging
+RUN /usr/bin/python3 -m pip install pyYAML
+
+# Install library dependencies
+WORKDIR /workspace/deps
+
+# FlashAttention
+RUN git clone https://github.com/ROCm/flash-attention/ flash_attention \
+    # latest stable commit of ck_tile/fa3 branch
+    && cd flash_attention && git checkout cace3592812640486b04196a209bb85d12267b4c \
+    && git submodule update --init --recursive \
+    && PYTORCH_ROCM_ARCH='gfx942' GPU_ARCHS="gfx942" MAX_JOBS=64 pip install --no-build-isolation -e .
+
+ADD patches /workspace/deps/patches
+
+# Megatron-core
+RUN git clone --recursive https://github.com/ROCm/Megatron-LM.git megatron_lm
+RUN pip uninstall -y megatron-core
+# dev branch commit
+RUN cd megatron_lm && git checkout megatron_190213a_mlperf   \
+    && pip install -e . && cd megatron/core/datasets && make
+
+ENV PYTHONPATH "${PYTHONPATH}:/workspace/deps/megatron_lm"
+
+# mambe dependency required for NeMo
+RUN git clone https://github.com/state-spaces/mamba.git mamba_ssm \
+    && cd mamba_ssm \
+    && git checkout v2.2.2 \
+    && export HIP_ARCHITECTURES="gfx942" \
+    && pip install --no-cache-dir --verbose .
+
+# NeMo
+RUN git clone https://github.com/NVIDIA/NeMo nemo \
+    && cd nemo && git checkout v2.1.0
+RUN cd /workspace/deps/nemo \
+    && git apply /workspace/deps/patches/nemo_v2_1_0.patch \
+    && pip install --no-build-isolation -e ".[nlp]"
+
+# NeMo-Run
+RUN pip install git+https://github.com/NVIDIA/NeMo-Run.git@v0.4.0
+
+# Python deps
+# Important this should be done after NeMo, otherwise the pinned transformers==4.40.2 version will be overwritten
+COPY requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt
+
+# Transformer Engine
+ARG TE_COMMIT=te_v1.9_mlperf_llama2
+RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git \
+    # dev branch commit
+    && cd TransformerEngine && git checkout $TE_COMMIT && git submodule update --init --recursive \
+    # Workaround logging debug info to the console
+    && sed -i 's/self.logger.info/self.logger.debug/g' /workspace/deps/TransformerEngine/transformer_engine/pytorch/attention.py \
+    && sed -i 's/warnings.warn/if False: warnings.warn/g' /workspace/deps/TransformerEngine/transformer_engine/pytorch/attention.py \
+    && sed -i '/.*\"window_size should be.*/d' /workspace/deps/TransformerEngine/transformer_engine/common/fused_attn_rocm/fused_attn.cpp \
+    && NVTE_FUSED_ATTN_AOTRITON=0 NVTE_ROCM_ARCH='gfx942' NVTE_FRAMEWORK='pytorch' NVTE_USE_HIPBLASLT=1 MAX_JOBS=128 PYTORCH_ROCM_ARCH='gfx942' GPU_ARCHS='gfx942' pip install  -e .
+
+# Install hipBLASLt (FP8 tuned gemms - second round)
+ RUN git clone https://github.com/ROCm/hipBLASLt.git \
+     && cd hipBLASLt && git checkout ebc770851dfb99a1bbb8ef2e5873c753f8011a47 \
+     && sudo apt-get update \
+     && apt install -y python3.10-venv \
+     && ./install.sh -idc -a gfx942
+
+# RPD
+RUN sudo apt-get update && \
+    apt --fix-broken install -y && \
+    apt-get install -y\
+    sqlite3 libsqlite3-dev \
+    libfmt-dev
+
+RUN git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData \
+    && cd rocmProfileData \
+    && cd rocpd_python \
+    && python3 setup.py bdist_wheel \
+    && pip install dist/*.whl \
+    && cd .. \
+    && cd rpd_tracer \
+    && python3 setup.py bdist_wheel \
+    && pip3 install dist/*.whl \
+    && cd .. \
+    && make; make install
+
+WORKDIR /workspace/code
+
+# Copy the current state of the code inside the image
+COPY . .