The openmind-accelerate is a product which allows you to use NVIDIA Megatron-LM in accelerate framework.
Project description
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格式数据
-
修改 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'
-
修改 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
-
在 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: µ_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,我们会尽快回复。感谢您的支持。
安全声明
为保障使用过程安全,推荐用户参考《安全声明》了解相关安全信息,进行必要的安全加固。
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distribution
Hashes for openmind_accelerate-0.5.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 938c2f64cd83806726dc9b0b9a073c465573a99b53902ec2ee4338cadcff7d3a |
|
MD5 | 607e8d5587dbb819c448ccce89cbab25 |
|
BLAKE2b-256 | f19cdd3478854b380d4eee3f24dc8816b039d1e934eadf9a032249be6ced6a23 |