# Megatron-LM

**Repository Path**: wangzw1022/Megatron-LM

## Basic Information

- **Project Name**: Megatron-LM
- **Description**: Megatron-LM NPU插件仓
- **Primary Language**: Python
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 14
- **Created**: 2023-07-14
- **Last Updated**: 2024-06-11

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Megatron-NPU

## 1 简介与特性介绍

Megatron-LM 是由 NVIDIA 的应用深度学习研究团队开发的一款功能强大的大型Transformer仓。而本仓库megatron_npu为昇腾基于[Megatron-LM原始仓](https://github.com/NVIDIA/Megatron-LM)所开发的适配仓，已适配特性如下：

- 数据并行（Data parallel）
- 模型并行（Tensor parallel）
- 序列并行（Sequence parallel）
- 流水并行（Pipeline parallel）
- 分布式优化器（Distributed optimizer）

## 2 环境准备
> 基于安全性考虑，建议您以非root的安全账户执行脚本。
### 2.1 Pytorch框架训练环境准备
请参考昇腾官方文档《[Pytorch框架训练环境准备](https://www.hiascend.com/document/detail/zh/ModelZoo/pytorchframework/ptes)》。建议您在准备好模型训练环境以后，将umask调整为`027`或以上。

### 2.2 克隆原始仓
  ```shell
  git clone https://github.com/NVIDIA/Megatron-LM.git
  cd Megatron-LM
  git checkout 285068c8108e0e8e6538f54fe27c3ee86c5217a2
  ```
### 2.3 安装 megatron_npu

在Megatron-LM目录下执行：

  ```shell
  git clone https://gitee.com/ascend/Megatron-LM.git megatron_npu_adaptor
  cd megatron_npu_adaptor
  pip install -e .
  ```
> 如需要保存安装megatron_npu的日志，可在pip install命令后面添加参数 `--log <PATH>`，并对您指定的路径`<PATH>`做好权限管控。
> 如需要代码二次开发，可在pip install命令后面加上参数 `-e`，并对目录做好权限管控。
### 2.4 安装昇腾发布的相关版本torch与apex包及其他依赖
  ```shell
  pip install torch-*.whl  
  pip install torch_npu-*.whl
  pip install apex-*.whl

  pip install -r requirements.txt
  ```
### 2.5 使用megatron_npu
在执行代码的最开始处`import megatron_npu`，这样脚本执行时原生的megatron才会被megatron_npu patch。

## 3 数据集

### 3.1 获取数据集。
我们以enwiki数据集为例，下载地址为Megatron官方提供，您可以根据自身需求选择数据集训练。

对于步骤1，我们以`./dataset`作为存储数据集的目录，并默认赋予`750`的权限控制，您可以自行选择存放目录，并参考[附录A 文件（夹）各场景权限管控推荐最大值](#A-文件（夹）各场景权限管控推荐最大值)进行权限设置。

```shell
# 步骤1: 在megatron_npu_adaptor目录下建立数据集目录dataset并进入，您可以选择
mkdir dataset && chmod 750 dataset && cd dataset

# 步骤2: 下载 enwiki 数据集
wget https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2
bzip2 -dk enwiki-latest-pages-articles.xml.bz2

# 步骤3: 下载bert-large Vocab表
wget https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt

# 步骤4: 安装 WikiExtractor 对xml格式的数据集进行初步处理
pip3 install wikiextractor
python3 -m wikiextractor.WikiExtractor enwiki-latest-pages-articles.xml --json
```

### 3.2 处理数据集
在megatron_npu_adaptor根目录执行脚本，对数据集和VOCAB词表进行处理。
   ```shell
   # 退出dataset目录，回到megatron_npu_adaptor根目录
   cd ..
   # 执行数据预处理脚本
   bash ./tests/dataset_preprocess_t5.sh
   ```
### 3.3 数据集目录结构
经过上述步骤处理后，数据集默认放置在megatron_npu_adaptor根目录的`./dataset/en_wiki`下，其目录结构与内容如下所示：

   ```
   ├── ./dataset/en_wiki             
         ├── my-t5_text_sentence.bin
         ├── my-t5_text_sentence.idx
   ```

## 4 训练
以下脚本执行模型的训练与评估，均通过调用torch.distributed.launch函数，函数参数的说明，请您参考[Pytorch官方文档](https://pytorch.org/docs/stable/distributed.html)。
脚本传入launch函数的入参仅作为参考，您可以根据需求调整参数，并建议您参考[附录A 文件（夹）各场景权限管控推荐最大值](#A-文件（夹）各场景权限管控推荐最大值)，对生成文件（如权重文件checkpoint）进行安全加固。
### 4.1 执行如下前置命令
   ```shell
   cd ./tests_gpt/
   mv pretrain_gpt.py ../../
   ```
### 4.2 运行训练脚本
以单机8卡训练为例：
```shell
bash pretrain_gpt_distributed.sh
```
训练完成后，权重文件默认保存在./checkpoint下，并输出模型训练精度和性能信息。
### 4.3 使能bf16、fp16的大kernel

#### 4.3.1 命令行参数说明

pretrain_gpt_distributed_fp16.sh脚本实际调用并透传以下三个参数到torch.distributed.launch函数调用，具体如下：

|   参数名称    |                           参数说明                           |
| :-----------: | :----------------------------------------------------------: |
|     --pre     | 透传到launch函数的--pre-tockens参数，用于控制flash attention三角的大小 |
|    --next     | 透传到launch函数的--next-tockens参数，用于控制flash attention三角的大小 |
| --shape_order |    透传到launch函数的--shape_order参数，用于输入排布格式     |

#### 4.3.2 执行训练

以单机8卡训练为例：

```shell
bash pretrain_gpt_distributed_fp16.sh --pre=65536 --next=65536 --shape_order=SBH #FP16 flash-attn SBH输入
bash pretrain_gpt_distributed_fp16.sh --pre=65536 --next=65536 --shape_order=BSH #FP16 flash-attn BSH输入
bash pretrain_gpt_distributed_fp16.sh --pre=65536 --next=0 --shape_order=SBH #FP16 sparse-attn SBH输入
bash pretrain_gpt_distributed_fp16.sh --pre=65536 --next=0 --shape_order=BSH #FP16 sparse-attn BSH输入
bash pretrain_gpt_distributed_bf16.sh --pre=65536 --next=65536 --shape_order=SBH #BF16 flash-attn SBH输入
bash pretrain_gpt_distributed_bf16.sh --pre=65536 --next=65536 --shape_order=BSH #BF16 flash-attn BSH输入
bash pretrain_gpt_distributed_bf16.sh --pre=65536 --next=0 --shape_order=SBH #BF16 sparse-attn SBH输入
bash pretrain_gpt_distributed_bf16.sh --pre=65536 --next=0 --shape_order=BSH #BF16 sparse-attn BSH输入
bash pretrain_x1_distributed_2_6B_32P.sh --pre=2048 --next=0 --shape_order=SBH #x1场景
bash pretrain_x1_distributed_13B_32P.sh --pre=2048 --next=0 --shape_order=SBH #x1场景
bash pretrain_x1_distributed_65B_64P.sh --pre=2048 --next=0 --shape_order=SBH #x1场景
bash pretrain_x1_distributed_175B_96P.sh --pre=2048 --next=0 --shape_order=SBH #x1场景
```
训练完成后，权重文件会默认保存在./checkpoint下，并输出模型训练精度和性能信息。

## 5 评估
### LAMBADA Cloze Accuracy

1. 在test_gpt下，执行如下前置命令
   ```shell
   mv main.py ../../
   mv ../../tasks/zeroshot_gpt ../../
   ```

2. 获取测评数据集
   进入链接：https://github.com/cybertronai/bflm/blob/master/lambada_test.jsonl ，下载lambada_test.jsonl，重命名为lambada_test.json，放置在test_gpt下

3. 运行原仓zeroshot评估样例（可跳过）
   进入链接：https://catalog.ngc.nvidia.com/orgs/nvidia/models/megatron_lm_345m/files ，选择File Brower，下载release/mp_rank_00/model_optim_rng.pt，放置在`checkpoint_dist_test/mp_rank_00/`，然后执行样例脚本：

   ```shell
   bash test_gpt_distributed_sample.sh
   ```

4. 非大kernel zeroshot评估脚本

   以单机8卡训练为例：
   ```shell
   bash test_gpt_distributed.sh
   ```

5. bf16、fp16的大kernel zeroshot评估脚本

   以单机8卡训练为例，命令行参数说明可参考 4.3.1 命令行参数说明。

   ```shell
   bash test_gpt_distributed_fp16.sh --pre=65536 --next=65536 --shape_order=SBH #FP16 flash-attn SBH输入
   bash test_gpt_distributed_fp16.sh --pre=65536 --next=65536 --shape_order=BSH #FP16 flash-attn BSH输入
   bash test_gpt_distributed_fp16.sh --pre=2048 --next=0 --shape_order=SBH #FP16 sparse-attn SBH输入
   bash test_gpt_distributed_fp16.sh --pre=2048 --next=0 --shape_order=BSH #FP16 sparse-attn BSH输入
   bash test_gpt_distributed_bf16.sh --pre=65536 --next=65536 --shape_order=SBH #BF16 flash-attn SBH输入
   bash test_gpt_distributed_bf16.sh --pre=65536 --next=65536 --shape_order=BSH #BF16 flash-attn BSH输入
   bash test_gpt_distributed_bf16.sh --pre=2048 --next=0 --shape_order=SBH #BF16 sparse-attn SBH输入
   bash test_gpt_distributed_bf16.sh --pre=2048 --next=0 --shape_order=BSH #BF16 sparse-attn BSH输入
   ```
## 6 关于卸载

1. Pytorch框架训练环境的卸载可以参考[昇腾官方文档](https://www.hiascend.com/document/detail/zh/ModelZoo/pytorchframework/ptes/ptes_00032.html)。

2. megatron_npu的卸载只需执行命令：

  > 如需要保存卸载megatron_npu的日志，可在pip uninstall命令后面添加参数 `--log <PATH>`，并对您指定的路径`<PATH>`做好权限管控。
  
  ```python
  pip uninstall megatron_npu
  ```
  
# 安全声明

[Megatron-NPU 安全声明](SECURITYNOTE.md)