预训练 BERT 使用 Hugging Face 和 PyTorch 在 AMD GPU 上

Pre-training BERT using Hugging Face & PyTorch on an AMD GPU — ROCm Blogs

2024年1月26日，作者：Vara Lakshmi Bayanagari.

这篇博客解释了如何从头开始使用 Hugging Face 库和 PyTorch 后端在 AMD GPU 上为英文语料(WikiText-103-raw-v1)预训练 BERT 基础模型的端到端过程。

你可以在 GitHub folder中找到与这篇博客相关的文件。

BERT简介

BERT是一种在2019年提出的语言表示模型。其模型架构基于一个transformer编码器，其中自注意力层对输入的每个token对进行注意力计算，整合了来自两个方向的上下文（因此称为BERT的“双向”特性）。在此之前，像ELMo和GPT这样的模型只使用从左到右的（单向）架构，这极大地限制了模型的表现力；模型性能依赖于微调。

本博客解释了BERT所采用的预训练任务，这些任务在通用语言理解评估（GLUE）基准测试中取得了最先进的成果。在接下来的章节中，我们将展示在PyTorch中的实现。

这篇BERT paper最先提出了一种新的预训练方法，称为掩码语言建模（MLM）。MLM随机掩盖输入的某些部分，并对一批输入进行训练以预测这些被掩盖的tokens。预训练期间，在对输入进行分词之后，15%的tokens被随机挑选。其中，80%被替换为一个`[MASK]`标记，10%被替换为一个随机标记，10%则保持不变。

在下面的示例中，MLM预处理方法如下：`dog`标记保持不变，`Golden`和`years`标记被掩盖，并且`and`标记被随机替换为`paper`标记。预训练的目标是使用`CategoricalCrossEntropy`损失来预测这些标记，以便模型学习语言的语法、模式和结构。

Input sentence: My dog is a Golden Retriever and his is 5 years old

After MLM: My dog is a [MASK] Retriever paper his is 5 [MASK] old

此外，为了捕捉句子之间的关系，超越掩码语言建模任务，论文提出了第二个预训练任务，称为下一个句子预测（NSP）。在不改变架构的情况下，论文证明了NSP有助于提升问答（QA）和自然语言推理（NLI）任务的结果。

这个任务不直接输入token流，而是输入一对句子的token，例如`A`和`B`，以及一个前置分类标记（`[CLS]`）。分类标记指示句对是随机组合的（label=0）还是`B`是`A`的下一个句子（label=1）。因此，NSP预训练是一种二元分类任务。

_IsNext_ Pair: [1] My dog is a Golden Retriever. He is five years old.

Not _IsNext_ Pair: [0] My dog is a Golden Retriever. The next chapter in the book is a biography.

总之，数据集首先进行预处理以形成一对句子，然后进行分词，并最终随机掩盖某些tokens。预处理后的输入批次要么*填充*（使用`[PAD]`标记）或*修剪*（到_max_seq_length_超参数），以便所有输入元素在加载到BERT模型中之前都统一为相同的长度。BERT模型配有两个分类头：一个用于MLM（`num_cls_heads = vocab_size），另一个用于NSP（num_cls_heads=2`）。来自两个预训练任务的分类损失之和用于训练BERT。

在多台 AMD GPU 上的实现

在开始之前，确保您已经满足以下要求：

1. 在搭载 AMD GPU 的设备上安装 ROCm 兼容的 PyTorch。本实验在 ROCm 5.7.0 和 PyTorch 2.0.1 上进行了测试。

2. 运行命令 pip install datasets transformers accelerate 以安装 Hugging Face 的相关库。

3.  运行 accelerate config 命令以设置分布式训练参数，详见此处。在本实验中，我们使用了单节点上的八块 GPU 并行计算，运用了 DistributedDataParallel。

实现

