梯度检查点(Gradient Checkpointing)
大模型的参数量巨大,即使将batch_size设置为1并使用梯度累积的方式更新,也仍然会OOM。原因是通常在计算梯度时,我们需要将所有前向传播时的激活值保存下来,这消耗大量显存。
还有另外一种延迟计算的思路,丢掉前向传播时的激活值,在计算梯度时需要哪部分的激活值就重新计算哪部分的激活值,这样做倒是解决了显存不足的问题,但加大了计算量同时也拖慢了训练。
梯度检查点(Gradient Checkpointing)在上述两种方式之间取了一个平衡,这种方法采用了一种策略选择了计算图上的一部分激活值保存下来,其余部分丢弃,这样被丢弃的那一部分激活值需要在计算梯度时重新计算。
下面这个动图展示了一种简单策略:前向传播过程中计算节点的激活值并保存,计算下一个节点完成后丢弃中间节点的激活值,反向传播时如果有保存下来的梯度就直接使用,如果没有就使用保存下来的前一个节点的梯度重新计算当前节点的梯度再使用。

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 方法保存分词器时,会生成多个文件。这些文件包含了分词器的各种配置和模型数据,以便之后能够准确地加载和使用分词器。以下是这些文件的具体作用说明:
**special_tokens_map.json**:**tokenizer.model**:**tokenizer_config.json**:**added_tokens.json**:**tokenization_internlm2.py**:
这些文件一起定义了分词器的行为,并且可以通过调用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**的作用:- 加载自定义分词器: