DeepSpeed Chat中監(jiān)督指令微調(diào)過程的源碼

文章目錄

系列文章
0x0. 前言
0x1.  Supervised finetuning (SFT) 教程翻譯
     如何訓練模型
     如何對SFT checkpoint進行評測?
     模型和數(shù)據(jù)
    來自OPT-1.3B及其SFT變體（使用不同微調(diào)數(shù)據(jù)）的提示示例
     一些參數(shù)解釋和可訓練的最大模型
     其它
0x2. 評測腳本解讀
0x3. 訓練腳本解讀
    0x3.1 頭文件相關解析
        create_prompt_dataset解析
        工具函數(shù)解析
            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函數(shù)解析
0x4. hybrid_engine的細節(jié) & log
0x5. 總結(jié)

系列文章

DeepSpeed-Chat 打造類ChatGPT全流程 筆記一

【DeepSpeed 教程翻譯】三，在 DeepSpeed中使用 PyTorch Profiler和Flops Profiler

DeepSpeed結(jié)合Megatron-LM訓練GPT2模型筆記（上）

【DeepSpeed 教程翻譯】二，Megatron-LM GPT2，Zero 和 ZeRO-Offload

【DeepSpeed 教程翻譯】開始，安裝細節(jié)和CIFAR-10 Tutorial

0x0. 前言

在 DeepSpeed-Chat 打造類ChatGPT全流程 筆記一 中跑通了DeepSpeed Chat的訓練和推理流程，DeepSpeed Chat的訓練流程包含監(jiān)督指令微調(diào)（SFT），Reward模型微調(diào)，基于人類反饋的強化學習（RLHF）三個步驟。接著上面文章的todo，這篇文章主要是解析一下監(jiān)督指令微調(diào)（SFT）階段的代碼實現(xiàn)。

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

監(jiān)督微調(diào)（SFT）與在自然語言任務（例如，WikiText-103）上的標準語言模型微調(diào)非常相似。主要的區(qū)別來自于數(shù)據(jù)集資源，SFT將收集高質(zhì)量的查詢-回答對來微調(diào)模型以達到人類更傾向的生成結(jié)果。

如何訓練模型

我們提供了多個腳本用于在單個GPU（例如，單個A6000-48G，V100-32G，A100-40G等），單節(jié)點（例如，8/16x V100-32G，8 A100-40G/80G）和多節(jié)點設置（例如，64x A100-80G）上進行訓練，這些可以在 training_scripts 目錄中找到。例如，如果你有一個單獨的A6000-48G，你可以簡單地運行對應的腳本

training_scripts/single_gpu/run_1.3b.sh

來訓練一個OPT-1.3b模型。我們的單節(jié)點腳本很容易擴展到多節(jié)點系統(tǒng)。

如何對SFT checkpoint進行評測?

一旦你使用上述代碼完成訓練，你可以簡單地執(zhí)行 bash evaluation_scripts/run_prompt.sh

它會要求用戶提供兩個模型的路徑：(a) 原始預訓練模型（即 --model_name_or_path_baseline facebook/opt-1.3b）和 (b) 微調(diào)后的模型（即 --model_name_or_path_finetune output/check_base）。"prompt_eval.py" 包含了幾個可以根據(jù)你的喜好進行更新的提示。

模型和數(shù)據(jù)

由于GPT3沒有開源的checkpoint，我們使用了Meta OPT家族的預訓練模型（即facebook/opt-1.3b）。你也可以使用其他預訓練模型（如GPT-Neo，Bloom等）。至于數(shù)據(jù)集，我們也使用了來自Huggingface數(shù)據(jù)集的開源數(shù)據(jù)集，具體如下：

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

感謝DeepSpeed RLHF的數(shù)據(jù)抽象和融合技術，我們現(xiàn)在可以將多個數(shù)據(jù)源合并用于訓練。然而，重要的是要注意，不同的數(shù)據(jù)集可能使用不同的提示詞（例如，Dohas/rm-static使用"Human:"表示查詢，"Assistant:"表示回答）。因此，用戶必須自行對齊這些提示。在我們的例子中，我們一致使用了Dohas/rm-static的格式。通過我們的評估，我們發(fā)現(xiàn)整合多樣化的數(shù)據(jù)集可以提高模型的質(zhì)量。請參考下一節(jié)以獲取不同查詢-答案對的示例。

來自OPT-1.3B及其SFT變體（使用不同微調(diào)數(shù)據(jù)）的提示示例

在這里插入圖片描述

一些參數(shù)解釋和可訓練的最大模型

main.py文件中使用的大多數(shù)參數(shù)都有清晰的解釋，如果你有解碼器模型微調(diào)的經(jīng)驗，通常很容易理解。然而，如果你對其中任何一個不清楚，請不要猶豫在GitHub問題上向我們求助。在這一部分，我們提供了一些具體的參數(shù)解釋和它們的使用方法。

對于用戶來說，一個重要的考慮是確定他們可以使用當前系統(tǒng)訓練的最大模型大小。在這里，我們提供了一個估計這個限制的方法。假設你不使用卸載功能，并啟用(i)零階段3（如果使用多個GPU），(ii)梯度檢查點，以及(iii)LoRA，那么你可以訓練的大致最大模型大小（以十億參數(shù)為單位）可以估計為"總GPU內(nèi)存（GB）除以3"。例如，如果你有一臺單一的A6000-48G GPU，你可能可以訓練最多16十億參數(shù)的模型。需要注意的是，這只是一個粗略的估計，你應該自己驗證。

其它

從InstructGPT的工作中，我們建議為了得到更好的人類偏好的答案，讓模型過度擬合（即更長的訓練周期）。通過我們的探索，我們發(fā)現(xiàn)這對于較小模型的微調(diào)，如OPT-1.3B，特別有幫助。值得注意的是，我們在腳本中提供的超參數(shù)并沒有經(jīng)過大量的調(diào)整。因此，我們鼓勵用戶和實踐者自己找到最優(yōu)的配置。此外，我們的系統(tǒng)可以很容易地擴展到其他語言，如中文和日語。為了證明這一點，我們在"training_scripts/other_language"目錄下包含了兩個例子。

0x2. 評測腳本解讀

DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/evaluation_scripts/run_prompt.sh這個腳本是對輸入的模型進行評測，我們看一下內(nèi)容：

# 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

這表示我們可以調(diào)用prompt_eval.py來對baseline模型和finetune模型進行對比評測。所以評測的實現(xiàn)就是這里的prompt_eval.py，接下來對prompt_eval.py做一個解讀。

首先是頭文件導入部分，除了一些常規(guī)的需要用到的包還從DeepSpeed Chat的utils下面導入了2個函數(shù)。

#DeepSpeedTeam
importargparse
importlogging
importtorch
importsys
importos

fromtransformersimport(
AutoModelForCausalLM,)

sys.path.append(
os.path.abspath(os.path.join(os.path.dirname(__file__),os.path.pardir)))
fromutils.model.model_utilsimportcreate_hf_model#從utils下面導入自定義函數(shù)
fromutils.utilsimportload_hf_tokenizer

logger=logging.getLogger(__name__)

接下來解析一下create_hf_model和load_hf_tokenizer兩個函數(shù)的實現(xiàn)。首先是create_hf_model函數(shù)：

#這段代碼定義了一個名為create_hf_model的函數(shù)，該函數(shù)的作用是創(chuàng)建或加載一個預訓練模型。該函數(shù)的主要參數(shù)包括：
#model_class：模型的類別，例如GPT-2、BERT等。
#tokenizer：用于模型的分詞器。
#ds_config:DeepSpeed的配置參數(shù)。
#rlhf_training：一個標志，用來表示是否正在進行RLHF（ReinforcementLearningfromHumanFeedback）訓練。
#disable_dropout：一個標志，用來表示是否禁用dropout。Dropout是一種防止過擬合的技術。
defcreate_hf_model(model_class,
model_name_or_path,
tokenizer,
ds_config=None,
rlhf_training=False,
disable_dropout=False):
#根據(jù)model_name_or_path從預訓練模型獲取模型配置model_config。
model_config=AutoConfig.from_pretrained(model_name_or_path)
#如果disable_dropout為真，則將模型配置中的dropout設為0.0。
ifdisable_dropout:
model_config.dropout=0.0
#Note:dschfisdefinedinfunctionscopetoavoidglobaleffects
#https://huggingface.co/docs/transformers/main_classes/deepspeed#nontrainer-deepspeed-integration
#根據(jù)ds_config中的設置，創(chuàng)建DeepSpeed的配置對象dschf，以便進行DeepSpeed優(yōu)化。
ifds_configisnotNoneandds_config["zero_optimization"]["stage"]==3:
dschf=HfDeepSpeedConfig(ds_config)
else:
dschf=None
#根據(jù)rlhf_training的值，確定是從配置中創(chuàng)建模型還是從預訓練模型中加載模型。如果rlhf_training為真，則根據(jù)模型配置創(chuàng)建模型；否則，從預訓練模型加載模型。
ifrlhf_training:
#theweightloadingishandledbycreatecriticmodel
model=model_class.from_config(model_config)
else:
model=model_class.from_pretrained(
model_name_or_path,
from_tf=bool(".ckpt"inmodel_name_or_path),
config=model_config)
#將模型的結(jié)束標記和填充標記設為分詞器的結(jié)束標記id。
model.config.end_token_id=tokenizer.eos_token_id
model.config.pad_token_id=model.config.eos_token_id
#調(diào)整模型的詞匯表大小，使其為8的倍數(shù)。這樣做的目的是為了在某些硬件（如GPU）上提高效率。
model.resize_token_embeddings(int(
8*
math.ceil(len(tokenizer)/8.0)))#makethevocabsizemultipleof8

returnmodel

然后是load_hf_tokenizer函數(shù)：

#這段代碼定義了一個名為load_hf_tokenizer的函數(shù)，該函數(shù)的作用是加載一個指定的分詞器。這個函數(shù)的參數(shù)有：
#model_name_or_path：模型的名稱或者預訓練模型的路徑。
#fast_tokenizer：是否使用FastTokenizer，這是一種更快的分詞器實現(xiàn)。
defload_hf_tokenizer(model_name_or_path,fast_tokenizer=True):
#檢查model_name_or_path是否存在。如果存在，可能意味著這是一個本地的模型路徑。
ifos.path.exists(model_name_or_path):
#Locallytokenizerloadinghassomeissue,soweneedtoforcedownload
#如果是本地模型，再檢查在該路徑下的"config.json"文件是否存在。如果存在，打開該文件，并從文件中讀取"_name_or_path"字段，將其賦值給model_name，然后通過AutoTokenizer.from_pretrained函數(shù)，使用該模型名稱和fast_tokenizer參數(shù)，從HuggingFace模型庫下載并加載分詞器。
model_json=os.path.join(model_name_or_path,"config.json")
ifos.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函數(shù)，使用model_name_or_path和fast_tokenizer參數(shù)，從HuggingFace模型庫下載并加載分詞器。
tokenizer=AutoTokenizer.from_pretrained(model_name_or_path,
fast_tokenizer=True)
returntokenizer

接下來是參數(shù)解析部分，解析如下：

