DeepSpeed-Chat 打造类ChatGPT全流程 笔记二之监督指令微调

文章目录

  • 系列文章
  • 0x0. 前言
  • 0x1. Supervised finetuning (SFT) 教程翻译
    • 如何训练模型
    • 如何对SFT checkpoint进行评测?
    • 模型和数据
    • ☀️来自OPT-1.3B及其SFT变体(使用不同微调数据)的提示示例
    • ☀️ 一些参数解释和可训练的最大模型
    • 其它
  • 0x2. 评测脚本解读
  • 0x3. 训练脚本解读
    • 0x3.1 头文件相关解析
      • create_prompt_dataset解析
      • 工具函数解析
        • print_rank_0
        • to_device
        • save_hf_format
        • set_random_seed
        • get_all_reduce_mean
        • get_optimizer_grouped_parameters
        • save_zero_three_model
        • load_hf_tokenizer
        • convert_linear_layer_to_lora
        • convert_lora_to_linear_layer
        • only_optimize_lora_parameters
        • create_hf_model
    • 0x3.2 main.py主体解析
      • parse_args解析
      • main函数解析
  • 0x4. hybrid_engine的细节 & log
  • 0x5. 总结

系列文章

  • DeepSpeed-Chat 打造类ChatGPT全流程 笔记一
  • 【DeepSpeed 教程翻译】三,在 DeepSpeed中使用 PyTorch Profiler和Flops Profiler
  • DeepSpeed结合Megatron-LM训练GPT2模型笔记(上)
  • 【DeepSpeed 教程翻译】二,Megatron-LM GPT2,Zero 和 ZeRO-Offload
  • 【DeepSpeed 教程翻译】开始,安装细节和CIFAR-10 Tutorial

0x0. 前言

在 DeepSpeed-Chat 打造类ChatGPT全流程 笔记一 中跑通了DeepSpeed Chat的训练和推理流程,DeepSpeed Chat的训练流程包含监督指令微调(SFT),Reward模型微调,基于人类反馈的强化学习(RLHF)三个步骤。接着上面文章的todo,这篇文章主要是解析一下监督指令微调(SFT)阶段的代码实现。

0x1. Supervised finetuning (SFT) 教程翻译

监督微调(SFT)与在自然语言任务(例如,WikiText-103)上的标准语言模型微调非常相似。主要的区别来自于数据集资源,SFT将收集高质量的查询-回答对来微调模型以达到人类更倾向的生成结果。

如何训练模型

我们提供了多个脚本用于在单个GPU(例如,单个A6000-48G,V100-32G,A100-40G等),单节点(例如,8/16x V100-32G,8 A100-40G/80G)和多节点设置(例如,64x A100-80G)上进行训练,这些可以在 training_scripts 目录中找到。例如,如果你有一个单独的A6000-48G,你可以简单地运行对应的脚本

 training_scripts/single_gpu/run_1.3b.sh

来训练一个OPT-1.3b模型。我们的单节点脚本很容易扩展到多节点系统。

如何对SFT checkpoint进行评测?

一旦你使用上述代码完成训练,你可以简单地执行 bash evaluation_scripts/run_prompt.sh

它会要求用户提供两个模型的路径:(a) 原始预训练模型(即 --model_name_or_path_baseline facebook/opt-1.3b)和 (b) 微调后的模型(即 --model_name_or_path_finetune output/check_base)。“prompt_eval.py” 包含了几个可以根据你的喜好进行更新的提示。

模型和数据

由于GPT3没有开源的checkpoint,我们使用了Meta OPT家族的预训练模型(即facebook/opt-1.3b)。你也可以使用其他预训练模型(如GPT-Neo,Bloom等)。至于数据集,我们也使用了来自Huggingface数据集的开源数据集,具体如下:

Dahoas/rm-static
Dahoas/full-hh-rlhf
Dahoas/synthetic-instruct-gptj-pairwise
yitingxie/rlhf-reward-datasets
openai/webgpt_comparisons 
stanfordnlp/SHP

感谢DeepSpeed RLHF的数据抽象和融合技术,我们现在可以将多个数据源合并用于训练。然而,重要的是要注意,不同的数据集可能使用不同的提示词(例如,Dohas/rm-static使用"Human:"表示查询,"Assistant:"表示回答)。因此,用户必须自行对齐这些提示。在我们的例子中,我们一致使用了Dohas/rm-static的格式。通过我们的评估,我们发现整合多样化的数据集可以提高模型的质量。请参考下一节以获取不同查询-答案对的示例。

☀️来自OPT-1.3B及其SFT变体(使用不同微调数据)的提示示例

DeepSpeed-Chat 打造类ChatGPT全流程 笔记二之监督指令微调_第1张图片

☀️ 一些参数解释和可训练的最大模型

main.py文件中使用的大多数参数都有清晰的解释,如果你有解码器模型微调的经验,通常很容易理解。然而,如果你对其中任何一个不清楚,请不要犹豫在GitHub问题上向我们求助。在这一部分,我们提供了一些具体的参数解释和它们的使用方法。

参数 解释 注意事项
–data_path 用于微调模型的数据 你可以指定多个数据资源来训练模型,例如:Dahoas/rm-static Dahoas/full-hh-rlhf
–data_split 为三步训练切分数据 根据InstructGPT,我们提供了切分数据集的能力,使得每个分区只在一个步骤中使用。设置为"2,4,4"意味着我们分别使用20%,40%,40%的数据在每个步骤中。如果你只做SFT,或者你发现在不同步骤中使用重叠数据是可以的/有帮助的,你可以将它改为"10,0,0"。
–sft_only_data_path 用于微调模型的单响应数据 对于只在步骤1中使用的单响应数据,你应该将它们作为这个参数的一部分,而不是上面的data_path参数。这个参数中的数据集将不会被切分,而只在步骤1中全面使用。
–gradient_checkpoint 为模型启用梯度检查点(也称为激活检查点) 这可以显著降低训练内存成本
–offload DeepSpeed特定功能。将模型卸载到CPT/NVME以节省内存 这可以在内存消耗较少的情况下训练更大的模型。但是它会减慢训练的速度。
–zero_stage DeepSpeed特定功能,适用于多GPU系统 这可以帮助将模型/优化器分布在多个GPU上。请参见https://www.deepspeed.ai/tutorials/zero/
–lora_dim 当它大于0时,将启用LoRA 通常,LoRA需要更大的学习率才能更好地收敛
–lora_module_name 启用LoRA模块的范围。
–only_optimize_lora 冻结所有其他参数,只优化LoRA相关参数
–gradient_checkpoint, --lora_dim, only_optimize_lora 当启用LoRA和梯度检查点时,不能启用只优化LoRA 如果全部启用,将影响梯度流(也就是由PyTorch支持的auto-grad系统后端)

对于用户来说,一个重要的考虑是确定他们可以使用当前系统训练的最大模型大小。在这里,我们提供了一个估计这个限制的方法。假设你不使用卸载功能,并启用(i)零阶段3(如果使用多个GPU),(ii)梯度检查点,以及(iii)LoRA,那么你可以训练的大致最大模型大小(以十亿参数为单位)可以估计为"总GPU内存(GB)除以3"。例如,如果你有一台单一的A6000-48G GPU,你可能可以训练最多16十亿参数的模型。需要注意的是,这只是一个粗略的估计,你应该自己验证。

其它

从InstructGPT的工作中,我们建议为了得到更好的人类偏好的答案,让模型过度拟合(即更长的训练周期)。通过我们的探索,我们发现这对于较小模型的微调,如OPT-1.3B,特别有帮助。值得注意的是,我们在脚本中提供的超参数并没有经过大量的调整。因此,我们鼓励用户和实践者自己找到最优的配置。此外,我们的系统可以很容易地扩展到其他语言,如中文和日语。为了证明这一点,我们在"training_scripts/other_language"目录下包含了两个例子。

0x2. 评测脚本解读

DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/evaluation_scripts/run_prompt.sh这个脚本是对输入的模型进行评测,我们看一下内容:

# You can provide two models to compare the performance of the baseline and the finetuned model
export CUDA_VISIBLE_DEVICES=0
python prompt_eval.py \
    --model_name_or_path_baseline XXX \
    --model_name_or_path_finetune XXX

这表示我们可以调用prompt_eval.py来对baseline模型和finetune模型进行对比评测。所以评测的实现就是这里的prompt_eval.py,接下来对prompt_eval.py做一个解读。

  1. 首先是头文件导入部分,除了一些常规的需要用到的包还从DeepSpeed Chat的utils下面导入了2个函数。
# DeepSpeed Team
import argparse
import logging
import torch
import sys
import os

from transformers import (
    AutoModelForCausalLM, )

sys.path.append(
    os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
from utils.model.model_utils import create_hf_model # 从utils下面导入自定义函数
from utils.utils import load_hf_tokenizer

logger = logging.getLogger(__name__)

接下来解析一下create_hf_modelload_hf_tokenizer两个函数的实现。首先是create_hf_model函数:

# 这段代码定义了一个名为create_hf_model的函数,该函数的作用是创建或加载一个预训练模型。该函数的主要参数包括:
# model_class:模型的类别,例如GPT-2、BERT等。
# tokenizer:用于模型的分词器。
# ds_config: DeepSpeed的配置参数。
# rlhf_training:一个标志,用来表示是否正在进行RLHF(Reinforcement Learning from Human Feedback)训练。
# disable_dropout:一个标志,用来表示是否禁用dropout。Dropout是一种防止过拟合的技术。
def create_hf_model(model_class,
                    model_name_or_path,
                    tokenizer,
                    ds_config=None,
                    rlhf_training=False,
                    disable_dropout=False):
    # 根据model_name_or_path从预训练模型获取模型配置model_config。
    model_config = AutoConfig.from_pretrained(model_name_or_path)
    # 如果disable_dropout为真,则将模型配置中的dropout设为0.0。
    if disable_dropout:
        model_config.dropout = 0.0
    # Note: dschf is defined in function scope to avoid global effects
    # https://huggingface.co/docs/transformers/main_classes/deepspeed#nontrainer-deepspeed-integration
    # 根据ds_config中的设置,创建DeepSpeed的配置对象dschf,以便进行DeepSpeed优化。
    if ds_config is not None and ds_config["zero_optimization"]["stage"] == 3:
        dschf = HfDeepSpeedConfig(ds_config)
    else:
        dschf = None
    # 根据rlhf_training的值,确定是从配置中创建模型还是从预训练模型中加载模型。如果rlhf_training为真,则根据模型配置创建模型;否则,从预训练模型加载模型。
    if rlhf_training:
        # the weight loading is handled by create critic model
        model = model_class.from_config(model_config)
    else:
        model = model_class.from_pretrained(
            model_name_or_path,
            from_tf=bool(".ckpt" in model_name_or_path),
            config=model_config)
    # 将模型的结束标记和填充标记设为分词器的结束标记id。
    model.config.end_token_id = tokenizer.eos_token_id
    model.config.pad_token_id = model.config.eos_token_id
    # 调整模型的词汇表大小,使其为8的倍数。这样做的目的是为了在某些硬件(如GPU)上提高效率。
    model.resize_token_embeddings(int(
        8 *
        math.ceil(len(tokenizer) / 8.0)))  # make the vocab size multiple of 8

    return model

然后是load_hf_tokenizer函数:

# 这段代码定义了一个名为load_hf_tokenizer的函数,该函数的作用是加载一个指定的分词器。这个函数的参数有:
# model_name_or_path:模型的名称或者预训练模型的路径。
# fast_tokenizer:是否使用Fast Tokenizer,这是一种更快的分词器实现。
def load_hf_tokenizer(model_name_or_path, fast_tokenizer=True): 
    # 检查model_name_or_path是否存在。如果存在,可能意味着这是一个本地的模型路径。
    if os.path.exists(model_name_or_path):
        # Locally tokenizer loading has some issue, so we need to force download
        # 如果是本地模型,再检查在该路径下的"config.json"文件是否存在。如果存在,打开该文件,并从文件中读取"_name_or_path"字段,将其赋值给model_name,然后通过AutoTokenizer.from_pretrained函数,使用该模型名称和fast_tokenizer参数,从HuggingFace模型库下载并加载分词器。
        model_json = os.path.join(model_name_or_path, "config.json")
        if os.path.exists(model_json):
            model_json_file = json.load(open(model_json))
            model_name = model_json_file["_name_or_path"]
            tokenizer = AutoTokenizer.from_pretrained(model_name,
                                                      fast_tokenizer=True)
    else:
        # 如果model_name_or_path不存在,直接使用AutoTokenizer.from_pretrained函数,使用model_name_or_path和fast_tokenizer参数,从HuggingFace模型库下载并加载分词器。
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
                                                  fast_tokenizer=True)
    return tokenizer
  1. 接下来是参数解析部分,解析如下:
