如何使用Intel Extension for PyTorch(IPEX)优化Huggingface SWAG多项选择任务

你之前仅修改了模型挂载的设备指向，没有启用IPEX的算子融合、指令集优化等核心能力，因此无法达到预期的性能提升效果，完整适配方案如下：

代码层面的完整修改 #

• 模型与优化器的IPEX封装
  在model.to(device)执行完成、优化器初始化完成后，添加IPEX优化逻辑，代码片段如下：

  import torch
  import intel_pytorch_extension as ipex
  
  # 启用oneDNN算子融合
  ipex.enable_onednn_fusion(True)
  # 封装模型和优化器，开启IPEX核心优化能力
  model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.float32)
  

• 混合精度适配（CPU支持AVX512/AMX指令集时推荐开启）
  在训练循环的前向传播逻辑外，添加CPU混合精度上下文，代码片段如下：

  with torch.cpu.amp.autocast(dtype=torch.bfloat16):
      outputs = model(**batch)
      loss = outputs.loss
  

启动参数与环境变量配置 #

启动任务前先配置CPU多线程优化环境变量，根据你的CPU物理核心数调整OMP_NUM_THREADS取值，完整启动命令如下：

export DATASET_NAME=swag
# 线程优化配置，<物理核心数>替换为你的CPU实际物理核心数量
export OMP_NUM_THREADS=<物理核心数>
export KMP_AFFINITY=granularity=fine,compact,1,0
export KMP_BLOCKTIME=1

accelerate launch run_swag_no_trainer.py \
  --model_name_or_path bert-base-cased \
  --dataset_name $DATASET_NAME \
  --max_seq_length 128 \
  --per_device_train_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir /tmp/$DATASET_NAME/ \
  --bf16 True # 支持BF16的设备添加该参数

性能参考 #

• 支持AMX指令集的4代及以上Xeon CPU，开启BF16混合精度后训练速度可提升2-3倍
• 支持AVX512指令集的CPU，开启优化后速度可提升1.5倍以上

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