#這段代碼定義了一個名為parse_args的函數(shù)，該函數(shù)的作用是解析命令行參數(shù)。它使用Python的argparse庫來完成這個工作。下面是每個參數(shù)的詳解：
defparse_args():
parser=argparse.ArgumentParser(description="EvalthefinetuedSFTmodel")
#model_name_or_path_baseline：基線模型的路徑，這是一個必須提供的參數(shù)（required=True）。
parser.add_argument(
"--model_name_or_path_baseline",
type=str,
help="Pathtobaselinemodel",
required=True,
)
#model_name_or_path_finetune：微調(diào)后模型的路徑，這也是一個必須提供的參數(shù)。
parser.add_argument(
"--model_name_or_path_finetune",
type=str,
help="Pathtopretrainedmodel",
required=True,
)
#num_beams：用于指定集束搜索的集束寬度，其默認值為1。
parser.add_argument(
"--num_beams",
type=int,
default=1,
help='Specifynumofbeams',
)
#num_beam_groups：用于指定集束搜索的組數(shù)，其默認值為1。
parser.add_argument(
"--num_beam_groups",
type=int,
default=1,
help='Specifynumofbeams',
)
#top_k：用于指定在Top-K采樣中考慮的最高可能性詞匯的數(shù)量，其默認值為4。
parser.add_argument(
"--top_k",
type=int,
default=4,
help='Specifynumofbeams',
)
#penalty_alpha：懲罰因子，其默認值為0.6。
parser.add_argument(
"--penalty_alpha",
type=float,
default=0.6,
help='Specifynumofbeams',
)
#num_return_sequences：生成序列的數(shù)量，其默認值為1。
parser.add_argument(
"--num_return_sequences",
type=int,
default=1,
help='Specifynumofreturnsequences',
)
#max_new_tokens：生成的最大新token數(shù)，其默認值為100。
parser.add_argument(
"--max_new_tokens",
type=int,
default=100,
help='Specifynumofreturnsequences',
)
#language：語言類型，可以是"English"、"Chinese"或"Japanese"，默認為"English"。
parser.add_argument("--language",
type=str,
default="English",
choices=["English","Chinese","Japanese"])
#parser.parse_args()這個函數(shù)將解析命令行參數(shù)，并將結(jié)果保存在一個Namespace對象中。這個對象被返回，可以在其他地方使用這些參數(shù)。
args=parser.parse_args()

returnargs

接下來是generate函數(shù)的解析：

#這個函數(shù)是用來利用訓練好的模型生成文本的，它接受以下參數(shù)
#model：已經(jīng)訓練好的模型。
#tokenizer：用于將文本轉(zhuǎn)換為模型可理解的輸入的工具。
#inputs：模型的輸入數(shù)據(jù)。
#num_beams：在使用束搜索算法時的束寬，其默認值為1。
#num_beam_groups：在使用分組束搜索時的組數(shù)，默認為1。
#do_sample：是否進行隨機采樣。如果設為True，則在生成過程中會隨機選擇下一個單詞，而不是僅選擇最可能的單詞。默認為False。
#num_return_sequences：模型返回的序列數(shù)，默認為1。
#max_new_tokens：模型生成的最大新token數(shù)，即最大生成文本的長度，默認為100。
defgenerate(model,
tokenizer,
inputs,
num_beams=1,
num_beam_groups=1,
do_sample=False,
num_return_sequences=1,
max_new_tokens=100):
#函數(shù)首先使用模型的generate方法，根據(jù)提供的參數(shù)生成文本。
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解碼為可讀的文本。注意，這里跳過了特殊的令牌（如填充和開始/結(jié)束令牌），并且不會清理tokenize產(chǎn)生的額外空格。
result=tokenizer.batch_decode(generate_ids,
skip_special_tokens=True,
clean_up_tokenization_spaces=False)
returnresult

接下來是generate_constrastive_search函數(shù)的解析：

#這個函數(shù)叫做generate_constrastive_search，它是用于利用訓練好的模型進行對比搜索生成文本的。這個函數(shù)接受以下參數(shù)：
#model：已經(jīng)訓練好的模型。
#tokenizer：用于將文本轉(zhuǎn)換為模型可理解的輸入的工具。
#inputs：模型的輸入數(shù)據(jù)。
#top_k：在每一步生成時，只考慮概率最高的top_k個候選項，然后進行隨機抽樣。默認為4。
#penalty_alpha：用于懲罰新生成的token與原始輸入之間的差異，默認為0.6。
#num_return_sequences：模型返回的序列數(shù)，默認為1。
#max_new_tokens：模型生成的最大新token數(shù)，即最大生成文本的長度，默認為100。
defgenerate_constrastive_search(model,
tokenizer,
inputs,
top_k=4,
penalty_alpha=0.6,
num_return_sequences=1,
max_new_tokens=100):
#函數(shù)首先使用模型的generate方法，根據(jù)提供的參數(shù)生成文本。注意這里使用了模型的一個特殊的生成方式，這種方式在每一步生成時，只考慮概率最高的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方法將生成的tokenID解碼為可讀的文本。注意，這里跳過了特殊的token（如填充和開始/結(jié)束token），并且不會清理token化產(chǎn)生的額外空格。
result=tokenizer.batch_decode(generate_ids,
skip_special_tokens=True,
clean_up_tokenization_spaces=False)
returnresult

接下來是一個簡單的打印工具函數(shù)：

#gen_output：這是一個列表，其中包含了我們希望打印的內(nèi)容，每一項都是一段文本。
defprint_utils(gen_output):
#函數(shù)會遍歷gen_output列表中的每一項，然后將每一項都打印出來。為了在不同項之間增加一些可視化的分隔，函數(shù)在每一項前后都額外打印了一個空行。
foriinrange(len(gen_output)):
print()
print(gen_output[i])
print()

然后是prompt_eval這個函數(shù)，這個函數(shù)prompt_eval的目的是評估和比較基線模型（model_baseline）和微調(diào)過的模型（model_fintuned）對于一組提示（prompts）的生成性能。讓我們逐行進行解析：

#輸入?yún)?shù)包括：args（命令行參數(shù)）、model_baseline（基線模型）、model_fintuned（微調(diào)模型）、tokenizer（用于編碼和解碼的分詞器）、device（指定運行模型的設備）、prompts（一組要評估的提示）。
defprompt_eval(args,model_baseline,model_fintuned,tokenizer,device,
prompts):
#對于prompts中的每一個提示，我們都做以下操作：
forpromptinprompts:
#使用分詞器將提示轉(zhuǎn)換為模型所需的輸入格式，并將其移至指定的設備上。
inputs=tokenizer(prompt,return_tensors="pt").to(device)
#打印一條消息表示我們現(xiàn)在正在進行基線模型的生成。
print("==========Baseline:Greedy=========")
#然后，我們調(diào)用之前定義的generate函數(shù)使用貪婪搜索方法生成文本，并使用print_utils函數(shù)打印生成的結(jié)果。
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)
#打印一條消息表示我們現(xiàn)在正在進行微調(diào)模型的生成。
print("==========finetune:Greedy=========")
#同樣地，我們調(diào)用generate函數(shù)使用貪婪搜索方法生成文本，并使用print_utils函數(shù)打印生成的結(jié)果。
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)
#注意：在此函數(shù)中，貪婪搜索被用作基線方法。然而，該函數(shù)還提供了其他幾種搜索策略的例子，包括多項式采樣、束搜索、束搜索多項式采樣、多樣性束搜索和對比搜索。這些策略在此函數(shù)中都被注釋掉了，但你可以根據(jù)需要去掉注釋，使用這些策略。

#print("==========finetune:Multinomialsampling=========")
#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:BeamSearch=========")
#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-searchmultinomialsampling=========")
#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:DiverseBeamSearch=========")
#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:ConstrastiveSearch=========")
#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)
#最后，打印一條消息表示這個提示的處理已經(jīng)結(jié)束。然后打印兩個空行作為分隔。
print("====================promptend=============================")
print()
print()

解析main函數(shù)：

#main函數(shù)負責解析命令行參數(shù)、準備模型和分詞器、定義提示，然后使用這些來評估和比較基線模型和微調(diào)模型。
defmain():
#這個main函數(shù)是整個腳本的入口點。它首先通過parse_args函數(shù)解析命令行參數(shù)。然后它設置了運行模型的設備為第一個GPU。
args=parse_args()

device=torch.device("cuda:0")
#接著，它使用load_hf_tokenizer函數(shù)加載分詞器，然后使用create_hf_model函數(shù)創(chuàng)建基線模型（model_baseline）和微調(diào)模型（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)

#在接下來的部分，函數(shù)定義了一組用于評估的提示。注意，這里特別指出，如果提示以空格結(jié)束，那么沒有經(jīng)過微調(diào)的原始模型有可能會陷入停滯并無法產(chǎn)生響應。微調(diào)過的模型在這方面表現(xiàn)得更好。因此，這里所有的提示都以冒號":"結(jié)束，以使得比較更有意義。
#這個腳本支持英文、中文和日文的評估，它通過args.language參數(shù)判斷用戶選擇的語言，并根據(jù)此選擇加載對應的提示。
ifargs.language=="English":
prompts=[
"Human:PleasetellmeaboutMicrosoftinafewsentence?Assistant:",
"Human:Explainthemoonlandingtoa6yearoldinafewsentences.Assistant:",
"Human:Writeashortpoemaboutawisefrog.Assistant:",
"Human:WhowaspresidentoftheUnitedStatesin1955?Assistant:",
"Human:Howdoesatelescopework?Assistant:",
"Human:Whydobirdsmigratesouthforthewinter?Assistant:"
]
elifargs.language=="Chinese":
prompts=[
"Human:請用幾句話介紹一下微軟?Assistant:",
"Human:用幾句話向6歲的孩子解釋登月。Assistant:",
"Human:寫一首關于一只聰明的青蛙的短詩。Assistant:",
"Human:誰是1955年的美國總統(tǒng)?Assistant:","Human:望遠鏡是如何工作的?Assistant:",
"Human:鳥類為什么要南遷過冬?Assistant:"
]
elifargs.language=="Japanese":
prompts=[
"Human:マイクロソフトについて簡単に教えてください。Assistant:",
"Human:6歳児に月面著陸を短い文で説明する。Assistant:",
"Human:賢いカエルについて短い詩を書いてください。Assistant:",
"Human:1955年のアメリカ合衆(zhòng)國大統(tǒng)領は誰?Assistant:",
"Human:望遠鏡はどのように機能しますか?Assistant:",
"Human:鳥が冬に南に移動するのはなぜですか?Assistant:"
]
#最后，它調(diào)用prompt_eval函數(shù)，將所有的參數(shù)和提示傳遞給它，以進行模型評估。
prompt_eval(args,model_baseline,model_fintuned,tokenizer,device,
prompts)

0x3. 訓練腳本解讀

訓練過程都集中在DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py這個腳本，我們逐行解析一下這個腳本。

0x3.1 頭文件相關解析

#DeepSpeedTeam
#首先，它導入了Python的標準庫，如argparse（用于解析命令行參數(shù)），os和math。
#然后，它導入了PyTorch庫，這是一個用于深度學習的開源庫，同時也導入了一些PyTorch的輔助模塊，如DataLoader（用于加載數(shù)據(jù)）、
#RandomSampler和SequentialSampler（用于數(shù)據(jù)抽樣）以及DistributedSampler（用于在分布式設置中進行數(shù)據(jù)抽樣）。
importargparse
importos
importmath
importsys