# 这段代码定义了一个名为parse_args的函数,该函数的作用是解析命令行参数。它使用Python的argparse库来完成这个工作。下面是每个参数的详解:
def parse_args():
    parser = argparse.ArgumentParser(description="Eval the finetued SFT model")
    # model_name_or_path_baseline:基线模型的路径,这是一个必须提供的参数(required=True)。
    parser.add_argument(
        "--model_name_or_path_baseline",
        type=str,
        help="Path to baseline model",
        required=True,
    )
    # model_name_or_path_finetune:微调后模型的路径,这也是一个必须提供的参数。
    parser.add_argument(
        "--model_name_or_path_finetune",
        type=str,
        help="Path to pretrained model",
        required=True,
    )
    # num_beams:用于指定集束搜索的集束宽度,其默认值为1。
    parser.add_argument(
        "--num_beams",
        type=int,
        default=1,
        help='Specify num of beams',
    )
    # num_beam_groups:用于指定集束搜索的组数,其默认值为1。
    parser.add_argument(
        "--num_beam_groups",
        type=int,
        default=1,
        help='Specify num of beams',
    )
    # top_k:用于指定在Top-K采样中考虑的最高可能性词汇的数量,其默认值为4。
    parser.add_argument(
        "--top_k",
        type=int,
        default=4,
        help='Specify num of beams',
    )
    # penalty_alpha:惩罚因子,其默认值为0.6。
    parser.add_argument(
        "--penalty_alpha",
        type=float,
        default=0.6,
        help='Specify num of beams',
    )
    # num_return_sequences:生成序列的数量,其默认值为1。
    parser.add_argument(
        "--num_return_sequences",
        type=int,
        default=1,
        help='Specify num of return sequences',
    )
    # max_new_tokens:生成的最大新token数,其默认值为100。
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=100,
        help='Specify num of return sequences',
    )
    # language:语言类型,可以是"English"、"Chinese"或"Japanese",默认为"English"。
    parser.add_argument("--language",
                        type=str,
                        default="English",
                        choices=["English", "Chinese", "Japanese"])
		# parser.parse_args()这个函数将解析命令行参数,并将结果保存在一个Namespace对象中。这个对象被返回,可以在其他地方使用这些参数。
    args = parser.parse_args()

    return args

  1. 接下来是generate函数的解析:
# 这个函数是用来利用训练好的模型生成文本的,它接受以下参数
# model:已经训练好的模型。
# tokenizer:用于将文本转换为模型可理解的输入的工具。
# inputs:模型的输入数据。
# num_beams:在使用束搜索算法时的束宽,其默认值为1。
# num_beam_groups:在使用分组束搜索时的组数,默认为1。
# do_sample:是否进行随机采样。如果设为True,则在生成过程中会随机选择下一个单词,而不是仅选择最可能的单词。默认为False。
# num_return_sequences:模型返回的序列数,默认为1。
# max_new_tokens:模型生成的最大新token数,即最大生成文本的长度,默认为100。
def generate(model,
             tokenizer,
             inputs,
             num_beams=1,
             num_beam_groups=1,
             do_sample=False,
             num_return_sequences=1,
             max_new_tokens=100):
		# 函数首先使用模型的generate方法,根据提供的参数生成文本。
    generate_ids = model.generate(inputs.input_ids,
                                  num_beams=num_beams,
                                  num_beam_groups=num_beam_groups,
                                  do_sample=do_sample,
                                  num_return_sequences=num_return_sequences,
                                  max_new_tokens=max_new_tokens)
		# 使用tokenizer的batch_decode方法将生成的令牌ID解码为可读的文本。注意,这里跳过了特殊的令牌(如填充和开始/结束令牌),并且不会清理tokenize产生的额外空格。
    result = tokenizer.batch_decode(generate_ids,
                                    skip_special_tokens=True,
                                    clean_up_tokenization_spaces=False)
    return result
  1. 接下来是generate_constrastive_search函数的解析:

# 这个函数叫做 generate_constrastive_search,它是用于利用训练好的模型进行对比搜索生成文本的。这个函数接受以下参数:
# model:已经训练好的模型。
# tokenizer:用于将文本转换为模型可理解的输入的工具。
# inputs:模型的输入数据。
# top_k:在每一步生成时,只考虑概率最高的top_k个候选项,然后进行随机抽样。默认为4。
# penalty_alpha:用于惩罚新生成的token与原始输入之间的差异,默认为0.6。
# num_return_sequences:模型返回的序列数,默认为1。
# max_new_tokens:模型生成的最大新token数,即最大生成文本的长度,默认为100。
def generate_constrastive_search(model,
                                 tokenizer,
                                 inputs,
                                 top_k=4,
                                 penalty_alpha=0.6,
                                 num_return_sequences=1,
                                 max_new_tokens=100):
		# 函数首先使用模型的generate方法,根据提供的参数生成文本。注意这里使用了模型的一个特殊的生成方式,这种方式在每一步生成时,只考虑概率最高的top_k个候选项,然后进行随机抽样,同时使用了一个惩罚因子penalty_alpha来惩罚新生成的token与原始输入之间的差异。
    generate_ids = model.generate(inputs.input_ids,
                                  top_k=top_k,
                                  penalty_alpha=penalty_alpha,
                                  num_return_sequences=num_return_sequences,
                                  max_new_tokens=max_new_tokens)
		# 然后,使用tokenizer的batch_decode方法将生成的token ID解码为可读的文本。注意,这里跳过了特殊的token(如填充和开始/结束token),并且不会清理token化产生的额外空格。
    result = tokenizer.batch_decode(generate_ids,
                                    skip_special_tokens=True,
                                    clean_up_tokenization_spaces=False)
    return result
  1. 接下来是一个简单的打印工具函数:
# gen_output:这是一个列表,其中包含了我们希望打印的内容,每一项都是一段文本。
def print_utils(gen_output):
    # 函数会遍历gen_output列表中的每一项,然后将每一项都打印出来。为了在不同项之间增加一些可视化的分隔,函数在每一项前后都额外打印了一个空行。
    for i in range(len(gen_output)):
        print()
        print(gen_output[i])
        print()
  1. 然后是prompt_eval这个函数,这个函数prompt_eval的目的是评估和比较基线模型(model_baseline)和微调过的模型(model_fintuned)对于一组提示(prompts)的生成性能。让我们逐行进行解析:
# 输入参数包括:args(命令行参数)、model_baseline(基线模型)、model_fintuned(微调模型)、tokenizer(用于编码和解码的分词器)、device(指定运行模型的设备)、prompts(一组要评估的提示)。
def prompt_eval(args, model_baseline, model_fintuned, tokenizer, device,
                prompts):
    # 对于prompts中的每一个提示,我们都做以下操作:
    for prompt in prompts:
        # 使用分词器将提示转换为模型所需的输入格式,并将其移至指定的设备上。
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        # 打印一条消息表示我们现在正在进行基线模型的生成。
        print("==========Baseline: Greedy=========")
        # 然后,我们调用之前定义的generate函数使用贪婪搜索方法生成文本,并使用print_utils函数打印生成的结果。
        r_base = generate(model_baseline,
                          tokenizer,
                          inputs,
                          num_beams=1,
                          num_return_sequences=args.num_return_sequences,
                          max_new_tokens=args.max_new_tokens)
        print_utils(r_base)
        # 打印一条消息表示我们现在正在进行微调模型的生成。
        print("==========finetune: Greedy=========")
        # 同样地,我们调用generate函数使用贪婪搜索方法生成文本,并使用print_utils函数打印生成的结果。
        r_finetune_g = generate(model_fintuned,
                                tokenizer,
                                inputs,
                                num_beams=1,
                                num_return_sequences=args.num_return_sequences,
                                max_new_tokens=args.max_new_tokens)
        print_utils(r_finetune_g)
        # 注意:在此函数中,贪婪搜索被用作基线方法。然而,该函数还提供了其他几种搜索策略的例子,包括多项式采样、束搜索、束搜索多项式采样、多样性束搜索和对比搜索。这些策略在此函数中都被注释掉了,但你可以根据需要去掉注释,使用这些策略。

        # print("==========finetune: Multinomial sampling=========")
        # r_finetune_m = generate(model_fintuned, tokenizer, inputs,
        #                         num_beams=1,
        #                         do_sample=True,
        #                         num_return_sequences=args.num_return_sequences,
        #                         max_new_tokens=args.max_new_tokens)
        # print_utils(r_finetune_m)
        # print("==========finetune: Beam Search=========")
        # r_finetune_b = generate(model_fintuned, tokenizer, inputs,
        #                         num_beams=args.num_beams,
        #                         num_return_sequences=args.num_return_sequences,
        #                         max_new_tokens=args.max_new_tokens)
        # print_utils(r_finetune_b)
        # print("==========finetune: Beam-search multinomial sampling=========")
        # r_finetune_s = generate(model_fintuned, tokenizer, inputs,
        #                         num_beams=args.num_beams,
        #                         do_sample=True,
        #                         num_return_sequences=args.num_return_sequences,
        #                         max_new_tokens=args.max_new_tokens)
        # print_utils(r_finetune_s)
        # print("==========finetune: Diverse Beam Search=========")
        # r_finetune_d = generate(model_fintuned, tokenizer, inputs,
        #                         num_beams=args.num_beams,
        #                         num_beam_groups=args.num_beam_groups,
        #                         num_return_sequences=args.num_return_sequences,
        #                         max_new_tokens=args.max_new_tokens)
        # print_utils(r_finetune_d)
        # print("==========finetune: Constrastive Search=========")
        # r_finetune_c = generate_constrastive_search(model_fintuned, tokenizer, inputs,
        #                                             top_k=args.top_k,
        #                                             penalty_alpha=args.penalty_alpha,
        #                                             num_return_sequences=args.num_return_sequences,
        #                                             max_new_tokens=args.max_new_tokens)
        # print_utils(r_finetune_c)
        # 最后,打印一条消息表示这个提示的处理已经结束。然后打印两个空行作为分隔。
        print("====================prompt end=============================")
        print()
        print()
  1. 解析main函数:
# main函数负责解析命令行参数、准备模型和分词器、定义提示,然后使用这些来评估和比较基线模型和微调模型。
def main():
	  # 这个main函数是整个脚本的入口点。它首先通过parse_args函数解析命令行参数。然后它设置了运行模型的设备为第一个GPU。
    args = parse_args()

    device = torch.device("cuda:0")
    # 接着,它使用load_hf_tokenizer函数加载分词器,然后使用create_hf_model函数创建基线模型(model_baseline)和微调模型(model_fintuned)
    tokenizer = load_hf_tokenizer(args.model_name_or_path_baseline,
                                  fast_tokenizer=True)

    model_baseline = create_hf_model(AutoModelForCausalLM,
                                     args.model_name_or_path_baseline,
                                     tokenizer, None)
    model_fintuned = create_hf_model(AutoModelForCausalLM,
                                     args.model_name_or_path_finetune,
                                     tokenizer, None)
    # 然后,这些模型被移动到指定的设备上。
    model_baseline.to(device)
    model_fintuned.to(device)

    # 在接下来的部分,函数定义了一组用于评估的提示。注意,这里特别指出,如果提示以空格结束,那么没有经过微调的原始模型有可能会陷入停滞并无法产生响应。微调过的模型在这方面表现得更好。因此,这里所有的提示都以冒号":"结束,以使得比较更有意义。
    # 这个脚本支持英文、中文和日文的评估,它通过args.language参数判断用户选择的语言,并根据此选择加载对应的提示。
    if args.language == "English":
        prompts = [
            "Human: Please tell me about Microsoft in a few sentence? Assistant:",
            "Human: Explain the moon landing to a 6 year old in a few sentences. Assistant:",
            "Human: Write a short poem about a wise frog. Assistant:",
            "Human: Who was president of the United States in 1955? Assistant:",
            "Human: How does a telescope work? Assistant:",
            "Human: Why do birds migrate south for the winter? Assistant:"
        ]
    elif args.language == "Chinese":
        prompts = [
            "Human: 请用几句话介绍一下微软? Assistant:",
            "Human: 用几句话向6岁的孩子解释登月。 Assistant:",
            "Human: 写一首关于一只聪明的青蛙的短诗。 Assistant:",
            "Human: 谁是1955年的美国总统? Assistant:", "Human: 望远镜是如何工作的? Assistant:",
            "Human: 鸟类为什么要南迁过冬? Assistant:"
        ]
    elif args.language == "Japanese":
        prompts = [
            "Human: マイクロソフトについて簡単に教えてください。 Assistant:",
            "Human: 6歳児に月面着陸を短い文で説明する。 Assistant:",
            "Human: 賢いカエルについて短い詩を書いてください。 Assistant:",
            "Human: 1955年のアメリカ合衆国大統領は誰? Assistant:",
            "Human: 望遠鏡はどのように機能しますか? Assistant:",
            "Human: 鳥が冬に南に移動するのはなぜですか? Assistant:"
        ]
    # 最后,它调用prompt_eval函数,将所有的参数和提示传递给它,以进行模型评估。
    prompt_eval(args, model_baseline, model_fintuned, tokenizer, device,
                prompts)

