Accelerating PyTorch Transformers with Intel Sapphire Rapids – part 1

TL;DR

Hugging Face demonstrates that training PyTorch Transformers on a four‑node cluster of Intel Sapphire Rapids CPUs with the Intel Extension for PyTorch (IPEX) and oneCCL yields an 8× speed‑up over comparable Ice Lake instances and scales almost linearly, making CPU‑based training a cost‑effective alternative to GPUs.

Why Train on CPUs?

Training on Intel Xeon CPUs can be cheaper and more scalable than GPU‑based training, especially for small‑ to medium‑sized models and datasets. Xeon CPUs provide AVX‑512, Hyper‑Threading, and now Advanced Matrix Extensions (AMX), which accelerate matrix multiplication. CPUs are widely available, can be repurposed for other workloads, and cloud spot instances can reduce costs by up to 90% compared to on‑demand instances.

Advanced Matrix Extensions (AMX)

Sapphire Rapids introduces AMX, a set of new instructions that speed up matrix multiplication—the core operation in deep‑learning training. AMX supports BF16 and INT8 data types and adds 2‑dimensional tile registers. To use AMX, the Linux kernel must be version 5.16 or newer; however, Intel and AWS have back‑ported the required support to kernel 5.15 on the bare‑metal r7iz.metal-16xl instances used in this guide.

Building a Sapphire Rapids Cluster on AWS

The simplest way to access Sapphire Rapids hardware is via the Amazon EC2 R7iz bare‑metal instances (e.g., r7iz.metal-16xl). Because the preview does not yet support AMX in virtual machines, the guide uses bare‑metal instances with 64 vCPUs and 512 GB RAM.

Network Setup

  • Open SSH (port 22) on all nodes.
  • Configure password‑less SSH from the master node to every node, including the master itself.
  • Create a security group that allows unrestricted intra‑cluster traffic and attach it to all instances.

Setting Up the Master Node

  1. Launch an r7iz.metal-16xl instance with Ubuntu 20.04 AMI ami-07cd3e6c4915b2d18.
  2. Verify AMX support with lscpu (flags amx_bf16 amx_tile amx_int8).
  3. Install dependencies:
    sudo apt-get update
sudo apt install libgoogle-perftools-dev -y
sudo apt-get install python3-pip -y
pip install --upgrade pip
pip install virtualenv
virtualenv cluster_env
source cluster_env/bin/activate
pip install torch==1.13.0 -f https://download.pytorch.org/whl/cpu
pip install intel_extension_for_pytorch==1.13.0 -f https://developer.intel.com/ipex-whl-stable-cpu
pip install -f https://developer.intel.com/ipex-whl-stable-cpu
pip install transformers==4.24.0
git clone https://github.com/huggingface/transformers.git
cd transformers && git checkout v4.24.0
  4. Generate an SSH key pair (ssh-keygen) and store it as ~/.ssh/cluster.
  5. Create an AMI from this configured instance for reuse.

Setting Up the Remaining Nodes

  1. Launch three additional r7iz.metal-16xl instances using the AMI created above.
  2. On the master, edit ~/.ssh/config to define hosts node1, node2, and node3 with their private IPs and the cluster key.
  3. Verify password‑less access with ssh node[1-3].
  4. Create a ~/hosts file listing all nodes, including localhost for the master:
    localhost
node1
node2
node3

Launching Distributed Training

The example fine‑tunes DistilBERT on the SQuAD dataset.

Single‑Node Baseline

source ~/cluster_env/bin/activate
cd ~/transformers/examples/pytorch/question-answering
pip install -r requirements.txt
export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so"
python run_qa.py \
  --model_name_or_path distilbert-base-uncased \
  --dataset_name squad \
  --do_train --do_eval \
  --per_device_train_batch_size 32 \
  --num_train_epochs 1 \
  --output_dir /tmp/debug_squad/ \
  --use_ipex --bf16 --no_cuda

A single epoch completed in ≈26 minutes, compared with ≈3 h 30 min on an Ice Lake c6i.16xlarge instance—a 8× speed‑up.

Distributed Run on Four Nodes

Environment variables for oneCCL and thread placement:

oneccl_bindings_for_pytorch_path=$(python -c "from oneccl_bindings_for_pytorch import cwd; print(cwd)")
source $oneccl_bindings_for_pytorch_path/env/setvars.sh
export MASTER_ADDR=172.31.3.190
export NUM_PROCESSES=8               # 2 processes per node × 4 nodes
export NUM_PROCESSES_PER_NODE=2
export CCL_WORKER_COUNT=2
export CCL_WORKER_AFFINITY=auto
export KMP_HW_SUBSET=1T
export OMP_NUM_THREADS=24           # training threads per process

Launch with mpirun:

mpirun -f ~/hosts \
  -n $NUM_PROCESSES -ppn $NUM_PROCESSES_PER_NODE \
  -genv OMP_NUM_THREADS=24 \
  -genv LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so" \
  python3 run_qa.py \
    --model_name_or_path distilbert-base-uncased \
    --dataset_name squad \
    --do_train --do_eval \
    --per_device_train_batch_size 32 \
    --num_train_epochs 1 \
    --output_dir /tmp/debug_squad/ \
    --overwrite_output_dir \
    --no_cuda \
    --xpu_backend ccl \
    --bf16

The epoch time dropped to 7 min 30 s, only 1 minute above the ideal linear scaling target of 6 min 30 s (26 min / 4). The screenshot in the original post shows two training processes per node and the master process orchestrating the run.

Implications and Takeaways

  • Cost Efficiency – CPU spot instances can be dramatically cheaper than GPU instances while delivering comparable training speed for many workloads.
  • Scalability – Near‑linear scaling across four Sapphire Rapids nodes demonstrates that the combination of IPEX and oneCCL effectively distributes both compute and communication.
  • Ease of Use – No code changes are required; enabling --use_ipex and --bf16 activates the new AMX instructions automatically.
  • Future Work – The authors plan a follow‑up post covering inference performance on Sapphire Rapids, which will complete the picture of end‑to‑end CPU acceleration.

Further Resources

Sources