模型预训练 ​

基础概念 ​

预训练是一种深度学习模型训练的策略，通常在大规模的数据集上进行。预训练的目标是通过在一个相关但较大的任务上训练模型，使得模型学习到通用的特征表示。但是随着大模型参数和所需训练数据量的急剧增长，单个机器的资源上限已无法满足训练要求，于是就引出了分布式训练的概念。

分布式训练指的是将深度学习模型任务分解为多个子任务，并在多个计算设备上并行的进行训练。分布式训练极大地提升了大模型的训练速度，可以大幅降低模型训练的总体时间。

本文档中的PreTrainer是基于Accelerate实现了多框架（Megatron、DeepSpeed以及FSDP）的分布式能力，并提供了通用的预训练流程管理功能。

环境准备 ​

torch: 2.1.0
transformers: 4.45.2
accelerate: 0.28.0
deepspeed: 0.15.2
megatron_core: 0.4.0rc0

安装Megatron-LM分布式框架 ​

若用户需要使用Megatron-LM分布式框架，则还需执行以下步骤。

1. 安装Megatron（参考MindSpeed的Megatron安装方式）

   shell

   git clone https://github.com/NVIDIA/Megatron-LM.git
   cd Megatron-LM
   git checkout bcce6f54e075e3c3374ea67adefe54f3f2da2b07
   pip install --no-use-pep517 -e .  # 使用"--no-use-pep517 -e"安装megatron全部文件
   

2. 安装MindSpeed

   shell

   git clone  https://gitee.com/ascend/MindSpeed.git
   cd MindSpeed
   git checkout origin/1.0.RC1
   pip install -r requirements.txt
   pip install -e .
   

3. 使用pip安装魔乐社区openmind_accelerate插件

   shell

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

4. 安装accelerate与deepspeed

   shell

   pip install deepspeed==0.15.2
   pip install accelerate==0.28.0
   

openMind Library环境准备 ​

#aarch64环境下安装
pip install openmind[pt] 

#x86环境下安装
pip install openmind[pt] --extra-index-url https://download.pytorch.org/whl/cpu 

openMind Library依赖环境安装请参考openMind Library安装指南。 安装完成后请使用pip list检查版本依赖，如果在安装上述依赖的时候，accelerate或transformers版本被刷新，请重新刷回指定版本。

快速使用 ​

我们提供了样例配置文件和启动脚本，方便用户一键使用。

PreTrainer的使用步骤如下所示 ​

准备数据 ​

用户需要准备好自己的预训练数据，例如alpaca_en数据。 如果用户需要使用Megatron-LM分布式框架，可参考Megatron的数据处理方法 进行处理。

准备模型 ​

用户需要准备好模型文件，例如llama2模型。 如果用户需要使用Megatron-LM分布式框架，则只需要准备config.json和tokenizer相关文件即可。

准备预训练参数 ​

预训练参数可以通过加载 llama2_config/llama2-megatron-json-dataset.yaml 文件自动生成，用户可参考此处基于json格式微调数据集的样例配置文件：

启动 ​

• Accelerate配置文件可参考：accelerate_config/accelerate_megatron_config.yaml

  yaml

  compute_environment: LOCAL_MACHINE
  debug: false
  distributed_type: MEGATRON_LM
  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
  
  

