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

文章目錄

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

系列文章

DeepSpeed-Chat 打造類(lèi)ChatGPT全流程 筆記一

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

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

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

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

0x0. 前言

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

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

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

如何訓(xùn)練模型

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

training_scripts/single_gpu/run_1.3b.sh

來(lái)訓(xùn)練一個(gè)OPT-1.3b模型。我們的單節(jié)點(diǎn)腳本很容易擴(kuò)展到多節(jié)點(diǎn)系統(tǒng)。

如何對(duì)SFT checkpoint進(jìn)行評(píng)測(cè)?

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

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

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

由于GPT3沒(méi)有開(kāi)源的checkpoint，我們使用了Meta OPT家族的預(yù)訓(xùn)練模型（即facebook/opt-1.3b）。你也可以使用其他預(yù)訓(xùn)練模型（如GPT-Neo，Bloom等）。至于數(shù)據(jù)集，我們也使用了來(lái)自Huggingface數(shù)據(jù)集的開(kāi)源數(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ù)抽象和融合技術(shù)，我們現(xiàn)在可以將多個(gè)數(shù)據(jù)源合并用于訓(xùn)練。然而，重要的是要注意，不同的數(shù)據(jù)集可能使用不同的提示詞（例如，Dohas/rm-static使用"Human:"表示查詢(xún)，"Assistant:"表示回答）。因此，用戶(hù)必須自行對(duì)齊這些提示。在我們的例子中，我們一致使用了Dohas/rm-static的格式。通過(guò)我們的評(píng)估，我們發(fā)現(xiàn)整合多樣化的數(shù)據(jù)集可以提高模型的質(zhì)量。請(qǐng)參考下一節(jié)以獲取不同查詢(xún)-答案對(duì)的示例。

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

在這里插入圖片描述

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

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

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

其它

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

0x2. 評(píng)測(cè)腳本解讀

DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/evaluation_scripts/run_prompt.sh這個(gè)腳本是對(duì)輸入的模型進(jìn)行評(píng)測(cè)，我們看一下內(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來(lái)對(duì)baseline模型和finetune模型進(jìn)行對(duì)比評(píng)測(cè)。所以評(píng)測(cè)的實(shí)現(xiàn)就是這里的prompt_eval.py，接下來(lái)對(duì)prompt_eval.py做一個(gè)解讀。

首先是頭文件導(dǎo)入部分，除了一些常規(guī)的需要用到的包還從DeepSpeed Chat的utils下面導(dǎo)入了2個(gè)函數(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下面導(dǎo)入自定義函數(shù)
fromutils.utilsimportload_hf_tokenizer

logger=logging.getLogger(__name__)

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

#這段代碼定義了一個(gè)名為create_hf_model的函數(shù)，該函數(shù)的作用是創(chuàng)建或加載一個(gè)預(yù)訓(xùn)練模型。該函數(shù)的主要參數(shù)包括：
#model_class：模型的類(lèi)別，例如GPT-2、BERT等。
#tokenizer：用于模型的分詞器。
#ds_config:DeepSpeed的配置參數(shù)。
#rlhf_training：一個(gè)標(biāo)志，用來(lái)表示是否正在進(jìn)行RLHF（ReinforcementLearningfromHumanFeedback）訓(xùn)練。
#disable_dropout：一個(gè)標(biāo)志，用來(lái)表示是否禁用dropout。Dropout是一種防止過(guò)擬合的技術(shù)。
defcreate_hf_model(model_class,
model_name_or_path,
tokenizer,
ds_config=None,
rlhf_training=False,
disable_dropout=False):
#根據(jù)model_name_or_path從預(yù)訓(xùn)練模型獲取模型配置model_config。
model_config=AutoConfig.from_pretrained(model_name_or_path)
#如果disable_dropout為真，則將模型配置中的dropout設(shè)為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中的設(shè)置，創(chuàng)建DeepSpeed的配置對(duì)象dschf，以便進(jìn)行DeepSpeed優(yōu)化。
ifds_configisnotNoneandds_config["zero_optimization"]["stage"]==3:
dschf=HfDeepSpeedConfig(ds_config)
else:
dschf=None
#根據(jù)rlhf_training的值，確定是從配置中創(chuàng)建模型還是從預(yù)訓(xùn)練模型中加載模型。如果rlhf_training為真，則根據(jù)模型配置創(chuàng)建模型；否則，從預(yù)訓(xùn)練模型加載模型。
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é)束標(biāo)記和填充標(biāo)記設(shè)為分詞器的結(jié)束標(biāo)記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ù)：

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

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

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

returnargs

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

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

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

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

接下來(lái)是一個(gè)簡(jiǎn)單的打印工具函數(shù)：

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

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

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

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

解析main函數(shù)：

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

device=torch.device("cuda:0")
#接著，它使用load_hf_tokenizer函數(shù)加載分詞器，然后使用create_hf_model函數(shù)創(chuàng)建基線(xiàn)模型（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)
#然后，這些模型被移動(dòng)到指定的設(shè)備上。
model_baseline.to(device)
model_fintuned.to(device)

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

0x3. 訓(xùn)練腳本解讀

訓(xùn)練過(guò)程都集中在DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py這個(gè)腳本，我們逐行解析一下這個(gè)腳本。

0x3.1 頭文件相關(guān)解析

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

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

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

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

