对接Megatron框架

前提条件

下载Megatron框架和准备数据集请参见Megatron-LM工程网站，完成“4.1 执行如下前置命令”步骤。
进入Megatron-LM目录下。

Megatron框架的下载准备以及数据集的下载处理由使用者完成。
Megatron支持的版本为CommitID：285068c8108e0e8e6538f54fe27c3ee86c5217a2。

操作步骤

Megatron框架和MindIO TTP SDK安装完成之后，需要修改代码才可在框架中启动TTP模块，在训练中同步优化器数据更新状态到TTP模块，故障时TTP能够及时生成临终CheckPoint。

修改pretrain_gpt.py。
1. 打开“pretrain_gpt.py”文件。
```
vim pretrain_gpt.py
```
2. 按“i”进入编辑模式，新增以下内容，即导入megatron_npu模块。
```
import megatron_npu
```
3. 按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

修改training.py。

打开“megatron/training.py”文件。
```
vim megatron/training.py
```

按“i”进入编辑模式，修改以下内容。

新增以下内容，即导入mindio_ttp模块。

import mindio_ttp.framework_ttp as ttp
import mindio_ttp.megatron_ttp as mega

并修改添加以下代码。

# 在setup_model_and_optimizer方法中做修改：
args.iteration = mega.load_checkpoint(model, optimizer, opt_param_scheduler)

# 在train_step方法中的optimizer.reduce_model_grads(args, timers)后添加：
# all rank need to barrier
torch.distributed.barrier()

# 在timers('optimizer', log_level=1).start(barrier=args.barrier_with_L1_time)后添加：
ttp.ttp_start_updating_os(-1)

# 在update_successful, grad_norm, num_zeros_in_grad = optimizer.step(args, timers)后添加：
if hasattr(optimizer, 'increase_step'):
    optimizer.increase_step()
ttp.ttp_end_updating_os(optimizer.current_step)

# 在train方法前一行添加：
@ttp.ttp_to_persist

# 在total_loss_dict = {}后添加：
# begin to init ttp processor & controller
cur_rank = torch.distributed.get_rank()
world_size: int = torch.distributed.get_world_size()
import os
masterIp = os.getenv('MASTER_ADDR', '127.0.0.1')
port = 8000
dp_ranks = mpu.get_data_parallel_all_rank()
if cur_rank == 0:
    ttp.ttp_init_controller(cur_rank, world_size, len(dp_ranks)//2, False)
    ttp.ttp_start_controller(masterIp, port)
ttp.ttp_init_processor(cur_rank, dp_ranks, len(dp_ranks), world_size, len(dp_ranks)//2, False)
ttp.ttp_start_processor(masterIp, port)
ttp.ttp_register_save_ckpt_handler(mega.ttp_save_checkpoint)
ttp.ttp_register_rename_handler(mega.rename_save_result)

# 在iteration += 1后添加：
ttp.ttp_set_ckpt_args((iteration, model, optimizer, opt_param_scheduler))

按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

__init__.py。
1. 打开“megatron/optimizer/__init__.py”文件。
```
vim megatron/optimizer/__init__.py
```
2. 按“i”进入编辑模式，修改以下内容。
  1. 新增以下内容，即导入ReplicatedOptimizer。
```
from mindio_ttp.megatron_ttp import ReplicatedOptimizer
```
  2. 修改“opt_ty”参数值。
```
opt_ty = ReplicatedOptimizer \
```
3. 按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

修改parallel_state.py。

打开“megatron/core/parallel_state.py”文件。
```
vim megatron/core/parallel_state.py
```

按“i”进入编辑模式，修改代码。

# 在_GLOBAL_MEMORY_BUFFER = None之后添加：
all_data_parallel_group_ranks = None

# 在120行左右global _DATA_PARALLEL_GLOBAL_RANKS后添加
global all_data_parallel_group_ranks

# 在211行左右model_parallel_is_initialized方法前添加：
def get_data_parallel_all_rank():
    assert _DATA_PARALLEL_GLOBAL_RANKS is not None, \
        "Data parallel group is not initialized"
    return _DATA_PARALLEL_GLOBAL_RANKS

def get_all_data_parallel_ranks():
    return all_data_parallel_group_ranks

按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

修改pretrain_gpt_distributed.sh。
1. 拷贝“pretrain_gpt_distributed.sh”到Megatron-LM目录下。
```
cd megatron_npu_adaptor/tests_gpt/
cp pretrain_gpt_distributed.sh ../../
```
2. 进入Megatron-LM目录，打开拷贝后的pretrain_gpt_distributed.sh文件。
```
cd ../../
vim ./pretrain_gpt_distributed.sh
```
3. 按“i”进入编辑模式，修改脚本。
  1. 添加如下命令。
```
source /实际路径/megatron_npu adaptor/tests_gpt/env_npu.sh
```
  2. 导出主节点IP地址变量。
```
export MASTER_ADDR=xxx.xxx.xxx.xxx
# 修改启动方式
ttprun $DISTRIBUTED_ARGS
```
    卸载MindIO.TTP后，启动方式需要还原为torchrun。
  3. 导出变量。
```
# 集群故障后使用TTP保存数据等待框架退出的时间，不做此动作默认120s，视情况修改。
export TTP_OT=360 
```
  4. 修改NNODES为实际节点数量，各个节点NODE_RANK编号从主节点=0开始依次增加，其余CheckPoint保存路径根据实际情况修改。
4. 按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

父主题： 使用指导