Hugging Face 使用 Torch 作为大多数模型的默认后端，从而实现了这两个框架的良好结合。为了简化常规训练步骤并避免样板代码，Hugging Face 提供了一个名为 Trainer 的类，该类模仿了 PyTorch 的功能。类似地，Lightning AI 提供了 Trainer 类。此外，对于分布式训练，Hugging Face 可能更方便，因为代码中没有额外的配置设置，系统会根据 accelerate config 自动检测并利用所有 GPU 设备。然而，如果你希望进一步自定义你的模型并对加载预训练检查点做出额外修改，原生的 PyTorch 是更好的选择。这篇博客解释了使用 Hugging Face 的 transformers 库对 BERT 进行端到端预训练，同时提供了简化的数据预处理管道。

使用 Hugging Face 的 Trainer 进行 BERT 预训练可以用几行代码来总结。transformer 编码器、MLM 分类头和 NSP 分类头都打包在 Hugging Face 的 BertForPreTraining 模型中，该模型返回一个累积分类损失，如我们在 介绍 中所解释的。模型使用默认的 BERT base 配置参数（`NUM_LAYERS`、`ACT_FUNC`、`BATCH_SIZE`、`HIDDEN_SIZE`、`EMBED_DIM` 等）进行初始化。你可以从 Hugging Face 的 BertConfig 中导入这些参数。

那就是全部了吗？几乎。训练最关键的部分是数据预处理。预处理分为三个步骤：

1.  将你的数据集重新组织为每个文档的句子字典。这对于从随机文档中选取随机句子以进行 NSP 任务非常有用。为此，可以对整个数据集使用简单的for循环。

2. 使用 Hugging Face 的 AutoTokenizer 来对所有句子进行标记化。

3. 使用另一个 for 循环，创建 50% 随机对和 50% 顺序对的句子对。

我已经对 WikiText-103-raw-v1 语料库（2,500 M单词）进行了上述的预处理步骤，并将生成的验证集放在这里。预处理的训练集已上传到 Hugging Face Hub。

接下来，导入 DataCollatorForLanguageModeling 收集器以运行 MLM 预处理，并获取掩码和句子分类标签。在使用 Trainer 类时，我们只需要访问 torch.utils.data.Dataset 和一个收集函数。与 TensorFlow 不同，Hugging Face 的 Trainer 会从数据集和收集器函数中创建数据加载器。为了演示，我们使用了有 3,000+ 句对的 Wikitext-103-raw-v1 验证集。

tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
collater = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)
# tokenized_dataset = datasets.load_from_disk(args.dataset_file)
tokenized_dataset_valid = datasets.load_from_disk('./wikiTokenizedValid.hf')

创建一个 TrainerArguments 实例，并传递所有必需的参数，如以下代码所示。这部分代码有助于在训练模型时抽象样板代码。该类很灵活，因为它提供了 100 多个参数来适应不同的训练模式；有关更多信息，请参阅 Hugging face transformers 页面。

你现在可以使用 t.train() 来训练模型了。你还可以通过将 resume_from_checkpoint=True 参数传递给 t.train() 来恢复训练。trainer 类会提取 output_dir 文件夹中的最新检查点，并继续训练直到达到总共 num_train_epochs。

train_args = TrainingArguments(output_dir=args.output_dir, overwrite_output_dir =True, per_device_train_batch_size =args.BATCH_SIZE, logging_first_step=True,
                                   logging_strategy='epoch', evaluation_strategy = 'epoch', save_strategy ='epoch', num_train_epochs=args.EPOCHS,save_total_limit=50)
t = Trainer(model, args = train_args, data_collator=collater, train_dataset = tokenized_dataset, optimizers=(optimizer, None), eval_dataset = tokenized_dataset_valid)
t.train()#resume_from_checkpoint=True)

上述模型使用Adam优化器（`learning_rate=2e-5`）和`per_device_train_batch_size=8`进行了大约400个epoch的训练。在一块AMD GPU（MI210，ROCm 5.7.0，PyTorch 2.0.1）上，使用3,000+句对的验证集进行预训练仅需几个小时。训练曲线如图1所示。可以使用最佳模型检查点微调不同的数据集，并在各种NLP任务上测试其表现。

完整的代码如下：