• 模型配置文件可参考：llama2_config/llama2-megatron-json-dataset.yaml

  yaml

  num_training_steps: 1000
  micro_batch_size: &micro_batch_size 4
  dp: 1
  gradient_accumulation_steps: &gradient_accumulation_steps 8
  ### seq_length需要小于或等于模型权重配置文件config.json中，"max_position_embeddings"字段的值
  seq_length: &seq_length 4096
  megatron_dataset_flag: False
  ### data_path请传入本地微调数据集所在路径
  data_path: &data_path '/path/to/alpaca_en/alpaca_data_en_52k.json'
  ### 微调模型权重保存路径
  save_dir: './saves'
  save_interval: 10000
  eval_interval: 10000
  ### openmind_model_path请传入本地模型权重文件夹所在路径
  openmind_model_path: '/path/to/llama2-7b-hf'
  dtype: 'bf16'
  
  plugin_args:
    tp_degree: 8
    pp_degree: 1
    num_micro_batches: *gradient_accumulation_steps
    gradient_clipping: 1.0
    use_distributed_optimizer: False
    sequence_parallelism: False
    other_megatron_args:
      ### tokenizer_model请传入本地模型权重文件中，tokenizer.model文件所在路径
      tokenizer_model: &tokenizer_model '/path/to/llama2-7b-hf/tokenizer.model'
      tokenizer_type: &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: False
      attention_dropout: 0.0
      attention_softmax_in_fp32: False
      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: False
      weight_decay: 0.1
      no_load_optim: True
      no_load_rng: True
      eval_iters: &eval_iters 10
      position_embedding_type: "rope"
  
  dataloader_config:
    return_tensors: 'pt'
    padding: 'max_length'
    pad_to_multiple_of: *seq_length
    max_length: *seq_length
  
  

• 预训练程序文件可参考train_with_megatron_json_dataset.py，此python脚本不能直接运行，如需运行，请自行下载如下仓库获取utils相关代码，然后将accelerate_examples/examples/utils复制到此脚本同目录下。

  shell

  git clone https://modelers.cn/AI-Research/accelerate_examples.git
  cp -r accelerate_examples/examples/utils ./   # 自行替换目的路径为train_with_megatron_json_dataset.py所在路径
  

  python

  import os
  
  import openmind_accelerate
  from openmind import PreTrainingArguments, PreTrainer
  
  from utils.config import get_pretrain_config_file
  from utils.accelerator import make_accelerator
  from utils.data import make_train_and_eval_dataloader
  from utils.tokenizer import get_tokenizer
  
  pretrain_args = PreTrainingArguments.from_yaml(get_pretrain_config_file())
  
  os.makedirs(pretrain_args.save_dir, exist_ok=True)
  
  accelerator = make_accelerator(pretrain_args=pretrain_args)
  
  tokenizer = get_tokenizer(tokenizer_path=pretrain_args.openmind_model_path, use_fast=False)
  transformer_dataloader_config = pretrain_args.get_dataloader_config()
  train_dataloader, eval_dataloader = make_train_and_eval_dataloader(
      dataloader_config=transformer_dataloader_config,
      micro_batch_size=pretrain_args.micro_batch_size,
      data_files=pretrain_args.data_path,
      max_length=pretrain_args.seq_length,
      tokenizer=tokenizer,
      accelerator=accelerator
  )
  
  pretrainer = PreTrainer(pretrain_args=pretrain_args,
                          train_dataloader=train_dataloader,
                          eval_dataloader=eval_dataloader,
                          )
  pretrainer.train()
  

在完成上述环境配置以及配置文件准备后，即可通过如下命令启动微调，请确保其中的训练脚本和配置文件为本地实际路径。

accelerate launch --config_file accelerate_config/accelerate_megatron_config.yaml train_with_megatron_json_dataset.py --pretrain_config_file llama2_config/llama2-megatron-json-dataset.yaml

进阶使用 ​

定义预训练参数 ​

在我们定义PreTrainer之前首先需要定义一个PreTrainingArguments类，它将包含PreTrainer用于训练和评估的所有超参数。用户可以通过配置文件或者直接传参初始化预训练参数。

使用配置文件 ​

预训练参数可以通过加载yaml文件自动生成，更多yaml样例可参考：样例链接。

from openmind import PreTrainingArguments

# 路径需要替换为本地路径
pretrain_args = PreTrainingArguments.from_yaml(
    "openmind-accelerate/examples/llama2_config/llama2-megatron.yaml"
)

直接传参 ​

预训练参数也可以通过传参的方式实例化。使用Megatron模型训练Megatron数据集的预训练器初始化流程如下。

参数链接请点击：PreTrainingArguments说明。

from openmind import PreTrainingArguments

# 路径需要替换为本地路径
pretrain_args = PreTrainingArguments(
    megatron_dataset_flag=True,
    data_path="HaM/alpaca_en",
    num_training_steps=1000,
    micro_batch_size=4,
    dp=1,
    gradient_accumulation_steps=8,
    seq_length=2048,
)