#之后，它將當(dāng)前腳本的父目錄添加到系統(tǒng)路徑中，以便可以從該目錄下的utils目錄導(dǎo)入一些自定義函數(shù)和模塊。
sys.path.append(
os.path.abspath(os.path.join(os.path.dirname(__file__),os.path.pardir)))
#最后，它從utils目錄中導(dǎo)入了一些自定義模塊和函數(shù)，包括數(shù)據(jù)處理函數(shù)（create_prompt_dataset），打印和設(shè)備轉(zhuǎn)換函數(shù)（print_rank_0和to_device），
#模型保存函數(shù)（save_hf_format），隨機(jī)種子設(shè)置函數(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這個(gè)函數(shù)實(shí)際上直接或者間接的用到了utils/data中raw_dataset.py和data_utils.py，為了搞清楚這個(gè)函數(shù)，我們需要對(duì)這兩個(gè)文件做一個(gè)解析。

首先解析一下raw_dataset.py。這里先定義了一個(gè)PromptRawDataset類(lèi)：

#DeepSpeedTeam
fromdatasetsimportload_dataset
fromtorch.utils.dataimportSubset
importre


#這段代碼定義了一個(gè)名為PromptRawDataset的類(lèi)，這個(gè)類(lèi)是一個(gè)模板類(lèi)，用于處理和組織模型輸入數(shù)據(jù)的格式。
#如果有新的數(shù)據(jù)集需要進(jìn)行處理，可以繼承這個(gè)類(lèi)并實(shí)現(xiàn)相應(yīng)的方法來(lái)確保數(shù)據(jù)的統(tǒng)一格式和接口。
classPromptRawDataset(object):
#首先，這個(gè)類(lèi)的構(gòu)造函數(shù)__init__接收四個(gè)參數(shù)：output_path（輸出路徑），seed（隨機(jī)種子），
#local_rank（本地等級(jí)）和dataset_name（數(shù)據(jù)集名稱(chēng)）。
#在構(gòu)造函數(shù)中，如果數(shù)據(jù)集名稱(chēng)不是'local/jsonfile'，
#那么會(huì)使用HuggingFace的datasets庫(kù)的load_dataset函數(shù)來(lái)加載數(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)
#然后，這個(gè)類(lèi)定義了一些方法，這些方法在默認(rèn)情況下并沒(méi)有實(shí)現(xiàn)（只是返回None或者空操作），
#這是因?yàn)檫@個(gè)類(lèi)只是一個(gè)模板類(lèi)，這些方法需要在實(shí)際使用時(shí)在子類(lèi)中具體實(shí)現(xiàn)。
defget_train_data(self):#獲取訓(xùn)練數(shù)據(jù)
return

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

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

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

#Therejectedresponseshouldbeintheformatof:""+actual_response_sentence
#Ifthedatasetdoesnothaverejectedresponse,returnNone
#get_rejected方法用于獲取樣本中的rejected（被拒絕的回應(yīng)，這可能用于一些特定的訓(xùn)練場(chǎng)景，比如在對(duì)抗訓(xùn)練中，但如果數(shù)據(jù)集中沒(méi)有這樣的數(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

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

#Englishdataset
#這個(gè)類(lèi)OpenaiWebgptcomparisonsDataset繼承自PromptRawDataset類(lèi)，
#針對(duì)"openai/webgpt_comparisons"這個(gè)具體的數(shù)據(jù)集進(jìn)行了特化。
classOpenaiWebgptcomparisonsDataset(PromptRawDataset):
#在構(gòu)造函數(shù)__init__中，調(diào)用了父類(lèi)的構(gòu)造函數(shù)，并設(shè)定了dataset_name和dataset_name_clean兩個(gè)屬性，
#分別為"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中獲取訓(xùn)練數(shù)據(jù)和測(cè)試數(shù)據(jù)。
#它們與之前的DahoasRmstaticDataset類(lèi)不同之處在于，它們使用get_raw_dataset_split_index
#方法對(duì)訓(xùn)練數(shù)據(jù)進(jìn)行了劃分，將其劃分為訓(xùn)練集和驗(yàn)證集，并返回對(duì)應(yīng)的數(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方法分別從樣本中獲取提示，已選回應(yīng)和被拒絕的回應(yīng)。
#這里假定樣本是一個(gè)字典，其中包含了名為'question'，'score_0'，'score_1'，'answer_0'和'answer_1'的字段。
#其中，'question'字段是一個(gè)字典，包含了'full_text'字段。這個(gè)字段包含了人類(lèi)提出的問(wèn)題。
#'score_0'和'score_1'字段是字符串，表示對(duì)'answer_0'和'answer_1'的評(píng)分。
#如果'score_0'大于等于'score_1'，那么'answer_0'就是已選回應(yīng)，'answer_1'就是被拒絕的回應(yīng)，反之亦然。
#在獲取已選回應(yīng)和被拒絕的回應(yīng)時(shí)，還對(duì)回應(yīng)進(jìn)行了處理，
#去除了所有形如"[...]"或"(...)"的文本，然后在回應(yīng)前添加了一個(gè)空格。
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)置庫(kù)，提供了一系列散列函數(shù)，如MD5、SHA1等。
fromitertoolsimportchain#Python的內(nèi)置庫(kù)，提供了一系列用于操作迭代器的函數(shù)。
from.importraw_datasets