set_seed(42)
parser = argparse.ArgumentParser()
parser.add_argument('--BATCH_SIZE', type=int, default = 8) # 32 is the global batch size, since I use 8 GPUs
parser.add_argument('--EPOCHS', type=int, default=200)
parser.add_argument('--train', action='store_true')
parser.add_argument('--dataset_file', type=str, default= './wikiTokenizedValid.hf')
parser.add_argument('--lr', default = 0.00005, type=float)
parser.add_argument('--output_dir', default = './acc_valid/')
args = parser.parse_args()

accelerator = Accelerator()

if args.train:
    args.dataset_file = './wikiTokenizedTrain.hf'
    args.output_dir = './acc/'
print(args)

tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
collater = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)
tokenized_dataset = datasets.load_from_disk(args.dataset_file)
tokenized_dataset_valid = datasets.load_from_disk('./wikiTokenizedValid.hf')

model = BertForPreTraining(BertConfig.from_pretrained("bert-base-cased"))
optimizer = torch.optim.Adam(model.parameters(), lr =args.lr)

device = accelerator.device
model.to(accelerator.device)
train_args = TrainingArguments(output_dir=args.output_dir, overwrite_output_dir =True, per_device_train_batch_size =args.BATCH_SIZE, logging_first_step=True,
                               logging_strategy='epoch', evaluation_strategy = 'epoch', save_strategy ='epoch', num_train_epochs=args.EPOCHS,save_total_limit=50)#, lr_scheduler_type=None)
t = Trainer(model, args = train_args, data_collator=collater, train_dataset = tokenized_dataset, optimizers=(optimizer, None), eval_dataset = tokenized_dataset_valid)
t.train()#resume_from_checkpoint=True)

推理

以一个示例文本为例，使用分词器将其转换为输入tokens，并通过collator生成一个掩码输入。

collater = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15, pad_to_multiple_of=128)
text="The author takes his own advice when it comes to writing: he seeks to ground his claims in clear, concrete examples. He shows specific examples of bad writing to help readers better grasp exactly what he’s critiquing"
tokens = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(text))
inp = collater([tokens])
inp['attention_mask'] = torch.where(inp['input_ids']==0,0,1)

使用预训练的权重初始化模型并进行推理。你将看到模型生成的随机tokens没有上下文意义。

config = BertConfig.from_pretrained('bert-base-cased')
model = BertForPreTraining.from_pretrained('./acc_valid/checkpoint-19600/')
model.eval()
out = model(inp['input_ids'], inp['attention_mask'], labels=inp['labels'])

print('Input: ', tokenizer.decode(inp['input_ids'][0][:30]), '\n')
print('Output: ', tokenizer.decode(torch.argmax(out[0], -1)[0][:30]))

输入和输出如下所示。该模型在一个非常小的数据集（3,000多句子）上进行了训练；你可以通过在更大的数据集上训练，例如`wikiText-103-raw-v1`的训练切分数据，来提高性能。

The author takes his own advice when it comes to writing : he [MASK] to ground his claims in clear, concrete examples. He shows specific examples of bad
The Churchill takes his own, when it comes to writing : he continued to ground his claims in clear, this examples. He shows is examples of bad

源代码存储在这个 GitHub 文件夹。

结论

我们所描述的预训练BERT基础模型的过程可以很容易地扩展到不同大小的BERT版本以及不同的数据集。我们使用Hugging Face Trainer和PyTorch后端在AMD GPU上训练了我们的模型。对于训练，我们使用了`wikiText-103-raw-v1`数据集的验证集，但这可以很容易地替换为训练集，只需下载我们在Hugging Face Hub上的仓库中托管的预处理和标记化的训练文件Hugging Face Hub.

在本文中，我们通过MLM和NSP预训练任务复制了BERT的预训练过程，这与许多公共平台上仅使用MLM的方法不同。此外，我们没有使用数据集的小部分，而是预处理并上传了整个数据集到Hub上供您方便使用。在未来的文章中，我们将讨论在多个AMD GPU上使用数据并行和分布式策略来训练各种机器学习应用。