0x3. 训练脚本解读

训练过程都集中在DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py这个脚本,我们逐行解析一下这个脚本。

0x3.1 头文件相关解析

# DeepSpeed Team
# 首先,它导入了Python的标准库,如argparse(用于解析命令行参数),os和math。
# 然后,它导入了PyTorch库,这是一个用于深度学习的开源库,同时也导入了一些PyTorch的辅助模块,如DataLoader(用于加载数据)、
# RandomSampler和SequentialSampler(用于数据抽样)以及DistributedSampler(用于在分布式设置中进行数据抽样)。
import argparse
import os
import math
import sys

import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler

# 接下来,它导入了Hugging Face的transformers库的一些模块,包括用于因果语言建模的模型(AutoModelForCausalLM),优化器调度类型(SchedulerType),默认的数据整理函数(default_data_collator)和获取优化器调度器的函数(get_scheduler)。
from transformers import (
    AutoModelForCausalLM,
    SchedulerType,
    default_data_collator,
    get_scheduler,
)

# 然后,它导入了deepspeed库,这是一个为大规模模型训练优化的库。它也导入了deepspeed库中的一些模块,包括优化器类(DeepSpeedCPUAdam和FusedAdam)
import deepspeed
from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam

# 之后,它将当前脚本的父目录添加到系统路径中,以便可以从该目录下的utils目录导入一些自定义函数和模块。
sys.path.append(
    os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
# 最后,它从utils目录中导入了一些自定义模块和函数,包括数据处理函数(create_prompt_dataset),打印和设备转换函数(print_rank_0和to_device),
# 模型保存函数(save_hf_format),随机种子设置函数(set_random_seed),求平均函数(get_all_reduce_mean),
# 获取优化器参数组的函数(get_optimizer_grouped_parameters),保存和加载模型的函数(save_zero_three_model和load_hf_tokenizer),
# 以及创建模型和处理模型的函数(create_hf_model)。这些函数在脚本中的后续部分都将被使用。
from utils.data.data_utils import create_prompt_dataset
from utils.utils import print_rank_0, to_device, save_hf_format, set_random_seed, get_all_reduce_mean, get_optimizer_grouped_parameters, save_zero_three_model, load_hf_tokenizer
from utils.ds_utils import get_train_ds_config
from utils.module.lora import convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters
from utils.model.model_utils import create_hf_model

create_prompt_dataset解析

create_prompt_dataset这个函数实际上直接或者间接的用到了utils/dataraw_dataset.pydata_utils.py,为了搞清楚这个函数,我们需要对这两个文件做一个解析。

首先解析一下raw_dataset.py。这里先定义了一个PromptRawDataset类:

# DeepSpeed Team
from datasets import load_dataset
from torch.utils.data import Subset
import re


# 这段代码定义了一个名为PromptRawDataset的类,这个类是一个模板类,用于处理和组织模型输入数据的格式。
# 如果有新的数据集需要进行处理,可以继承这个类并实现相应的方法来确保数据的统一格式和接口。
class PromptRawDataset(object):
		# 首先,这个类的构造函数__init__接收四个参数:output_path(输出路径),seed(随机种子),
		# local_rank(本地等级)和dataset_name(数据集名称)。
		# 在构造函数中,如果数据集名称不是'local/jsonfile',
		# 那么会使用Hugging Face的datasets库的load_dataset函数来加载数据集。


    def __init__(self, output_path, seed, local_rank, dataset_name):
        self.output_path = output_path
        self.seed = seed
        self.local_rank = local_rank
        if not dataset_name == 'local/jsonfile':
            self.raw_datasets = load_dataset(dataset_name)
    # 然后,这个类定义了一些方法,这些方法在默认情况下并没有实现(只是返回None或者空操作),
    # 这是因为这个类只是一个模板类,这些方法需要在实际使用时在子类中具体实现。
    def get_train_data(self):     # 获取训练数据
        return

    def get_eval_data(self):     # 获取评估数据
        return

    # The prompt should be in the format of: " Human: " + actual_prompt_sentence + " Assistant:"
    # get_prompt方法用于获取样本中的prompt(提示,这是模型的输入)。
    def get_prompt(self, sample):
        return

    # The chosen response should be in the format of: " " + actual_response_sentence
    # get_chosen方法用于获取样本中的chosen(已选的回应,这是模型需要生成的目标输出)。
    def get_chosen(self, sample):
        return

    # The rejected response should be in the format of: " " + actual_response_sentence
    # If the dataset does not have rejected response, return None
    # get_rejected方法用于获取样本中的rejected(被拒绝的回应,这可能用于一些特定的训练场景,比如在对抗训练中,但如果数据集中没有这样的数据,可以返回None)。
    def get_rejected(self, sample):
        return
    # 获取样本中的prompt和chosen
    def get_prompt_and_chosen(self, sample):
        return
    # 获取样本中的prompt和rejected
    def get_prompt_and_rejected(self, sample):
        return

接下来就是每个具体数据集的定义,我这里以 OpenaiWebgptcomparisonsDataset 为例解析一下,剩下的读者又需要可以自行理解:

# English dataset
# 这个类OpenaiWebgptcomparisonsDataset继承自PromptRawDataset类,
# 针对"openai/webgpt_comparisons"这个具体的数据集进行了特化。
class OpenaiWebgptcomparisonsDataset(PromptRawDataset):
    # 在构造函数__init__中,调用了父类的构造函数,并设定了dataset_name和dataset_name_clean两个属性,
    # 分别为"openai/webgpt_comparisons"和"openai_webgpt_comparisons"。
    def __init__(self, output_path, seed, local_rank, dataset_name):
        super().__init__(output_path, seed, local_rank, dataset_name)
        self.dataset_name = "openai/webgpt_comparisons"
        self.dataset_name_clean = "openai_webgpt_comparisons"
    # get_train_data和get_eval_data方法分别从raw_datasets中获取训练数据和测试数据。
    # 它们与之前的DahoasRmstaticDataset类不同之处在于,它们使用get_raw_dataset_split_index
    # 方法对训练数据进行了划分,将其划分为训练集和验证集,并返回对应的数据子集。
    def get_train_data(self):
        from .data_utils import get_raw_dataset_split_index
        dataset = self.raw_datasets["train"]
        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
                                            self.dataset_name_clean,
                                            self.seed, "train_eval", "9,1", 0,
                                            len(dataset))
        dataset = Subset(dataset, index)
        return dataset

    def get_eval_data(self):
        from .data_utils import get_raw_dataset_split_index
        dataset = self.raw_datasets["train"]
        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
                                            self.dataset_name_clean,
                                            self.seed, "train_eval", "9,1", 1,
                                            len(dataset))
        dataset = Subset(dataset, index)
        return dataset
     # get_prompt,get_chosen和get_rejected方法分别从样本中获取提示,已选回应和被拒绝的回应。
     # 这里假定样本是一个字典,其中包含了名为'question','score_0','score_1','answer_0'和'answer_1'的字段。
		# 其中,'question'字段是一个字典,包含了'full_text'字段。这个字段包含了人类提出的问题。
		# 'score_0'和'score_1'字段是字符串,表示对'answer_0'和'answer_1'的评分。
		# 如果'score_0'大于等于'score_1',那么'answer_0'就是已选回应,'answer_1'就是被拒绝的回应,反之亦然。
		# 在获取已选回应和被拒绝的回应时,还对回应进行了处理,
		# 去除了所有形如"[...]"或"(...)"的文本,然后在回应前添加了一个空格。
    def get_prompt(self, sample):
        return " Human: " + sample['question']['full_text'] + " Assistant:"

    def get_chosen(self, sample):
        if float(sample['score_0']) >= float(sample['score_1']):
            response = sample['answer_0']
        else:
            response = sample['answer_1']
        # This data has citation square brackets and numbers (e.g., "[1]").
        # Right now we are not doing browser-assisted finetuning, thus we
        # remove these citations to avoid confusing the model.
        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
        return " " + response

    def get_rejected(self, sample):
        if float(sample['score_0']) < float(sample['score_1']):
            response = sample['answer_0']
        else:
            response = sample['answer_1']
        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
        return " " + response
    # get_prompt_and_chosen和get_prompt_and_rejected方法则分别返回样本中的'prompt'和'chosen',
    # 以及'prompt'和'rejected'的组合。这两个方法的返回值可以直接作为模型的输入和目标输出。
    # 在返回这两个组合时,也进行了类似的处理,去除了所有形如"[...]"或"(...)"的文本。
    def get_prompt_and_chosen(self, sample):
        if float(sample['score_0']) >= float(sample['score_1']):
            response = sample['answer_0']
        else:
            response = sample['answer_1']
        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
        return " Human: " + sample['question'][
            'full_text'] + " Assistant: " + response

    def get_prompt_and_rejected(self, sample):
        if float(sample['score_0']) < float(sample['score_1']):
            response = sample['answer_0']
        else:
            response = sample['answer_1']
        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
        return " Human: " + sample['question'][
            'full_text'] + " Assistant: " + response

接着解析 data_utils.py

# DeepSpeed Team
"""
Part of the code was adopted from https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/data/dataset_utils.py
"""
import torch
from torch.utils.data import Dataset, Subset, ConcatDataset
from torch.nn.utils.rnn import pad_sequence
import torch.nn.functional as F
from datasets import load_dataset
import numpy as np
import os
import hashlib # Python的内置库,提供了一系列散列函数,如MD5、SHA1等。
from itertools import chain # Python的内置库,提供了一系列用于操作迭代器的函数。
from . import raw_datasets

# 这段代码定义了一个名为get_raw_dataset的函数,其主要作用是根据传入的数据集名称dataset_name
# 返回一个适当的PromptRawDataset子类的实例。
# 这个函数有四个参数:dataset_name,output_path,seed和local_rank。
def get_raw_dataset(dataset_name, output_path, seed, local_rank):
    # 在函数中,根据dataset_name的具体值来创建不同的PromptRawDataset子类的实例。
    # 例如,如果dataset_name为"Dahoas/rm-static",那么就创建一个DahoasRmstaticDataset的实例;
    # 如果dataset_name为"Dahoas/full-hh-rlhf",那么就创建一个DahoasFullhhrlhfDataset的实例,以此类推。
    if "Dahoas/rm-static" in dataset_name:
        return raw_datasets.DahoasRmstaticDataset(output_path, seed,
                                                  local_rank, dataset_name)
    elif "Dahoas/full-hh-rlhf" in dataset_name:
        return raw_datasets.DahoasFullhhrlhfDataset(output_path, seed,
                                                    local_rank, dataset_name)
    elif "Dahoas/synthetic-instruct-gptj-pairwise" in dataset_name:
        return raw_datasets.DahoasSyntheticinstructgptjpairwiseDataset(
            output_path, seed, local_rank, dataset_name)
    elif "yitingxie/rlhf-reward-datasets" in dataset_name:
        return raw_datasets.YitingxieRlhfrewarddatasetsDataset(
            output_path, seed, local_rank, dataset_name)
    elif "openai/webgpt_comparisons" in dataset_name:
        return raw_datasets.OpenaiWebgptcomparisonsDataset(
            output_path, seed, local_rank, dataset_name)
    elif "stanfordnlp/SHP" in dataset_name:
        return raw_datasets.StanfordnlpSHPDataset(output_path, seed,
                                                  local_rank, dataset_name)
    elif "pvduy/sharegpt_alpaca_oa_vicuna_format" in dataset_name:
        return raw_datasets.PvduySharegptalpacaoavicunaformatDataset(
            output_path, seed, local_rank, dataset_name)
    elif "wangrui6/Zhihu-KOL" in dataset_name:
        return raw_datasets.Wangrui6ZhihuKOLDataset(output_path, seed,
                                                    local_rank, dataset_name)
    elif "Cohere/miracl-zh-queries-22-12" in dataset_name:
        return raw_datasets.CohereMiraclzhqueries2212Dataset(
            output_path, seed, local_rank, dataset_name)
    elif "Hello-SimpleAI/HC3-Chinese" in dataset_name:
        return raw_datasets.HelloSimpleAIHC3ChineseDataset(
            output_path, seed, local_rank, dataset_name)
    elif "mkqa-Chinese" in dataset_name:
        return raw_datasets.MkqaChineseDataset(output_path, seed, local_rank,
                                               "mkqa")
    elif "mkqa-Japanese" in dataset_name:
        return raw_datasets.MkqaJapaneseDataset(output_path, seed, local_rank,
                                                "mkqa")
    elif "Cohere/miracl-ja-queries-22-12" in dataset_name:
        return raw_datasets.CohereMiracljaqueries2212Dataset(
            output_path, seed, local_rank, dataset_name)
    elif "lmqg/qg_jaquad" in dataset_name:
        return raw_datasets.LmqgQgjaquadDataset(output_path, seed, local_rank,
                                                dataset_name)
    elif "lmqg/qag_jaquad" in dataset_name:
        return raw_datasets.LmqgQagjaquadDataset(output_path, seed, local_rank,
                                                 dataset_name)
    # 如果dataset_name是"local/jsonfile",则会检查在路径chat_path + '/data/train.json'
    # 和chat_path + '/data/eval.json'下是否存在文件。如果存在,则创建一个LocalJsonFileDataset的实例;
    # 如果不存在,则抛出一个RuntimeError异常。
    elif "local/jsonfile" in dataset_name:
        chat_path = os.path.abspath(
            os.path.join(os.path.dirname(__file__), os.path.pardir,
                         os.path.pardir, os.path.pardir))
        if not (os.path.isfile(chat_path + '/data/train.json')
                and os.path.isfile(chat_path + '/data/eval.json')):
            raise RuntimeError(
                f"Please check both the train.json and eval.json files in your applications/DeepSpeed-Chat/data directory."
            )
        return raw_datasets.LocalJsonFileDataset(output_path, seed, local_rank,
                                                 dataset_name, chat_path)
    else:
        # 如果dataset_name没有在以上的所有条件中匹配到,那么函数也会抛出一个RuntimeError异常,表示没有为这个数据集的配置。
        raise RuntimeError(
            f"We do not have configs for dataset {dataset_name}, but you can add it by yourself in raw_datasets.py."
        )