#這段代碼定義了一個(gè)名為get_raw_dataset的函數(shù)，其主要作用是根據(jù)傳入的數(shù)據(jù)集名稱(chēng)dataset_name
#返回一個(gè)適當(dāng)?shù)腜romptRawDataset子類(lèi)的實(shí)例。
#這個(gè)函數(shù)有四個(gè)參數(shù)：dataset_name，output_path，seed和local_rank。
defget_raw_dataset(dataset_name,output_path,seed,local_rank):
#在函數(shù)中，根據(jù)dataset_name的具體值來(lái)創(chuàng)建不同的PromptRawDataset子類(lèi)的實(shí)例。
#例如，如果dataset_name為"Dahoas/rm-static"，那么就創(chuàng)建一個(gè)DahoasRmstaticDataset的實(shí)例；
#如果dataset_name為"Dahoas/full-hh-rlhf"，那么就創(chuàng)建一個(gè)DahoasFullhhrlhfDataset的實(shí)例，以此類(lèi)推。
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"，則會(huì)檢查在路徑chat_path+'/data/train.json'
#和chat_path+'/data/eval.json'下是否存在文件。如果存在，則創(chuàng)建一個(gè)LocalJsonFileDataset的實(shí)例；
#如果不存在，則拋出一個(gè)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沒(méi)有在以上的所有條件中匹配到，那么函數(shù)也會(huì)拋出一個(gè)RuntimeError異常，表示沒(méi)有為這個(gè)數(shù)據(jù)集的配置。
raiseRuntimeError(
f"Wedonothaveconfigsfordataset{dataset_name},butyoucanadditbyyourselfinraw_datasets.py."
)


	

	再看下 get_shuffle_idx函數(shù)：

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


	

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

	
#這個(gè)函數(shù)主要是根據(jù)提供的參數(shù)分割數(shù)據(jù)集，并生成一個(gè)分割索引。
#它首先檢查索引文件是否存在，如果不存在，則生成分割索引，并保存到文件。
#然后，它從文件中加載索引，并返回索引列表。
#它接受8個(gè)參數(shù)：local_rank、輸出路徑、數(shù)據(jù)集名稱(chēng)、種子、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（以逗號(hào)分隔的字符串）分割為浮點(diǎn)數(shù)列表。
splits_sum=sum(splits)#計(jì)算所有分割的總和。
splits=[split/splits_sumforsplitinsplits]#將每個(gè)分割值歸一化，使其和為1。
splits_index=[0]#初始化分割索引列表，起始值為0
forindex,splitinenumerate(splits):#遍歷分割列表。
#將新的索引（當(dāng)前索引加上歸一化的分割值與數(shù)據(jù)大小的乘積）添加到索引列表。
splits_index.append(splits_index[index]+
int(round(split*float(data_size))))
#計(jì)算最后一個(gè)索引與數(shù)據(jù)大小的差值。
diff=splits_index[-1]-data_size
#遍歷除第一個(gè)外的所有索引。
forindexinrange(1,len(splits_index)):
#將差值從每個(gè)索引中減去，以確保最后一個(gè)索引等于數(shù)據(jù)大小。
splits_index[index]-=diff
#斷言最后一個(gè)索引等于數(shù)據(jù)大小
assertsplits_index[-1]==data_size

#生成一個(gè)亂序的索引。
shuffle_idx=get_shuffle_idx(seed,data_size)
forsplit_iinrange(len(splits)):#遍歷每個(gè)分割。
#根據(jù)給定參數(shù)生成亂序索引分割文件名。
shuffle_idx_split_file_name=f"{output_path}/{dataset_name}_seed{seed}_{split_name}_{data_split}_{split_i}.npy"
#提取亂序索引的一個(gè)分割。
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()


	

	接下來(lái)解析一下繼承自Dataset的PromptDataset類(lèi)：

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

#定義類(lèi)的__getitem__方法，它接受一個(gè)參數(shù)idx，返回索引idx處的數(shù)據(jù)。這是PyTorch數(shù)據(jù)集的必要方法。
def__getitem__(self,idx):
#如果訓(xùn)練階段為1，則返回一個(gè)字典，包含input_ids、attention_mask和labels，它們都來(lái)自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"]
}
#如果訓(xùn)練階段為2，則返回來(lái)自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"]
#如果訓(xùn)練階段為3，則返回來(lái)自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ù)：

	
#這是一個(gè)名為create_dataset_split的函數(shù)，它的功能是根據(jù)給定的訓(xùn)練階段（train_phase），創(chuàng)建并返回相應(yīng)的數(shù)據(jù)集分割。
#具體來(lái)說(shuō)，它為每個(gè)訓(xùn)練階段生成不同的數(shù)據(jù)集列表，并將它們放入PromptDataset對(duì)象中。
#函數(shù)接受6個(gè)參數(shù)：當(dāng)前數(shù)據(jù)集(current_dataset)、原始數(shù)據(jù)集(raw_dataset)、訓(xùn)練階段(train_phase)、
#分詞器(tokenizer)、會(huì)話(huà)結(jié)束標(biāo)記(end_of_conversation_token)和最大序列長(zhǎng)度(max_seq_len)。
defcreate_dataset_split(current_dataset,raw_dataset,train_phase,tokenizer,
end_of_conversation_token,max_seq_len):
#創(chuàng)建三個(gè)空的列表，用于存儲(chǔ)對(duì)話(huà)提示（prompt_dataset）、選定的對(duì)話(huà)（chosen_dataset）和被拒絕的對(duì)話(huà)（reject_dataset）。
prompt_dataset=[]
chosen_dataset=[]
reject_dataset=[]
#如果訓(xùn)練階段為1，則將接受的對(duì)話(huà)進(jìn)行分詞并添加到chosen_dataset中。
iftrain_phase==1:
#遍歷當(dāng)前數(shù)據(jù)集。
fori,tmp_datainenumerate(current_dataset):
#tokenizethetext
#從原始數(shù)據(jù)集中獲取對(duì)話(huà)提示和接受的對(duì)話(huà)。
chosen_sentence=raw_dataset.get_prompt_and_chosen(
tmp_data)#theacceptresponse
#如果接受的對(duì)話(huà)不為空，則將其分詞并添加到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)
#如果訓(xùn)練階段為2，則將接受和被拒絕的對(duì)話(huà)都進(jìn)行分詞并分別添加到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)
#如果訓(xùn)練階段為3，則將對(duì)話(huà)提示進(jìn)行分詞并添加到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)建一個(gè)新的PromptDataset對(duì)象，并返回。這個(gè)對(duì)象包含了對(duì)話(huà)提示、接受的對(duì)話(huà)和被拒絕的對(duì)話(huà)的數(shù)據(jù)集，以及分詞器的填充標(biāo)記ID和訓(xùn)練階段。
returnPromptDataset(prompt_dataset,chosen_dataset,reject_dataset,
tokenizer.pad_token_id,train_phase)


	

	接下來(lái)再解析一下create_dataset這個(gè)函數(shù)：

	