importtorch
fromtorch.utils.dataimportDataLoader,RandomSampler,SequentialSampler
fromtorch.utils.data.distributedimportDistributedSampler

#接下來，它導入了HuggingFace的transformers庫的一些模塊，包括用于因果語言建模的模型（AutoModelForCausalLM），優(yōu)化器調(diào)度類型（SchedulerType），默認的數(shù)據(jù)整理函數(shù)（default_data_collator）和獲取優(yōu)化器調(diào)度器的函數(shù)（get_scheduler）。
fromtransformersimport(
AutoModelForCausalLM,
SchedulerType,
default_data_collator,
get_scheduler,
)

#然后，它導入了deepspeed庫，這是一個為大規(guī)模模型訓練優(yōu)化的庫。它也導入了deepspeed庫中的一些模塊，包括優(yōu)化器類（DeepSpeedCPUAdam和FusedAdam）
importdeepspeed
fromdeepspeed.ops.adamimportDeepSpeedCPUAdam,FusedAdam

#之后，它將當前腳本的父目錄添加到系統(tǒng)路徑中，以便可以從該目錄下的utils目錄導入一些自定義函數(shù)和模塊。
sys.path.append(
os.path.abspath(os.path.join(os.path.dirname(__file__),os.path.pardir)))
#最后，它從utils目錄中導入了一些自定義模塊和函數(shù)，包括數(shù)據(jù)處理函數(shù)（create_prompt_dataset），打印和設備轉(zhuǎn)換函數(shù)（print_rank_0和to_device），
#模型保存函數(shù)（save_hf_format），隨機種子設置函數(shù)（set_random_seed），求平均函數(shù)（get_all_reduce_mean），
#獲取優(yōu)化器參數(shù)組的函數(shù)（get_optimizer_grouped_parameters），保存和加載模型的函數(shù)（save_zero_three_model和load_hf_tokenizer），
#以及創(chuàng)建模型和處理模型的函數(shù)（create_hf_model）。這些函數(shù)在腳本中的后續(xù)部分都將被使用。
fromutils.data.data_utilsimportcreate_prompt_dataset
fromutils.utilsimportprint_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
fromutils.ds_utilsimportget_train_ds_config
fromutils.module.loraimportconvert_linear_layer_to_lora,convert_lora_to_linear_layer,only_optimize_lora_parameters
fromutils.model.model_utilsimportcreate_hf_model

create_prompt_dataset解析

create_prompt_dataset這個函數(shù)實際上直接或者間接的用到了utils/data中raw_dataset.py和data_utils.py，為了搞清楚這個函數(shù)，我們需要對這兩個文件做一個解析。

首先解析一下raw_dataset.py。這里先定義了一個PromptRawDataset類：

#DeepSpeedTeam
fromdatasetsimportload_dataset
fromtorch.utils.dataimportSubset
importre


#這段代碼定義了一個名為PromptRawDataset的類，這個類是一個模板類，用于處理和組織模型輸入數(shù)據(jù)的格式。
#如果有新的數(shù)據(jù)集需要進行處理，可以繼承這個類并實現(xiàn)相應的方法來確保數(shù)據(jù)的統(tǒng)一格式和接口。
classPromptRawDataset(object):
#首先，這個類的構(gòu)造函數(shù)__init__接收四個參數(shù)：output_path（輸出路徑），seed（隨機種子），
#local_rank（本地等級）和dataset_name（數(shù)據(jù)集名稱）。
#在構(gòu)造函數(shù)中，如果數(shù)據(jù)集名稱不是'local/jsonfile'，
#那么會使用HuggingFace的datasets庫的load_dataset函數(shù)來加載數(shù)據(jù)集。


def__init__(self,output_path,seed,local_rank,dataset_name):
self.output_path=output_path
self.seed=seed
self.local_rank=local_rank
ifnotdataset_name=='local/jsonfile':
self.raw_datasets=load_dataset(dataset_name)
#然后，這個類定義了一些方法，這些方法在默認情況下并沒有實現(xiàn)（只是返回None或者空操作），
#這是因為這個類只是一個模板類，這些方法需要在實際使用時在子類中具體實現(xiàn)。
defget_train_data(self):#獲取訓練數(shù)據(jù)
return

defget_eval_data(self):#獲取評估數(shù)據(jù)
return

#Thepromptshouldbeintheformatof:"Human:"+actual_prompt_sentence+"Assistant:"
#get_prompt方法用于獲取樣本中的prompt（提示，這是模型的輸入）。
defget_prompt(self,sample):
return

#Thechosenresponseshouldbeintheformatof:""+actual_response_sentence
#get_chosen方法用于獲取樣本中的chosen（已選的回應，這是模型需要生成的目標輸出）。
defget_chosen(self,sample):
return

#Therejectedresponseshouldbeintheformatof:""+actual_response_sentence
#Ifthedatasetdoesnothaverejectedresponse,returnNone
#get_rejected方法用于獲取樣本中的rejected（被拒絕的回應，這可能用于一些特定的訓練場景，比如在對抗訓練中，但如果數(shù)據(jù)集中沒有這樣的數(shù)據(jù)，可以返回None）。
defget_rejected(self,sample):
return
#獲取樣本中的prompt和chosen
defget_prompt_and_chosen(self,sample):
return
#獲取樣本中的prompt和rejected
defget_prompt_and_rejected(self,sample):
return

接下來就是每個具體數(shù)據(jù)集的定義，我這里以 OpenaiWebgptcomparisonsDataset 為例解析一下，剩下的讀者又需要可以自行理解：

#Englishdataset
#這個類OpenaiWebgptcomparisonsDataset繼承自PromptRawDataset類，
#針對"openai/webgpt_comparisons"這個具體的數(shù)據(jù)集進行了特化。
classOpenaiWebgptcomparisonsDataset(PromptRawDataset):
#在構(gòu)造函數(shù)__init__中，調(diào)用了父類的構(gòu)造函數(shù)，并設定了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中獲取訓練數(shù)據(jù)和測試數(shù)據(jù)。
#它們與之前的DahoasRmstaticDataset類不同之處在于，它們使用get_raw_dataset_split_index
#方法對訓練數(shù)據(jù)進行了劃分，將其劃分為訓練集和驗證集，并返回對應的數(shù)據(jù)子集。
defget_train_data(self):
from.data_utilsimportget_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)
returndataset

defget_eval_data(self):
from.data_utilsimportget_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)
returndataset
#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'就是被拒絕的回應，反之亦然。
#在獲取已選回應和被拒絕的回應時，還對回應進行了處理，
#去除了所有形如"[...]"或"(...)"的文本，然后在回應前添加了一個空格。
defget_prompt(self,sample):
return"Human:"+sample['question']['full_text']+"Assistant:"

defget_chosen(self,sample):
iffloat(sample['score_0'])>=float(sample['score_1']):
response=sample['answer_0']
else:
response=sample['answer_1']
#Thisdatahascitationsquarebracketsandnumbers(e.g.,"[1]").
#Rightnowwearenotdoingbrowser-assistedfinetuning,thuswe
#removethesecitationstoavoidconfusingthemodel.
response=re.sub(r"[([].*?[)]]","",response)
response=re.sub(r"[([].*?[)]]","",response)
return""+response