再看下 get_shuffle_idx函数:

# 这个函数的作用是生成一个大小为size的乱序索引数组,它接受两个参数:seed和size。
def get_shuffle_idx(seed, size):
    np_rng = np.random.RandomState(seed=seed) # 创建一个NumPy的随机状态生成器对象np_rng,seed是随机种子,确定了随机数的生成序列。
    dtype_ = np.uint32 # 设置其为NumPy的uint32类型,这是一个无符号32位整数类型。
    if size >= (np.iinfo(np.uint32).max - 1): # 如果size大于或等于uint32的最大值减一,这里减一是为了防止可能的溢出。
        dtype_ = np.int64 # 则将dtype_改为int64,这是一个64位的有符号整数类型。
    shuffle_idx = np.arange(start=0, stop=size, step=1, dtype=dtype_) # 创建一个由0开始,步长为1,到size结束(不包含size),并且数据类型为dtype_的等差数列,将其赋值给shuffle_idx。
    np_rng.shuffle(shuffle_idx) # 使用np_rng随机状态生成器对shuffle_idx进行随机排列,这样就打乱了shuffle_idx的顺序。
    return shuffle_idx # 返回乱序后的shuffle_idx。

接着解析get_raw_dataset_split_index函数:

# 这个函数主要是根据提供的参数分割数据集,并生成一个分割索引。
# 它首先检查索引文件是否存在,如果不存在,则生成分割索引,并保存到文件。
# 然后,它从文件中加载索引,并返回索引列表。
# 它接受8个参数:local_rank、输出路径、数据集名称、种子、split_name、data_split、split_index和数据大小。
def get_raw_dataset_split_index(local_rank, output_path, dataset_name, seed,
                                split_name, data_split, split_index,
                                data_size):
    # 根据给定参数生成索引文件名。
    index_file_name = f"{output_path}/{dataset_name}_seed{seed}_{split_name}_{data_split}_{split_index}.npy"
    # reindex each time when using local jsonfile since it's more likely to get modified
    # 如果索引文件不存在,或者数据集名为'jsonfile',则执行下面的操作。
    if (not os.path.isfile(index_file_name)) or (dataset_name == 'jsonfile'):
        splits = [float(s) for s in data_split.split(',')] # 将data_split(以逗号分隔的字符串)分割为浮点数列表。
        splits_sum = sum(splits) # 计算所有分割的总和。
        splits = [split / splits_sum for split in splits] # 将每个分割值归一化,使其和为1。
        splits_index = [0] # 初始化分割索引列表,起始值为0
        for index, split in enumerate(splits): # 遍历分割列表。
            # 将新的索引(当前索引加上归一化的分割值与数据大小的乘积)添加到索引列表。
            splits_index.append(splits_index[index] +
                                int(round(split * float(data_size))))
        # 计算最后一个索引与数据大小的差值。
        diff = splits_index[-1] - data_size
        # 遍历除第一个外的所有索引。
        for index in range(1, len(splits_index)):
            # 将差值从每个索引中减去,以确保最后一个索引等于数据大小。
            splits_index[index] -= diff
        # 断言最后一个索引等于数据大小
        assert splits_index[-1] == data_size
        
        # 生成一个乱序的索引。
        shuffle_idx = get_shuffle_idx(seed, data_size)
        for split_i in range(len(splits)): # 遍历每个分割。
            # 根据给定参数生成乱序索引分割文件名。
            shuffle_idx_split_file_name = f"{output_path}/{dataset_name}_seed{seed}_{split_name}_{data_split}_{split_i}.npy" 
            # 提取乱序索引的一个分割。
            shuffle_idx_split = shuffle_idx[
                splits_index[split_i]:splits_index[split_i + 1]]
            # 将乱序索引分割保存到文件。
            np.save(shuffle_idx_split_file_name,
                    shuffle_idx_split,
                    allow_pickle=True)
    # 加载索引文件。
    index = np.load(index_file_name, allow_pickle=True)
    # 将索引数组转换为列表并返回。
    return index.tolist()

接下来解析一下继承自DatasetPromptDataset类:

# 这是一个自定义的PromptDataset类,它继承自torch.utils.data.Dataset。
# 这是一个数据集类,通常被用于PyTorch中数据的加载和预处理。
class PromptDataset(Dataset):
    # 类的构造函数,它接受五个参数:prompt_dataset、chosen_dataset、reject_dataset、pad_token_id和train_phase。
    def __init__(self, prompt_dataset, chosen_dataset, reject_dataset,
                 pad_token_id, train_phase) -> None:
        super().__init__() # 调用父类torch.utils.data.Dataset的构造函数。
        self.prompt_dataset = prompt_dataset # 将传入的参数赋值给类的成员变量。
        self.chosen_dataset = chosen_dataset
        self.reject_dataset = reject_dataset
        self.pad_token_id = pad_token_id
        self.train_phase = train_phase
   
    def __len__(self): # 定义类的__len__方法,它返回数据集的长度。这是PyTorch数据集的必要方法。
        length = len(self.chosen_dataset) # 初始设定数据集长度为chosen_dataset的长度。
        if self.train_phase == 3: 
            length = len(self.prompt_dataset) # 如果训练阶段为3,则数据集长度设定为prompt_dataset的长度。
        return length # 返回计算得出的数据集长度。
    
    # 定义类的__getitem__方法,它接受一个参数idx,返回索引idx处的数据。这是PyTorch数据集的必要方法。
    def __getitem__(self, idx): 
        # 如果训练阶段为1,则返回一个字典,包含input_ids、attention_mask和labels,它们都来自chosen_dataset的索引idx处。
        if self.train_phase == 1:
            return {
                "input_ids": self.chosen_dataset[idx]["input_ids"],
                "attention_mask": self.chosen_dataset[idx]["attention_mask"],
                "labels": self.chosen_dataset[idx]["input_ids"]
            }
        # 如果训练阶段为2,则返回来自chosen_dataset和reject_dataset的input_ids和attention_mask。
        elif self.train_phase == 2:
            return self.chosen_dataset[idx]["input_ids"], self.chosen_dataset[idx]["attention_mask"], \
                self.reject_dataset[idx]["input_ids"], self.reject_dataset[idx]["attention_mask"]
        # 如果训练阶段为3,则返回来自prompt_dataset的input_ids、attention_mask和pad_token_id
        elif self.train_phase == 3:
            return self.prompt_dataset[idx]["input_ids"],self.prompt_dataset[idx]["attention_mask"], \
                self.pad_token_id

接着再解析一下create_dataset_split函数:

# 这是一个名为create_dataset_split的函数,它的功能是根据给定的训练阶段(train_phase),创建并返回相应的数据集分割。
# 具体来说,它为每个训练阶段生成不同的数据集列表,并将它们放入PromptDataset对象中。
# 函数接受6个参数:当前数据集(current_dataset)、原始数据集(raw_dataset)、训练阶段(train_phase)、
# 分词器(tokenizer)、会话结束标记(end_of_conversation_token)和最大序列长度(max_seq_len)。
def create_dataset_split(current_dataset, raw_dataset, train_phase, tokenizer,
                         end_of_conversation_token, max_seq_len):
    # 创建三个空的列表,用于存储对话提示(prompt_dataset)、选定的对话(chosen_dataset)和被拒绝的对话(reject_dataset)。
    prompt_dataset = []
    chosen_dataset = []
    reject_dataset = []
    # 如果训练阶段为1,则将接受的对话进行分词并添加到chosen_dataset中。
    if train_phase == 1:
        # 遍历当前数据集。
        for i, tmp_data in enumerate(current_dataset):
            # tokenize the text
            # 从原始数据集中获取对话提示和接受的对话。
            chosen_sentence = raw_dataset.get_prompt_and_chosen(
                tmp_data)  # the accept response
            # 如果接受的对话不为空,则将其分词并添加到chosen_dataset中。
            if chosen_sentence is not None:
                chosen_sentence += end_of_conversation_token
                chosen_token = tokenizer(chosen_sentence,
                                         max_length=max_seq_len,
                                         padding="max_length",
                                         truncation=True,
                                         return_tensors="pt")
                chosen_token["input_ids"] = chosen_token["input_ids"].squeeze(
                    0)
                chosen_token["attention_mask"] = chosen_token[
                    "attention_mask"].squeeze(0)
                chosen_dataset.append(chosen_token)
    # 如果训练阶段为2,则将接受和被拒绝的对话都进行分词并分别添加到chosen_dataset和reject_dataset中。
    elif train_phase == 2:
        for i, tmp_data in enumerate(current_dataset):
            # tokenize the text
            chosen_sentence = raw_dataset.get_prompt_and_chosen(
                tmp_data)  # the accept response
            reject_sentence = raw_dataset.get_prompt_and_rejected(
                tmp_data)  # the accept response
            if chosen_sentence is not None and reject_sentence is not None:
                chosen_sentence += end_of_conversation_token  # the accept response
                reject_sentence += end_of_conversation_token
                chosen_token = tokenizer(chosen_sentence,
                                         max_length=max_seq_len,
                                         padding="max_length",
                                         truncation=True,
                                         return_tensors="pt")
                reject_token = tokenizer(reject_sentence,
                                         max_length=max_seq_len,
                                         padding="max_length",
                                         truncation=True,
                                         return_tensors="pt")
                chosen_token["input_ids"] = chosen_token["input_ids"]
                chosen_token["attention_mask"] = chosen_token["attention_mask"]
                chosen_dataset.append(chosen_token)

                reject_token["input_ids"] = reject_token["input_ids"]
                reject_token["attention_mask"] = reject_token["attention_mask"]
                reject_dataset.append(reject_token)
    # 如果训练阶段为3,则将对话提示进行分词并添加到prompt_dataset中。
    elif train_phase == 3:
        for i, tmp_data in enumerate(current_dataset):
            # tokenize the text
            prompt = raw_dataset.get_prompt(tmp_data)
            if prompt is not None:
                prompt_token = tokenizer(prompt, return_tensors="pt")
                prompt_token["input_ids"] = prompt_token["input_ids"]
                prompt_token["attention_mask"] = prompt_token["attention_mask"]
                for key_word in ["input_ids", "attention_mask"]:
                    length = prompt_token[key_word].size()[-1]
                    if length > max_seq_len:
                        y = prompt_token[key_word].squeeze(0)[length -
                                                              (max_seq_len -
                                                               1):].flip(0)
                    else:
                        y = prompt_token[key_word].squeeze(0).flip(0)
                    prompt_token[key_word] = y
                prompt_dataset.append(prompt_token)
    # 创建一个新的PromptDataset对象,并返回。这个对象包含了对话提示、接受的对话和被拒绝的对话的数据集,以及分词器的填充标记ID和训练阶段。
    return PromptDataset(prompt_dataset, chosen_dataset, reject_dataset,
                         tokenizer.pad_token_id, train_phase)

