Reading

大模型训练基础

梯度检查点(Gradient Checkpointing)

大模型的参数量巨大,即使将batch_size设置为1并使用梯度累积的方式更新,也仍然会OOM。原因是通常在计算梯度时,我们需要将所有前向传播时的激活值保存下来,这消耗大量显存。

还有另外一种延迟计算的思路,丢掉前向传播时的激活值,在计算梯度时需要哪部分的激活值就重新计算哪部分的激活值,这样做倒是解决了显存不足的问题,但加大了计算量同时也拖慢了训练

梯度检查点(Gradient Checkpointing)在上述两种方式之间取了一个平衡,这种方法采用了一种策略选择了计算图上的一部分激活值保存下来,其余部分丢弃,这样被丢弃的那一部分激活值需要在计算梯度时重新计算

下面这个动图展示了一种简单策略:前向传播过程中计算节点的激活值并保存,计算下一个节点完成后丢弃中间节点的激活值,反向传播时如果有保存下来的梯度就直接使用,如果没有就使用保存下来的前一个节点的梯度重新计算当前节点的梯度再使用。

image

Transformer框架开启梯度检查点非常简单,仅需在TrainingArguments中指定gradient checkpoint为True即可:

training_args = TrainingArguments(
    per_device_train_batch_size=1, gradient_accumulation_steps=4, gradient_checkpointing=True, **default_args
)

trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()

Tokenizer

Transformers加载和保存权重

使用 transformers 库的 tokenizer.save_pretrained 方法保存分词器时,会生成多个文件。这些文件包含了分词器的各种配置和模型数据,以便之后能够准确地加载和使用分词器。以下是这些文件的具体作用说明:

  1. special_tokens_map.json
    • 记录了分词器中的特殊标记(例如 [CLS], [SEP], [PAD], [UNK], [MASK] 等)与它们在词汇表中的映射关系。这个文件确保在加载分词器时,这些特殊标记能正确地映射和使用。
  2. tokenizer.model
    • 分词器的核心模型文件,通常是一个二进制文件,包含了分词器的词汇表和编码逻辑。这是分词器用来将文本转换为 token id 的关键数据文件。
  3. tokenizer_config.json
    • 包含了分词器的配置信息,比如分词器类型、特殊标记、预处理规则等。这是一个 JSON 文件,用于在加载分词器时配置其行为。
  4. added_tokens.json
    • 保存了在预训练模型中未包含的、在训练或微调过程中额外添加的 token。例如,一些特定于任务的标记或符号。
  5. tokenization_internlm2.py
    • 这是分词器的实现代码文件。这个文件包含了与 InternLM2 模型相关的分词器逻辑。当加载分词器时,这个文件会被用来重建分词器对象。此文件可能是特定于某些自定义或扩展的分词器实现。

这些文件一起定义了分词器的行为,并且可以通过调用 tokenizer.from_pretrained 方法加载回来。如果你想分享或重新使用这个分词器,保留这些文件是必要的。

tokenizer 是如何知道要通过 tokenization_internlm2.py 文件加载分词器的,这是由 tokenizer_config.json 文件中的配置决定的。具体而言,tokenizer_config.json 文件中会包含一个键(tokenizer_class)来指定分词器类的名称。加载分词器时,transformers 库会检查该配置文件,根据这个键的值来确定要加载的分词器类。

  "bos_token": "<s>",
  "chat_template": "{{ bos_token }}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}",
  "clean_up_tokenization_spaces": false,
  "eos_token": "</s>",
  "model_max_length": 4096,
  "pad_token": "</s>",
  "tokenizer_class": "InternLM2Tokenizer",
  "unk_token": "<unk>"
  • tokenizer_class 的作用
    • tokenizer_class 键指定了分词器类的名称。在上面的例子中,它的值是 "InternLM2Tokenizer"。当你调用 transformers.AutoTokenizer.from_pretrained() 时,transformers 库会根据这个值去寻找名为 InternLM2Tokenizer 的类。
  • 加载自定义分词器
    • 当加载自定义分词器时,transformers 库会搜索 transformers 包中的分词器类。如果没有找到(比如自定义分词器类不在库中),则会使用当前目录或指定路径中的 tokenization_internlm2.py 文件来找到并加载该类。