ColossalAI GPT2分布式训练调试配置—GPT系列训练与部署

        本文为博主原创文章，未经博主允许不得转载。

        本文为专栏《Python从零开始进行AIGC大模型训练与推理》系列文章，地址为“https://blog.csdn.net/suiyingy/article/details/130169592”。

        本专栏之前文章详细介绍了Colossal-AI环境搭建及GPT2数据处理与训练。Colossal-AI框架的主要优势在于分布式训练，进而提高训练效率。但是这种启动方式无法使用Pycharm等IDE直接进行调试。本节将重点介绍如何使用Pycharm来调试。Colossal-AI分布式GPT训练程序，文中所述方法不仅支持colossal run方法启动的程序，也支持torchrun或python -m torch.distributed.launch所启动的程序。这是因为它们本质上都是torch分布式启动方法。

        另外，本专栏具体更新可关注文章下方公众号，也可关注本专栏。所有相关文章会在《Python从零开始进行AIGC大模型训练与推理》中进行更新，地址为“https://blog.csdn.net/suiyingy/article/details/130169592”。所有AIGC类模型部署的体验效果将在RdFast小程序中同步上线。

1 分布式训练启动方法

        由于Colossal-AI框架GPT2训练的默认启动方法为分布式模式，它的启动命令可以使用colossalai run，或torchrun，或python -m torch.distributed.launch。具体示例如下所示。

# 1、colossalai run方式启动GPT2训练
colossalai run --nproc_per_node=2 train_gpt.py --config=gpt2_configs/gpt2_vanilla.py --from_torch
# 2、torch run方式启动GPT2训练
torchrun --nproc_per_node=2 train_gpt.py --config=gpt2_configs/gpt2_vanilla.py --from_torch
# 3、python -m torch.distributed.launch方式启动GPT2训练
python -m torch.distributed.launch --nproc_per_node=2 train_gpt.py --config=gpt2_configs/gpt2_vanilla.py --from_torch

        虽然分布式训练启动方法有多种，但是却不支持Python直接启动。例如，“python train_gpt.py --config=gpt2_configs/gpt2_vanilla.py --from_torch”会报错误“RuntimeError: Could not find 'RANK' in the torch environment, visit https://www.colossalai.org/ for more information on launching with torch”。

2 Python启动方式配置

        上述分布式启动方式之一是python -m torch.distributed.launch，这说明python程序入口是launch.py文件。该文件通常位于python环境库site-packages/torch.distributed目录下，例如“/root/miniconda3/envs/clai/lib/python3.8/site-packages/torch/distributed/launch.py”。该文件的启动参数则为“--nproc_per_node=2 train_gpt.py --config=gpt2_configs/gpt2_vanilla.py --from_torch”。

        在PyCharm中，我们可以按照如上内容来配置启动程序和参数。配置页面路径为Run->Edit Configurations...，如下图所示。

图1 PyCharm配置启动程序入口

        点击“Edit Configurations...”后，PyCharm会弹出“Run/Debug Configurations”配置页面，如下图所示。需要注意，页面中分为Python和Python tests两类。我们需要在Python类别中进行配置。Python tests中类别没有启动程序路径和参数配置选项。

图2 Run/Debug Configurations配置页面

        具体配置步骤如下所示。远程服务器调试时，只需将启动程序和配餐参数中路径设置成服务器相应路径即可。远程服务器配置过程请参考第4部分。

        （1）点击左上方“+”号，在弹出页面选择Python后，页面会主动创建一个Unamed配置。Unamed为配置名称，用户可以在Name字段中进行自定义，比如修改为“gpt2”。

        （2）在Scripts path中输入launch.py的完整绝对路劲，例如“/root/miniconda3/envs/clai/lib/python3.8/site-packages/torch/distributed/launch.py”。

        （3）在Parameters中输入启动参数，例如“--nproc_per_node=1 /root/project/ColossalAI-Examples/language/gpt/train_gpt.py --config=/root/project/ColossalAI-Examples/language/gpt/gpt2_configs/gpt2_vanilla.py --from_torch”。注意，这里为了方便后续调试而将nproc_per_node设置为1，即并行节点数为1。

        步骤（3）设置完成后，运行程序会调试找不到运行程序文件或配置文件，一般设置成完整绝对路径即可解决问题。

        （4）点击运行按钮即可开始训练。

        配置完成界面如下所示。

 图3 配置完成页面

 图4 程序运行按钮

3 程序调试

        启动程序完成后，程序调试只要以Debug方式运行即可调试。Debug运行之前，可在实际主程序中设置断点，以便进行后续调试，如下图所示。

 图5 主程序断点设置

        通过不断步入等调试操作，我们可以定位到以下几个关键断点位置：

        （1）Colossal-AI框架训练入口：colossalai/trainer/_trainer.py第319行“for epoch in range(last_epoch, epochs):”。

        （2）Colossal-AI框架训练损失计算结果：colossalai/trainer/_trainer.py第181行logits, label, loss = self.engine.execute_schedule(”。

        （3）Colossal-AI框架模型计算入口：colossalai/amp/naive_amp/naive_amp.py第152行“out = self.model(*args, **kwargs)”。

        （4）Pytorch模型计算入口：torch/nn/modules/module.py第1130行“return forward_call(*input, **kwargs)”。

        （5）GPT模型计算入口：colossalai/titans/model/gpt/gpt.py第105行“x = self.embed(input_ids)”。

4 远程调试

4.1 服务器配置

        由于模型经常部署在服务器上，所以我们需要进行远程运行与调试。PyCharm远程配置页面为“Tools->Deployment->Configuration”，并在该页面点击“+”后选择“SFTP”。此时，页面会弹出新建服务器对话框，用户给服务器命名即可，例如test。

 图6 远程服务器配置

        在上述远程服务器配置页面的Connection中完成ssh远程连接配置，并在Mappings中完成路径配置。Mappings中的Local path为本机路径，Deployment path为服务器对应路径。Web path不用设置。程序运行时会将Deployment path的文件同步到本机的Local path。

4.2 python解释器

        PyCharm远程Python解释器配置页面为“Files->Settings->Project: python ->Python Interpreter->SSH Interpreter”，如下图所示。

图7 Python解释器配置

4.3 远程调试

        PyCharm远程调试与上述第3部分调试完全一致。调试之前，我们需要通过Tools->Deployment将运行程序同步下载到本地，如train_gpt.py等，然后在该文件中设置断点，最后以Debug方式运行程序即可开始调试。

        本文为博主原创文章，未经博主允许不得转载。

        本文为专栏《Python从零开始进行AIGC大模型训练与推理》系列文章，地址为“https://blog.csdn.net/suiyingy/article/details/130169592”。

        另外，本专栏具体更新可关注文章下方公众号，也可关注本专栏。所有相关文章会在《Python从零开始进行AIGC大模型训练与推理》中进行更新，地址为“https://blog.csdn.net/suiyingy/article/details/130169592”。所有AIGC类模型部署的体验效果将在RdFast小程序中同步上线。