接下来再解析一下create_dataset这个函数:

这段代码定义了一个函数 create_dataset,主要负责创建训练数据集和评估数据集,具体的功能细节如下:
def create_dataset(local_rank, dataset_name, data_split, output_path,
                   train_phase, seed, tokenizer, end_of_conversation_token,
                   max_seq_len):
    # 调用 get_raw_dataset 函数,该函数根据提供的数据集名称、输出路径、随机种子和local_rank等参数,从各种预定义的数据集中获取所需的原始数据集。
    raw_dataset = get_raw_dataset(dataset_name, output_path, seed, local_rank)
    train_dataset = raw_dataset.get_train_data() # 从原始数据集中获取训练数据。
    #  获取训练数据集的索引,涉及数据的分割。
    train_index = get_raw_dataset_split_index(local_rank, output_path,
                                              raw_dataset.dataset_name_clean,
                                              seed, "train", data_split,
                                              train_phase - 1,
                                              len(train_dataset))
    # 根据上一步获取的索引,创建训练数据的子集。
    train_dataset = Subset(train_dataset, train_index)
    # 调用 create_dataset_split 函数对上一步获得的数据子集进行进一步处理,
    # 这可能包括对文本的标记化(tokenization),并且创建一个PromptDataset 对象。
    train_dataset = create_dataset_split(train_dataset, raw_dataset,
                                         train_phase, tokenizer,
                                         end_of_conversation_token,
                                         max_seq_len)
    # 是用于创建评估数据集的,步骤与训练数据集的创建基本相同。
    eval_dataset = raw_dataset.get_eval_data()
    eval_index = get_raw_dataset_split_index(local_rank, output_path,
                                             raw_dataset.dataset_name_clean,
                                             seed, "eval",
                                             data_split, train_phase - 1,
                                             len(eval_dataset))
    eval_dataset = Subset(eval_dataset, eval_index)
    eval_dataset = create_dataset_split(eval_dataset, raw_dataset, train_phase,
                                        tokenizer, end_of_conversation_token,
                                        max_seq_len)
    return train_dataset, eval_dataset

这里需要额外注意一点,create_dataset_split 这个函数传入的参数是 train_phase,而不是像 get_raw_dataset_split_index 函数那样传入 train_phase-1,这是为什么?这是因为train_phase用来标识在模型训练过程中的哪个阶段,是一个直接的指示符,与数据分割或索引无关。它用来在create_dataset_split函数中确定当前处于哪个训练阶段,并根据不同的训练阶段对数据进行不同的处理。而在调用get_raw_dataset_split_index函数时,传入的是train_phase - 1,这是因为在这个函数中,我们需要根据当前训练阶段的前一个阶段(由train_phase - 1表示)的数据分割或索引情况,来决定如何对当前阶段的数据进行分割或索引。create_dataset_split主要关注如何根据训练阶段对数据进行处理,而get_raw_dataset_split_index主要关注如何根据前一个训练阶段的数据分割或索引情况,对当前阶段的数据进行分割或索引。

回收本节的开头,我们解析create_prompt_dataset函数:

# 这个函数的主要目的是创建一个包含训练和评估数据集的“提示”数据集,并将这两个数据集保存在指定的文件中。具体来说:
def create_prompt_dataset(local_rank,
                          data_path,
                          data_split,
                          output_path,
                          train_phase,
                          seed,
                          tokenizer,
                          max_seq_len,
                          end_of_conversation_token="<|endoftext|>",
                          sft_only_data_path=[],
                          reload=False):
    """
    Creates the prompt dataset
    """
    # os.makedirs(output_path, exist_ok=True): 创建输出目录,如果目录已经存在则不会引发异常。
    os.makedirs(output_path, exist_ok=True)
    # 构造文件名,这个文件名包含了很多有关数据集和模型的信息,如数据路径、数据分割、训练阶段、
    # 随机种子、tokenizer的名称、最大序列长度等。然后将这个文件名哈希化,以避免文件名过长。
    fname = "_".join(data_path)
    sft_cache_key = "_".join(sft_only_data_path)
    tokenizer_name = tokenizer.init_kwargs["name_or_path"].replace("/", "_")
    fname = f"{fname}_split{data_split}_phase{train_phase}_seed{seed}_tokenizer{tokenizer_name}_seqlen{max_seq_len}_sft{sft_cache_key}"
    fname = "_".join(fname.split("/"))
    fname = hashlib.sha256(fname.encode()).hexdigest(
    )  # hash the file name to avoid too long file name
    # 构造训练数据集和评估数据集的文件路径。
    train_fname = f"{output_path}/traindata_{fname}.pt"
    eval_fname = f"{output_path}/evaldata_{fname}.pt"

    # 检查训练数据集和评估数据集的文件是否都已经存在,如果存在,则表示缓存已经找到,否则表示需要创建缓存。
    cache_found = os.path.isfile(train_fname) and os.path.isfile(eval_fname)
    buf_create_cache = torch.ByteTensor([not cache_found]).cuda()
    torch.distributed.all_reduce(buf_create_cache)
    
    # 如果当前进程是主进程(local_rank <= 0)并且需要创建缓存,就执行以下操作。
    if local_rank <= 0 and (buf_create_cache.item() != 0 or reload):
        # 如果只有一个数据集,直接调用create_dataset函数创建训练数据集和评估数据集。
        if len(data_path) == 1:  # Single dataset.
            train_dataset, eval_dataset = create_dataset(
                local_rank, data_path[0], data_split, output_path, train_phase,
                seed, tokenizer, end_of_conversation_token, max_seq_len)
        else:  # Blending datasets.
            # 如果有多个数据集,对每个数据集都调用create_dataset函数,并把得到的训练数据集和评估数据集添加到对应的列表中,

            train_datasets = []
            eval_datasets = []
            train_size = 0
            eval_size = 0
            for d_path in data_path:
                train_dataset, eval_dataset = create_dataset(
                    local_rank, d_path, data_split, output_path, train_phase,
                    seed, tokenizer, end_of_conversation_token, max_seq_len)
                train_datasets.append(train_dataset)
                eval_datasets.append(eval_dataset)
                train_size += len(train_dataset)
                eval_size += len(eval_dataset)
            # 然后使用ConcatDataset和Subset函数合并数据集。
            train_dataset = ConcatDataset(train_datasets)
            shuffle_idx = get_shuffle_idx(seed, train_size)
            train_dataset = Subset(train_dataset, shuffle_idx.tolist())
            eval_dataset = ConcatDataset(eval_datasets)
            shuffle_idx = get_shuffle_idx(seed, eval_size)
            eval_dataset = Subset(eval_dataset, shuffle_idx.tolist())

        # Append the SFT-only dataset if it exists, and current phase is 1(SFT).
        # 如果当前是第一阶段的训练(SFT)并且指定了仅用于SFT的数据集,那么对这些数据集执行类似的操作,
        # 然后把得到的训练数据集和评估数据集添加到原有的数据集中。
        if train_phase == 1 and sft_only_data_path:
            sft_train_datasets = []
            sft_eval_datasets = []
            sft_train_size = 0
            sft_eval_size = 0
            for sft_path in sft_only_data_path:
                sft_train_dataset, sft_eval_dataset = create_dataset(
                    local_rank,
                    sft_path,
                    "10,0,0",
                    output_path,
                    train_phase,
                    seed,
                    tokenizer,
                    end_of_conversation_token,
                    max_seq_len,
                )
                sft_train_datasets.append(sft_train_dataset)
                sft_eval_datasets.append(sft_eval_dataset)
                sft_train_size += len(sft_train_dataset)
                sft_eval_size += len(sft_eval_dataset)
            if sft_train_datasets:  # Check if sft_train_datasets is not empty
                sft_train_dataset = ConcatDataset(sft_train_datasets)
                train_dataset = ConcatDataset(
                    [train_dataset, sft_train_dataset])
                shuffle_idx = get_shuffle_idx(seed, len(train_dataset))
                train_dataset = Subset(train_dataset, shuffle_idx.tolist())
            if sft_eval_datasets:  # Check if sft_eval_datasets is not empty
                sft_eval_dataset = ConcatDataset(sft_eval_datasets)
                eval_dataset = ConcatDataset([eval_dataset, sft_eval_dataset])
                shuffle_idx = get_shuffle_idx(seed, len(eval_dataset))
                eval_dataset = Subset(eval_dataset, shuffle_idx.tolist())
        # 把训练数据集和评估数据集保存到对应的文件中。
        torch.save(train_dataset, train_fname)
        torch.save(eval_dataset, eval_fname)
    # 在多进程环境中,确保所有进程都完成了数据集的保存操作。
    torch.distributed.barrier()
    return torch.load(train_fname), torch.load(eval_fname)

工具函数解析

这里解析一下下面导入的代码,在DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/utils.py这个路径下:

from utils.utils import print_rank_0, to_device, save_hf_format, set_random_seed, get_all_reduce_mean, get_optimizer_grouped_parameters, save_zero_three_model, load_hf_tokenizer

print_rank_0

# 在rank0也就是master rank打印信息,防止每个机器或GPU都打印消息造成大量重复信息
def print_rank_0(msg, rank=0):
    if rank <= 0:
        print(msg)

to_device

# 这个函数的作用是把一个数据批次移动到指定的设备上。在PyTorch中,数据默认是在CPU上的,
# 如果要在GPU上进行运算,就需要先把数据移到GPU上。这个函数通过遍历批次中的所有元素并
# 调用to(device)方法来实现这一点。如果某个元素不能被移到指定的设备上
#(例如,这个元素是一个字符串或者其他不能在GPU上使用的类型),那么就直接保留这个元素,不进行任何操作。
def to_device(batch, device):
    output = {}
    for k, v in batch.items():
        try:
            output[k] = v.to(device)
        except:
            output[k] = v
    return output

save_hf_format

# 这段代码的功能是将训练好的模型以Hugging Face格式保存,这样之后就可以使用Hugging Face库的from_pretrained方法加载了。
def save_hf_format(model, tokenizer, args, sub_folder=""):
    # used to save huggingface format, so we can use it for hf.from_pretrained
    # 首先,这行代码检查模型是否有'module'这个属性。这在PyTorch中是一个常见的模式,
    # 当模型被封装在torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel时,
    # 模型的所有属性都会被存储在'module'属性中。所以这行代码的目的是确保我们总是在原始模型上进行操作,而不是并行化的包装器。
    model_to_save = model.module if hasattr(model, 'module') else model
    # 定义模型的配置和权重的文件名。
    CONFIG_NAME = "config.json"
    WEIGHTS_NAME = "pytorch_model.bin"
    # 定义模型保存的目录路径。
    output_dir = os.path.join(args.output_dir, sub_folder)
    # 如果模型保存的目录不存在,那么创建这个目录。
    os.makedirs(output_dir, exist_ok=True)
    # 定义模型权重和配置的保存路径。
    output_model_file = os.path.join(output_dir, WEIGHTS_NAME)
    output_config_file = os.path.join(output_dir, CONFIG_NAME)
    # 获取模型的状态字典,这个字典包含了模型的所有权重。
    save_dict = model_to_save.state_dict()
    # 遍历状态字典中的所有键。这里使用list(save_dict.keys())而不是直接使用
    # save_dict.keys()是因为我们在遍历过程中会修改字典,直接遍历字典的键在Python中会导致错误。
    for key in list(save_dict.keys()):
        # if "lora" in key::如果键中包含"lora",那么删除这个键值对。
        # "lora"可能是某种特定的模型组件或参数,这里将其排除在保存的模型权重之外。
        if "lora" in key:
            del save_dict[key]
    # 将处理过的状态字典保存到文件中。
    torch.save(save_dict, output_model_file)
    # 将模型的配置保存为JSON文件。
    model_to_save.config.to_json_file(output_config_file)
    # 将分词器的词汇表保存到指定的目录。
    tokenizer.save_vocabulary(output_dir)

set_random_seed