使用Megatron框架预训练模型 ​

用户完成预训练参数配置后即可启动Megatron模型预训练。

• Accelerate对接Megatron的配置文件可参考：accelerate_config/accelerate_megatron_config.yaml
• 使用Megatron框架训练Json数据运行示例可参考：train_with_megatron_json_dataset.py。
• Json格式数据预训练配置文件示例可参考：llama2_config/llama2-megatron-json-dataset.yaml。

用户只需要将准备好的train_dataloader（eval_dataloader非必选），传给PreTrainer，即可使用用户自定义的dataloader预训练模型。

accelerate launch --config_file accelerate_config/accelerate_megatron_config.yaml train_with_megatron_json_dataset.py --pretrain_config_file llama2_config/llama2-megatron-json-dataset.yaml

自定义Megatron框架处理流程（可选） ​

自定义处理函数 ​

如下代码所示，PreTrainer接口在使用Megatron预训练时，支持用户根据实际场景按需自定义datasets_provider、model_provider、get_batch和loss_function中的任意函数，并将函数指针赋值到如下属性中。自定义函数的实现可参考官方样例pretrain_gpt.py。

• custom_megatron_datasets_provider_function：用于提供Megatron的训练和验证数据集。
• custom_get_batch_function：用于生成批次数据。
• custom_model_provider_function：用于构建模型。
• custom_loss_function：返回损失函数。

import openmind_accelerate
from openmind import PreTrainingArguments
from pretrain_gpt import (
    train_valid_test_datasets_provider,
    get_batch as megatron_gpt_get_batch,
    model_provider as megatron_gpt_model_provider,
    loss_func as megatron_gpt_loss_func,
)

# 路径需要替换为本地路径
pretrain_args = PreTrainingArguments.from_yaml(
    "openmind-accelerate/examples/llama2_config/llama2-megatron-json-dataset.yaml"
)
train_valid_test_datasets_provider.is_distributed = True
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,
    }
)

自定义解析模型配置文件 ​

用户可依据Accelerate解析模型配置的格式，自定义模型配置文件解析函数。以下为PreTrainer内置的llama模型配置文件解析函数，用户可以根据实际情况参考。

import openmind_accelerate
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

使用其他框架预训练模型 ​

PreTrainer是基于Accelerate实现的多框架分布式能力，所以PreTrainer除了支持Megatron框架，还支持DeepSpeed和FSDP分布式框架。如下以DeepSpeed分布式框架为例： 用户完成Json格式预训练参数配置后即可启动DeepSpeed模型预训练。

• Accelerate对接DeepSpeed的配置文件示例可参考：accelerate_config/accelerate_deepspeed_config.yaml。
• 使用DeepSpeed框架训练Json数据运行示例可参考：train_with_deepspeed.py。
• Json格式数据预训练配置文件示例可参考：llama2_config/llama2-deepspeed.yaml。

num_training_steps: 1000
micro_batch_size: 1
dp: 8
gradient_accumulation_steps: 8
seq_length: 4096
megatron_dataset_flag: False
data_path: '/path/to/alpaca_en/alpaca_data_en_52k.json'
save_dir: './saves'
save_interval: 10000
eval_interval: 10000
openmind_model_path: '/path/to/llama2-7b-hf'
dtype: 'bf16'

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

### seq_length、max_length以及padding的值均需要小于或等于模型权重配置文件config.json中，"max_position_embeddings"字段的值

accelerate launch --config_file accelerate_config/accelerate_deepspeed_config.yaml train_with_deepspeed.py --pretrain_config_file llama2_config/llama2-deepspeed.yaml

PreTrainingArguments说明 ​

PreTrainer说明 ​

PreTrainer接口会根据Accelerate是否使用Megatron-LM分布式加速库（以环境变量ACCELERATE_USE_MEGATRON_LM=="true"为依据），来选择创建Megatron预训练器或其他预训练器。

Megatron预训练器 ​

其他预训练器 ​

感谢社区贡献的 llama2 模型以及 alpaca_en 数据集