這段代碼定義了一個(gè)函數(shù)create_dataset，主要負(fù)責(zé)創(chuàng)建訓(xùn)練數(shù)據(jù)集和評(píng)估數(shù)據(jù)集，具體的功能細(xì)節(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ù)集名稱(chēng)、輸出路徑、隨機(jī)種子和local_rank等參數(shù)，從各種預(yù)定義的數(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ù)集中獲取訓(xùn)練數(shù)據(jù)。
#獲取訓(xùn)練數(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)建訓(xùn)練數(shù)據(jù)的子集。
train_dataset=Subset(train_dataset,train_index)
#調(diào)用create_dataset_split函數(shù)對(duì)上一步獲得的數(shù)據(jù)子集進(jìn)行進(jìn)一步處理，
#這可能包括對(duì)文本的標(biāo)記化(tokenization)，并且創(chuàng)建一個(gè)PromptDataset對(duì)象。
train_dataset=create_dataset_split(train_dataset,raw_dataset,
train_phase,tokenizer,
end_of_conversation_token,
max_seq_len)
#是用于創(chuàng)建評(píng)估數(shù)據(jù)集的，步驟與訓(xùn)練數(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


	

	這里需要額外注意一點(diǎn)，create_dataset_split 這個(gè)函數(shù)傳入的參數(shù)是 train_phase，而不是像 get_raw_dataset_split_index 函數(shù)那樣傳入 train_phase-1，這是為什么？這是因?yàn)閠rain_phase用來(lái)標(biāo)識(shí)在模型訓(xùn)練過(guò)程中的哪個(gè)階段，是一個(gè)直接的指示符，與數(shù)據(jù)分割或索引無(wú)關(guān)。它用來(lái)在create_dataset_split函數(shù)中確定當(dāng)前處于哪個(gè)訓(xùn)練階段，并根據(jù)不同的訓(xùn)練階段對(duì)數(shù)據(jù)進(jìn)行不同的處理。而在調(diào)用get_raw_dataset_split_index函數(shù)時(shí)，傳入的是train_phase - 1，這是因?yàn)樵谶@個(gè)函數(shù)中，我們需要根據(jù)當(dāng)前訓(xùn)練階段的前一個(gè)階段（由train_phase - 1表示）的數(shù)據(jù)分割或索引情況，來(lái)決定如何對(duì)當(dāng)前階段的數(shù)據(jù)進(jìn)行分割或索引。create_dataset_split主要關(guān)注如何根據(jù)訓(xùn)練階段對(duì)數(shù)據(jù)進(jìn)行處理，而get_raw_dataset_split_index主要關(guān)注如何根據(jù)前一個(gè)訓(xùn)練階段的數(shù)據(jù)分割或索引情況，對(duì)當(dāng)前階段的數(shù)據(jù)進(jìn)行分割或索引。

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

	
#這個(gè)函數(shù)的主要目的是創(chuàng)建一個(gè)包含訓(xùn)練和評(píng)估數(shù)據(jù)集的“提示”數(shù)據(jù)集，并將這兩個(gè)數(shù)據(jù)集保存在指定的文件中。具體來(lái)說(shuō)：
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)存在則不會(huì)引發(fā)異常。
os.makedirs(output_path,exist_ok=True)
#構(gòu)造文件名，這個(gè)文件名包含了很多有關(guān)數(shù)據(jù)集和模型的信息，如數(shù)據(jù)路徑、數(shù)據(jù)分割、訓(xùn)練階段、
#隨機(jī)種子、tokenizer的名稱(chēng)、最大序列長(zhǎng)度等。然后將這個(gè)文件名哈?；?，以避免文件名過(guò)長(zhǎng)。
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)造訓(xùn)練數(shù)據(jù)集和評(píng)估數(shù)據(jù)集的文件路徑。
train_fname=f"{output_path}/traindata_{fname}.pt"
eval_fname=f"{output_path}/evaldata_{fname}.pt"

#檢查訓(xùn)練數(shù)據(jù)集和評(píng)估數(shù)據(jù)集的文件是否都已經(jīng)存在，如果存在，則表示緩存已經(jīng)找到，否則表示需要?jiǎ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)

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