# 这段代码的目的是为所有可能用到的随机数生成器设定种子,以确保实验结果的可复现性。
# 随机数生成器的种子是生成随机数序列的初始值,当我们用同一个种子初始化随机数生成器时,它会生成完全一样的随机数序列。
def set_random_seed(seed):
    # 首先检查种子是否是None。如果是None,那么就跳过这段代码,随机数生成器将会用一个随机的种子。
    if seed is not None: 
        set_seed(seed)# 这一行调用transformers库的set_seed的函数,将种子设定为指定的值。
        random.seed(seed) # 将Python内置的随机数生成器的种子设定为指定的值。
        np.random.seed(seed) # 将NumPy的随机数生成器的种子设定为指定的值。
        torch.manual_seed(seed) # 将PyTorch的随机数生成器的种子设定为指定的值。
        torch.cuda.manual_seed_all(seed) # 将PyTorch的所有GPU随机数生成器的种子设定为指定的值。

get_all_reduce_mean

# 这段代码是在分布式训练环境中进行平均值计算的函数,通过这段代码,
# 所有的处理器(或者叫节点)上的同一个tensor都会被加和起来,然后除以总的处理器数,得到平均值。
def get_all_reduce_mean(tensor):
    # 这行代码执行一个分布式的reduce操作。reduce操作是指所有处理器中的同一个tensor都被某种方式结合起来。
    # 在这个例子中,torch.distributed.ReduceOp.SUM表示所有处理器上的tensor将被加和起来。
    # 加和的结果会在所有处理器上都可用。
    torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
    # 这行代码将前一步得到的加和结果除以处理器的数量(也叫作 world size)。
    # 这样,tensor就变成了所有处理器上原始tensor的平均值。
    tensor = tensor / torch.distributed.get_world_size()
    # 最后,这个平均值tensor被返回。在所有处理器上,这个函数返回的tensor都是相同的,
    # 等于所有处理器上原始tensor的平均值。
    return tensor

get_optimizer_grouped_parameters

# 这段代码的作用是将模型中的参数分组以便于在优化器中使用。它将模型参数分为两组:
# 一组需要进行权重衰减(L2正则化)的参数,另一组不需要进行权重衰减的参数。
def get_optimizer_grouped_parameters(model,
                                     weight_decay,
                                     no_decay_name_list=[
                                         "bias", "LayerNorm.weight"
                                     ]):
    # 它定义了一个列表 optimizer_grouped_parameters,其中包含两个字典。每个字典都对应一个参数组,包含 "params" 和 "weight_decay" 这两个关键字。
    optimizer_grouped_parameters = [
        # 在第一个字典中,它从模型参数中选出那些名称不包含 "bias" 或 "LayerNorm.weight" 
        # 且需要求梯度的参数。这些参数在优化过程中会应用 weight_decay 作为权重衰减项。
        {
            "params": [
                p for n, p in model.named_parameters()
                if (not any(nd in n
                            for nd in no_decay_name_list) and p.requires_grad)
            ],
            "weight_decay":
            weight_decay,
        },
        # 在第二个字典中,它选出那些名称包含 "bias" 或 "LayerNorm.weight" 且需要求梯度的参数。
        # 这些参数在优化过程中不会应用权重衰减,即其 "weight_decay" 值为0。

        {
            "params": [
                p for n, p in model.named_parameters()
                if (any(nd in n
                        for nd in no_decay_name_list) and p.requires_grad)
            ],
            "weight_decay":
            0.0,
        },
    ]
    return optimizer_grouped_parameters

这种参数的分组策略是很常见的。比如在训练Transformer模型时,通常会为权重和偏置项设定不同的学习策略。这是因为权重衰减对于防止过拟合很有帮助,但对于某些参数(如偏置项或者层归一化的权重)可能会导致性能下降,因此常常会排除这些参数不进行权重衰减。

save_zero_three_model

# 这个函数的主要功能是筛选出那些在DeepSpeed Zero 3优化中被离线存储,但在当前还未获取的参数。
# 在DeepSpeed Zero 3优化中,一些模型参数在使用过后会被离线存储,以此释放GPU显存。
# 当这些参数需要再次被使用时,需要先获取到本地。
def _z3_params_to_fetch(param_list):
    # 这个条件语句判断一个参数是否是被DeepSpeed Zero 3优化过的,且其状态为"未获取"(NOT_AVAILABLE)。
    # 对于被DeepSpeed Zero 3优化过的参数,它们有一个ds_id属性和一个ds_status属性,其中ds_status表示参数的当前状态。
    return [
        p for p in param_list
        if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE
    ]

# 这个函数的主要作用是保存一个使用了DeepSpeed Zero优化(可能为stage 3)的模型。
# DeepSpeed的Zero优化技术是为了解决模型参数、优化器状态和梯度等内存占用问题,
# 通过这种方式,可以训练比当前GPU内存更大的模型。
def save_zero_three_model(model_ema, global_rank, save_dir, zero_stage=0):
    # 首先,检查输入的zero_stage是否为3,确定是否使用了DeepSpeed Zero阶段3优化。
    zero_stage_3 = (zero_stage == 3)
    # 然后,确保保存模型的目录存在。
    os.makedirs(save_dir, exist_ok=True)
    # 定义模型权重保存的完整路径。
    WEIGHTS_NAME = "pytorch_model.bin"
    output_model_file = os.path.join(save_dir, WEIGHTS_NAME)

    # 如果模型是被包裹在其它结构(如DataParallel或DistributedDataParallel)中的,我们需要取出真实的模型实例。
    model_to_save = model_ema.module if hasattr(model_ema,
                                                'module') else model_ema
    if not zero_stage_3:
    	 # 如果没有使用Zero阶段3优化,直接使用PyTorch的torch.save函数保存模型状态。 
        if global_rank == 0:
            torch.save(model_to_save.state_dict(), output_model_file)
    else:
        # 如果使用了Zero阶段3优化,因为模型的部分参数和优化器状态在不同的设备上,所以需要先将它们收集起来。
        output_state_dict = {}
        for k, v in model_to_save.named_parameters():

            if hasattr(v, 'ds_id'):
                # deepspeed.zero.GatheredParameters是DeepSpeed提供的一个上下文管理器,
                # 它可以将分布在多个设备上的参数收集到一起。这部分参数保存在CPU上。
                with deepspeed.zero.GatheredParameters(_z3_params_to_fetch([v
                                                                            ]),
                                                       enabled=zero_stage_3):
                    v_p = v.data.cpu()
            else:
                v_p = v.cpu()
            # 然后,将收集好的参数(并且不包含“lora”关键字的参数)添加到输出状态字典中。
            if global_rank == 0 and "lora" not in k:
                output_state_dict[k] = v_p
        # 最后,再使用torch.save函数保存模型状态。
        if global_rank == 0:
            torch.save(output_state_dict, output_model_file)
        # 同时为了节省内存,使用del关键字删除了存储参数的字典。
        del output_state_dict

load_hf_tokenizer

# 这个函数的目标是加载一个Hugging Face tokenizer,这个tokenizer是用来将文本数据转化为模型可以接受的形式的。
# model_name_or_path: 这个参数可以是一个模型名字或者一个本地路径,指示从哪里获取tokenizer。
# fast_tokenizer=True: 这个参数指示是否使用Hugging Face的"fast" tokenizer。"fast" tokenizer是用Rust编写的,运行速度更快。
def load_hf_tokenizer(model_name_or_path, fast_tokenizer=True):
    # 这行代码检查model_name_or_path是否是一个存在的路径。如果是,那么就尝试从本地文件系统中加载tokenizer。
    if os.path.exists(model_name_or_path):
        # Locally tokenizer loading has some issue, so we need to force download
        # 这行代码拼接路径,尝试找到模型的config文件。
        model_json = os.path.join(model_name_or_path, "config.json")
        # 这行代码检查config文件是否存在。如果存在,那么就尝试从config文件中获取模型名字。
        if os.path.exists(model_json):
            model_json_file = json.load(open(model_json)) # 这行代码加载config文件。
            model_name = model_json_file["_name_or_path"] # 这行代码从config文件中获取模型名字。
            tokenizer = AutoTokenizer.from_pretrained(model_name,
                                                      fast_tokenizer=True) # 这行代码使用模型名字从Hugging Face的模型库中下载并加载tokenizer。
    else: # 如果model_name_or_path不是一个存在的路径,那么就认为它是一个模型名字。
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
                                                  fast_tokenizer=True) # 这行代码使用模型名字从Hugging Face的模型库中下载并加载tokenizer。
    return tokenizer

再解析一下DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/ds_utils.py导入的get_train_ds_config函数:

# DeepSpeed Team
GLOBAL_BATCH_SIZE = 32
MICRO_BATCH_SIZE = 4

# 这个函数的主要目标是生成一个字典,用来配置DeepSpeed训练模型的各种参数。
# DeepSpeed是一个用于加速大规模模型训练的库,其中包括一种叫做ZeRO(Zero Redundancy Optimizer)的优化技术,
# 可以大大减少模型训练过程中的显存和内存占用。
def get_train_ds_config(offload,
                        stage=2,
                        enable_hybrid_engine=False,
                        inference_tp_size=1,
                        release_inference_cache=False,
                        pin_parameters=True,
                        tp_gather_partition_size=8,
                        max_out_tokens=512):
    # 这行代码首先定义了一个变量device。如果参数offload为True,那么就将device设置为"cpu",
    # 表示模型的参数和优化器状态将被"offload"到CPU内存中,以减少GPU显存的占用。
    # 如果offload为False,那么device就被设置为"none",表示不进行offload。
    device = "cpu" if offload else "none"
    # 这行代码定义了一个字典zero_opt_dict,包含了ZeRO优化器的各种配置参数。
    zero_opt_dict = {
        "stage": stage, # ZeRO的阶段,可以为0、1、2、3,阶段越高,显存优化效果越好,但计算复杂度也越高。
        "offload_param": { # offload_param和offload_optimizer这两个字段控制是否将模型的参数和优化器状态offload到CPU内存。
            "device": device
        },
        "offload_optimizer": {
            "device": device
        }, 
        # 这三个字段只有在ZeRO stage 3时有效,用于控制参数offloading的细节。
        "stage3_param_persistence_threshold": 1e4,
        "stage3_max_live_parameters": 3e7,
        "stage3_prefetch_bucket_size": 3e7,
        # 这个字段控制是否使用显存更高效的线性层实现。由于当前的显存高效实现可能会导致数值不稳定,所以这里默认设为False。
        "memory_efficient_linear": False
    }
    # 返回一个配置字典,这个字典包含了训练模型的各种参数。
    return {
        "train_batch_size": GLOBAL_BATCH_SIZE, # 全局的批次大小
        "train_micro_batch_size_per_gpu": MICRO_BATCH_SIZE, # 每个GPU的微批次大小。
        "steps_per_print": 10, # 每隔10步打印一次训练信息。
        "zero_optimization": zero_opt_dict, # 之前定义的ZeRO配置字典。
        "fp16": { # 控制是否使用混合精度训练,以及loss scaling的窗口大小。
            "enabled": True,
            "loss_scale_window": 100
        },
        "gradient_clipping": 1.0, # 梯度裁剪的阈值。
        "prescale_gradients": False, # 是否在梯度累计之前就进行梯度缩放,通常用于防止梯度下溢。
        "wall_clock_breakdown": False, # 是否进行每步训练时间的详细分析。
        "hybrid_engine": {
            "enabled": enable_hybrid_engine,
            "max_out_tokens": max_out_tokens,
            "inference_tp_size": inference_tp_size,
            "release_inference_cache": release_inference_cache,
            "pin_parameters": pin_parameters,
            "tp_gather_partition_size": tp_gather_partition_size,
        }
    }

再解析一下DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/module/lora.py导入的convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters 这几个函数。

在解析这几个函数之前需要先解析一下 LinearLayer_LoRA 这个类:

# LinearLayer_LoRA 是一个实现了 LoRA(Low-Rank Adaptation)策略的 PyTorch 模块。
# LoRA 是一种在大型预训练模型(如 GPT-3 或者 BERT)微调中使用的策略,可以提高微调的效果。
# 这个类是一个简单的 LoRA 实现,只支持线性层。
class LinearLayer_LoRA(nn.Module):
    # an simple implementation of LoRA
    # for now only support Linear Layer
    # 初始化函数,定义了一些参数,如权重(weight)、LoRA 维度(lora_dim)、
    # LoRA 缩放因子(lora_scaling)、LoRA dropout (lora_droppout)以及偏置(bias)。
    def __init__(self,
                 weight,
                 lora_dim=0,
                 lora_scaling=1,
                 lora_droppout=0,
                 bias=None):
        super(LinearLayer_LoRA, self).__init__()
        self.weight = weight
        self.bias = bias
        # 在初始化过程中,首先检查 lora_dim 是否大于0,如果不是,就抛出异常。
        if lora_dim <= 0:
            raise ValueError(
                "You are training to use LoRA, whose reduced dim should be larger than 1"
            )
        # 然后尝试获取权重的形状,如果权重是在 ZeRO 阶段3中的参数,那么通过 ds_shape 属性获取形状;
        # 否则直接使用 shape 属性。
        try:
            # for zero stage 3
            rows, columns = weight.ds_shape
        except:
            rows, columns = weight.shape
        # 然后初始化 LoRA 的左右两个权重矩阵,并设置它们为可学习的参数。
        self.lora_right_weight = nn.Parameter(torch.zeros(
            columns,
            lora_dim))  # apply transpose so in forward we do not need to
        self.lora_left_weight = nn.Parameter(torch.zeros(lora_dim, rows))
        self.lora_scaling = lora_scaling / lora_dim
        # 如果 lora_droppout 大于0,则创建一个 Dropout 层;否则创建一个 Identity 层。
        if lora_droppout > 0:
            self.lora_dropout = nn.Dropout(lora_droppout)
        else:
            self.lora_dropout = nn.Identity()
        # 然后调用 reset_parameters 方法对 LoRA 权重进行初始化。
        self.reset_parameters()
        # disable the original weight gradient
        # 最后,关闭原始权重的梯度,设置 LoRA 融合标志位为 False。
        self.weight.requires_grad = False
        # fuse LoRA to the original weight
        self.fuse_lora = False
    
    def eval(self):
        self.lora_dropout.eval() # 将模型设置为评估模式,这时候 Dropout 层会停止工作。

    #   self.fuse_lora_weight()

    def train(self, mode=True):
        self.lora_dropout.train(mode) # 将模型设置为训练模式,这时候 Dropout 层会开始工作。
        # self.unfuse_lora_weight()

    def reset_parameters(self): # 初始化 LoRA 权重的方法。右权重使用 kaiming 均匀分布进行初始化,左权重初始化为全0。
        nn.init.kaiming_uniform_(self.lora_right_weight, a=math.sqrt(5))
        nn.init.zeros_(self.lora_left_weight)
    
    # fuse_lora_weight(self) 和 unfuse_lora_weight(self):
    # 这两个方法用于将 LoRA 权重融合到原始权重中,或者从原始权重中解融合。
    # 融合操作实质上是将原始权重与 LoRA 权重的乘积(缩放后)相加。
    def fuse_lora_weight(self):
        if not self.fuse_lora:
            self.weight.data += self.lora_scaling * torch.matmul(
                self.lora_left_weight.t(), self.lora_right_weight.t())
        self.fuse_lora = True

    def unfuse_lora_weight(self):
        if self.fuse_lora:
            self.weight.data -= self.lora_scaling * torch.matmul(
                self.lora_left_weight.t(), self.lora_right_weight.t())
        self.fuse_lora = False
    # 前向传播函数。如果 LoRA 权重已融合,则直接对输入进行线性变换;
    # 否则,会额外计算一个 LoRA 项,该项是输入通过 Dropout 层,然后与 LoRA 权重相乘得到的。
    def forward(self, input):
        if self.fuse_lora:
            return F.linear(input, self.weight, self.bias)
        else:
            return F.linear(
                input, self.weight,
                self.bias) + (self.lora_dropout(input) @ self.lora_right_weight
                              @ self.lora_left_weight) * self.lora_scaling

convert_linear_layer_to_lora

# 这个函数 convert_linear_layer_to_lora 是用来将模型中的线性层转换为 LoRA 层的。
# 在训练深度学习模型时,这种方法能够在保持预训练模型参数不变的同时,通过添加额外的参数来微调模型。
# convert the linear layer to LoRA
def convert_linear_layer_to_lora(model,
                                 part_module_name,
                                 lora_dim=0,
                                 lora_scaling=1,
                                 lora_droppout=0):
    repalce_name = []
    # 函数首先遍历模型中的所有模块(model.named_modules()),找出名称中包含 part_module_name 的线性层(nn.Linear),
    # 并将这些层的名称添加到 repalce_name 列表中。
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and part_module_name in name:
            repalce_name.append(name)
    # 然后,函数遍历 repalce_name 列表,使用 recursive_getattr 函数获取模型中对应名称的模块。
    # 这些模块是需要被替换成 LoRA 层的线性层。
    for name in repalce_name:
        module = recursive_getattr(model, name)
        # 对于每一个需要被替换的模块,函数创建一个 LinearLayer_LoRA 实例 tmp,
        # 并将其传输到与原始模块相同的设备和数据类型上。创建 LinearLayer_LoRA 实例时,
        # 需要传入原始模块的权重、偏置以及 LoRA 层的一些参数,如 lora_dim、lora_scaling 和 lora_droppout。
        tmp = LinearLayer_LoRA(
            module.weight, lora_dim, lora_scaling, lora_droppout,
            module.bias).to(module.weight.device).to(module.weight.dtype)
        # 创建完 LinearLayer_LoRA 实例后,函数使用 recursive_setattr 函数将原始模块替换为 LinearLayer_LoRA 实例。
        recursive_setattr(model, name, tmp)
    return model

convert_lora_to_linear_layer

# 这个函数的主要功能是筛选出那些在DeepSpeed Zero 3优化中被离线存储,但在当前还未获取的参数。
# 在DeepSpeed Zero 3优化中,一些模型参数在使用过后会被离线存储,以此释放GPU显存。
# 当这些参数需要再次被使用时,需要先获取到本地。
def _z3_params_to_fetch(param_list):
    # 这个条件语句判断一个参数是否是被DeepSpeed Zero 3优化过的,且其状态为"未获取"(NOT_AVAILABLE)。
    # 对于被DeepSpeed Zero 3优化过的参数,它们有一个ds_id属性和一个ds_status属性,其中ds_status表示参数的当前状态。
    return [
        p for p in param_list
        if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE
    ]

# 这个函数 convert_lora_to_linear_layer 是用来将模型中的 LoRA 层转换回线性层的。
# 在训练深度学习模型时,这个操作可以用于在训练完 LoRA 层后,将模型恢复到原始的状态,
# 以便进行下一步的操作,如模型的保存、加载等。
# convert the LoRA layer to linear layer
def convert_lora_to_linear_layer(model):
    repalce_name = []
    # 函数首先遍历模型中的所有模块(model.named_modules()),找出所有的 LoRA 层(LinearLayer_LoRA),
    # 并将这些层的名称添加到 repalce_name 列表中。
    for name, module in model.named_modules():
        if isinstance(module, LinearLayer_LoRA):
            repalce_name.append(name)
    # 然后,函数遍历 repalce_name 列表,使用 recursive_getattr 函数获取模型中对应名称的 LoRA 层。
    for name in repalce_name:
        module = recursive_getattr(model, name)
        # 对于每一个 LoRA 层,函数首先检查是否处于 zero stage 3(DeepSpeed 的一个特性,用于在多GPU训练中减少内存占用)。
        # 如果是,则设置 zero_stage_3 为 True。
        zero_stage_3 = hasattr(module.weight, 'ds_id')
        with deepspeed.zero.GatheredParameters(_z3_params_to_fetch([
                module.weight, module.bias, module.lora_left_weight,
                module.lora_right_weight
        ]),
                                               modifier_rank=0,
                                               enabled=zero_stage_3):
            module.fuse_lora_weight()
    return model

only_optimize_lora_parameters

# 这个函数的作用是关闭模型中除LoRA参数之外的所有参数的梯度。这意味着在训练过程中,只有LoRA参数会被优化,其他参数保持不变。
def only_optimize_lora_parameters(model):
    # turn off the gradient of all the parameters except the LoRA parameters
    # 遍历模型的所有参数。每个参数都有一个唯一的名称name和对应的参数值param。
    for name, param in model.named_parameters():
        # 查当前参数的名称是否包含lora_right_weight或lora_left_weight。
        # 这是因为在LoRA(Low-Rank Adaptation)中,只有这两种参数是需要优化的。
        if "lora_right_weight" in name or "lora_left_weight" in name:
            # param.requires_grad = True 如果参数名包含lora_right_weight或lora_left_weight,
            # 则设置参数的requires_grad属性为True,表示需要对此参数进行梯度下降优化。
            param.requires_grad = True
        else:
            param.requires_grad = False
    return model

create_hf_model

# 这段代码定义了一个名为create_hf_model的函数,该函数的作用是创建或加载一个预训练模型。该函数的主要参数包括:
# model_class:模型的类别,例如GPT-2、BERT等。
# tokenizer:用于模型的分词器。
# ds_config: DeepSpeed的配置参数。
# rlhf_training:一个标志,用来表示是否正在进行RLHF(Reinforcement Learning from Human Feedback)训练。
# disable_dropout:一个标志,用来表示是否禁用dropout。Dropout是一种防止过拟合的技术。
def create_hf_model(model_class,
                    model_name_or_path,
                    tokenizer,
                    ds_config=None,
                    rlhf_training=False,
                    disable_dropout=False):
    # 根据model_name_or_path从预训练模型获取模型配置model_config。
    model_config = AutoConfig.from_pretrained(model_name_or_path)
    # 如果disable_dropout为真,则将模型配置中的dropout设为0.0。
    if disable_dropout:
        model_config.dropout = 0.0
    # Note: dschf is defined in function scope to avoid global effects
    # https://huggingface.co/docs/transformers/main_classes/deepspeed#nontrainer-deepspeed-integration
    # 根据ds_config中的设置,创建DeepSpeed的配置对象dschf,以便进行DeepSpeed优化。
    if ds_config is not None and ds_config["zero_optimization"]["stage"] == 3:
        dschf = HfDeepSpeedConfig(ds_config)
    else:
        dschf = None
    # 根据rlhf_training的值,确定是从配置中创建模型还是从预训练模型中加载模型。如果rlhf_training为真,则根据模型配置创建模型;否则,从预训练模型加载模型。
    if rlhf_training:
        # the weight loading is handled by create critic model
        model = model_class.from_config(model_config)
    else:
        model = model_class.from_pretrained(
            model_name_or_path,
            from_tf=bool(".ckpt" in model_name_or_path),
            config=model_config)
    # 将模型的结束标记和填充标记设为分词器的结束标记id。
    model.config.end_token_id = tokenizer.eos_token_id
    model.config.pad_token_id = model.config.eos_token_id
    # 调整模型的词汇表大小,使其为8的倍数。这样做的目的是为了在某些硬件(如GPU)上提高效率。
    model.resize_token_embeddings(int(
        8 *
        math.ceil(len(tokenizer) / 8.0)))  # make the vocab size multiple of 8

    return model

0x3.2 main.py主体解析

parse_args解析

