Ascend megatron script

benchmarks/compatibility/scripts/megatron/README.md

@@ -0,0 +1,31 @@
+### How to run Megatron-LM on Ascend platform?
+
+1. Install MindSpeed by following the [Ascend official guide](https://gitcode.com/Ascend/MindSpeed/blob/master/docs/user-guide/installation.md). Recommend to use Python 3.10.x and PyTorch 2.1.0, some code in Mindspeed may not be compatible with Python 3.9.x.
+
+2. With MindSpeed Core, you can run Megatron-LM on Ascend training devices by adding just a single line of code. 
+    ##### Example usage
+
+    Using the GPT model as an example, modify the `pretrain_gpt.py` file located in your Megatron-LM directory. Add the new line `import mindspeed.megatron_adaptor` right after import torch.
+
+    See the example modification below:
+    ```
+    import torch
+    import mindspeed.megatron_adaptor # new code added
+    from functools import partial
+    from contextlib import nullcontext
+    import inspect
+    ```
+    After making this change, you can launch your Megatron-LM training tasks on Ascend devices.
+    For detailed instructions, please refer to the [Quick Start Guide](https://gitcode.com/Ascend/MindSpeed/blob/master/docs/user-guide/getting_started.md).
+
+
+### How to test this repo Megatron-LM scripts on Ascend platform?
+1. With MindSpeed installed, init Ascend related env variable.
+    ```
+    source /usr/local/Ascend/nnal/atb/set_env.sh
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh
+    ```
+2. Go to $INFINIPERF_ROOT/benchmarks/compatibility/scripts/megatron folder and run test script by the following command
+    ```
+    bash run_megatron_test.sh
+    ```

benchmarks/compatibility/scripts/megatron/build.sh

@@ -5,17 +5,36 @@ if [ -z "$INFINIPERF_ROOT" ]; then
     exit 1
 fi
 
-export OSCAR_DATA_PATH=/workspace/bolunz/Data/oscar-en-10k.jsonl
-export LLAMA2_7B_MODEL_PATH=/workspace/bolunz/Models/Llama-2-7b-hf
+export OSCAR_DATA_PATH=/home/libaoming/workplace/oscar-en-10k.jsonl
+export LLAMA2_7B_MODEL_PATH=/home/libaoming/workplace/Llama-2-7b-hf
 
 export MEGATRON_ROOT=$INFINIPERF_ROOT/benchmarks/compatibility/Megatron-LM
 
+if [ "$INFINIPERF_PLATFORM" = "ASCEND_NPU" ]; then
+    echo "INFO: ASCEND_NPU platform detected. Change Magetron-LM to core_v0.12.1 branch..."
+    cd $MEGATRON_ROOT
+    # Clean change not staged
+    git reset --hard HEAD
+    git clean -fd
+    git checkout main
+    git pull --rebase
+    git checkout core_v0.12.1
+    cd -
+    echo "INFO: Magetron-LM repo changed to core_v0.12.1 branch successfully."
+
+    echo "INFO: ASCEND_NPU platform detected. Copying preprocess_data files..."
+    cp -v $INFINIPERF_ROOT/benchmarks/compatibility/scripts/megatron/preprocess_data.py $MEGATRON_ROOT/tools/
+    echo "INFO: preprocess_data file for ASCEND_NPU platform copied successfully."
+else
+    echo "INFO: ASCEND_NPU platform not detected (PLATFORM is '$INFINIPERF_PLATFORM'). Skipping file copy."
+fi
+
 cp $INFINIPERF_ROOT/benchmarks/compatibility/scripts/megatron/training.py $MEGATRON_ROOT/megatron/training
-cp $MEGATRON_ROOT/pretrain_gpt.py $MEGATRON_ROOT/pretrain_llama.py
+cp $INFINIPERF_ROOT/benchmarks/compatibility/scripts/megatron/pretrain_llama.py $MEGATRON_ROOT/pretrain_llama.py
 
 if ! python -c "import sentencepiece" &> /dev/null; then
     echo "sentencepiece 未安装，正在安装..."
     pip install sentencepiece
 else
     echo "sentencepiece 已安装。"
-fi
+fi