CPU

现代CPU能够通过利用底层硬件内置的优化和在 fp16 或 bf16 数据类型上进行训练来高效地训练大型模型。

本指南重点介绍如何使用混合精度在 Intel CPU 上训练大型模型。PyTorch 训练的 CPU 后端已启用 AMP。

Trainer 通过添加 --use_cpu 和 --bf16 参数支持 CPU 上的 AMP 训练。以下示例演示了 run_qa.py 脚本。

python run_qa.py \
 --model_name_or_path google-bert/bert-base-uncased \
 --dataset_name squad \
 --do_train \
 --do_eval \
 --per_device_train_batch_size 12 \
 --learning_rate 3e-5 \
 --num_train_epochs 2 \
 --max_seq_length 384 \
 --doc_stride 128 \
 --output_dir /tmp/debug_squad/ \
 --bf16 \
 --use_cpu

这些参数也可以添加到 TrainingArguments，如下所示。

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
    use_cpu=True,
)

资源

在用 Intel Sapphire Rapids 加速 PyTorch Transformers 博客文章中了解更多关于在 Intel CPU 上训练的信息。

< > 在 GitHub 上更新