def parse_args():
    # 创建一个argparse的解析器对象,这个对象可以添加命令行参数和处理它们。description参数提供了一个对程序的简单描述。
    parser = argparse.ArgumentParser(
        description=
        "Finetune a transformers model on a causal language modeling task")
    parser.add_argument('--data_path',
                        nargs='*',
                        default=['Dahoas/rm-static'],
                        help='Path to the training dataset. Accepted format:'
                        '1) a single data path, 2) multiple datasets in the'
                        'form: dataset1-path dataset2-path ...')
    parser.add_argument('--data_split',
                        type=str,
                        default='2,4,4',
                        help='Comma-separated list of proportions for training'
                        'phase 1, 2, and 3 data. For example the split `6,2,2`'
                        'will use 60% of data for phase 1, 20% for phase 2'
                        'and 20% for phase 3.')
    parser.add_argument(
        '--sft_only_data_path',
        nargs='*',
        default=[],
        help='Path to the dataset for only using in SFT phase.')
    parser.add_argument(
        '--data_output_path',
        type=str,
        default='/data_turbo/home/zhangxiaoyu/data_files/',
        help=
        'Where to store the data-related files such as shuffle index. This needs to be on a local storage of a node (not on a shared storage)'
    )
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help=
        "Path to pretrained model or model identifier from huggingface.co/models.",
        required=True,
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=16,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=16,
        help="Batch size (per device) for the evaluation dataloader.",
    )
    parser.add_argument(
        "--max_seq_len",
        type=int,
        default=512,
        help="The maximum sequence length.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=1e-3,
        help=
        "Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument("--weight_decay",
                        type=float,
                        default=0.,
                        help="Weight decay to use.")
    parser.add_argument("--num_train_epochs",
                        type=int,
                        default=1,
                        help="Total number of training epochs to perform.")
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help=
        "Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--lr_scheduler_type",
        type=SchedulerType,
        default="cosine",
        help="The scheduler type to use.",
        choices=[
            "linear", "cosine", "cosine_with_restarts", "polynomial",
            "constant", "constant_with_warmup"
        ],
    )
    parser.add_argument(
        "--num_warmup_steps",
        type=int,
        default=0,
        help="Number of steps for the warmup in the lr scheduler.")
    parser.add_argument("--output_dir",
                        type=str,
                        default=None,
                        help="Where to store the model.")
    parser.add_argument("--seed",
                        type=int,
                        default=1234,
                        help="A seed for reproducible training.")
    parser.add_argument("--local_rank",
                        type=int,
                        default=-1,
                        help="local_rank for distributed training on gpus")
    parser.add_argument('--gradient_checkpointing',
                        action='store_true',
                        help='Enable HF gradient checkpointing for model.')
    parser.add_argument('--disable_dropout',
                        action='store_true',
                        help='Disable the dropout of the model.')
    # deepspeed features
    parser.add_argument('--offload',
                        action='store_true',
                        help='Enable ZeRO Offload techniques.')
    parser.add_argument(
        '--zero_stage',
        type=int,
        default=0,
        help='ZeRO optimization stage for Actor model (and clones).')
    ## LoRA for efficient training setting
    parser.add_argument("--lora_dim",
                        type=int,
                        default=0,
                        help="If > 0, use LoRA for efficient training.")
    parser.add_argument("--lora_module_name",
                        type=str,
                        default="decoder.layers.",
                        help="The scope of LoRA.")
    parser.add_argument('--only_optimize_lora',
                        action='store_true',
                        help='Only optimize the LoRA parameters.')
    # 这一行将DeepSpeed的配置参数添加到解析器中。
    parser = deepspeed.add_config_arguments(parser)
    # 这一行解析命令行参数并将它们存储在args对象中
    args = parser.parse_args()

    # Validate settings
    # 在这个代码块中,验证一些特定的参数设置是否合法。
    # 例如,如果同时启用了gradient_checkpointing和仅优化LoRA参数,那么将会抛出一个错误。
    if args.gradient_checkpointing and args.lora_dim > 0:
        assert (
            not args.only_optimize_lora
        ), "--gradient_checkpointing and --only_optimize_lora cannot be enabled at the same time."

    return args

main函数解析

接下来是训练部分的核心函数,也是全文最后一个函数main的解析。

# 这个函数是主函数,是训练语言模型的主流程,主要步骤包括解析命令行参数、
# 设置设备、准备数据、定义模型、配置优化器和学习率调度器、进行训练和评估等。
def main():
    # 解析命令行参数。
    args = parse_args()
    
    # 如果本地排名为-1,说明不在分布式训练环境下,设备设置为"cuda";
    # 否则根据args.local_rank设置对应的cuda设备,并初始化分布式训练。
    if args.local_rank == -1:
        device = torch.device("cuda")
    else:
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        # torch.distributed.init_process_group(backend='nccl')
        deepspeed.init_distributed()
    
    # 获取全局rank。
    args.global_rank = torch.distributed.get_rank()
    # 获取deepspeed的训练配置。
    ds_config = get_train_ds_config(offload=args.offload,
                                    stage=args.zero_stage)
    # 在配置中设置训练时每个GPU的微批次大小和总的批次大小。
    ds_config[
        'train_micro_batch_size_per_gpu'] = args.per_device_train_batch_size
    ds_config[
        'train_batch_size'] = args.per_device_train_batch_size * torch.distributed.get_world_size(
        ) * args.gradient_accumulation_steps

    # If passed along, set the training seed now.
    # 设置随机种子以保证结果的可复现性。
    set_random_seed(args.seed)
    
    # 在此处添加一个barrier操作,确保所有进程都执行到这一点后再继续执行后续操作。
    torch.distributed.barrier()
    # 加载预训练模型对应的分词器。
    tokenizer = load_hf_tokenizer(args.model_name_or_path, fast_tokenizer=True)
    tokenizer.pad_token = tokenizer.eos_token
    # make sure tokenizer is right pad in our logic
    tokenizer.padding_side = 'right'
    # 创建预训练模型。
    model = create_hf_model(AutoModelForCausalLM,
                            args.model_name_or_path,
                            tokenizer,
                            ds_config,
                            disable_dropout=args.disable_dropout)

    # 如果参数lora_dim大于0,将模型的线性层转换为LoRa层;如果只优化LoRa参数,关闭其他参数的梯度。
    if args.lora_dim > 0:
        model = convert_linear_layer_to_lora(model, args.lora_module_name,
                                             args.lora_dim)
        if args.only_optimize_lora:
            model = only_optimize_lora_parameters(model)

    # Prepare the data
    # 创建数据集和数据加载器:包括训练集和验证集,以及对应的采样器和数据加载器。
    train_phase = 1
    train_dataset, eval_dataset = create_prompt_dataset(
        args.local_rank,
        args.data_path,
        args.data_split,
        args.data_output_path,
        train_phase,
        args.seed,
        tokenizer,
        args.max_seq_len,
        sft_only_data_path=args.sft_only_data_path)
    # DataLoaders creation:
    if args.local_rank == -1:
        train_sampler = RandomSampler(train_dataset)
        eval_sampler = SequentialSampler(eval_dataset)
    else:
        train_sampler = DistributedSampler(train_dataset)
        eval_sampler = DistributedSampler(eval_dataset)
    train_dataloader = DataLoader(train_dataset,
                                  collate_fn=default_data_collator,
                                  sampler=train_sampler,
                                  batch_size=args.per_device_train_batch_size)
    eval_dataloader = DataLoader(eval_dataset,
                                 collate_fn=default_data_collator,
                                 sampler=eval_sampler,
                                 batch_size=args.per_device_eval_batch_size)
    # 定义模型评估函数,用于计算模型在验证集上的困惑度。
    def evaluation(model, eval_dataloader):
        model.eval() # 将模型切换为评估模式。
        losses = 0 # 初始化loss。
        for step, batch in enumerate(eval_dataloader): # 对于评估数据集中的每一个batch。
            batch = to_device(batch, device) # 将batch数据移到对应的设备上。
            with torch.no_grad(): # 在此上下文管理器中,不计算梯度,这样可以节省存储和计算资源。
                # 将batch数据输入模型,进行前向计算。
                outputs = model(**batch)

            loss = outputs.loss # 取出模型的输出中的loss。
            losses += loss.float() # 将当前的loss累加到总的losses中。
        losses = losses / (step + 1) # 计算平均的loss。
        try:
            perplexity = torch.exp(losses) # 尝试计算模型的困惑度,如果捕捉到溢出错误,将困惑度设置为无穷大。
        except OverflowError:
            perplexity = float("inf")
        # 尝试在所有设备上计算困惑度的平均值,如果发生任何错误,就忽略。
        try:
            perplexity = get_all_reduce_mean(perplexity).item()
        except:
            pass
        return perplexity
    
    # 根据是否使用权重衰减将模型参数分为两组。
    # Split weights in two groups, one with weight decay and the other not.
    optimizer_grouped_parameters = get_optimizer_grouped_parameters(
        model, args.weight_decay)
    # 选择优化器类型,如果启用了梯度Offload,使用DeepSpeedCPUAdam,否则使用FusedAdam。
    AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
    # 创建优化器。
    optimizer = AdamOptimizer(optimizer_grouped_parameters,
                              lr=args.learning_rate,
                              betas=(0.9, 0.95))
    # 计算每个epoch的更新步数。
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps)
    # 创建学习率调度器。
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps,
        num_training_steps=args.num_train_epochs * num_update_steps_per_epoch,
    )
    # 使用deepspeed初始化模型、优化器和学习率调度器。
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
        model=model,
        optimizer=optimizer,
        args=args,
        config=ds_config,
        lr_scheduler=lr_scheduler,
        dist_init_required=True)
    # 如果启用了梯度检查点,那么在模型中也启用梯度检查点。
    if args.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    # Train!
    # 使用 print_rank_0 函数在主节点(global_rank为0的节点)打印开始训练的信息。
    print_rank_0("***** Running training *****", args.global_rank)
    # 在主节点打印在第0个epoch(训练开始前)进行模型评估的信息。
    print_rank_0(
        f"***** Evaluating perplexity, Epoch {0}/{args.num_train_epochs} *****",
        args.global_rank)
    # 调用 evaluation 函数对模型进行评估,得到模型的困惑度。
    perplexity = evaluation(model, eval_dataloader)
    # 在主节点打印模型的困惑度。
    print_rank_0(f"ppl: {perplexity}", args.global_rank)
    
    # 循环 args.num_train_epochs 轮进行训练。
    for epoch in range(args.num_train_epochs):
        print_rank_0(
            f"Beginning of Epoch {epoch+1}/{args.num_train_epochs}, Total Micro Batches {len(train_dataloader)}",
            args.global_rank) # 在每个epoch开始时,在主节点打印开始新的训练周期的信息。
        model.train() # 将模型设置为训练模式。
        for step, batch in enumerate(train_dataloader): # 对于训练数据集中的每一个batch。
            batch = to_device(batch, device) # 将batch数据移到对应的设备上。
            outputs = model(**batch, use_cache=False) # 将batch数据输入模型,进行前向计算。
            loss = outputs.loss # 取出模型的输出中的loss。
            model.backward(loss) # 进行反向传播,计算梯度。
            model.step() # 更新模型的参数。

        # Evaluate perplexity on the validation set.
        # 在每个epoch结束后,在主节点打印开始评估的信息。
        print_rank_0(
            f"***** Evaluating perplexity, Epoch {epoch+1}/{args.num_train_epochs} *****",
            args.global_rank)
        # 对模型进行评估,得到模型的困惑度。
        perplexity = evaluation(model, eval_dataloader)
        # 在主节点打印模型的困惑度。
        print_rank_0(f"ppl: {perplexity}", args.global_rank)
        # 更新模型中的epoch计数。
        model.tput_timer.update_epoch_count()
    
    # 如果设置了输出目录,进行以下操作。
    if args.output_dir is not None:
    	 # 在主节点打印开始保存模型的信息。
        print_rank_0('saving the final model ...', args.global_rank)
        # 将模型中的LoRA层转换为线性层。
        model = convert_lora_to_linear_layer(model)
        
        # 如果是主节点,进行以下操作。
        if args.global_rank == 0:
            # 以Hugging Face的模型格式保存模型。
            save_hf_format(model, tokenizer, args)
        # 如果使用了Zero Redundancy Optimizer(Zero)的第三阶段,进行以下操作。
        if args.zero_stage == 3:
            # For zero stage 3, each gpu only has a part of the model, so we need a special save function
            # 使用特殊的保存函数保存模型。在Zero的第三阶段,每个GPU只有模型的一部分,所以需要特殊的保存函数。
            save_zero_three_model(model,
                                  args.global_rank,
                                  args.output_dir,
                                  zero_stage=args.zero_stage)

0x4. hybrid_engine的细节 & log

从训练过程的输出日志来看hybrid_engine是默认关闭的,DeepSpeed-Chat 打造类ChatGPT全流程 笔记一 里面提到DeepSpeed Hybrid Engine是用在加速 RLHF 流程中最耗时的部分也就是第三步,而本文介绍的监督指令微调是第一步,所以即使开启hybrid_engine加速效果应该也比较有限,所以这里默认关闭。

DeepSpeed-Chat 打造类ChatGPT全流程 笔记二之监督指令微调_第2张图片
hybrid_engine的优化方法和原理在后续文章中继续探索。

这里分享一下我复现官方sample训练的第一阶段的log:https://paste.ubuntu.com/p/vcG49hQmCW/

0x5. 总结

这篇文章解析了DeepSpeed Chat中监督指令微调这个过程的源码,这个过程和一般的PyTorch DDP分布式训练区别不是特别大,主要是自定义prompt数据集以及将普通的训练流程中的组件如模型,优化器,学习率调度器等等,使用DeepSpeed来warp一下,来用上DeepSpeed提供的Zero,Gradient Checkpoint(注意这个其实就是activation checkpoint)等特性。本文是完全按照训练流程顺序阅读代码,并补全了训练过程中所有涉及到的工具函数或者新的特性如LoRA微调的代码解析。DeepSpeed Chat这部分代码写得比较清晰易懂,因为是在接口层面来使用DeepSpeed,相当于基于DeepSpeed做应用所以代码中不会涉及到DeepSpeed的底层代码,只需要关注算法流程。但这个代码在LoRA微调这部分感觉设计的耦合性有一点高,如果要新增新的微调方式比如QLoRA可能写法就不太优雅了。

你可能感兴趣的:(chatgpt,笔记)