Intel Sapphire Rapids와 함께 PyTorch Transformer 가속 – 파트 1

TL;DR

Hugging Face는 Intel Extension for PyTorch (IPEX)와 oneCCL을 사용하여 Intel Sapphire Rapids CPU 4노드 클러스터에서 PyTorch Transformer를 학습하면, 유사한 Ice Lake 인스턴스 대비 8배 속도 향상을 제공하고 거의 선형적으로 확장되어 CPU 기반 학습이 GPU에 대한 비용 효율적인 대안이 됨을 보여줍니다.

CPU에서 학습하는 이유

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는 행렬 곱셈(딥러닝 학습의 핵심 연산)을 가속화하는 새로운 명령어 집합인 AMX를 도입합니다. AMX는 BF16 및 INT8 데이터 타입을 지원하고 2차원 타일 레지스터를 추가합니다. AMX를 사용하려면 Linux 커널 버전이 5.16 이상이어야 하지만, Intel과 AWS는 이 가이드에서 사용되는 베어메탈 r7iz.metal-16xl 인스턴스의 커널 5.15에 필요한 지원을 백포팅했습니다.

AWS에서 Sapphire Rapids 클러스터 구축

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.

네트워크 설정

  • 모든 노드에서 SSH(포트 22)를 엽니다.
  • 마스터 노드에서 모든 노드(마스터 자체 포함)로 비밀번호 없이 SSH를 설정합니다.
  • 제한 없는 클러스터 내부 트래픽을 허용하는 보안 그룹을 생성하고 모든 인스턴스에 연결합니다.

마스터 노드 설정

  1. Ubuntu 20.04 AMI ami-07cd3e6c4915b2d18를 사용하여 r7iz.metal-16xl 인스턴스를 시작합니다.
  2. lscpu(플래그 amx_bf16 amx_tile amx_int8)로 AMX 지원을 확인합니다.
  3. 필요한 패키지를 설치합니다:
    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. ssh-keygen을 사용해 SSH 키 쌍을 생성하고 ~/.ssh/cluster에 저장합니다.
  5. 이 설정된 인스턴스로부터 재사용 가능한 AMI를 생성합니다.

나머지 노드 설정

  1. 위에서 만든 AMI를 사용하여 추가로 r7iz.metal-16xl 인스턴스 3대를 시작합니다.
  2. 마스터에서 ~/.ssh/config를 편집하여 호스트 node1, node2, node3을 각각의 사설 IP와 cluster 키로 정의합니다.
  3. ssh node[1-3] 명령으로 비밀번호 없는 접속을 확인합니다.
  4. 모든 노드와 마스터(localhost)를 나열한 ~/hosts 파일을 생성합니다:
    localhost
node1
node2
node3

분산 학습 실행

예제에서는 DistilBERT를 SQuAD 데이터셋에 대해 파인튜닝합니다.

단일 노드 기준

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

단일 에포크가 ≈26 분에 완료되었으며, Ice Lake c6i.16xlarge 인스턴스에서 ≈3 시간 30 분이 걸린 것에 비해 8배 속도 향상을 보였습니다.

4노드 분산 실행

oneCCL 및 스레드 배치를 위한 환경 변수:

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

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

에포크 시간이 7분 30초로 감소했으며, 이상적인 선형 확장 목표인 6분 30초(26분 / 4)보다 1분만 초과했습니다. 원본 게시물의 스크린샷에는 노드당 두 개의 학습 프로세스와 마스터 프로세스가 실행을 조정하는 모습이 표시됩니다.

시사점 및 요약

  • 비용 효율성 – CPU 스팟 인스턴스는 GPU 인스턴스보다 비용이 크게 저렴하면서도 많은 워크로드에 대해 비슷한 학습 속도를 제공합니다.
  • 확장성 – 4개의 Sapphire Rapids 노드에서 거의 선형에 가까운 확장을 보여주며, IPEX와 oneCCL의 조합이 계산과 통신을 효과적으로 분산함을 증명합니다.
  • 사용 편의성 – 코드 변경이 필요 없으며, --use_ipex와 --bf16 옵션을 활성화하면 새로운 AMX 명령어가 자동으로 적용됩니다.
  • 향후 작업 – 저자들은 Sapphire Rapids에서의 추론 성능을 다루는 후속 게시물을 계획하고 있으며, 이를 통해 엔드‑투‑엔드 CPU 가속의 전체 그림을 완성할 예정입니다.

추가 자료

Sources