GPU集群上分布式训练大模型

总结一下如何在超算系统上进行预训练大模型的分布式训练 / 微调，文中代码已上传至 github

实验环境

集群1：国家广州超算 星逸A800智能AI集群
GPU：8 * Nvdia Tesla-A800 80G显存
CPU：2 * 28核 Intel Xeon Gold 6348
内存：1024GB

集群2：并行科技 中国国家网格 N12 区（cngrid12）
GPU：4 * Nvdia Tesla-V100 16G显存
CPU：20 核 Intel® Xeon® CPU E5-2640 v4
内存：128GB

在超算分布式环境上和本地训练有几点不同：

1. 超算环境无法科学上网，需要手动下载并上传：数据、tokenizer、模型和模型参数，并在代码中作相应修改。
2. 通过 slurm 进行作业管理，编写并提交 sbatch 脚本来运行作业。
3. 每个集群的环境各不相同，移植时需要注意配置环境和预加载相关的库。
4. 训练超大模型时，单个GPU显存有限，仅使用 torch.nn.parallel 数据并行常常无法加载完整模型无法完成训练，或只能小批次训练。因此需要用到分布式训练框架，常见的分布式训练框架有Horovod，Megatron-LM，DeepSpeed等。

1 举例：bert-large

1.1 本地单卡训练bert-large（假设GPU显存足够大）

Step1 编写训练代码 run_bert.py ：

import torch
from transformers import BertTokenizer, BertForMaskedLM, DataCollatorForLanguageModeling
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm

# 加载预训练的tokenizer和模型
tokenizer = BertTokenizer.from_pretrained("bert-large-uncased")
model = BertForMaskedLM.from_pretrained("bert-large-uncased")

# 加载WikiText数据集
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

# 数据预处理
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# 设置数据加载器
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)
train_loader = DataLoader(tokenized_dataset, batch_size=4, shuffle=True, collate_fn=data_collator)  # 根据显存调整batch_size

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 设置优化器
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)