openmind-accelerate

简介

accelerate

accelerate 是一个用于加速深度学习训练的工具库。它提供了一些优化技术和工具，能够提高训练速度、性能和效率。支持混合精度训练、模型并行和数据并行、分布式训练等功能。

openmind_accelerate

openmind-accelerate 为 accelerate 的插件仓，通过复用 accelerate 的 plugin 机制，增加其对于 nvidia 官方 megatron 框架的支持。

其打包的 whl 包名为 openmind_accelerate。

如何使用openmind_accelerate

import openmind_accelerate

用户只需简单的导入 openmind_accelerate，即可自动获取通过 accelerate 调用 nvidia 官方 megatron 的能力，并且不改变原生 accelerate 的使用方式。

同时，openmind_accelerate 还会自动根据环境信息，导入适配 npu 的相关能力。

而这一切，都不需要用户额外的关心。

环境准备

python3.8版本及以上

1. 安装依赖

请安装最新昇腾软件栈：https://www.hiascend.com/zh/

依赖软件
Driver
Firmware
CANN
Kernel
PyTorch
torch_npu
apex
MindSpeed-1.0

2. 获取 Megatron-LM

如有旧版本 Megatron-LM 请先卸载，再进行安装操作。

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout bcce6f54e075e3c3374ea67adefe54f3f2da2b07
pip install --no-use-pep517 -e .

3. 安装 openmind

git clone https://gitee.com/openmind-ai/openmind.git
cd openmind
pip install -e .

4. 安装 openmind_accelerate

git clone https://gitee.com/openmind-ai/openmind-accelerate.git
cd openmind-accelerate

#aarch64平台
pip install -e .

#x86平台
pip install -e . --extra-index-url https://download.pytorch.org/whl/cpu 

快速开始预训练

准备数据

用户可以准备好自己的预训练数据，例如falcon，slimpajama，alpaca数据集。

如果用户需要预训练 Megatron 模型，则还要参考Megatron的数据处理方法进行处理。

准备好数据集后在llama2-megatron.yaml配置文件中设置data_path，传入数据集路径。

准备模型

用户可以准备好模型文件，例如llama2模型。

如果用户需要预训练 Megatron 模型，则只需要准备 config.json 和 tokenizer 相关文件即可。

准备好模型后在llama2-megatron.yaml配置文件中设置openmind_model_path和tokenizer_model，传入模型和分词器路径。

启动

## 默认使用以下配置文件 ##
# examples/train_with_megatron.py
# examples/accelerate_config/accelerate_megatron_config.yaml
# examples/llama2_config/llama2-megatron.yaml

cd openmind-accelerate
bash examples/train_launch.sh

进阶使用

使用Megatron框架预训练模型

场景1：使用megatron格式数据

1. 修改 examples/llama2_config/llama2-megatron.yaml 配置文件，设置项：

   data_path: 'data/falcon-slimpajama-merged-dataset/llama2-mt_text_document'
   save_dir: 'model/llama-2-7b-hf_save'
   save_interval: 100
   openmind_model_path: 'model/llama-2-7b-hf'
   plugin_args:
     tp_degree: 4
     other_megatron_args:
       tokenizer_model: 'model/llama-2-7b-hf/tokenizer.model'
   

2. 修改 examples/llama2/train_launch.sh 启动脚本，设置项：

   export CUDA_VISIBLE_DEVICES_=0,1,2,3
   source /usr/loacal/Ascend/ascend-toolkit/set_env.sh
   
   # 设置对应配置文件及脚本
   SCRIPT_PATH=examples/train_with_megatron.py
   ACCELERATE_CONF_PATH=examples/accelerate_config/accelerate_megatron_config.yaml
   PRETRAIN_CONF_PATH=examples/llama2_config/llama2-megatron.yaml
   

3. 在 openmind-accelerate 根目录下执行 examples/llama2/train_launch.sh 启动脚本。

场景2：使用json格式数据

