[RFC] AMD integration for Megablocks

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "third_party/Megatron-LM"]
 	path = third_party/Megatron-LM
-	url = https://github.com/stanford-futuredata/Megatron-LM.git
+	url = https://github.com/ROCm/Stanford-Megatron-LM
+[submodule "third_party/Stanford-Megatron-LM"]
+	path = third_party/Stanford-Megatron-LM
+	url = https://github.com/ROCm/Stanford-Megatron-LM

Dockerfile

@@ -1,9 +1,18 @@
-FROM nvcr.io/nvidia/pytorch:23.09-py3
+FROM rocm/pytorch:rocm6.3_ubuntu24.04_py3.12_pytorch_release_2.4.0
+WORKDIR /root/
 
-RUN pip install stanford-stk==0.0.6
+# Install the application dependencies
+RUN pip install regex
+RUN pip install nltk
+RUN pip install pybind11
 
-RUN pip install flash-attn
+RUN git clone --recurse-submodules -j8 https://github.com/ROCm/megablocks && \
+    cd megablocks && \
+    ./patch_torch.sh && \
+    python setup.py install
 
-ENV PYTHONPATH="/mount/megablocks/third_party/Megatron-LM:${PYTHONPATH}"
+RUN cd megablocks/third_party/Stanford-Megatron-LM && \
+    git checkout rocm_6_3_patch && \
+    ./apply_patch.sh
 
-WORKDIR /mount/megablocks
+CMD ["/bin/bash"]

README.md

@@ -40,4 +40,4 @@ We provide scripts for pre-training Transformer MoE and dMoE language models und
   volume={5},
   year={2023}
 }
-```
+```

csrc/replicate.h

@@ -1,3 +1,5 @@
+// Modifications:  Copyright Advanced Micro Devices, Inc.  SPDX License:  MIT.
+
 #undef CUB_WRAPPED_NAMESPACE
 #define CUB_WRAPPED_NAMESPACE megablocks
 
@@ -8,6 +10,12 @@
 #include <c10/cuda/CUDAStream.h>
 #include <torch/extension.h>
 