defget_rejected(self,sample):
iffloat(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

defget_prompt_and_rejected(self,sample):
iffloat(sample['score_0'])

	

	接著解析 data_utils.py：

	
#DeepSpeedTeam
"""
Partofthecodewasadoptedfromhttps://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/data/dataset_utils.py
"""
importtorch
fromtorch.utils.dataimportDataset,Subset,ConcatDataset
fromtorch.nn.utils.rnnimportpad_sequence
importtorch.nn.functionalasF
fromdatasetsimportload_dataset
importnumpyasnp
importos
importhashlib#Python的內(nèi)置庫，提供了一系列散列函數(shù)，如MD5、SHA1等。
fromitertoolsimportchain#Python的內(nèi)置庫，提供了一系列用于操作迭代器的函數(shù)。
from.importraw_datasets

#這段代碼定義了一個名為get_raw_dataset的函數(shù)，其主要作用是根據(jù)傳入的數(shù)據(jù)集名稱dataset_name
#返回一個適當?shù)腜romptRawDataset子類的實例。
#這個函數(shù)有四個參數(shù)：dataset_name，output_path，seed和local_rank。
defget_raw_dataset(dataset_name,output_path,seed,local_rank):
#在函數(shù)中，根據(jù)dataset_name的具體值來創(chuàng)建不同的PromptRawDataset子類的實例。
#例如，如果dataset_name為"Dahoas/rm-static"，那么就創(chuàng)建一個DahoasRmstaticDataset的實例；
#如果dataset_name為"Dahoas/full-hh-rlhf"，那么就創(chuàng)建一個DahoasFullhhrlhfDataset的實例，以此類推。
if"Dahoas/rm-static"indataset_name:
returnraw_datasets.DahoasRmstaticDataset(output_path,seed,
local_rank,dataset_name)
elif"Dahoas/full-hh-rlhf"indataset_name:
returnraw_datasets.DahoasFullhhrlhfDataset(output_path,seed,
local_rank,dataset_name)
elif"Dahoas/synthetic-instruct-gptj-pairwise"indataset_name:
returnraw_datasets.DahoasSyntheticinstructgptjpairwiseDataset(
output_path,seed,local_rank,dataset_name)
elif"yitingxie/rlhf-reward-datasets"indataset_name:
returnraw_datasets.YitingxieRlhfrewarddatasetsDataset(
output_path,seed,local_rank,dataset_name)
elif"openai/webgpt_comparisons"indataset_name:
returnraw_datasets.OpenaiWebgptcomparisonsDataset(
output_path,seed,local_rank,dataset_name)
elif"stanfordnlp/SHP"indataset_name:
returnraw_datasets.StanfordnlpSHPDataset(output_path,seed,
local_rank,dataset_name)
elif"pvduy/sharegpt_alpaca_oa_vicuna_format"indataset_name:
returnraw_datasets.PvduySharegptalpacaoavicunaformatDataset(
output_path,seed,local_rank,dataset_name)
elif"wangrui6/Zhihu-KOL"indataset_name:
returnraw_datasets.Wangrui6ZhihuKOLDataset(output_path,seed,
local_rank,dataset_name)
elif"Cohere/miracl-zh-queries-22-12"indataset_name:
returnraw_datasets.CohereMiraclzhqueries2212Dataset(
output_path,seed,local_rank,dataset_name)
elif"Hello-SimpleAI/HC3-Chinese"indataset_name:
returnraw_datasets.HelloSimpleAIHC3ChineseDataset(
output_path,seed,local_rank,dataset_name)
elif"mkqa-Chinese"indataset_name:
returnraw_datasets.MkqaChineseDataset(output_path,seed,local_rank,
"mkqa")
elif"mkqa-Japanese"indataset_name:
returnraw_datasets.MkqaJapaneseDataset(output_path,seed,local_rank,
"mkqa")
elif"Cohere/miracl-ja-queries-22-12"indataset_name:
returnraw_datasets.CohereMiracljaqueries2212Dataset(
output_path,seed,local_rank,dataset_name)
elif"lmqg/qg_jaquad"indataset_name:
returnraw_datasets.LmqgQgjaquadDataset(output_path,seed,local_rank,
dataset_name)
elif"lmqg/qag_jaquad"indataset_name:
returnraw_datasets.LmqgQagjaquadDataset(output_path,seed,local_rank,
dataset_name)
#如果dataset_name是"local/jsonfile"，則會檢查在路徑chat_path+'/data/train.json'
#和chat_path+'/data/eval.json'下是否存在文件。如果存在，則創(chuàng)建一個LocalJsonFileDataset的實例；
#如果不存在，則拋出一個RuntimeError異常。
elif"local/jsonfile"indataset_name:
chat_path=os.path.abspath(
os.path.join(os.path.dirname(__file__),os.path.pardir,
os.path.pardir,os.path.pardir))
ifnot(os.path.isfile(chat_path+'/data/train.json')
andos.path.isfile(chat_path+'/data/eval.json')):
raiseRuntimeError(
f"Pleasecheckboththetrain.jsonandeval.jsonfilesinyourapplications/DeepSpeed-Chat/datadirectory."
)
returnraw_datasets.LocalJsonFileDataset(output_path,seed,local_rank,
dataset_name,chat_path)
else:
#如果dataset_name沒有在以上的所有條件中匹配到，那么函數(shù)也會拋出一個RuntimeError異常，表示沒有為這個數(shù)據(jù)集的配置。
raiseRuntimeError(
f"Wedonothaveconfigsfordataset{dataset_name},butyoucanadditbyyourselfinraw_datasets.py."
)


	

	再看下 get_shuffle_idx函數(shù)：

	
#這個函數(shù)的作用是生成一個大小為size的亂序索引數(shù)組，它接受兩個參數(shù)：seed和size。
defget_shuffle_idx(seed,size):
np_rng=np.random.RandomState(seed=seed)#創(chuàng)建一個NumPy的隨機狀態(tài)生成器對象np_rng，seed是隨機種子，確定了隨機數(shù)的生成序列。
dtype_=np.uint32#設置其為NumPy的uint32類型，這是一個無符號32位整數(shù)類型。
ifsize>=(np.iinfo(np.uint32).max-1):#如果size大于或等于uint32的最大值減一，這里減一是為了防止可能的溢出。
dtype_=np.int64#則將dtype_改為int64，這是一個64位的有符號整數(shù)類型。
shuffle_idx=np.arange(start=0,stop=size,step=1,dtype=dtype_)#創(chuàng)建一個由0開始，步長為1，到size結(jié)束（不包含size），并且數(shù)據(jù)類型為dtype_的等差數(shù)列，將其賦值給shuffle_idx。
np_rng.shuffle(shuffle_idx)#使用np_rng隨機狀態(tài)生成器對shuffle_idx進行隨機排列，這樣就打亂了shuffle_idx的順序。
returnshuffle_idx#返回亂序后的shuffle_idx。


	

	接著解析get_raw_dataset_split_index函數(shù)：

	
#這個函數(shù)主要是根據(jù)提供的參數(shù)分割數(shù)據(jù)集，并生成一個分割索引。
#它首先檢查索引文件是否存在，如果不存在，則生成分割索引，并保存到文件。
#然后，它從文件中加載索引，并返回索引列表。
#它接受8個參數(shù)：local_rank、輸出路徑、數(shù)據(jù)集名稱、種子、split_name、data_split、split_index和數(shù)據(jù)大小。
defget_raw_dataset_split_index(local_rank,output_path,dataset_name,seed,
split_name,data_split,split_index,
data_size):
#根據(jù)給定參數(shù)生成索引文件名。
index_file_name=f"{output_path}/{dataset_name}_seed{seed}_{split_name}_{data_split}_{split_index}.npy"
#reindexeachtimewhenusinglocaljsonfilesinceit'smorelikelytogetmodified
#如果索引文件不存在，或者數(shù)據(jù)集名為'jsonfile'，則執(zhí)行下面的操作。
if(notos.path.isfile(index_file_name))or(dataset_name=='jsonfile'):
splits=[float(s)forsindata_split.split(',')]#將data_split（以逗號分隔的字符串）分割為浮點數(shù)列表。
splits_sum=sum(splits)#計算所有分割的總和。
splits=[split/splits_sumforsplitinsplits]#將每個分割值歸一化，使其和為1。
splits_index=[0]#初始化分割索引列表，起始值為0
forindex,splitinenumerate(splits):#遍歷分割列表。
#將新的索引（當前索引加上歸一化的分割值與數(shù)據(jù)大小的乘積）添加到索引列表。
splits_index.append(splits_index[index]+
int(round(split*float(data_size))))
#計算最后一個索引與數(shù)據(jù)大小的差值。
diff=splits_index[-1]-data_size
#遍歷除第一個外的所有索引。
forindexinrange(1,len(splits_index)):
#將差值從每個索引中減去，以確保最后一個索引等于數(shù)據(jù)大小。
splits_index[index]-=diff
#斷言最后一個索引等于數(shù)據(jù)大小
assertsplits_index[-1]==data_size

#生成一個亂序的索引。
shuffle_idx=get_shuffle_idx(seed,data_size)
forsplit_iinrange(len(splits)):#遍歷每個分割。
#根據(jù)給定參數(shù)生成亂序索引分割文件名。
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)
#將索引數(shù)組轉(zhuǎn)換為列表并返回。
returnindex.tolist()


	

	接下來解析一下繼承自Dataset的PromptDataset類：

	
#這是一個自定義的PromptDataset類，它繼承自torch.utils.data.Dataset。
#這是一個數(shù)據(jù)集類，通常被用于PyTorch中數(shù)據(jù)的加載和預處理。
classPromptDataset(Dataset):
#類的構(gòu)造函數(shù)，它接受五個參數(shù)：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__()#調(diào)用父類torch.utils.data.Dataset的構(gòu)造函數(shù)。
self.prompt_dataset=prompt_dataset#將傳入的參數(shù)賦值給類的成員變量。
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__方法，它返回數(shù)據(jù)集的長度。這是PyTorch數(shù)據(jù)集的必要方法。
length=len(self.chosen_dataset)#初始設定數(shù)據(jù)集長度為chosen_dataset的長度。
ifself.train_phase==3:
length=len(self.prompt_dataset)#如果訓練階段為3，則數(shù)據(jù)集長度設定為prompt_dataset的長度。
returnlength#返回計算得出的數(shù)據(jù)集長度。

#定義類的__getitem__方法，它接受一個參數(shù)idx，返回索引idx處的數(shù)據(jù)。這是PyTorch數(shù)據(jù)集的必要方法。
def__getitem__(self,idx):
#如果訓練階段為1，則返回一個字典，包含input_ids、attention_mask和labels，它們都來自chosen_dataset的索引idx處。
ifself.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。
elifself.train_phase==2:
returnself.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
elifself.train_phase==3:
returnself.prompt_dataset[idx]["input_ids"],self.prompt_dataset[idx]["attention_mask"],
self.pad_token_id


	

	接著再解析一下create_dataset_split函數(shù)：

	
#這是一個名為create_dataset_split的函數(shù)，它的功能是根據(jù)給定的訓練階段（train_phase），創(chuàng)建并返回相應的數(shù)據(jù)集分割。
#具體來說，它為每個訓練階段生成不同的數(shù)據(jù)集列表，并將它們放入PromptDataset對象中。
#函數(shù)接受6個參數(shù)：當前數(shù)據(jù)集(current_dataset)、原始數(shù)據(jù)集(raw_dataset)、訓練階段(train_phase)、
#分詞器(tokenizer)、會話結(jié)束標記(end_of_conversation_token)和最大序列長度(max_seq_len)。
defcreate_dataset_split(current_dataset,raw_dataset,train_phase,tokenizer,
end_of_conversation_token,max_seq_len):
#創(chuàng)建三個空的列表，用于存儲對話提示（prompt_dataset）、選定的對話（chosen_dataset）和被拒絕的對話（reject_dataset）。
prompt_dataset=[]
chosen_dataset=[]
reject_dataset=[]
#如果訓練階段為1，則將接受的對話進行分詞并添加到chosen_dataset中。
iftrain_phase==1:
#遍歷當前數(shù)據(jù)集。
fori,tmp_datainenumerate(current_dataset):
#tokenizethetext
#從原始數(shù)據(jù)集中獲取對話提示和接受的對話。
chosen_sentence=raw_dataset.get_prompt_and_chosen(
tmp_data)#theacceptresponse
#如果接受的對話不為空，則將其分詞并添加到chosen_dataset中。
ifchosen_sentenceisnotNone:
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中。
eliftrain_phase==2:
fori,tmp_datainenumerate(current_dataset):
#tokenizethetext
chosen_sentence=raw_dataset.get_prompt_and_chosen(
tmp_data)#theacceptresponse
reject_sentence=raw_dataset.get_prompt_and_rejected(
tmp_data)#theacceptresponse
ifchosen_sentenceisnotNoneandreject_sentenceisnotNone:
chosen_sentence+=end_of_conversation_token#theacceptresponse
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中。
eliftrain_phase==3:
fori,tmp_datainenumerate(current_dataset):
#tokenizethetext
prompt=raw_dataset.get_prompt(tmp_data)
ifpromptisnotNone:
prompt_token=tokenizer(prompt,return_tensors="pt")
prompt_token["input_ids"]=prompt_token["input_ids"]
prompt_token["attention_mask"]=prompt_token["attention_mask"]
forkey_wordin["input_ids","attention_mask"]:
length=prompt_token[key_word].size()[-1]
iflength>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)
#創(chuàng)建一個新的PromptDataset對象，并返回。這個對象包含了對話提示、接受的對話和被拒絕的對話的數(shù)據(jù)集，以及分詞器的填充標記ID和訓練階段。
returnPromptDataset(prompt_dataset,chosen_dataset,reject_dataset,
tokenizer.pad_token_id,train_phase)


	

	接下來再解析一下create_dataset這個函數(shù)：

	
