NUMA-Torch: 멀티소켓 CPU에서 컴파일러 기반 NUMA 인식 LLM 프리필
arXiv:2611.07105v1 발표 유형: 신규 초록: 온칩 캐시에 매개변수를 저장할 수 없는 충분히 큰 변환기 기반 대규모 언어 모델(LLM)의 경우, 멀티 소켓 CPU에서의 프리필 처리량은 최대 부동소수점 처리량이 아닌 지속 가능한 DRAM 대역폭에 의해 종종 제한됩니다. 표준 PyTorch 컴파일(즉, TorchInductor) 경로에서는 메모리 뱅크와 연산 코어 간의 연관성을 고려하지 않고 모델 가중치를 할당합니다. 이로 인해 가중치 접근의 상당 부분이 소켓 간 상호 연결을 통과하게 되어 각 소켓에서 사용할 수 있는 유효 메모리 대역폭이 크게 저하됩니다. 우리는 TorchInductor를 NUMA-aware로 만드는 컴파일러 수준의 솔루션을 통해 이러한 병목 현상을 해결합니다. 각 NUMA 소켓을 독립적인 텐서 병렬 파티션으로 취급하여 NUMA-aware Tensor Parallelism(NUMA-TP)을 도입합니다. 이 방법은 소켓 간에 가중치를 투명하게 분할하고 연산 커널을 스케줄링하여 스레드 풀과 메모리 할당이 CPU 연관성을 활용해 해당 소켓 내에서 엄격히 로컬로 유지되도록 합니다. 우리는 GPT-2(117M, 345M, 762M) 및 BERT(Small, Large)와 같은 표준 다중 헤드 주의와 MLP를 테스트하며, Multi-Query Attention과 GeGLU 게이트된 MLP를 결합한 Gemma 3 270M에서도 이 접근법을 평가합니다. 데이터 지역성을 엄격히 준수함으로써 우리의 접근법은 소켓 간 트래픽을 성공적으로 완화하여 듀얼 소켓 Haswell 플랫폼에서 최대 2.63배, 듀얼 소켓 Cascade Lake 플랫폼에서 최대 1.87배, 그리고 사중 소켓 Haswell-EX 플랫폼에서 최대 2.27배의 프리필 속도 향상을 달성했습니다.
