Intel이 Sapphire Rapids CPU의 AMX 명령어 세트를 IPEX 라이브러리로 자동 활용해 PyTorch Transformer 분산 학습을 4개 노드에서 선형 스케일링 달성

Context

GPU는 모델 학습에 필수로 여겨져 왔으나 조달 어려움, 높은 비용, 유휴 시간 활용 제한이 있다. CPU 기반 학습은 비용 효율적 대안이지만 기존 Ice Lake(3세대) Xeon CPU의 행렬 연산 성능이 제한적이었다.

Technical Solution

• AMX(Advanced Matrix Extensions) 명령어 도입: Sapphire Rapids(4세대) Xeon CPU에 2차원 타일 레지스터 추가로 BF16 및 INT8 행렬 곱셈 가속화
• Intel Extension for PyTorch(IPEX) 라이브러리 통합: 코드 변경 없이 AMX 명령어 자동 활용, Hugging Face Transformers와 사전 통합
• Intel oneAPI CCL(Collective Communications Library)을 통한 분산 학습 구현: mpirun으로 4개 노드에 총 8개 Python 프로세스 배치(노드당 2개)
• 환경 변수 기반 CPU 자원 할당: OMP_NUM_THREADS=24(학습용), CCL_WORKER_COUNT=2(통신용), 6개 코어는 커널 예약
• AWS EC2 R7iz.metal-16xl 인스턴스 클러스터 구성: 64 vCPU, 512GB RAM 베어메탈 인스턴스로 AMX 지원(가상화 환경 미지원으로 우회)

Impact

• 1 에포크 학습 시간: 26분 → 7분 30초(4개 노드 기준)
• 선형 스케일링 달성도: 이상적 6분 30초 대비 7분 30초(선형 스케일링에 매우 근접)
• AWS Spot 인스턴스 활용 시 추가 비용 절감: 온디맨드 대비 최대 90% 할인

Key Takeaway

Small/Medium 데이터셋 기반 Transformer 학습은 GPU 조달 의존을 제거하고 Intel CPU의 벡터화 명령어(AVX-512, AMX)를 활용하면 비용 대비 선형 스케일링을 달성할 수 있으며, 프레임워크 수준의 라이브러리 통합(IPEX, CCL)으로 사용자 코드 변경 없이 하드웨어 최적화를 투명하게 제공 가능하다.