????????????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).
????????#?如果當(dāng)前是第一階段的訓(xùn)練（SFT）并且指定了僅用于SFT的數(shù)據(jù)集，那么對(duì)這些數(shù)據(jù)集執(zhí)行類(lèi)似的操作，
????????#?然后把得到的訓(xùn)練數(shù)據(jù)集和評(píng)估數(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())
????????#?把訓(xùn)練數(shù)據(jù)集和評(píng)估數(shù)據(jù)集保存到對(duì)應(yīng)的文件中。
????????torch.save(train_dataset,?train_fname)
????????torch.save(eval_dataset,?eval_fname)
????#?在多進(jìn)程環(huán)境中，確保所有進(jìn)程都完成了數(shù)據(jù)集的保存操作。
????torch.distributed.barrier()
????return?torch.load(train_fname),?torch.load(eval_fname)



	

	工具函數(shù)解析

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

	
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打印信息，防止每個(gè)機(jī)器或GPU都打印消息造成大量重復(fù)信息
defprint_rank_0(msg,rank=0):
ifrank<=?0:
????????print(msg)


	

	to_device

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


	

	save_hf_format

	
#這段代碼的功能是將訓(xùn)練好的模型以HuggingFace格式保存，這樣之后就可以使用HuggingFace庫(kù)的from_pretrained方法加載了。
defsave_hf_format(model,tokenizer,args,sub_folder=""):
#usedtosavehuggingfaceformat,sowecanuseitforhf.from_pretrained
#首先，這行代碼檢查模型是否有'module'這個(gè)屬性。這在PyTorch中是一個(gè)常見(jiàn)的模式，
#當(dāng)模型被封裝在torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel時(shí)，
#模型的所有屬性都會(huì)被存儲(chǔ)在'module'屬性中。所以這行代碼的目的是確保我們總是在原始模型上進(jìn)行操作，而不是并行化的包裝器。
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)建這個(gè)目錄。
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)字典，這個(gè)字典包含了模型的所有權(quán)重。
save_dict=model_to_save.state_dict()
#遍歷狀態(tài)字典中的所有鍵。這里使用list(save_dict.keys())而不是直接使用
#save_dict.keys()是因?yàn)槲覀冊(cè)诒闅v過(guò)程中會(huì)修改字典，直接遍歷字典的鍵在Python中會(huì)導(dǎo)致錯(cuò)誤。
forkeyinlist(save_dict.keys()):
#if"lora"inkey:：如果鍵中包含"lora"，那么刪除這個(gè)鍵值對(duì)。
#"lora"可能是某種特定的模型組件或參數(shù)，這里將其排除在保存的模型權(quán)重之外。
if"lora"inkey:
delsave_dict[key]
#將處理過(guò)的狀態(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

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


	

	get_all_reduce_mean

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


	

	get_optimizer_grouped_parameters

	
#這段代碼的作用是將模型中的參數(shù)分組以便于在優(yōu)化器中使用。它將模型參數(shù)分為兩組：
#一組需要進(jìn)行權(quán)重衰減（L2正則化）的參數(shù)，另一組不需要進(jìn)行權(quán)重衰減的參數(shù)。
defget_optimizer_grouped_parameters(model,
weight_decay,
no_decay_name_list=[
"bias","LayerNorm.weight"
]):
#它定義了一個(gè)列表optimizer_grouped_parameters，其中包含兩個(gè)字典。每個(gè)字典都對(duì)應(yīng)一個(gè)參數(shù)組，包含"params"和"weight_decay"這兩個(gè)關(guān)鍵字。
optimizer_grouped_parameters=[
#在第一個(gè)字典中，它從模型參數(shù)中選出那些名稱(chēng)不包含"bias"或"LayerNorm.weight"
#且需要求梯度的參數(shù)。這些參數(shù)在優(yōu)化過(guò)程中會(huì)應(yīng)用weight_decay作為權(quán)重衰減項(xiàng)。
{
"params":[
pforn,pinmodel.named_parameters()
if(notany(ndinn
forndinno_decay_name_list)andp.requires_grad)
],
"weight_decay":
weight_decay,
},
#在第二個(gè)字典中，它選出那些名稱(chēng)包含"bias"或"LayerNorm.weight"且需要求梯度的參數(shù)。
#這些參數(shù)在優(yōu)化過(guò)程中不會(huì)應(yīng)用權(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ù)的分組策略是很常見(jiàn)的。比如在訓(xùn)練Transformer模型時(shí)，通常會(huì)為權(quán)重和偏置項(xiàng)設(shè)定不同的學(xué)習(xí)策略。這是因?yàn)闄?quán)重衰減對(duì)于防止過(guò)擬合很有幫助，但對(duì)于某些參數(shù)（如偏置項(xiàng)或者層歸一化的權(quán)重）可能會(huì)導(dǎo)致性能下降，因此常常會(huì)排除這些參數(shù)不進(jìn)行權(quán)重衰減。

	save_zero_three_model

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

#這個(gè)函數(shù)的主要作用是保存一個(gè)使用了DeepSpeedZero優(yōu)化（可能為stage3）的模型。
#DeepSpeed的Zero優(yōu)化技術(shù)是為了解決模型參數(shù)、優(yōu)化器狀態(tài)和梯度等內(nèi)存占用問(wèn)題，
#通過(guò)這種方式，可以訓(xùn)練比當(dāng)前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）中的，我們需要取出真實(shí)的模型實(shí)例。
model_to_save=model_ema.moduleifhasattr(model_ema,
'module')elsemodel_ema
ifnotzero_stage_3:
#如果沒(méi)有使用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)化，因?yàn)槟Ｐ偷牟糠謪?shù)和優(yōu)化器狀態(tài)在不同的設(shè)備上，所以需要先將它們收集起來(lái)。
output_state_dict={}
fork,vinmodel_to_save.named_parameters():

ifhasattr(v,'ds_id'):
#deepspeed.zero.GatheredParameters是DeepSpeed提供的一個(gè)上下文管理器，
#它可以將分布在多個(gè)設(shè)備上的參數(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”關(guān)鍵字的參數(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)
#同時(shí)為了節(jié)省內(nèi)存，使用del關(guān)鍵字刪除了存儲(chǔ)參數(shù)的字典。
deloutput_state_dict



	

	load_hf_tokenizer

	
#這個(gè)函數(shù)的目標(biāo)是加載一個(gè)HuggingFacetokenizer，這個(gè)tokenizer是用來(lái)將文本數(shù)據(jù)轉(zhuǎn)化為模型可以接受的形式的。
#model_name_or_path:這個(gè)參數(shù)可以是一個(gè)模型名字或者一個(gè)本地路徑，指示從哪里獲取tokenizer。
#fast_tokenizer=True:這個(gè)參數(shù)指示是否使用HuggingFace的"fast"tokenizer。"fast"tokenizer是用Rust編寫(xiě)的，運(yùn)行速度更快。
defload_hf_tokenizer(model_name_or_path,fast_tokenizer=True):
#這行代碼檢查model_name_or_path是否是一個(gè)存在的路徑。如果是，那么就嘗試從本地文件系統(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的模型庫(kù)中下載并加載tokenizer。
else:#如果model_name_or_path不是一個(gè)存在的路徑，那么就認(rèn)為它是一個(gè)模型名字。
tokenizer=AutoTokenizer.from_pretrained(model_name_or_path,
fast_tokenizer=True)#這行代碼使用模型名字從HuggingFace的模型庫(kù)中下載并加載tokenizer。
returntokenizer


	

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

	
#DeepSpeedTeam
GLOBAL_BATCH_SIZE=32
MICRO_BATCH_SIZE=4

#這個(gè)函數(shù)的主要目標(biāo)是生成一個(gè)字典，用來(lái)配置DeepSpeed訓(xùn)練模型的各種參數(shù)。
#DeepSpeed是一個(gè)用于加速大規(guī)模模型訓(xùn)練的庫(kù)，其中包括一種叫做ZeRO（ZeroRedundancyOptimizer）的優(yōu)化技術(shù)，
#可以大大減少模型訓(xùn)練過(guò)程中的顯存和內(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):
#這行代碼首先定義了一個(gè)變量device。如果參數(shù)offload為T(mén)rue，那么就將device設(shè)置為"cpu"，
#表示模型的參數(shù)和優(yōu)化器狀態(tài)將被"offload"到CPU內(nèi)存中，以減少GPU顯存的占用。
#如果offload為False，那么device就被設(shè)置為"none"，表示不進(jìn)行offload。
device="cpu"ifoffloadelse"none"
#這行代碼定義了一個(gè)字典zero_opt_dict，包含了ZeRO優(yōu)化器的各種配置參數(shù)。
zero_opt_dict={
"stage":stage,#ZeRO的階段，可以為0、1、2、3，階段越高，顯存優(yōu)化效果越好，但計(jì)算復(fù)雜度也越高。
"offload_param":{#offload_param和offload_optimizer這兩個(gè)字段控制是否將模型的參數(shù)和優(yōu)化器狀態(tài)offload到CPU內(nèi)存。
"device":device
},
"offload_optimizer":{
"device":device
},
#這三個(gè)字段只有在ZeROstage3時(shí)有效，用于控制參數(shù)offloading的細(xì)節(jié)。
"stage3_param_persistence_threshold":1e4,
"stage3_max_live_parameters":3e7,
"stage3_prefetch_bucket_size":3e7,
#這個(gè)字段控制是否使用顯存更高效的線(xiàn)性層實(shí)現(xiàn)。由于當(dāng)前的顯存高效實(shí)現(xiàn)可能會(huì)導(dǎo)致數(shù)值不穩(wěn)定，所以這里默認(rèn)設(shè)為False。
"memory_efficient_linear":False
}
#返回一個(gè)配置字典，這個(gè)字典包含了訓(xùn)練模型的各種參數(shù)。
return{
"train_batch_size":GLOBAL_BATCH_SIZE,#全局的批次大小
"train_micro_batch_size_per_gpu":MICRO_BATCH_SIZE,#每個(gè)GPU的微批次大小。
"steps_per_print":10,#每隔10步打印一次訓(xùn)練信息。
"zero_optimization":zero_opt_dict,#之前定義的ZeRO配置字典。
"fp16":{#控制是否使用混合精度訓(xùn)練，以及l(fā)ossscaling的窗口大小。
"enabled":True,
"loss_scale_window":100
},
"gradient_clipping":1.0,#梯度裁剪的閾值。
"prescale_gradients":False,#是否在梯度累計(jì)之前就進(jìn)行梯度縮放，通常用于防止梯度下溢。
"wall_clock_breakdown":False,#是否進(jìn)行每步訓(xùn)練時(shí)間的詳細(xì)分析。
"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導(dǎo)入的convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters 這幾個(gè)函數(shù)。

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

	
#LinearLayer_LoRA是一個(gè)實(shí)現(xiàn)了LoRA（Low-RankAdaptation）策略的PyTorch模塊。
#LoRA是一種在大型預(yù)訓(xùn)練模型（如GPT-3或者BERT）微調(diào)中使用的策略，可以提高微調(diào)的效果。
#這個(gè)類(lèi)是一個(gè)簡(jiǎn)單的LoRA實(shí)現(xiàn)，只支持線(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
#在初始化過(guò)程中，首先檢查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ù)，那么通過(guò)?ds_shape?屬性獲取形狀；
????????#?否則直接使用?shape?屬性。
????????try:
????????????#?for?zero?stage?3
????????????rows,?columns?=?weight.ds_shape
????????except:
????????????rows,?columns?=?weight.shape
????????#?然后初始化?LoRA?的左右兩個(gè)權(quán)重矩陣，并設(shè)置它們?yōu)榭蓪W(xué)習(xí)的參數(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)建一個(gè)?Dropout?層；否則創(chuàng)建一個(gè)?Identity?層。
????????if?lora_droppout?>0:
self.lora_dropout=nn.Dropout(lora_droppout)
else:
self.lora_dropout=nn.Identity()
#然后調(diào)用reset_parameters方法對(duì)LoRA權(quán)重進(jìn)行初始化。
self.reset_parameters()
#disabletheoriginalweightgradient
#最后，關(guān)閉原始權(quán)重的梯度，設(shè)置LoRA融合標(biāo)志位為False。
self.weight.requires_grad=False
#fuseLoRAtotheoriginalweight
self.fuse_lora=False

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

#self.fuse_lora_weight()

deftrain(self,mode=True):
self.lora_dropout.train(mode)#將模型設(shè)置為訓(xùn)練模式，這時(shí)候Dropout層會(huì)開(kāi)始工作。
#self.unfuse_lora_weight()

defreset_parameters(self):#初始化LoRA權(quán)重的方法。右權(quán)重使用kaiming均勻分布進(jìn)行初始化，左權(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)：
#這兩個(gè)方法用于將LoRA權(quán)重融合到原始權(quán)重中，或者從原始權(quán)重中解融合。
#融合操作實(shí)質(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)重已融合，則直接對(duì)輸入進(jìn)行線(xiàn)性變換；
#否則，會(huì)額外計(jì)算一個(gè)LoRA項(xiàng)，該項(xiàng)是輸入通過(guò)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

	
#這個(gè)函數(shù)convert_linear_layer_to_lora是用來(lái)將模型中的線(xiàn)性層轉(zhuǎn)換為L(zhǎng)oRA層的。
#在訓(xùn)練深度學(xué)習(xí)模型時(shí)，這種方法能夠在保持預(yù)訓(xùn)練模型參數(shù)不變的同時(shí)，通過(guò)添加額外的參數(shù)來(lái)微調(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()），找出名稱(chēng)中包含part_module_name的線(xiàn)性層（nn.Linear），
#并將這些層的名稱(chēng)添加到repalce_name列表中。
forname,moduleinmodel.named_modules():
ifisinstance(module,nn.Linear)andpart_module_nameinname:
repalce_name.append(name)
#然后，函數(shù)遍歷repalce_name列表，使用recursive_getattr函數(shù)獲取模型中對(duì)應(yīng)名稱(chēng)的模塊。
#這些模塊是需要被替換成LoRA層的線(xiàn)性層。
fornameinrepalce_name:
module=recursive_getattr(model,name)
#對(duì)于每一個(gè)需要被替換的模塊，函數(shù)創(chuàng)建一個(gè)LinearLayer_LoRA實(shí)例tmp，
#并將其傳輸?shù)脚c原始模塊相同的設(shè)備和數(shù)據(jù)類(lèi)型上。創(chuàng)建LinearLayer_LoRA實(shí)例時(shí)，
#需要傳入原始模塊的權(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í)例后，函數(shù)使用recursive_setattr函數(shù)將原始模塊替換為L(zhǎng)inearLayer_LoRA實(shí)例。
recursive_setattr(model,name,tmp)
returnmodel


	

	convert_lora_to_linear_layer

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

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

	
#這個(gè)函數(shù)的作用是關(guān)閉模型中除LoRA參數(shù)之外的所有參數(shù)的梯度。這意味著在訓(xùn)練過(guò)程中，只有LoRA參數(shù)會(huì)被優(yōu)化，其他參數(shù)保持不變。
defonly_optimize_lora_parameters(model):
#turnoffthegradientofalltheparametersexcepttheLoRAparameters
#遍歷模型的所有參數(shù)。每個(gè)參數(shù)都有一個(gè)唯一的名稱(chēng)name和對(duì)應(yīng)的參數(shù)值param。
forname,paraminmodel.named_parameters():
#查當(dāng)前參數(shù)的名稱(chēng)是否包含lora_right_weight或lora_left_weight。
#這是因?yàn)樵贚oRA（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è)置參數(shù)的requires_grad屬性為T(mén)rue，表示需要對(duì)此參數(shù)進(jìn)行梯度下降優(yōu)化。
param.requires_grad=True
else:
param.requires_grad=False
returnmodel


	

	create_hf_model

	
#這段代碼定義了一個(gè)名為create_hf_model的函數(shù)，該函數(shù)的作用是創(chuàng)建或加載一個(gè)預(yù)訓(xùn)練模型。該函數(shù)的主要參數(shù)包括：
#model_class：模型的類(lèi)別，例如GPT-2、BERT等。
#tokenizer：用于模型的分詞器。
#ds_config:DeepSpeed的配置參數(shù)。
#rlhf_training：一個(gè)標(biāo)志，用來(lái)表示是否正在進(jìn)行RLHF（ReinforcementLearningfromHumanFeedback）訓(xùn)練。
#disable_dropout：一個(gè)標(biāo)志，用來(lái)表示是否禁用dropout。Dropout是一種防止過(guò)擬合的技術(shù)。
defcreate_hf_model(model_class,
model_name_or_path,
tokenizer,
ds_config=None,
rlhf_training=False,
disable_dropout=False):
#根據(jù)model_name_or_path從預(yù)訓(xùn)練模型獲取模型配置model_config。
model_config=AutoConfig.from_pretrained(model_name_or_path)
#如果disable_dropout為真，則將模型配置中的dropout設(shè)為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中的設(shè)置，創(chuàng)建DeepSpeed的配置對(duì)象dschf，以便進(jìn)行DeepSpeed優(yōu)化。
ifds_configisnotNoneandds_config["zero_optimization"]["stage"]==3:
dschf=HfDeepSpeedConfig(ds_config)
else:
dschf=None
#根據(jù)rlhf_training的值，確定是從配置中創(chuàng)建模型還是從預(yù)訓(xùn)練模型中加載模型。如果rlhf_training為真，則根據(jù)模型配置創(chuàng)建模型；否則，從預(yù)訓(xùn)練模型加載模型。
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é)束標(biāo)記和填充標(biāo)記設(shè)為分詞器的結(jié)束標(biāo)記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)建一個(gè)argparse的解析器對(duì)象，這個(gè)對(duì)象可以添加命令行參數(shù)和處理它們。description參數(shù)提供了一個(gè)對(duì)程序的簡(jiǎn)單描述。
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ù)并將它們存儲(chǔ)在args對(duì)象中
args=parser.parse_args()

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