這段代碼定義了一個函數(shù)create_dataset，主要負責創(chuàng)建訓練數(shù)據(jù)集和評估數(shù)據(jù)集，具體的功能細節(jié)如下：
defcreate_dataset(local_rank,dataset_name,data_split,output_path,
train_phase,seed,tokenizer,end_of_conversation_token,
max_seq_len):
#調(diào)用get_raw_dataset函數(shù)，該函數(shù)根據(jù)提供的數(shù)據(jù)集名稱、輸出路徑、隨機種子和local_rank等參數(shù)，從各種預定義的數(shù)據(jù)集中獲取所需的原始數(shù)據(jù)集。
raw_dataset=get_raw_dataset(dataset_name,output_path,seed,local_rank)
train_dataset=raw_dataset.get_train_data()#從原始數(shù)據(jù)集中獲取訓練數(shù)據(jù)。
#獲取訓練數(shù)據(jù)集的索引，涉及數(shù)據(jù)的分割。
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))
#根據(jù)上一步獲取的索引，創(chuàng)建訓練數(shù)據(jù)的子集。
train_dataset=Subset(train_dataset,train_index)
#調(diào)用create_dataset_split函數(shù)對上一步獲得的數(shù)據(jù)子集進行進一步處理，
#這可能包括對文本的標記化(tokenization)，并且創(chuàng)建一個PromptDataset對象。
train_dataset=create_dataset_split(train_dataset,raw_dataset,
train_phase,tokenizer,
end_of_conversation_token,
max_seq_len)
#是用于創(chuàng)建評估數(shù)據(jù)集的，步驟與訓練數(shù)據(jù)集的創(chuàng)建基本相同。
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)
returntrain_dataset,eval_dataset


	

	這里需要額外注意一點，create_dataset_split 這個函數(shù)傳入的參數(shù)是 train_phase，而不是像 get_raw_dataset_split_index 函數(shù)那樣傳入 train_phase-1，這是為什么？這是因為train_phase用來標識在模型訓練過程中的哪個階段，是一個直接的指示符，與數(shù)據(jù)分割或索引無關。它用來在create_dataset_split函數(shù)中確定當前處于哪個訓練階段，并根據(jù)不同的訓練階段對數(shù)據(jù)進行不同的處理。而在調(diào)用get_raw_dataset_split_index函數(shù)時，傳入的是train_phase - 1，這是因為在這個函數(shù)中，我們需要根據(jù)當前訓練階段的前一個階段（由train_phase - 1表示）的數(shù)據(jù)分割或索引情況，來決定如何對當前階段的數(shù)據(jù)進行分割或索引。create_dataset_split主要關注如何根據(jù)訓練階段對數(shù)據(jù)進行處理，而get_raw_dataset_split_index主要關注如何根據(jù)前一個訓練階段的數(shù)據(jù)分割或索引情況，對當前階段的數(shù)據(jù)進行分割或索引。

	回收本節(jié)的開頭，我們解析create_prompt_dataset函數(shù)：

	
#這個函數(shù)的主要目的是創(chuàng)建一個包含訓練和評估數(shù)據(jù)集的“提示”數(shù)據(jù)集，并將這兩個數(shù)據(jù)集保存在指定的文件中。具體來說：
defcreate_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):
"""
Createsthepromptdataset
"""
#os.makedirs(output_path,exist_ok=True):創(chuàng)建輸出目錄，如果目錄已經(jīng)存在則不會引發(fā)異常。
os.makedirs(output_path,exist_ok=True)
#構(gòu)造文件名，這個文件名包含了很多有關數(shù)據(jù)集和模型的信息，如數(shù)據(jù)路徑、數(shù)據(jù)分割、訓練階段、
#隨機種子、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(
)#hashthefilenametoavoidtoolongfilename
#構(gòu)造訓練數(shù)據(jù)集和評估數(shù)據(jù)集的文件路徑。
train_fname=f"{output_path}/traindata_{fname}.pt"
eval_fname=f"{output_path}/evaldata_{fname}.pt"

#檢查訓練數(shù)據(jù)集和評估數(shù)據(jù)集的文件是否都已經(jīng)存在，如果存在，則表示緩存已經(jīng)找到，否則表示需要創(chuàng)建緩存。
cache_found=os.path.isfile(train_fname)andos.path.isfile(eval_fname)
buf_create_cache=torch.ByteTensor([notcache_found]).cuda()
torch.distributed.all_reduce(buf_create_cache)

#如果當前進程是主進程（local_rank<=?0）并且需要創(chuàng)建緩存，就執(zhí)行以下操作。
????if?local_rank?<=?0?and?(buf_create_cache.item()?!=?0?or?reload):
????????#?如果只有一個數(shù)據(jù)集，直接調(diào)用create_dataset函數(shù)創(chuàng)建訓練數(shù)據(jù)集和評估數(shù)據(jù)集。
????????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.
????????????#?如果有多個數(shù)據(jù)集，對每個數(shù)據(jù)集都調(diào)用create_dataset函數(shù)，并把得到的訓練數(shù)據(jù)集和評估數(shù)據(jù)集添加到對應的列表中，

????????????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函數(shù)合并數(shù)據(jù)集。
????????????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的數(shù)據(jù)集，那么對這些數(shù)據(jù)集執(zhí)行類似的操作，
????????#?然后把得到的訓練數(shù)據(jù)集和評估數(shù)據(jù)集添加到原有的數(shù)據(jù)集中。
????????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())
????????#?把訓練數(shù)據(jù)集和評估數(shù)據(jù)集保存到對應的文件中。
????????torch.save(train_dataset,?train_fname)
????????torch.save(eval_dataset,?eval_fname)
????#?在多進程環(huán)境中，確保所有進程都完成了數(shù)據(jù)集的保存操作。
????torch.distributed.barrier()
????return?torch.load(train_fname),?torch.load(eval_fname)



	

	工具函數(shù)解析

	這里解析一下下面導入的代碼，在DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/utils.py這個路徑下：

	
fromutils.utilsimportprint_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也就是masterrank打印信息，防止每個機器或GPU都打印消息造成大量重復信息
defprint_rank_0(msg,rank=0):
ifrank<=?0:
????????print(msg)


	

	to_device

	
#這個函數(shù)的作用是把一個數(shù)據(jù)批次移動到指定的設備上。在PyTorch中，數(shù)據(jù)默認是在CPU上的，
#如果要在GPU上進行運算，就需要先把數(shù)據(jù)移到GPU上。這個函數(shù)通過遍歷批次中的所有元素并
#調(diào)用to(device)方法來實現(xiàn)這一點。如果某個元素不能被移到指定的設備上
#（例如，這個元素是一個字符串或者其他不能在GPU上使用的類型），那么就直接保留這個元素，不進行任何操作。
defto_device(batch,device):
output={}
fork,vinbatch.items():
try:
output[k]=v.to(device)
except:
output[k]=v
returnoutput


	

	save_hf_format

	
#這段代碼的功能是將訓練好的模型以HuggingFace格式保存，這樣之后就可以使用HuggingFace庫的from_pretrained方法加載了。
defsave_hf_format(model,tokenizer,args,sub_folder=""):
#usedtosavehuggingfaceformat,sowecanuseitforhf.from_pretrained
#首先，這行代碼檢查模型是否有'module'這個屬性。這在PyTorch中是一個常見的模式，
#當模型被封裝在torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel時，
#模型的所有屬性都會被存儲在'module'屬性中。所以這行代碼的目的是確保我們總是在原始模型上進行操作，而不是并行化的包裝器。
model_to_save=model.moduleifhasattr(model,'module')elsemodel
#定義模型的配置和權(quán)重的文件名。
CONFIG_NAME="config.json"
WEIGHTS_NAME="pytorch_model.bin"
#定義模型保存的目錄路徑。
output_dir=os.path.join(args.output_dir,sub_folder)
#如果模型保存的目錄不存在，那么創(chuàng)建這個目錄。
os.makedirs(output_dir,exist_ok=True)
#定義模型權(quán)重和配置的保存路徑。
output_model_file=os.path.join(output_dir,WEIGHTS_NAME)
output_config_file=os.path.join(output_dir,CONFIG_NAME)
#獲取模型的狀態(tài)字典，這個字典包含了模型的所有權(quán)重。
save_dict=model_to_save.state_dict()
#遍歷狀態(tài)字典中的所有鍵。這里使用list(save_dict.keys())而不是直接使用
#save_dict.keys()是因為我們在遍歷過程中會修改字典，直接遍歷字典的鍵在Python中會導致錯誤。
forkeyinlist(save_dict.keys()):
#if"lora"inkey:：如果鍵中包含"lora"，那么刪除這個鍵值對。
#"lora"可能是某種特定的模型組件或參數(shù)，這里將其排除在保存的模型權(quán)重之外。
if"lora"inkey:
delsave_dict[key]
#將處理過的狀態(tài)字典保存到文件中。
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

	
#這段代碼的目的是為所有可能用到的隨機數(shù)生成器設定種子，以確保實驗結(jié)果的可復現(xiàn)性。
#隨機數(shù)生成器的種子是生成隨機數(shù)序列的初始值，當我們用同一個種子初始化隨機數(shù)生成器時，它會生成完全一樣的隨機數(shù)序列。
defset_random_seed(seed):
#首先檢查種子是否是None。如果是None，那么就跳過這段代碼，隨機數(shù)生成器將會用一個隨機的種子。
ifseedisnotNone:
set_seed(seed)#這一行調(diào)用transformers庫的set_seed的函數(shù)，將種子設定為指定的值。
random.seed(seed)#將Python內(nèi)置的隨機數(shù)生成器的種子設定為指定的值。
np.random.seed(seed)#將NumPy的隨機數(shù)生成器的種子設定為指定的值。
torch.manual_seed(seed)#將PyTorch的隨機數(shù)生成器的種子設定為指定的值。
torch.cuda.manual_seed_all(seed)#將PyTorch的所有GPU隨機數(shù)生成器的種子設定為指定的值。


	

	get_all_reduce_mean

	
#這段代碼是在分布式訓練環(huán)境中進行平均值計算的函數(shù)，通過這段代碼，
#所有的處理器（或者叫節(jié)點）上的同一個tensor都會被加和起來，然后除以總的處理器數(shù)，得到平均值。
defget_all_reduce_mean(tensor):
#這行代碼執(zhí)行一個分布式的reduce操作。reduce操作是指所有處理器中的同一個tensor都被某種方式結(jié)合起來。
#在這個例子中，torch.distributed.ReduceOp.SUM表示所有處理器上的tensor將被加和起來。
#加和的結(jié)果會在所有處理器上都可用。
torch.distributed.all_reduce(tensor,op=torch.distributed.ReduceOp.SUM)
#這行代碼將前一步得到的加和結(jié)果除以處理器的數(shù)量（也叫作worldsize）。
#這樣，tensor就變成了所有處理器上原始tensor的平均值。
tensor=tensor/torch.distributed.get_world_size()
#最后，這個平均值tensor被返回。在所有處理器上，這個函數(shù)返回的tensor都是相同的，
#等于所有處理器上原始tensor的平均值。
returntensor


	

	get_optimizer_grouped_parameters

	
#這段代碼的作用是將模型中的參數(shù)分組以便于在優(yōu)化器中使用。它將模型參數(shù)分為兩組：
#一組需要進行權(quán)重衰減（L2正則化）的參數(shù)，另一組不需要進行權(quán)重衰減的參數(shù)。
defget_optimizer_grouped_parameters(model,
weight_decay,
no_decay_name_list=[
"bias","LayerNorm.weight"
]):
#它定義了一個列表optimizer_grouped_parameters，其中包含兩個字典。每個字典都對應一個參數(shù)組，包含"params"和"weight_decay"這兩個關鍵字。
optimizer_grouped_parameters=[
#在第一個字典中，它從模型參數(shù)中選出那些名稱不包含"bias"或"LayerNorm.weight"
#且需要求梯度的參數(shù)。這些參數(shù)在優(yōu)化過程中會應用weight_decay作為權(quán)重衰減項。
{
"params":[
pforn,pinmodel.named_parameters()
if(notany(ndinn
forndinno_decay_name_list)andp.requires_grad)
],
"weight_decay":
weight_decay,
},
#在第二個字典中，它選出那些名稱包含"bias"或"LayerNorm.weight"且需要求梯度的參數(shù)。
#這些參數(shù)在優(yōu)化過程中不會應用權(quán)重衰減，即其"weight_decay"值為0。

{
"params":[
pforn,pinmodel.named_parameters()
if(any(ndinn
forndinno_decay_name_list)andp.requires_grad)
],
"weight_decay":
0.0,
},
]
returnoptimizer_grouped_parameters


	

	這種參數(shù)的分組策略是很常見的。比如在訓練Transformer模型時，通常會為權(quán)重和偏置項設定不同的學習策略。這是因為權(quán)重衰減對于防止過擬合很有幫助，但對于某些參數(shù)（如偏置項或者層歸一化的權(quán)重）可能會導致性能下降，因此常常會排除這些參數(shù)不進行權(quán)重衰減。

	save_zero_three_model

	