+#ifndef USE_ROCM
+  #define _LDG(arg) __ldg(arg)
+#else
+  #define _LDG(arg) *(arg)
+#endif
+
 #define CUDA_CALL(code)					    \
   do {                                                      \
     cudaError_t status = code;                              \
@@ -37,11 +45,11 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
   // Load the start/end for this bin.
   int bin_idx = blockIdx.x;
   int start = 0;
-  if (bin_idx > 0) start = __ldg(bins + bin_idx - 1);
-  int end = __ldg(bins + bin_idx);
+  if (bin_idx > 0) start = _LDG(bins + bin_idx - 1);
+  int end = _LDG(bins + bin_idx);
 
   // Load the value to replicate.
-  T value = __ldg((T*)x + bin_idx);
+  T value = _LDG((T*)x + bin_idx);
 
   // Offset to this threadblocks bin and this threads
   // offset within the bin.

exp/dmoe/dmoe_125m_8gpu.sh

@@ -131,12 +131,17 @@ PILE_DATASET="\
 1.0 \
 /mount/pile_gpt2/29_text_document"
 
+VOCAB_FILE=<Enter location here>
+MERGE_FILE=<Enter location here>
+DATA_PATH=<Enter location here>
+CHECKPOINT_PATH=<Enter location here>
+
 # NOTE: We don't train for enough tokens for the
 # split to matter.
 DATA_ARGUMENTS="\
---data-path ${PILE_DATASET} \
---vocab-file /mount/gpt2-vocab.json \
---merge-file /mount/gpt2-merges.txt \
+--data-path ${DATA_PATH} \
+--vocab-file ${VOCAB_FILE} \
+--merge-file ${MERGE_FILE} \
 --make-vocab-size-divisible-by 1024 \
 --split 969,30,1"
 
@@ -145,7 +150,7 @@ COMPUTE_ARGUMENTS="\
 --DDP-impl local \
 --moe-expert-model-parallelism \
 --no-async-tensor-model-parallel-allreduce \
---use-flash-attn"
+--no-gradient-accumulation-fusion"
 
 CHECKPOINT_ARGUMENTS="\
 --save-interval 2000 \
@@ -157,7 +162,7 @@ EVALUATION_ARGUMENTS="\
 --eval-interval 1000"
 
 torchrun ${DISTRIBUTED_ARGUMENTS} \
-       third_party/Megatron-LM/pretrain_gpt.py \
+       third_party/Stanford-Megatron-LM/pretrain_gpt.py \
        ${MOE_ARGUMENTS} \
        ${MODEL_ARGUMENTS} \
        ${TRAINING_ARGUMENTS} \

exp/dmoe/dmoe_356m_8gpu.sh

@@ -131,12 +131,17 @@ PILE_DATASET="\
 1.0 \
 /mount/pile_gpt2/29_text_document"
 
+VOCAB_FILE=<Enter location here>
+MERGE_FILE=<Enter location here>
+DATA_PATH=<Enter location here>
+CHECKPOINT_PATH=<Enter location here>
+
 # NOTE: We don't train for enough tokens for the
 # split to matter.
 DATA_ARGUMENTS="\
---data-path ${PILE_DATASET} \
---vocab-file /mount/gpt2-vocab.json \
---merge-file /mount/gpt2-merges.txt \
+--data-path ${DATA_PATH} \
+--vocab-file ${VOCAB_FILE} \
+--merge-file ${MERGE_FILE} \
 --make-vocab-size-divisible-by 1024 \
 --split 969,30,1"
 
@@ -145,7 +150,7 @@ COMPUTE_ARGUMENTS="\
 --DDP-impl local \
 --moe-expert-model-parallelism \
 --no-async-tensor-model-parallel-allreduce \
---use-flash-attn"
+--no-gradient-accumulation-fusion"
 
 CHECKPOINT_ARGUMENTS="\
 --save-interval 2000 \
@@ -158,7 +163,7 @@ EVALUATION_ARGUMENTS="\
 
 
 torchrun ${DISTRIBUTED_ARGUMENTS} \
-       third_party/Megatron-LM/pretrain_gpt.py \
+       third_party/Stanford-Megatron-LM/pretrain_gpt.py \
        ${MOE_ARGUMENTS} \
        ${MODEL_ARGUMENTS} \
        ${TRAINING_ARGUMENTS} \

exp/dmoe/dmoe_46m_8gpu.sh

@@ -131,12 +131,17 @@ PILE_DATASET="\
 1.0 \
 /mount/pile_gpt2/29_text_document"
 
+VOCAB_FILE=<Enter location here>
+MERGE_FILE=<Enter location here>
+DATA_PATH=<Enter location here>
+CHECKPOINT_PATH=<Enter location here>
+
 # NOTE: We don't train for enough tokens for the
 # split to matter.
 DATA_ARGUMENTS="\
---data-path ${PILE_DATASET} \
---vocab-file /mount/gpt2-vocab.json \
---merge-file /mount/gpt2-merges.txt \
+--data-path ${DATA_PATH} \
+--vocab-file ${VOCAB_FILE} \
+--merge-file ${MERGE_FILE} \
 --make-vocab-size-divisible-by 1024 \
 --split 969,30,1"
 
@@ -145,7 +150,7 @@ COMPUTE_ARGUMENTS="\
 --DDP-impl local \
 --moe-expert-model-parallelism \
 --no-async-tensor-model-parallel-allreduce \
---use-flash-attn"
+--no-gradient-accumulation-fusion"
 
 CHECKPOINT_ARGUMENTS="\
 --save-interval 2000 \
@@ -157,7 +162,7 @@ EVALUATION_ARGUMENTS="\
 --eval-interval 1000"
 
 torchrun ${DISTRIBUTED_ARGUMENTS} \
-       third_party/Megatron-LM/pretrain_gpt.py \
+       third_party/Stanford-Megatron-LM/pretrain_gpt.py \
        ${MOE_ARGUMENTS} \
        ${MODEL_ARGUMENTS} \
        ${TRAINING_ARGUMENTS} \

exp/dmoe/dmoe_760m_8gpu.sh

@@ -131,12 +131,17 @@ PILE_DATASET="\
 1.0 \
 /mount/pile_gpt2/29_text_document"
 
+VOCAB_FILE=<Enter location here>
+MERGE_FILE=<Enter location here>
+DATA_PATH=<Enter location here>
+CHECKPOINT_PATH=<Enter location here>
+
 # NOTE: We don't train for enough tokens for the
 # split to matter.
 DATA_ARGUMENTS="\
---data-path ${PILE_DATASET} \
---vocab-file /mount/gpt2-vocab.json \
---merge-file /mount/gpt2-merges.txt \
+--data-path ${DATA_PATH} \
+--vocab-file ${VOCAB_FILE} \
+--merge-file ${MERGE_FILE} \
 --make-vocab-size-divisible-by 1024 \
 --split 969,30,1"
 
@@ -145,7 +150,7 @@ COMPUTE_ARGUMENTS="\
 --DDP-impl local \
 --moe-expert-model-parallelism \
 --no-async-tensor-model-parallel-allreduce \
---use-flash-attn"
+--no-gradient-accumulation-fusion"
 
 CHECKPOINT_ARGUMENTS="\
 --save-interval 2000 \
@@ -157,7 +162,7 @@ EVALUATION_ARGUMENTS="\
 --eval-interval 1000"
 
 torchrun ${DISTRIBUTED_ARGUMENTS} \
-       third_party/Megatron-LM/pretrain_gpt.py \
+       third_party/Stanford-Megatron-LM/pretrain_gpt.py \
        ${MOE_ARGUMENTS} \
        ${MODEL_ARGUMENTS} \
        ${TRAINING_ARGUMENTS} \

exp/gpt2/gpt2_125m_1gpu.sh

@@ -104,18 +104,25 @@ PILE_DATASET="\
 1.0 \
 ../pile_gpt2/29_text_document"
 
+VOCAB_FILE=<Enter location here>
+MERGE_FILE=<Enter location here>
+DATA_PATH=<Enter location here>
+CHECKPOINT_PATH=<Enter location here>
+
 # NOTE: We don't train for enough tokens for the
 # split to matter.
 DATA_ARGUMENTS="\
---data-path ${PILE_DATASET} \
---vocab-file ../gpt2-vocab.json \
---merge-file ../gpt2-merges.txt \
+--data-path ${DATA_PATH} \
+--vocab-file ${VOCAB_FILE} \
+--merge-file ${MERGE_FILE} \
 --make-vocab-size-divisible-by 1024 \
 --split 969,30,1"
 
 COMPUTE_ARGUMENTS="\
 --fp16 \
---DDP-impl local"
+--DDP-impl local \
+--no-async-tensor-model-parallel-allreduce \
+--no-gradient-accumulation-fusion"
 
 CHECKPOINT_ARGUMENTS="\
 --save-interval 2000 \
@@ -127,7 +134,7 @@ EVALUATION_ARGUMENTS="\
 --eval-interval 1000"
 
 torchrun ${DISTRIBUTED_ARGUMENTS} \
-       third_party/Megatron-LM/pretrain_gpt.py \
+       third_party/Stanford-Megatron-LM/pretrain_gpt.py \
        ${MODEL_ARGUMENTS} \
        ${TRAINING_ARGUMENTS} \
        ${DATA_ARGUMENTS} \

exp/gpt2/gpt2_125m_8gpu.sh

@@ -1,4 +1,5 @@
 #!/bin/bash
+# Modifications:  Copyright Advanced Micro Devices, Inc.  SPDX License:  MIT.
 
 EXP_DIR=$1
 
@@ -106,35 +107,40 @@ PILE_DATASET="\
 1.0 \
 /mount/pile_gpt2/29_text_document"
 
+VOCAB_FILE=<Enter location here>
+MERGE_FILE=<Enter location here>
+DATA_PATH=<Enter location here>
+CHECKPOINT_PATH=<Enter location here>
+
 # NOTE: We don't train for enough tokens for the
 # split to matter.
 DATA_ARGUMENTS="\
---data-path ${PILE_DATASET} \
---vocab-file /mount/gpt2-vocab.json \
---merge-file /mount/gpt2-merges.txt \
+--data-path ${DATA_PATH} \
+--vocab-file ${VOCAB_FILE} \
+--merge-file ${MERGE_FILE} \
 --make-vocab-size-divisible-by 1024 \
 --split 969,30,1"
 
 COMPUTE_ARGUMENTS="\
 --bf16 \
 --DDP-impl local \
 --no-async-tensor-model-parallel-allreduce \
---use-flash-attn"
+--no-gradient-accumulation-fusion"
 
 CHECKPOINT_ARGUMENTS="\
 --save-interval 2000 \
---save ./${EXP_DIR}"
+--save ${CHECKPOINT_PATH}"
 
 EVALUATION_ARGUMENTS="\
 --eval-iters 100 \
 --log-interval 100 \
 --eval-interval 1000"
 
 torchrun ${DISTRIBUTED_ARGUMENTS} \
-       third_party/Megatron-LM/pretrain_gpt.py \
+       third_party/Stanford-Megatron-LM/pretrain_gpt.py \
        ${MODEL_ARGUMENTS} \
        ${TRAINING_ARGUMENTS} \
        ${DATA_ARGUMENTS} \
        ${COMPUTE_ARGUMENTS} \
        ${CHECKPOINT_ARGUMENTS} \
-       ${EVALUATION_ARGUMENTS} |& tee ./${EXP_DIR}/train.log
+       ${EVALUATION_ARGUMENTS} |& tee ${CHECKPOINT_PATH}/train.log

exp/gpt2/gpt2_1315m_1gpu.sh

@@ -103,18 +103,25 @@ PILE_DATASET="\
 1.0 \
 ../pile_gpt2/29_text_document"
 
+VOCAB_FILE=<Enter location here>
+MERGE_FILE=<Enter location here>
+DATA_PATH=<Enter location here>
+CHECKPOINT_PATH=<Enter location here>
+
 # NOTE: We don't train for enough tokens for the
 # split to matter.
 DATA_ARGUMENTS="\
---data-path ${PILE_DATASET} \
---vocab-file ../gpt2-vocab.json \
---merge-file ../gpt2-merges.txt \
+--data-path ${DATA_PATH} \
+--vocab-file ${VOCAB_FILE} \
+--merge-file ${MERGE_FILE} \
 --make-vocab-size-divisible-by 1024 \
 --split 969,30,1"
 
 COMPUTE_ARGUMENTS="\
 --fp16 \
---DDP-impl local"
+--DDP-impl local \
+--no-async-tensor-model-parallel-allreduce \
+--no-gradient-accumulation-fusion"
 
 CHECKPOINT_ARGUMENTS="\
 --save-interval 2000 \
@@ -126,7 +133,7 @@ EVALUATION_ARGUMENTS="\
 --eval-interval 1000"
 
 torchrun ${DISTRIBUTED_ARGUMENTS} \
-       third_party/Megatron-LM/pretrain_gpt.py \
+       third_party/Stanford-Megatron-LM/pretrain_gpt.py \
        ${MODEL_ARGUMENTS} \
        ${TRAINING_ARGUMENTS} \
        ${DATA_ARGUMENTS} \