returnargs


	

	main函數(shù)解析

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

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

#如果本地排名為-1，說(shuō)明不在分布式訓(xùn)練環(huán)境下，設(shè)備設(shè)置為"cuda"；
#否則根據(jù)args.local_rank設(shè)置對(duì)應(yīng)的cuda設(shè)備，并初始化分布式訓(xùn)練。
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的訓(xùn)練配置。
ds_config=get_train_ds_config(offload=args.offload,
stage=args.zero_stage)
#在配置中設(shè)置訓(xùn)練時(shí)每個(gè)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.
#設(shè)置隨機(jī)種子以保證結(jié)果的可復(fù)現(xiàn)性。
set_random_seed(args.seed)

#在此處添加一個(gè)barrier操作，確保所有進(jìn)程都執(zhí)行到這一點(diǎn)后再繼續(xù)執(zhí)行后續(xù)操作。
torch.distributed.barrier()
#加載預(yù)訓(xùn)練模型對(duì)應(yīng)的分詞器。
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)建預(yù)訓(xùn)練模型。
model=create_hf_model(AutoModelForCausalLM,
args.model_name_or_path,
tokenizer,
ds_config,
disable_dropout=args.disable_dropout)

#如果參數(shù)lora_dim大于0，將模型的線(xiàn)性層轉(zhuǎn)換為L(zhǎng)oRa層；如果只優(yōu)化LoRa參數(shù)，關(guān)閉其他參數(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ù)加載器：包括訓(xùn)練集和驗(yàn)證集，以及對(duì)應(yīng)的采樣器和數(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)
#定義模型評(píng)估函數(shù)，用于計(jì)算模型在驗(yàn)證集上的困惑度。
defevaluation(model,eval_dataloader):
model.eval()#將模型切換為評(píng)估模式。
losses=0#初始化loss。
forstep,batchinenumerate(eval_dataloader):#對(duì)于評(píng)估數(shù)據(jù)集中的每一個(gè)batch。
batch=to_device(batch,device)#將batch數(shù)據(jù)移到對(duì)應(yīng)的設(shè)備上。
withtorch.no_grad():#在此上下文管理器中，不計(jì)算梯度，這樣可以節(jié)省存儲(chǔ)和計(jì)算資源。
#將batch數(shù)據(jù)輸入模型，進(jìn)行前向計(jì)算。
outputs=model(**batch)