#這個函數(shù)的主要功能是篩選出那些在DeepSpeedZero3優(yōu)化中被離線存儲，但在當前還未獲取的參數(shù)。
#在DeepSpeedZero3優(yōu)化中，一些模型參數(shù)在使用過后會被離線存儲，以此釋放GPU顯存。
#當這些參數(shù)需要再次被使用時，需要先獲取到本地。
def_z3_params_to_fetch(param_list):
#這個條件語句判斷一個參數(shù)是否是被DeepSpeedZero3優(yōu)化過的，且其狀態(tài)為"未獲取"（NOT_AVAILABLE）。
#對于被DeepSpeedZero3優(yōu)化過的參數(shù)，它們有一個ds_id屬性和一個ds_status屬性，其中ds_status表示參數(shù)的當前狀態(tài)。
return[
pforpinparam_list
ifhasattr(p,'ds_id')andp.ds_status==ZeroParamStatus.NOT_AVAILABLE
]

#這個函數(shù)的主要作用是保存一個使用了DeepSpeedZero優(yōu)化（可能為stage3）的模型。
#DeepSpeed的Zero優(yōu)化技術是為了解決模型參數(shù)、優(yōu)化器狀態(tài)和梯度等內(nèi)存占用問題，
#通過這種方式，可以訓練比當前GPU內(nèi)存更大的模型。
defsave_zero_three_model(model_ema,global_rank,save_dir,zero_stage=0):
#首先，檢查輸入的zero_stage是否為3，確定是否使用了DeepSpeedZero階段3優(yōu)化。
zero_stage_3=(zero_stage==3)
#然后，確保保存模型的目錄存在。
os.makedirs(save_dir,exist_ok=True)
#定義模型權(quán)重保存的完整路徑。
WEIGHTS_NAME="pytorch_model.bin"
output_model_file=os.path.join(save_dir,WEIGHTS_NAME)

#如果模型是被包裹在其它結(jié)構(gòu)（如DataParallel或DistributedDataParallel）中的，我們需要取出真實的模型實例。
model_to_save=model_ema.moduleifhasattr(model_ema,
'module')elsemodel_ema
ifnotzero_stage_3:
#如果沒有使用Zero階段3優(yōu)化，直接使用PyTorch的torch.save函數(shù)保存模型狀態(tài)。
ifglobal_rank==0:
torch.save(model_to_save.state_dict(),output_model_file)
else:
#如果使用了Zero階段3優(yōu)化，因為模型的部分參數(shù)和優(yōu)化器狀態(tài)在不同的設備上，所以需要先將它們收集起來。
output_state_dict={}
fork,vinmodel_to_save.named_parameters():

ifhasattr(v,'ds_id'):
#deepspeed.zero.GatheredParameters是DeepSpeed提供的一個上下文管理器，
#它可以將分布在多個設備上的參數(shù)收集到一起。這部分參數(shù)保存在CPU上。
withdeepspeed.zero.GatheredParameters(_z3_params_to_fetch([v
]),
enabled=zero_stage_3):
v_p=v.data.cpu()
else:
v_p=v.cpu()
#然后，將收集好的參數(shù)（并且不包含“l(fā)ora”關鍵字的參數(shù)）添加到輸出狀態(tài)字典中。
ifglobal_rank==0and"lora"notink:
output_state_dict[k]=v_p
#最后，再使用torch.save函數(shù)保存模型狀態(tài)。
ifglobal_rank==0:
torch.save(output_state_dict,output_model_file)
#同時為了節(jié)省內(nèi)存，使用del關鍵字刪除了存儲參數(shù)的字典。
deloutput_state_dict



	

	load_hf_tokenizer

	
#這個函數(shù)的目標是加載一個HuggingFacetokenizer，這個tokenizer是用來將文本數(shù)據(jù)轉(zhuǎn)化為模型可以接受的形式的。
#model_name_or_path:這個參數(shù)可以是一個模型名字或者一個本地路徑，指示從哪里獲取tokenizer。
#fast_tokenizer=True:這個參數(shù)指示是否使用HuggingFace的"fast"tokenizer。"fast"tokenizer是用Rust編寫的，運行速度更快。
defload_hf_tokenizer(model_name_or_path,fast_tokenizer=True):
#這行代碼檢查model_name_or_path是否是一個存在的路徑。如果是，那么就嘗試從本地文件系統(tǒng)中加載tokenizer。
ifos.path.exists(model_name_or_path):
#Locallytokenizerloadinghassomeissue,soweneedtoforcedownload
#這行代碼拼接路徑，嘗試找到模型的config文件。
model_json=os.path.join(model_name_or_path,"config.json")
#這行代碼檢查config文件是否存在。如果存在，那么就嘗試從config文件中獲取模型名字。
ifos.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)#這行代碼使用模型名字從HuggingFace的模型庫中下載并加載tokenizer。
else:#如果model_name_or_path不是一個存在的路徑，那么就認為它是一個模型名字。
tokenizer=AutoTokenizer.from_pretrained(model_name_or_path,
fast_tokenizer=True)#這行代碼使用模型名字從HuggingFace的模型庫中下載并加載tokenizer。
returntokenizer


	

	再解析一下DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/ds_utils.py導入的get_train_ds_config函數(shù)：

	
#DeepSpeedTeam
GLOBAL_BATCH_SIZE=32
MICRO_BATCH_SIZE=4

#這個函數(shù)的主要目標是生成一個字典，用來配置DeepSpeed訓練模型的各種參數(shù)。
#DeepSpeed是一個用于加速大規(guī)模模型訓練的庫，其中包括一種叫做ZeRO（ZeroRedundancyOptimizer）的優(yōu)化技術，
#可以大大減少模型訓練過程中的顯存和內(nèi)存占用。
defget_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。如果參數(shù)offload為True，那么就將device設置為"cpu"，
#表示模型的參數(shù)和優(yōu)化器狀態(tài)將被"offload"到CPU內(nèi)存中，以減少GPU顯存的占用。
#如果offload為False，那么device就被設置為"none"，表示不進行offload。
device="cpu"ifoffloadelse"none"
#這行代碼定義了一個字典zero_opt_dict，包含了ZeRO優(yōu)化器的各種配置參數(shù)。
zero_opt_dict={
"stage":stage,#ZeRO的階段，可以為0、1、2、3，階段越高，顯存優(yōu)化效果越好，但計算復雜度也越高。
"offload_param":{#offload_param和offload_optimizer這兩個字段控制是否將模型的參數(shù)和優(yōu)化器狀態(tài)offload到CPU內(nèi)存。
"device":device
},
"offload_optimizer":{
"device":device
},
#這三個字段只有在ZeROstage3時有效，用于控制參數(shù)offloading的細節(jié)。
"stage3_param_persistence_threshold":1e4,
"stage3_max_live_parameters":3e7,
"stage3_prefetch_bucket_size":3e7,
#這個字段控制是否使用顯存更高效的線性層實現(xiàn)。由于當前的顯存高效實現(xiàn)可能會導致數(shù)值不穩(wěn)定，所以這里默認設為False。
"memory_efficient_linear":False
}
#返回一個配置字典，這個字典包含了訓練模型的各種參數(shù)。
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":{#控制是否使用混合精度訓練，以及l(fā)ossscaling的窗口大小。
"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 這幾個函數(shù)。

	在解析這幾個函數(shù)之前需要先解析一下 LinearLayer_LoRA 這個類：

	
#LinearLayer_LoRA是一個實現(xiàn)了LoRA（Low-RankAdaptation）策略的PyTorch模塊。
#LoRA是一種在大型預訓練模型（如GPT-3或者BERT）微調(diào)中使用的策略，可以提高微調(diào)的效果。
#這個類是一個簡單的LoRA實現(xiàn)，只支持線性層。
classLinearLayer_LoRA(nn.Module):
#ansimpleimplementationofLoRA
#fornowonlysupportLinearLayer
#初始化函數(shù)，定義了一些參數(shù)，如權(quán)重（weight）、LoRA維度（lora_dim）、
#LoRA縮放因子（lora_scaling）、LoRAdropout（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，如果不是，就拋出異常。
iflora_dim<=?0:
????????????raise?ValueError(
????????????????"You?are?training?to?use?LoRA,?whose?reduced?dim?should?be?larger?than?1"
????????????)
????????#?然后嘗試獲取權(quán)重的形狀，如果權(quán)重是在?ZeRO?階段3中的參數(shù)，那么通過?ds_shape?屬性獲取形狀；
????????#?否則直接使用?shape?屬性。
????????try:
????????????#?for?zero?stage?3
????????????rows,?columns?=?weight.ds_shape
????????except:
????????????rows,?columns?=?weight.shape
????????#?然后初始化?LoRA?的左右兩個權(quán)重矩陣，并設置它們?yōu)榭蓪W習的參數(shù)。
????????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，則創(chuàng)建一個?Dropout?層；否則創(chuàng)建一個?Identity?層。
????????if?lora_droppout?>0:
self.lora_dropout=nn.Dropout(lora_droppout)
else:
self.lora_dropout=nn.Identity()
#然后調(diào)用reset_parameters方法對LoRA權(quán)重進行初始化。
self.reset_parameters()
#disabletheoriginalweightgradient
#最后，關閉原始權(quán)重的梯度，設置LoRA融合標志位為False。
self.weight.requires_grad=False
#fuseLoRAtotheoriginalweight
self.fuse_lora=False

defeval(self):
self.lora_dropout.eval()#將模型設置為評估模式，這時候Dropout層會停止工作。

#self.fuse_lora_weight()

deftrain(self,mode=True):
self.lora_dropout.train(mode)#將模型設置為訓練模式，這時候Dropout層會開始工作。
#self.unfuse_lora_weight()

defreset_parameters(self):#初始化LoRA權(quán)重的方法。右權(quán)重使用kaiming均勻分布進行初始化，左權(quán)重初始化為全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權(quán)重融合到原始權(quán)重中，或者從原始權(quán)重中解融合。
#融合操作實質(zhì)上是將原始權(quán)重與LoRA權(quán)重的乘積（縮放后）相加。
deffuse_lora_weight(self):
ifnotself.fuse_lora:
self.weight.data+=self.lora_scaling*torch.matmul(
self.lora_left_weight.t(),self.lora_right_weight.t())
self.fuse_lora=True