基于场景1修改对应配置文件:

SCRIPT_PATH=examples/train_with_megatron_json_dataset.py
ACCELERATE_CONF_PATH=examples/accelerate_config/accelerate_megatron_config.yaml
PRETRAIN_CONF_PATH=examples/llama2_config/llama2-megatron-json-dataset.yaml  

在 examples/llama2_config/llama2-megatron-json-dataset.yaml 配置文件中配置 dataloader 参数，设置项：

dataloader_config:
  return_tensors: 'pt'
  padding: 'max_length'
  pad_to_multiple_of: *seq_length
  max_length: *seq_length

可以发现在 examples/train_with_megatron_json_dataset.py 中已经额外构建 dataloader 并传入 PreTrainer。

场景3：自定义Megatron处理流程

基于场景1修改对应配置文件:

SCRIPT_PATH=examples/train_with_megatron_custom.py
ACCELERATE_CONF_PATH=examples/accelerate_config/accelerate_megatron_config.yaml
PRETRAIN_CONF_PATH=examples/llama2_config/llama2-megatron.yaml

1、自定义处理函数

def train_valid_test_datasets_provider():
   """自定义数据集获取函数"""
   pass

def megatron_gpt_get_batch():
   """自定义批获取函数"""
   pass

def megatron_gpt_model_provider():
   """自定义模型获取函数"""
   pass

def megatron_gpt_loss_func():
   """自定义损失函数"""
   pass

pretrain_args.update_distributed_train_args(
	extra_args={
   		"custom_megatron_datasets_provider_function": train_valid_test_datasets_provider,
    	"custom_get_batch_function": megatron_gpt_get_batch,
   		"custom_model_provider_function": megatron_gpt_model_provider,
  		"custom_loss_function": megatron_gpt_loss_func,
   }
)

可以发现在 examples/train_with_megatron_custom.py 中已经将用户自定义函数通过 update_distributed_train_args 接口传入 pretrain_args。

2、自定义解析模型配置文件

用户可以按照 accelerate 解析模型 config 文件的格式，编写解析函数。

以下用已经内置的 llama 模型配置解析函数为例：

from accelerate.utils import add_model_config_to_megatron_parser

@add_model_config_to_megatron_parser('llama')
def parse_llama_config(megatron_lm_plugin, model, batch_data):
    model_type_name = "gpt"
    num_layers = model.config.num_hidden_layers
    pretraining_flag = True
    hidden_size = model.config.hidden_size
    num_attention_heads = model.config.num_attention_heads
    orig_vocab_size = model.config.vocab_size

    max_position_embeddings = getattr(model.config, "max_position_embeddings")
    seq_length = getattr(model.config, "max_sequence_length", None)
    if megatron_lm_plugin.seq_length is None:
        if seq_length is not None:
            megatron_lm_plugin.seq_length = seq_length
        elif megatron_lm_plugin.decoder_seq_length is not None:
            megatron_lm_plugin.seq_length = megatron_lm_plugin.decoder_seq_length
        elif batch_data is not None:
            megatron_lm_plugin.seq_length = batch_data["input_ids"].shape[1]
        else:
            megatron_lm_plugin.seq_length = max_position_embeddings

    megatron_lm_plugin.megatron_lm_default_args["return_logits"] = megatron_lm_plugin.return_logits
    megatron_lm_plugin.megatron_lm_default_args["tokenizer_type"] = "Llama2Tokenizer"
    megatron_lm_plugin.megatron_lm_default_args["model_type_name"] = model_type_name
    megatron_lm_plugin.megatron_lm_default_args["num_layers"] = num_layers
    megatron_lm_plugin.megatron_lm_default_args["pretraining_flag"] = pretraining_flag
    megatron_lm_plugin.megatron_lm_default_args["hidden_size"] = hidden_size
    megatron_lm_plugin.megatron_lm_default_args["num_attention_heads"] = num_attention_heads
    megatron_lm_plugin.megatron_lm_default_args["orig_vocab_size"] = orig_vocab_size
    megatron_lm_plugin.megatron_lm_default_args["max_position_embeddings"] = max_position_embeddings
    megatron_lm_plugin.megatron_lm_default_args["seq_length"] = megatron_lm_plugin.seq_length
    megatron_lm_plugin.megatron_lm_default_args["model_return_dict"] = model.config.return_dict