loss=outputs.loss#取出模型的輸出中的loss。
losses+=loss.float()#將當(dāng)前的loss累加到總的losses中。
losses=losses/(step+1)#計(jì)算平均的loss。
try:
perplexity=torch.exp(losses)#嘗試計(jì)算模型的困惑度，如果捕捉到溢出錯(cuò)誤，將困惑度設(shè)置為無(wú)窮大。
exceptOverflowError:
perplexity=float("inf")
#嘗試在所有設(shè)備上計(jì)算困惑度的平均值，如果發(fā)生任何錯(cuò)誤，就忽略。
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)化器類(lèi)型，如果啟用了梯度Offload，使用DeepSpeedCPUAdam，否則使用FusedAdam。
AdamOptimizer=DeepSpeedCPUAdamifargs.offloadelseFusedAdam
#創(chuàng)建優(yōu)化器。
optimizer=AdamOptimizer(optimizer_grouped_parameters,
lr=args.learning_rate,
betas=(0.9,0.95))
#計(jì)算每個(gè)epoch的更新步數(shù)。
num_update_steps_per_epoch=math.ceil(
len(train_dataloader)/args.gradient_accumulation_steps)
#創(chuàng)建學(xué)習(xí)率調(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)化器和學(xué)習(xí)率調(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)
#如果啟用了梯度檢查點(diǎn)，那么在模型中也啟用梯度檢查點(diǎn)。
ifargs.gradient_checkpointing:
model.gradient_checkpointing_enable()

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

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

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

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

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


	

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

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

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

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

	0x5. 總結(jié)

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