defunfuse_lora_weight(self):
ifself.fuse_lora:
self.weight.data-=self.lora_scaling*torch.matmul(
self.lora_left_weight.t(),self.lora_right_weight.t())
self.fuse_lora=False
#前向傳播函數(shù)。如果LoRA權(quán)重已融合，則直接對輸入進行線性變換；
#否則，會額外計算一個LoRA項，該項是輸入通過Dropout層，然后與LoRA權(quán)重相乘得到的。
defforward(self,input):
ifself.fuse_lora:
returnF.linear(input,self.weight,self.bias)
else:
returnF.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

	
#這個函數(shù)convert_linear_layer_to_lora是用來將模型中的線性層轉(zhuǎn)換為LoRA層的。
#在訓練深度學習模型時，這種方法能夠在保持預訓練模型參數(shù)不變的同時，通過添加額外的參數(shù)來微調(diào)模型。
#convertthelinearlayertoLoRA
defconvert_linear_layer_to_lora(model,
part_module_name,
lora_dim=0,
lora_scaling=1,
lora_droppout=0):
repalce_name=[]
#函數(shù)首先遍歷模型中的所有模塊（model.named_modules()），找出名稱中包含part_module_name的線性層（nn.Linear），
#并將這些層的名稱添加到repalce_name列表中。
forname,moduleinmodel.named_modules():
ifisinstance(module,nn.Linear)andpart_module_nameinname:
repalce_name.append(name)
#然后，函數(shù)遍歷repalce_name列表，使用recursive_getattr函數(shù)獲取模型中對應名稱的模塊。
#這些模塊是需要被替換成LoRA層的線性層。
fornameinrepalce_name:
module=recursive_getattr(model,name)
#對于每一個需要被替換的模塊，函數(shù)創(chuàng)建一個LinearLayer_LoRA實例tmp，
#并將其傳輸?shù)脚c原始模塊相同的設備和數(shù)據(jù)類型上。創(chuàng)建LinearLayer_LoRA實例時，
#需要傳入原始模塊的權(quán)重、偏置以及LoRA層的一些參數(shù)，如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)
#創(chuàng)建完LinearLayer_LoRA實例后，函數(shù)使用recursive_setattr函數(shù)將原始模塊替換為LinearLayer_LoRA實例。
recursive_setattr(model,name,tmp)
returnmodel


	

	convert_lora_to_linear_layer

	
#這個函數(shù)的主要功能是篩選出那些在DeepSpeedZero3優(yōu)化中被離線存儲，但在當前還未獲取的參數(shù)。
#在DeepSpeedZero3優(yōu)化中，一些模型參數(shù)在使用過后會被離線存儲，以此釋放GPU顯存。
#當這些參數(shù)需要再次被使用時，需要先獲取到本地。
def_z3_params_to_fetch(param_list):
#這個條件語句判斷一個參數(shù)是否是被DeepSpeedZero3優(yōu)化過的，且其狀態(tài)為"未獲取"（NOT_AVAILABLE）。
#對于被DeepSpeedZero3優(yōu)化過的參數(shù)，它們有一個ds_id屬性和一個ds_status屬性，其中ds_status表示參數(shù)的當前狀態(tài)。
return[
pforpinparam_list
ifhasattr(p,'ds_id')andp.ds_status==ZeroParamStatus.NOT_AVAILABLE
]

#這個函數(shù)convert_lora_to_linear_layer是用來將模型中的LoRA層轉(zhuǎn)換回線性層的。
#在訓練深度學習模型時，這個操作可以用于在訓練完LoRA層后，將模型恢復到原始的狀態(tài)，
#以便進行下一步的操作，如模型的保存、加載等。
#converttheLoRAlayertolinearlayer
defconvert_lora_to_linear_layer(model):
repalce_name=[]
#函數(shù)首先遍歷模型中的所有模塊（model.named_modules()），找出所有的LoRA層（LinearLayer_LoRA），
#并將這些層的名稱添加到repalce_name列表中。
forname,moduleinmodel.named_modules():
ifisinstance(module,LinearLayer_LoRA):
repalce_name.append(name)
#然后，函數(shù)遍歷repalce_name列表，使用recursive_getattr函數(shù)獲取模型中對應名稱的LoRA層。
fornameinrepalce_name:
module=recursive_getattr(model,name)
#對于每一個LoRA層，函數(shù)首先檢查是否處于zerostage3（DeepSpeed的一個特性，用于在多GPU訓練中減少內(nèi)存占用）。
#如果是，則設置zero_stage_3為True。
zero_stage_3=hasattr(module.weight,'ds_id')
withdeepspeed.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()
returnmodel


	

	only_optimize_lora_parameters

	
#這個函數(shù)的作用是關閉模型中除LoRA參數(shù)之外的所有參數(shù)的梯度。這意味著在訓練過程中，只有LoRA參數(shù)會被優(yōu)化，其他參數(shù)保持不變。
defonly_optimize_lora_parameters(model):
#turnoffthegradientofalltheparametersexcepttheLoRAparameters
#遍歷模型的所有參數(shù)。每個參數(shù)都有一個唯一的名稱name和對應的參數(shù)值param。
forname,paraminmodel.named_parameters():
#查當前參數(shù)的名稱是否包含lora_right_weight或lora_left_weight。
#這是因為在LoRA（Low-RankAdaptation）中，只有這兩種參數(shù)是需要優(yōu)化的。
if"lora_right_weight"innameor"lora_left_weight"inname:
#param.requires_grad=True如果參數(shù)名包含lora_right_weight或lora_left_weight，
#則設置參數(shù)的requires_grad屬性為True，表示需要對此參數(shù)進行梯度下降優(yōu)化。
param.requires_grad=True
else:
param.requires_grad=False
returnmodel


	

	create_hf_model

	
#這段代碼定義了一個名為create_hf_model的函數(shù)，該函數(shù)的作用是創(chuàng)建或加載一個預訓練模型。該函數(shù)的主要參數(shù)包括：
#model_class：模型的類別，例如GPT-2、BERT等。
#tokenizer：用于模型的分詞器。
#ds_config:DeepSpeed的配置參數(shù)。
#rlhf_training：一個標志，用來表示是否正在進行RLHF（ReinforcementLearningfromHumanFeedback）訓練。
#disable_dropout：一個標志，用來表示是否禁用dropout。Dropout是一種防止過擬合的技術。
defcreate_hf_model(model_class,
model_name_or_path,
tokenizer,
ds_config=None,
rlhf_training=False,
disable_dropout=False):
#根據(jù)model_name_or_path從預訓練模型獲取模型配置model_config。
model_config=AutoConfig.from_pretrained(model_name_or_path)
#如果disable_dropout為真，則將模型配置中的dropout設為0.0。
ifdisable_dropout:
model_config.dropout=0.0
#Note:dschfisdefinedinfunctionscopetoavoidglobaleffects
#https://huggingface.co/docs/transformers/main_classes/deepspeed#nontrainer-deepspeed-integration
#根據(jù)ds_config中的設置，創(chuàng)建DeepSpeed的配置對象dschf，以便進行DeepSpeed優(yōu)化。
ifds_configisnotNoneandds_config["zero_optimization"]["stage"]==3:
dschf=HfDeepSpeedConfig(ds_config)
else:
dschf=None
#根據(jù)rlhf_training的值，確定是從配置中創(chuàng)建模型還是從預訓練模型中加載模型。如果rlhf_training為真，則根據(jù)模型配置創(chuàng)建模型；否則，從預訓練模型加載模型。
ifrlhf_training:
#theweightloadingishandledbycreatecriticmodel
model=model_class.from_config(model_config)
else:
model=model_class.from_pretrained(
model_name_or_path,
from_tf=bool(".ckpt"inmodel_name_or_path),
config=model_config)
#將模型的結(jié)束標記和填充標記設為分詞器的結(jié)束標記id。
model.config.end_token_id=tokenizer.eos_token_id
model.config.pad_token_id=model.config.eos_token_id
#調(diào)整模型的詞匯表大小，使其為8的倍數(shù)。這樣做的目的是為了在某些硬件（如GPU）上提高效率。
model.resize_token_embeddings(int(
8*
math.ceil(len(tokenizer)/8.0)))#makethevocabsizemultipleof8

returnmodel


	

	0x3.2 main.py主體解析

	parse_args解析

	
defparse_args():
#創(chuàng)建一個argparse的解析器對象，這個對象可以添加命令行參數(shù)和處理它們。description參數(shù)提供了一個對程序的簡單描述。
parser=argparse.ArgumentParser(
description=
"Finetuneatransformersmodelonacausallanguagemodelingtask")
parser.add_argument('--data_path',
nargs='*',
default=['Dahoas/rm-static'],
help='Pathtothetrainingdataset.Acceptedformat:'
'1)asingledatapath,2)multipledatasetsinthe'
'form:dataset1-pathdataset2-path...')
parser.add_argument('--data_split',
type=str,
default='2,4,4',
help='Comma-separatedlistofproportionsfortraining'
'phase1,2,and3data.Forexamplethesplit`6,2,2`'
'willuse60%ofdataforphase1,20%forphase2'
'and20%forphase3.')
parser.add_argument(
'--sft_only_data_path',
nargs='*',
default=[],
help='PathtothedatasetforonlyusinginSFTphase.')
parser.add_argument(
'--data_output_path',
type=str,
default='/data_turbo/home/zhangxiaoyu/data_files/',
help=
'Wheretostorethedata-relatedfilessuchasshuffleindex.Thisneedstobeonalocalstorageofanode(notonasharedstorage)'
)
parser.add_argument(
"--model_name_or_path",
type=str,
help=
"Pathtopretrainedmodelormodelidentifierfromhuggingface.co/models.",
required=True,
)
parser.add_argument(
"--per_device_train_batch_size",
type=int,
default=16,
help="Batchsize(perdevice)forthetrainingdataloader.",
)
parser.add_argument(
"--per_device_eval_batch_size",
type=int,
default=16,
help="Batchsize(perdevice)fortheevaluationdataloader.",
)
parser.add_argument(
"--max_seq_len",
type=int,
default=512,
help="Themaximumsequencelength.",
)
parser.add_argument(
"--learning_rate",
type=float,
default=1e-3,
help=
"Initiallearningrate(afterthepotentialwarmupperiod)touse.",
)
parser.add_argument("--weight_decay",
type=float,
default=0.,
help="Weightdecaytouse.")
parser.add_argument("--num_train_epochs",
type=int,
default=1,
help="Totalnumberoftrainingepochstoperform.")
parser.add_argument(
"--gradient_accumulation_steps",
type=int,
default=1,
help=
"Numberofupdatesstepstoaccumulatebeforeperformingabackward/updatepass.",
)
parser.add_argument(
"--lr_scheduler_type",
type=SchedulerType,
default="cosine",
help="Theschedulertypetouse.",
choices=[
"linear","cosine","cosine_with_restarts","polynomial",
"constant","constant_with_warmup"
],
)
parser.add_argument(
"--num_warmup_steps",
type=int,
default=0,
help="Numberofstepsforthewarmupinthelrscheduler.")
parser.add_argument("--output_dir",
type=str,
default=None,
help="Wheretostorethemodel.")
parser.add_argument("--seed",
type=int,
default=1234,
help="Aseedforreproducibletraining.")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rankfordistributedtrainingongpus")
parser.add_argument('--gradient_checkpointing',
action='store_true',
help='EnableHFgradientcheckpointingformodel.')
parser.add_argument('--disable_dropout',
action='store_true',
help='Disablethedropoutofthemodel.')
#deepspeedfeatures
parser.add_argument('--offload',
action='store_true',
help='EnableZeROOffloadtechniques.')
parser.add_argument(
'--zero_stage',
type=int,
default=0,
help='ZeROoptimizationstageforActormodel(andclones).')
##LoRAforefficienttrainingsetting
parser.add_argument("--lora_dim",
type=int,
default=0,
help="If>0,useLoRAforefficienttraining.")
parser.add_argument("--lora_module_name",
type=str,
default="decoder.layers.",
help="ThescopeofLoRA.")
parser.add_argument('--only_optimize_lora',
action='store_true',
help='OnlyoptimizetheLoRAparameters.')
#這一行將DeepSpeed的配置參數(shù)添加到解析器中。
parser=deepspeed.add_config_arguments(parser)
#這一行解析命令行參數(shù)并將它們存儲在args對象中
args=parser.parse_args()

