如何在使用Hugging Face Seq2SeqTrainer的PyTorch代码中启用IPEX

启用IPEX优化Seq2SeqTrainer的操作步骤 #

方法1：使用Transformers Trainer原生IPEX支持（推荐） #

新版本Hugging Face Transformers的Trainer/Seq2SeqTrainer已经内置了IPEX适配，只需要两步操作即可完成启用：

• 首先在代码开头导入IPEX库

import intel_extension_for_pytorch as ipex

• 在定义training_args的时候，直接添加use_ipex=True参数即可，不需要修改你现有的Trainer初始化代码，示例如下：

training_args = Seq2SeqTrainingArguments(
    # 保留你原本设置的所有参数，比如output_dir、per_device_train_batch_size、学习率等
    use_ipex=True, # 新增这一行即可启用IPEX优化
)

Trainer会自动完成模型、优化器的IPEX适配，以及算子替换、精度优化等操作。

方法2：手动适配IPEX（适用于旧版本Transformers不支持use_ipex参数的场景） #

如果你的Transformers版本比较旧，没有内置IPEX参数支持，就需要在初始化Trainer之前手动优化模型和优化器：

• 先导入IPEX库

import intel_extension_for_pytorch as ipex
import torch

• 在模型加载完成后、Trainer初始化之前，对模型和优化器做IPEX优化，示例如下：

# 先定义好优化器，用默认优化器的话可以直接从TrainingArguments生成
optimizer = AdamW(model.parameters(), lr=training_args.learning_rate)

# 用IPEX优化模型和优化器，如果你要开启BF16混合精度训练可以添加dtype=torch.bfloat16参数
model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.float32)

• 再用你原来的代码初始化Trainer即可，不需要修改Trainer的入参。

额外优化建议 #

• 如果你的CPU支持AVX512或者AMX指令集，可以在training_args里加bf16=True开启BF16混合精度训练，配合IPEX可以获得更明显的性能提升。
• 可以调整dataloader_num_workers参数匹配你的CPU核心数，进一步提升数据加载的效率。

内容的提问来源于stack exchange，提问作者Janani Chandran