3、传递Megatron参数

在配置说明章节的预训练YAML配置文件说明中的 other_megatron_args 中配置透传给 Megatron 的参数。

使用其他框架预训练模型

以 deepspeed 框架为例。

基于使用 Megatron 框架预训练模型中场景1修改对应配置文件：

SCRIPT_PATH=examples/train_with_deepspeed.py
ACCELERATE_CONF_PATH=examples/accelerate_config/accelerate_deepspeed_config.yaml 
PRETRAIN_CONF_PATH=examples/llama2_config/llama2-deepspeed.yaml  

在 examples/llama2_config/llama2-deepspeed.yaml 配置文件中配置 dataloader 参数，设置项：

dataloader_config:
  return_tensors: 'pt'
  padding: 'max_length'
  pad_to_multiple_of: 4096
  max_length: 4096

可以发现在 examples/train_with_deepspeed.py 中已经额外构建 dataloader 并传入 PreTrainer。

配置说明

Accelerate配置文件说明

Framework：

Accelerate通用配置

distributed_type: MEGATRON_LM         # 分布式模式:MEGATRON_LM/DEEPSPEED
compute_environment: LOCAL_MACHINE
debug: false
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
num_machines: 1                       # 机器数
num_processes: 8                      # 进程数 
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false                        # 是否使用cpu

Accelerate Plugin配置

• megatron

预训练 YAML 配置文件中的 megatron 的参数会覆盖 Accelerate 配置文件中的 megatron plugin ，为了避免歧义，我们推荐在 Accelerate 配置文件中不配 megatron plugin 的相关参数，转到预训练 YAML 配置文件中配置。

• deepspeed

deepspeed_config:
  gradient_accumulation_steps: 8      # 梯度累计步数
  gradient_clipping: 1.0
  zero3_init_flag: falseo
  zero_stage: 2

预训练YAML配置文件说明

配置文件：

examples/llama2_config/llama2-deepspeed.yaml
examples/llama2_config/llama2-megatron-json-dataset.yaml
examples/llama2_config/llama2-megatron-multi-machine.yaml
examples/llama2_config/llama2-megatron.yaml

说明：

### General Args ###
num_training_steps: 1000                                        # 训练步数
micro_batch_size: &micro_batch_size 4                           # 微批次大小         
dp: 1                                                           # 并行度
gradient_accumulation_steps: &gradient_accumulation_steps 8     # 梯度累计步数
seq_length: &seq_length 4096                                    # 能够处理的最大的序列长度
megatron_dataset_flag: True                                     # megatron数据集标识
data_path: &data_path 'datasets/alpaca'                         # 数据集路径
save_dir: 'models/llama-2-7b-hf_save'                           # 模型保存路径
save_interval: 10000                                            # 模型保存间隔
eval_interval: 10000                                            # 模型评估间隔
openmind_model_path: 'models/llama-2-7b-hf'                     # openmind模型路径
dtype: 'bf16'                                                   # 数据类型

### Plugin Args ###
plugin_args:                                                    # megatron插件参数
  tp_degree: 8                                                  # 张量并行度
  pp_degree: 1                                                  # 流水并行度
  num_micro_batches: *gradient_accumulation_steps               
  gradient_clipping: 1.0
  use_distributed_optimizer: False
  sequence_parallelism: True                                   # 是否启用序列并行
  other_megatron_args:                                          # 透传给megatron的参数
    tokenizer_model: 'models/llama-2-7b-hf/tokenizer.model'                    
    tokenizer_type: 'Llama2Tokenizer'
    finetune: False                                        
    recompute_granularity: "full"
    recompute_method: "block"
    recompute_num_layers: 32
    optimizer: "adam"                                         
    lr: 1e-5                                                  
    min_lr: 1e-6                                           
    adam_beta2: 0.95
    add_bias_linear: False
    async_tensor_model_parallel_allreduce: True
    attention_dropout: 0.0
    attention_softmax_in_fp32: True
    bias_gelu_fusion: False
    ffn_hidden_size: 11008
    hidden_dropout: 0.0
    init_method_std: 0.01
    initial_loss_scale: 65536.0
    lr_decay_style: "cosine"                                 
    lr_warmup_fraction: 0.01                                    
    masked_softmax_fusion: False
    normalization: "RMSNorm"                                   
    split: &split "100,0,0"                                  
    swiglu: True
    untie_embeddings_and_output_weights: True
    use_flash_attn: True                                        
    weight_decay: 0.1
    no_load_optim: True
    no_load_rng: True
    position_embedding_type: "rope"

### Dataloader Config ###
dataloader_config:                                              # 透传给数据加载器DataCollatorForSeq2Seq的参数
  return_tensors: 'pt'                                         
  padding: 'max_length'                                                                         
  pad_to_multiple_of: 4096
  max_length: 4096  

启动脚本说明

启动脚本

#!/bin/bash
# Copyright (c) Huawei Technologies Co., Ltd. 2024, All rights reserved.

export CUDA_VISIBLE_DEVICES_=0,1,2,3,4,5,6,7                                                           # 使用的NPU,卡号逗号分割
export ASCEND_RT_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES_}
export LD_LIBRARY_PATH=/usr/local/lib:/home/anaconda3/lib:$LD_LIBRARY_PATH
export HCCL_CONNECT_TIMEOUT=1200
export COMBINED_ENABLE=1
export CUDA_DEVICE_MAX_CONNECTIONS=1
source /usr/local/Ascend/ascend-toolkit/set_env.sh                                                     # CANN包路径

SCRIPT_DIR=$(cd "$(dirname "$0")"; pwd)
SCRIPT_PATH="${SCRIPT_DIR}"/train_with_megatron.py                                                     # 训练脚本
ACCELERATE_CONF_PATH="${SCRIPT_DIR}"/accelerate_config/accelerate_megatron_config.yaml                 # accelerate配置文件
PRETRAIN_CONF_PATH="${SCRIPT_DIR}"/llama2_config/llama2-megatron.yaml                                  # pretrain配置文件

accelerate launch --config_file "${ACCELERATE_CONF_PATH}" "${SCRIPT_PATH}" --pretrain_config_file "${PRETRAIN_CONF_PATH}" | tee "${SCRIPT_DIR}"/train.log

参数配置

• ACCELERATE_CONF_PATH

可选：

examples/accelerate_config/accelerate_deepspeed_config.yaml  # accelerate关于deepspeed 的配置文件
examples/accelerate_config/accelerate_megatron_config.yaml   # accelerate关于 megatron 的配置文件

• PRETRAIN_CONF_PATH

可选：

examples/llama2_config/llama2-deepspeed.yaml                   # 使用 deepspeed 启动预训练的配置文件
examples/llama2_config/llama2-megatron.yaml                    # 使用 megatron 启动预训练的配置文件
examples/llama2_config/llama2-megatron-json-dataset.yaml       # 使用 megatron 和 json 格式数据集启动预训练的配置文件
examples/llama2_config/llama2-megatron-multi-machine.yaml      # 使用多机启动 megatron 预训练的配置文件

公网地址声明

本代码仓包含公网地址，公开性声明请参考《公网地址声明》。

建议与交流

欢迎大家为社区做贡献。如果有任何疑问或建议，请提交gitee Issues，我们会尽快回复。感谢您的支持。

安全声明

为保障使用过程安全，推荐用户参考《安全声明》了解相关安全信息，进行必要的安全加固。