#Validatesettings
#在這個代碼塊中，驗證一些特定的參數(shù)設置是否合法。
#例如，如果同時啟用了gradient_checkpointing和僅優(yōu)化LoRA參數(shù)，那么將會拋出一個錯誤。
ifargs.gradient_checkpointingandargs.lora_dim>0:
assert(
notargs.only_optimize_lora
),"--gradient_checkpointingand--only_optimize_loracannotbeenabledatthesametime."

returnargs


	

	main函數(shù)解析

	接下來是訓練部分的核心函數(shù)，也是全文最后一個函數(shù)main的解析。

	
#這個函數(shù)是主函數(shù)，是訓練語言模型的主流程，主要步驟包括解析命令行參數(shù)、
#設置設備、準備數(shù)據(jù)、定義模型、配置優(yōu)化器和學習率調(diào)度器、進行訓練和評估等。
defmain():
#解析命令行參數(shù)。
args=parse_args()

#如果本地排名為-1，說明不在分布式訓練環(huán)境下，設備設置為"cuda"；
#否則根據(jù)args.local_rank設置對應的cuda設備，并初始化分布式訓練。
ifargs.local_rank==-1:
device=torch.device("cuda")
else:
torch.cuda.set_device(args.local_rank)
device=torch.device("cuda",args.local_rank)
#Initializesthedistributedbackendwhichwilltakecareofsychronizingnodes/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

#Ifpassedalong,setthetrainingseednow.
#設置隨機種子以保證結(jié)果的可復現(xiàn)性。
set_random_seed(args.seed)

#在此處添加一個barrier操作，確保所有進程都執(zhí)行到這一點后再繼續(xù)執(zhí)行后續(xù)操作。
torch.distributed.barrier()
#加載預訓練模型對應的分詞器。
tokenizer=load_hf_tokenizer(args.model_name_or_path,fast_tokenizer=True)
tokenizer.pad_token=tokenizer.eos_token
#makesuretokenizerisrightpadinourlogic
tokenizer.padding_side='right'
#創(chuàng)建預訓練模型。
model=create_hf_model(AutoModelForCausalLM,
args.model_name_or_path,
tokenizer,
ds_config,
disable_dropout=args.disable_dropout)

#如果參數(shù)lora_dim大于0，將模型的線性層轉(zhuǎn)換為LoRa層；如果只優(yōu)化LoRa參數(shù)，關閉其他參數(shù)的梯度。
ifargs.lora_dim>0:
model=convert_linear_layer_to_lora(model,args.lora_module_name,
args.lora_dim)
ifargs.only_optimize_lora:
model=only_optimize_lora_parameters(model)

#Preparethedata
#創(chuàng)建數(shù)據(jù)集和數(shù)據(jù)加載器：包括訓練集和驗證集，以及對應的采樣器和數(shù)據(jù)加載器。
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)
#DataLoaderscreation:
ifargs.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)
#定義模型評估函數(shù)，用于計算模型在驗證集上的困惑度。
defevaluation(model,eval_dataloader):
model.eval()#將模型切換為評估模式。
losses=0#初始化loss。
forstep,batchinenumerate(eval_dataloader):#對于評估數(shù)據(jù)集中的每一個batch。
batch=to_device(batch,device)#將batch數(shù)據(jù)移到對應的設備上。
withtorch.no_grad():#在此上下文管理器中，不計算梯度，這樣可以節(jié)省存儲和計算資源。
#將batch數(shù)據(jù)輸入模型，進行前向計算。
outputs=model(**batch)

loss=outputs.loss#取出模型的輸出中的loss。
losses+=loss.float()#將當前的loss累加到總的losses中。
losses=losses/(step+1)#計算平均的loss。
try:
perplexity=torch.exp(losses)#嘗試計算模型的困惑度，如果捕捉到溢出錯誤，將困惑度設置為無窮大。
exceptOverflowError:
perplexity=float("inf")
#嘗試在所有設備上計算困惑度的平均值，如果發(fā)生任何錯誤，就忽略。
try:
perplexity=get_all_reduce_mean(perplexity).item()
except:
pass
returnperplexity

#根據(jù)是否使用權(quán)重衰減將模型參數(shù)分為兩組。
#Splitweightsintwogroups,onewithweightdecayandtheothernot.
optimizer_grouped_parameters=get_optimizer_grouped_parameters(
model,args.weight_decay)
#選擇優(yōu)化器類型，如果啟用了梯度Offload，使用DeepSpeedCPUAdam，否則使用FusedAdam。
AdamOptimizer=DeepSpeedCPUAdamifargs.offloadelseFusedAdam
#創(chuàng)建優(yōu)化器。
optimizer=AdamOptimizer(optimizer_grouped_parameters,
lr=args.learning_rate,
betas=(0.9,0.95))
#計算每個epoch的更新步數(shù)。
num_update_steps_per_epoch=math.ceil(
len(train_dataloader)/args.gradient_accumulation_steps)
#創(chuàng)建學習率調(diào)度器。
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初始化模型、優(yōu)化器和學習率調(diào)度器。
model,optimizer,_,lr_scheduler=deepspeed.initialize(
model=model,
optimizer=optimizer,
args=args,
config=ds_config,
lr_scheduler=lr_scheduler,
dist_init_required=True)
#如果啟用了梯度檢查點，那么在模型中也啟用梯度檢查點。
ifargs.gradient_checkpointing:
model.gradient_checkpointing_enable()

#Train!
#使用print_rank_0函數(shù)在主節(jié)點（global_rank為0的節(jié)點）打印開始訓練的信息。
print_rank_0("*****Runningtraining*****",args.global_rank)
#在主節(jié)點打印在第0個epoch（訓練開始前）進行模型評估的信息。
print_rank_0(
f"*****Evaluatingperplexity,Epoch{0}/{args.num_train_epochs}*****",
args.global_rank)
#調(diào)用evaluation函數(shù)對模型進行評估，得到模型的困惑度。
perplexity=evaluation(model,eval_dataloader)
#在主節(jié)點打印模型的困惑度。
print_rank_0(f"ppl:{perplexity}",args.global_rank)

#循環(huán)args.num_train_epochs輪進行訓練。
forepochinrange(args.num_train_epochs):
print_rank_0(
f"BeginningofEpoch{epoch+1}/{args.num_train_epochs},TotalMicroBatches{len(train_dataloader)}",
args.global_rank)#在每個epoch開始時，在主節(jié)點打印開始新的訓練周期的信息。
model.train()#將模型設置為訓練模式。
forstep,batchinenumerate(train_dataloader):#對于訓練數(shù)據(jù)集中的每一個batch。
batch=to_device(batch,device)#將batch數(shù)據(jù)移到對應的設備上。
outputs=model(**batch,use_cache=False)#將batch數(shù)據(jù)輸入模型，進行前向計算。
loss=outputs.loss#取出模型的輸出中的loss。
model.backward(loss)#進行反向傳播，計算梯度。
model.step()#更新模型的參數(shù)。

#Evaluateperplexityonthevalidationset.
#在每個epoch結(jié)束后，在主節(jié)點打印開始評估的信息。
print_rank_0(
f"*****Evaluatingperplexity,Epoch{epoch+1}/{args.num_train_epochs}*****",
args.global_rank)
#對模型進行評估，得到模型的困惑度。
perplexity=evaluation(model,eval_dataloader)
#在主節(jié)點打印模型的困惑度。
print_rank_0(f"ppl:{perplexity}",args.global_rank)
#更新模型中的epoch計數(shù)。
model.tput_timer.update_epoch_count()

#如果設置了輸出目錄，進行以下操作。
ifargs.output_dirisnotNone:
#在主節(jié)點打印開始保存模型的信息。
print_rank_0('savingthefinalmodel...',args.global_rank)
#將模型中的LoRA層轉(zhuǎn)換為線性層。
model=convert_lora_to_linear_layer(model)

#如果是主節(jié)點，進行以下操作。
ifargs.global_rank==0:
#以HuggingFace的模型格式保存模型。
save_hf_format(model,tokenizer,args)
#如果使用了ZeroRedundancyOptimizer（Zero）的第三階段，進行以下操作。
ifargs.zero_stage==3:
#Forzerostage3,eachgpuonlyhasapartofthemodel,soweneedaspecialsavefunction
#使用特殊的保存函數(shù)保存模型。在Zero的第三階段，每個GPU只有模型的一部分，所以需要特殊的保存函數(shù)。
save_zero_three_model(model,
args.global_rank,
args.output_dir,
zero_stage=args.zero_stage)


	

	0x4. hybrid_engine的細節(jié) & log

	從訓練過程的輸出日志來看hybrid_engine是默認關閉的，DeepSpeed-Chat 打造類ChatGPT全流程 筆記一 里面提到DeepSpeed Hybrid Engine是用在加速 RLHF 流程中最耗時的部分也就是第三步，而本文介紹的監(jiān)督指令微調(diào)是第一步，所以即使開啟hybrid_engine加速效果應該也比較有限，所以這里默認關閉。

	hybrid_engine的優(yōu)化方法和原理在后續(xù)文章中繼續(xù)探索。

	這里分享一下我復現(xiàn)官方sample訓練的第一階段的log：https://paste.ubuntu.com/p/vcG49hQmCW/

	0x5. 總結(jié)

	這篇文章解析了DeepSpeed Chat中監(jiān)督指令微調(diào)這個過程的源碼，這個過程和一般的PyTorch DDP分布式訓練區(qū)別不是特別大，主要是自定義prompt數(shù)據(jù)集以及將普通的訓練流程中的組件如模型，優(yōu)化器，學習率調(diào)度器等等，使用DeepSpeed來warp一下，來用上DeepSpeed提供的Zero，Gradient Checkpoint（注意這個其實就是activation checkpoint）等特性。本文是完全按照訓練流程順序閱讀代碼，并補全了訓練過程中所有涉及到的工具函數(shù)或者新的特性如LoRA微調(diào)的代碼解析。DeepSpeed Chat這部分代碼寫得比較清晰易懂，因為是在接口層面來使用DeepSpeed，相當于基于DeepSpeed做應用所以代碼中不會涉及到DeepSpeed的底層代碼，只需要關注算法流程。但這個代碼在LoRA微調(diào)這部分感覺設計的耦合性有一點高，如果要新增新的微調(diào)方式比如QLoRA可能寫法就不太優(yōu)雅了。