학습 성능이 크게 향상되었습니다. Bytedance의 Zheng Size는 대규모 모델에 대한 효율적인 분산 통신 및 컴퓨팅 통합을 달성하기 위한 Triton 분산 프레임워크를 설명합니다.

HyperAI가 주최하는 Meet AI Complier Technology Salon이 2025년 7회를 맞이했습니다. 커뮤니티 파트너와 수많은 업계 전문가들의 지원을 바탕으로 베이징, 상하이, 선전 등 여러 지역에 거점을 마련하여 개발자와 개발자 애호가들을 위한 소통 플랫폼을 제공하고, 선구적인 기술의 비밀을 밝히고, 현장 개발자들의 애플리케이션 피드백을 수렴하며, 기술 구현에 대한 실질적인 경험을 공유하고, 다각적인 관점에서 혁신적인 사고를 경청합니다.

위챗 공개 계정 "HyperAI Super Neuro"를 팔로우하고 키워드 "0705 AI Compiler"에 답글을 달면 공인 강사의 발표 PPT를 받으실 수 있습니다.

기조 연설 "Triton-distributed: 고성능 통신을 위한 네이티브 Python 프로그래밍"에서ByteDance의 시드 연구 과학자 Zheng Size이 논문에서는 대규모 모델 학습에서 Triton 분산의 통신 효율성과 플랫폼 간 적응성 측면에서 획기적인 진전을 자세히 분석하고, Python 프로그래밍을 통해 통신과 컴퓨팅을 긴밀하게 통합하는 방법도 설명합니다.공유 후, 현장은 금세 질문의 도가니로 빠져들었습니다. FLUX 프레임워크, Tile 프로그래밍 모델, AllGather 및 ReduceScatter 최적화 등 세부적인 내용에 대한 끊임없는 토론이 이어졌습니다. 토론은 핵심적인 기술적 어려움과 실무 경험에 초점을 맞추었고, 이론과 응용의 조화를 효과적으로 이끌어냈습니다.

HyperAI는 정 사이즈 씨의 연설을 원래 의도를 훼손하지 않고 편집하고 요약했습니다. 다음은 연설 전문입니다.

분산 훈련의 실제 과제

대규모 모델의 급속한 진화 맥락에서 학습과 추론은 모두분산 시스템은 없어서는 안 될 부분이 되었습니다.우리는 또한 이 방향으로 컴파일러 수준의 탐색을 수행했고 이 프로젝트를 오픈 소스화하고 Triton-Distributed라는 이름을 붙였습니다.

현재 주류를 이루는 하드웨어 상호 연결 방식으로는 NVLink, PCIe, 그리고 노드 간 네트워크 통신이 있습니다. 이상적인 조건에서 H100의 NVLink 단방향 대역폭은 450GB/s에 도달할 수 있지만, 대부분의 국내 구축 환경에서는 H800이 더 널리 사용됩니다. H800의 단방향 대역폭은 약 200GB/s에 불과하며, 전반적인 통신 성능과 토폴로지 복잡성이 크게 감소합니다.이 프로젝트에서 우리가 직면한 명백한 과제는 대역폭 부족과 비대칭적 통신 토폴로지로 인해 발생하는 시스템 성능 병목 현상이었습니다.

이러한 이유로 초기 분산 최적화는 텐서 병렬 처리, 파이프라인 병렬 처리, 데이터 병렬 처리와 같은 전략을 포함하여 수동으로 구현된 수많은 통신 연산자에 의존하는 경우가 많았으며, 이러한 모든 전략은 기본 통신 로직을 신중하게 작성해야 했습니다. NCCL 및 ROCm CCL과 같은 통신 라이브러리를 호출하는 것이 일반적이지만, 이러한 솔루션은 유연성과 이식성이 부족하고 개발 및 유지 관리 비용이 높은 경우가 많습니다.

기존 시스템의 병목 현상을 분석한 결과, 우리는 3가지 핵심 사실을 요약했습니다.

사실 1: 하드웨어 대역폭이 제한되어 통신 지연이 병목 현상이 됩니다.

첫 번째는 기본 하드웨어 조건으로 인한 제약입니다. H100을 사용하여 대규모 모델을 학습하는 경우, 컴퓨팅 지연 시간이 통신 지연 시간보다 훨씬 더 길어지는 경우가 많으므로 컴퓨팅과 통신의 중복 스케줄링에 특별히 주의를 기울일 필요가 없습니다. 그러나 현재 H800 환경에서는 통신 지연 시간이 상당히 길어집니다. 일부 시나리오에서는 학습 시간의 거의 절반이 통신 지연에 소모되어 전체 MSU(Model Scale Utilization)가 크게 감소하는 것으로 평가되었습니다. 통신과 컴퓨팅의 중복이 최적화되지 않으면 시스템은 심각한 자원 낭비 문제에 직면하게 됩니다.

소규모 및 중규모 사례에서는 이러한 손실이 허용될 수 있습니다. 하지만 MegaScale이나 DeepSeek의 학습 방식처럼 모델이 수천 개의 카드로 확장되면 누적된 리소스 손실이 수백만 달러에서 수천만 달러에 달할 수 있습니다. 이는 기업에 매우 현실적인 비용 압박을 가합니다.

추론 시나리오에서도 마찬가지입니다. DeepSeek의 초기 추론 배포에는 최대 320개의 카드가 사용되었습니다. 이후 압축 및 최적화에도 불구하고 통신 지연은 여전히 분산 시스템의 피할 수 없는 핵심 문제입니다. 따라서 프로그램 수준에서 통신 및 컴퓨팅을 효과적으로 스케줄링하고 전반적인 효율성을 개선하는 방법은 우리가 정면으로 맞서야 할 핵심 과제가 되었습니다.

사실 2: 높은 통신 오버헤드는 MFU 성능에 직접적인 영향을 미칩니다.

현재 대규모 모델 학습 및 추론에서 통신 오버헤드는 항상 주요 병목 현상입니다. 기반 계층이 NVLink, PCIe 또는 여러 세대의 GPU(예: A100 및 H800)를 사용하든 통신 오버헤드가 매우 높은 것을 확인했습니다. 특히 실제 국내 배포에서는 대역폭 제한이 더욱 뚜렷하기 때문에 통신 지연이 전반적인 효율성을 직접적으로 저하시킵니다.

대규모 모델 학습의 경우, 이러한 고주파 크로스 카드 통신은 시스템의 MFU를 크게 감소시킵니다. 따라서 통신 오버헤드 최적화는 학습 및 추론 성능 향상을 위한 중요한 개선 사항이며, 저희의 핵심 집중 분야 중 하나이기도 합니다.

사실 3: 프로그래밍 가능성과 성능 간의 격차

현재 분산 시스템에서는 프로그래밍 가능성과 성능 간에 여전히 큰 격차가 존재합니다. 과거에는 단일 카드 컴파일러의 최적화 기능, 즉 단일 카드에서 우수한 성능을 달성하는 방법에 더 많은 관심을 기울였습니다. 하지만 여러 카드를 사용하는 단일 머신이나 여러 노드에 걸친 분산 시스템으로 확장하면 상황은 더욱 복잡해집니다.

분산 통신은 NCCL, MPI, 토폴로지와 같은 많은 기본 기술 세부 사항을 포함하며, 이러한 세부 사항들은 다양한 전용 라이브러리에 분산되어 있고 사용에 대한 높은 기준점을 가지고 있습니다. 많은 경우 개발자는 통신 로직을 수동으로 구현하고, 계산 및 동기화를 수동으로 스케줄링해야 하므로 높은 개발 비용과 오류율을 초래합니다. 반면, 분산 환경에서 복잡한 통신 스케줄링 및 연산자 최적화를 자동으로 처리할 수 있는 도구가 있다면 개발자는 개발 기준을 크게 낮추고 분산 시스템의 가용성과 유지 보수성을 향상시킬 수 있습니다. 이는 Triton-Distributed에서 해결하고자 하는 문제 중 하나입니다.

위에서 언급한 세 가지 실질적인 문제를 바탕으로 우리는 Triton-Distributed에서 세 가지 핵심 방향을 제안했습니다.

첫째, 의사소통과 컴퓨팅의 중첩 메커니즘을 촉진합니다.통신 오버헤드가 점점 더 두드러지는 분산 시나리오에서는 컴퓨팅과 통신의 병렬 창을 최대한 많이 예약하여 시스템의 전반적인 효율성을 개선하고자 합니다.

둘째, 대형 모델의 컴퓨팅 및 통신 모드를 심층적으로 통합하고 적응시키는 것이 필요합니다.예를 들어, 동기적 대기를 줄이고 실행 경로를 압축하기 위해 AllReduce 및 Broadcast와 같은 일반적인 통신 패턴을 컴퓨팅 패턴이 있는 모델에 통합하려고 합니다.

마지막으로, 우리는 이러한 최적화가 개발자가 고도로 맞춤화된 CUDA 구현을 직접 작성하는 것에 의존하는 대신 컴파일러에 의해 수행되어야 한다고 믿습니다.분산 시스템 개발을 보다 추상적이고 효율적으로 만드는 것이 우리가 추구하는 방향입니다.

Triton 분산 아키텍처 분석: 고성능 통신을 위한 네이티브 Python

분산 학습에서 중첩을 구현하고자 하지만, 구현이 쉽지 않습니다. 개념적으로 중첩은 통신 지연을 감추기 위해 여러 스트림을 통해 계산과 통신을 동시에 수행하는 것을 의미합니다. 연산자 간 종속성이 없는 시나리오에서는 이 작업이 더 쉽지만, Tensor Parallel(TP) 또는 Expert Parallel(EP)에서는 GEMM을 수행하기 전에 AllGather를 완료해야 합니다. 두 가지 모두 중요 경로에 있기 때문에 중첩은 매우 어렵습니다.

현재 일반적인 방법은 다음과 같습니다. 첫째, 작업을 여러 마이크로 배치로 나누고 배치의 독립성을 유지하면서 중복을 구현하는 것입니다. 둘째, 단일 배치 내에서 타일 단위와 같이 더 세밀한 단위로 분할하고 커널 융합을 통해 병렬 효과를 구현하는 것입니다. Flux에서도 이러한 분할 및 스케줄링 메커니즘을 살펴보았습니다. 동시에, 대규모 모델 학습 시 통신 모드는 매우 복잡합니다. 예를 들어, DeepSeek은 MoE를 수행할 때 대역폭과 부하 분산을 고려하여 전체 대 전체 통신을 맞춤 설정해야 합니다. 예를 들어, 저지연 추론 및 양자화 시나리오에서 NCCL과 같은 일반 라이브러리는 성능 요구 사항을 충족하기 어렵고, 종종 직접 작성한 통신 커널을 요구하기 때문에 맞춤 설정 비용이 증가합니다.

그러므로 우리는 믿는다통신-컴퓨팅 융합의 최적화 기능은 복잡한 모델 구조와 다양한 하드웨어 환경에 대처하고 반복적인 수동 구현으로 인한 개발 부담을 피하기 위해 컴파일러 계층에서 담당해야 합니다.

2계층 통신 기본 추상화

우리의 컴파일러 설계에서는 상위 계층 최적화 표현 능력과 기본 배포의 실행 가능성을 모두 고려하기 위해 2계층 통신 기본 추상 구조를 채택했습니다.

첫 번째 계층은 상대적으로 고수준의 기본 요소로, 주로 타일 단위에서 컴퓨팅 스케줄링을 완료하고 통신을 위한 추상적인 인터페이스를 제공합니다.이는 통신 추상화로 랭크 간의 푸시/겟 작업을 사용하고 태그 식별 메커니즘을 통해 각 통신 동작을 구별하여 스케줄러가 데이터 흐름과 종속성을 더 쉽게 추적할 수 있도록 해줍니다.

두 번째 계층은 기본 구현에 더 가깝고 OpenSHMEM(Open Shared Memory 표준)과 유사한 기본 시스템을 사용합니다.이 계층은 실제 통신 동작을 구현하기 위해 기존 통신 라이브러리나 하드웨어 백엔드에 매핑하는 데 주로 사용됩니다.

또한,다중 랭크 시나리오에서는 교차 랭크 동기화를 위한 장벽 및 신호 제어 메커니즘도 도입해야 합니다.예를 들어, 자신의 데이터가 기록되었다는 것을 다른 랭크에 알려야 할 때나, 특정 랭크의 데이터가 준비될 때까지 기다려야 할 때, 이러한 유형의 동기화 신호는 매우 중요합니다.

컴파일러 아키텍처 및 의미 모델링

컴파일 스택 측면에서, 저희의 전반적인 프로세스는 여전히 기존 Triton 컴파일 프레임워크를 기반으로 합니다. Triton은 소스 코드부터 시작하여 먼저 사용자 코드를 추상 구문 트리(AST)로 변환한 다음, 이를 Triton IR로 변환합니다. 저희가 구축한 Triton-Distributed에서는원래의 Triton IR이 확장되었고 분산 의미론을 위한 새로운 IR 계층이 추가되었습니다.이 분산 IR은 대기(wait) 및 알림(notify)과 같은 동기화 작업에 대한 시맨틱 모델링을 도입하여 랭크 간의 통신 종속성을 설명합니다. 동시에, OpenSHMEM이 저수준 통신 호출을 지원할 수 있도록 일련의 시맨틱 인터페이스를 설계합니다.

실제 코드 생성 단계에서 이러한 의미 체계는 기본 통신 라이브러리에 대한 외부 호출에 매핑될 수 있습니다. LLVM 중간 계층을 통해 OpenSHMEM에서 제공하는 라이브러리의 비트코드 버전(소스 코드가 아님)에 이러한 호출을 직접 연결하여 계층 간 효율적인 공유 메모리 통신을 구현합니다. 이 방법은 Triton이 소스 코드에서 외부 라이브러리에 대한 직접 액세스를 지원하지 않는다는 한계를 우회하여, 공유 메모리 관련 호출을 통해 심볼 확인을 완료하고 컴파일 과정에서 원활하게 연결할 수 있도록 합니다.

고수준 기본형과 저수준 실행 간의 매핑 메커니즘

Triton 분산 환경에서는 고수준 추상화와 저수준 제어를 포괄하는 통신 기본 시스템들을 설계했습니다.consumer_tile_wait를 예로 들면, 개발자는 기다릴 타일 ID만 선언하면 시스템이 현재 연산자 의미 체계(예: AllGather)를 기반으로 통신 대상의 특정 순위와 오프셋을 자동으로 추론하여 동기화 로직을 완성합니다. 고수준 기본 요소는 특정 데이터 소스 및 신호 전송의 세부 정보를 보호하여 개발 효율성을 향상시킵니다.

반면, 저수준 기본형은 더욱 세밀한 제어 기능을 제공합니다. 개발자는 신호 포인터, 범위(GPU 또는 시스템), 메모리 의미론(획득, 해제 등), 그리고 예상 값을 수동으로 지정해야 합니다. 이 메커니즘은 더 복잡하지만, 통신 지연 시간과 스케줄링 정확도에 대한 요구 사항이 매우 높은 시나리오에 적합합니다.

고수준 기본 요소는 크게 신호 제어와 데이터 제어라는 두 가지 범주로 나눌 수 있습니다. 신호 제어의 의미론에서,우리는 주로 생산자, 소비자, 동료의 세 가지 유형의 역할을 정의합니다.Triton 분산 통신은 읽기 및 쓰기 신호를 통해 동기화를 수행하는데, 이는 분산 통신의 핸드셰이크 메커니즘과 유사합니다. 데이터 전송을 위해 Triton 분산은 푸시(push)와 풀(pull)이라는 두 가지 기본 방식을 제공하는데, 이는 원격 카드로 데이터를 전송하거나 원격 카드에서 로컬 카드로 데이터를 가져오는 데 사용됩니다.

모든 저수준 통신 기본 요소는 OpenSHMEM 표준을 따르며, 현재 NVSHMEM과 ROCSHMEM을 지원합니다. 고수준 기본 요소와 저수준 기본 요소 사이에는 명확한 매핑 관계가 있으며, 컴파일러는 간결한 인터페이스를 저수준 동기화 및 전송 명령어로 자동 변환합니다. 이 메커니즘을 통해,트리톤 분산 방식은 통신 스케줄링의 고성능 기능을 유지할 뿐만 아니라 분산 프로그래밍의 복잡성도 크게 줄입니다.

Triton 분산 환경에서 고수준 통신 기본 요소(예: notify 및 wait)의 설계 목표는 간결한 의미론으로 카드 간 동기화 요구 사항을 설명하는 것이며, 컴파일러는 이를 해당 기본 실행 로직으로 변환하는 역할을 합니다. notify를 예로 들어보면, notify와 wait은 한 쌍의 동기화 의미론을 형성합니다. notify는 알림을 전송하는 데 사용되고 wait은 데이터 준비가 완료될 때까지 기다리는 데 사용됩니다. 개발자는 타일 ID만 지정하면 시스템은 운영자 유형 및 통신 토폴로지에 따라 통신 대상 및 신호 오프셋과 같은 기본 세부 정보를 자동으로 추론합니다.

구체적인 기본 구현은 배포 환경에 따라 달라집니다. 예를 들어, GPU가 8개인 시나리오에서는 이러한 유형의 동기화를 스레드 내에서 _syncthreads() 및 atomic_dd를 통해 구현할 수 있습니다. 크로스 머신 배포에서는 NVSHMEM 또는 ROCSHMEM에서 제공하는 signal_up과 같은 기본 요소를 사용하여 동일한 작업을 수행합니다. 이러한 메커니즘은 고수준 의미론과 저수준 기본 요소 간의 매핑 관계를 구성하며, 뛰어난 다용성과 확장성을 제공합니다.

GEMM ReduceScatter 통신 시나리오를 예로 들어 보겠습니다. 시스템에 GPU가 4개 있고 각 타일의 목표 위치가 미리 계산된 메타 정보(예: 타일 할당 및 각 랭크의 배리어 번호)에 의해 결정된다고 가정해 보겠습니다. 개발자는 Triton으로 작성된 GEMM 커널에 notify 문만 추가하면 되고, ReduceScatter 커널은 wait을 사용하여 동기적으로 데이터를 수신합니다.

전체 프로세스는 Python으로 표현할 수 있으며, 명확한 통신 로직과 손쉬운 스케줄링을 통해 듀얼 스트림 시작의 커널 모드도 지원합니다. 이 메커니즘은 크로스 카드 통신 프로그래밍의 표현력을 향상시킬 뿐만 아니라, 기본 구현의 복잡성을 크게 줄여 대규모 분산 모델의 효율적인 학습 및 추론을 위한 강력한 기본 기능을 제공합니다.

다차원 중첩 최적화: 스케줄링 메커니즘에서 토폴로지 인식까지

Triton 분산 시스템은 비교적 간결한 고수준 통신 원시 인터페이스를 제공하지만, 커널을 작성하고 최적화하는 실제 과정에는 여전히 몇 가지 기술적 장벽이 존재합니다. 원시 설계가 뛰어난 표현력을 가지고 있음에도 불구하고, 진정으로 유연하게 적용하고 심층적으로 최적화할 수 있는 사용자 수는 여전히 제한적입니다. 본질적으로 통신 최적화는 여전히 엔지니어링 경험과 스케줄링 이해에 크게 의존하는 작업이며, 현재로서는 개발자가 직접 제어해야 합니다. 이 문제를 해결하기 위해 몇 가지 주요 최적화 경로를 요약했습니다. 다음은 Triton 분산 시스템의 일반적인 구현 전략입니다.

푸시 vs. 풀: 데이터 흐름 방향 및 장벽 번호 제어

통신과 계산의 오버랩 최적화에서,Triton-distributed는 푸시와 풀의 두 가지 데이터 전송 방법을 제공합니다.비록 두 가지의 의미적 차이는 단지 "능동적 전송"과 "수동적 풀링"의 방향에 있을 뿐이지만, 실제 분산 실행에서는 성능과 스케줄링 제어 기능 면에서 확실한 차이가 있습니다.

배리어의 개수를 예로 들면, 풀 모드는 일반적으로 두 개의 배리어를 필요로 합니다. 하나는 상대방이 풀링하기 전에 로컬 데이터가 준비되었는지 확인하는 배리어이고, 다른 하나는 전체 통신 주기 동안 로컬 작업에 의해 데이터가 수정되지 않도록 보호하여 데이터 불일치나 읽기-쓰기 충돌을 방지하는 배리어입니다. 푸시 모드에서는 원격지에 데이터가 기록된 후 하나의 배리어만 설정하면 모든 장치를 동기화하여 전반적인 제어를 간소화합니다.

하지만 풀 모드에는 장점도 있습니다. 로컬 노드가 데이터 풀링 순서를 능동적으로 제어할 수 있으므로 통신 타이밍과 컴퓨팅 중복을 더욱 정확하게 스케줄링할 수 있습니다. 통신과 컴퓨팅 간의 중복 효과를 극대화하고 병렬성을 확보하려는 경우, 풀 모드는 더 큰 유연성을 제공합니다.

일반적으로 주요 목표가 중복을 개선하는 것이라면 풀 모드가 권장됩니다. 별도의 AllGather나 ReduceScatter 커널과 같은 일부 순수 통신 작업에서는 단순성과 낮은 오버헤드로 인해 푸시 모드가 더 일반적입니다.

스위즐링 스케줄링: 데이터 지역성에 따라 동적으로 순서 조정

통신과 연산의 중첩은 기본 요소 선택뿐만 아니라 스케줄링 전략에도 영향을 받습니다. 그중 스위즐링(Swizzling)은 토폴로지 인식 기반 스케줄링 최적화 방법으로, 크로스 카드 컴퓨팅 시 실행 유휴 시간을 줄이는 것을 목표로 합니다. 분산 환경에서는 각 GPU 카드를 독립적인 실행 단위로 볼 수 있습니다. 각 카드가 처음에 서로 다른 데이터 조각을 보유하기 때문에 모든 카드가 동일한 타일 인덱스에서 연산을 시작하면 일부 랭크는 데이터가 준비될 때까지 기다려야 하며, 이로 인해 실행 단계에서 긴 유휴 시간이 발생하여 전반적인 컴퓨팅 효율이 저하됩니다.

스위즐링의 핵심 아이디어는 다음과 같습니다.시작 계산 오프셋은 각 카드에 있는 기존 로컬 데이터의 위치에 따라 동적으로 조정됩니다.예를 들어, AllGather 시나리오에서 각 카드는 자신의 데이터 처리를 우선 순위로 지정하고 원격 타일에서 동시에 데이터를 가져오도록 하여 통신 및 연산의 동시 스케줄링을 구현할 수 있습니다. 모든 카드가 타일 0에서 처리를 시작하면 랭크 0만 즉시 연산을 시작할 수 있으며, 나머지 랭크는 데이터 대기로 인해 직렬 지연이 발생합니다.

교차 머신 ReduceScatter 시나리오와 같이 더 복잡한 상황에서는 Swizzling 전략이 네트워크 토폴로지와 함께 설계되어야 합니다. 두 노드를 예로 들어, 합리적인 스케줄링 방법은 다음과 같습니다. 상대 노드에서 필요한 데이터 계산을 우선시하고, 가능한 한 빨리 교차 머신 지점 간 통신을 트리거합니다. 그리고 전송 과정에서 로컬 노드에서 필요한 데이터를 병렬로 계산하여 통신과 컴퓨팅의 중첩 효과를 극대화합니다.

현재 이러한 유형의 스케줄링 최적화는 컴파일러가 일반 최적화에서 핵심 성능 경로를 희생하는 것을 방지하기 위해 여전히 프로그래머가 제어하고 있습니다. 또한 스위즐링과 같은 세부 사항을 이해하는 것이 개발자에게 일정한 기준이 된다는 것을 알고 있습니다. 앞으로 개발자들이 분산 연산자 개발 모델을 더욱 빠르게 숙달하고 개방적이고 효율적인 Triton 분산 프로그래밍 생태계를 점진적으로 구축할 수 있도록 더 많은 실제 사례와 템플릿 코드를 제공하고자 합니다.

불완전한 블록 스케줄링: 순위 타일에 따른 우선순위 처리

실제 대규모 모델 학습 및 추론 시나리오에서 연산자의 입력 모양은 종종 불규칙하며, 특히 토큰 길이가 고정되지 않은 경우 더욱 그렇습니다. 또한 타일 블록을 깔끔하고 균일하게 유지하기 어렵습니다.이러한 불완전한 타일링으로 인해 일부 타일이 여러 랭크에 걸쳐 있게 됩니다. 즉, 같은 타일의 데이터가 여러 장치에 분산되어 스케줄링과 동기화의 복잡성이 증가합니다.

AllGather GEMM을 예로 들어, 타일에 로컬 데이터와 원격 데이터가 모두 포함되어 있다고 가정해 보겠습니다. 이 타일에서 계산을 시작하면 원격 데이터가 먼저 전송될 때까지 기다려야 하며, 이로 인해 추가적인 버블이 발생하고 전체 계산의 병렬성에 영향을 미칩니다. 더 나은 방법은 이 교차 랭크 타일을 건너뛰고, 로컬에서 완전히 이용 가능한 데이터 처리를 우선시하며, 원격 입력이 실행될 때까지 기다리는 타일을 마지막으로 스케줄링하여 통신과 계산 간의 중복을 최대화하는 것입니다.

ReduceScatter 시나리오에서는 스케줄링 순서를 반대로 해야 합니다. 교차 순위 타일의 계산 결과는 가능한 한 빨리 원격 노드로 전송되어야 하므로, 원격 노드에서 사용하는 타일의 우선순위를 정하는 것이 가장 좋은 전략입니다. 이렇게 하면 교차 머신 데이터 전송을 가능한 한 빨리 완료하고 원격 종속성을 줄일 수 있습니다.

MoE 하의 동적 분류 전략

MoE(전문가 혼합) 모델에서는 라우팅 결과에 따라 여러 전문가에게 토큰을 분배해야 하며, 일반적으로 전체 대 전체 통신 및 그룹 GEMM 계산이 수반됩니다. 통신과 컴퓨팅의 중복 효율성을 개선하기 위해 Triton 분산은 동적 정렬(Dynamic Sorting)을 도입했습니다. 동적 정렬은 통신 데이터 의존도에 따라 컴퓨팅 작업을 단계적으로 스케줄링하고, 데이터 의존도가 낮은 작업을 우선적으로 처리합니다.

이러한 순서는 각 단계의 계산이 가능한 가장 낮은 통신 차단으로 시작할 수 있도록 보장하여 All-to-All과 Group GEMM 간의 더 나은 중복을 달성합니다.전반적인 스케줄링은 데이터 종속성이 가장 적은 타일에서 시작하여 점차 종속성이 복잡한 타일로 확장되어 실행 동시성을 극대화합니다.

하드웨어 기반 통신 가속

Triton 배포는 특정 하드웨어 기능과 결합하여 통신 최적화도 지원합니다.특히 NVSwitch 아키텍처를 사용할 경우, 내장된 SHARP Accelerator를 사용하여 저지연 통신 계산을 수행할 수 있습니다. 이 모듈은 스위치 칩에서 Broadcast 및 AllReduce와 같은 연산을 수행하여 전송 경로에서 데이터 집계 가속을 달성하고 지연 시간과 대역폭 소모를 줄입니다. 관련 명령어는 Triton 배포판에 통합되었으며, 해당 하드웨어를 사용하는 사용자는 이 명령어를 직접 호출하여 더욱 효율적인 통신 커널을 구축할 수 있습니다.

AOT 컴파일 최적화: 추론 지연 오버헤드 감소

Triton-distributed는 추론 시나리오에서 지연 시간에 매우 민감한 요구 사항에 맞춰 특별히 최적화된 AOT(Ahead-of-Time) 메커니즘을 도입했습니다. Triton은 기본적으로 JIT(Just-In-Time) 컴파일 방식을 사용하며, 함수가 처음 실행될 때 상당한 컴파일 및 캐시 오버헤드가 발생합니다.

AOT 메커니즘을 통해 사용자는 실행 전에 함수를 바이트코드로 미리 컴파일하고, 추론 단계에서 직접 로드하여 실행할 수 있습니다. 이를 통해 JIT 컴파일 과정을 생략하고 컴파일 및 캐싱으로 인한 지연 시간을 효과적으로 줄일 수 있습니다. 이를 기반으로 Triton-distributed는 AOT 메커니즘을 확장하여 이제 분산 환경에서 AOT 컴파일 및 배포를 지원하여 분산 추론의 성능을 더욱 향상시킵니다.

성능 측정 및 사례 재현

우리는 NVIDIA H800, AMD GPU, 8카드 GPU 및 크로스 머신 클러스터를 포함하여 다중 플랫폼 및 다중 작업 시나리오에서 Triton 분산형의 성능에 대한 포괄적인 테스트를 수행했으며, PyTorch 및 Flux와 같은 주류 분산 구현 솔루션을 비교했습니다.

8개의 GPU 카드에서Triton 분산형은 AG GEMM 및 GEMM RS 작업에서 PyTorch 구현에 비해 상당한 속도 향상을 달성합니다.수동으로 최적화된 Flux 솔루션과 비교했을 때, 스위즐링 스케줄링, 통신 오프로드, AOT 컴파일 등 여러 최적화를 통해 더 나은 성능을 달성합니다. 동시에, AMD 플랫폼에서 PyTorch + RCCL을 조합한 것과 비교했을 때, 전반적인 가속은 다소 낮지만 상당한 최적화를 달성합니다. 주요 한계는 테스트 하드웨어의 낮은 컴퓨팅 성능과 스위치가 없는 토폴로지입니다.

AllReduce 작업에서우리가 테스트한 하드웨어 구성에서 작은 메시지부터 큰 메시지까지 다양한 메시지 크기에 대해 Triton 분산 방식은 NCCL보다 상당히 빠른 속도를 보였으며, 평균적으로 약 1.6배의 속도 향상을 보였습니다.Attention 시나리오에서는 주로 gather-KV 유형의 Attention 연산을 테스트했습니다. PyTorch Touch의 네이티브 구현과 비교했을 때, 8개의 GPU 카드에 분산된 Triton의 성능은 약 5배 향상될 수 있습니다. 또한 오픈 소스 Ring Attention 구현보다 약 2배 향상된 성능을 보입니다.

교차 머신 테스트:AG GEMM은 1.3배 빠르고 GEMM RS는 1.4배 더 빠르며, 이는 Flux보다 약간 낮지만 형태의 유연성과 확장성 측면에서 더 많은 장점을 가지고 있습니다.고속 추론 시나리오에서 단일 토큰 디코딩도 테스트했습니다. 지연 시간은 1M 토큰 환경에서 20~30마이크로초 이내로 제어되었으며, NVLink 및 PCIe와 호환됩니다.

또한, DeepEP의 분산 스케줄링 로직을 재현하여 All-to-All 라우팅 및 컨텍스트 분산 전략을 주로 조정했습니다. 64개 미만의 카드가 있는 시나리오에서 Triton 분산 시스템의 성능은 DeepEP와 기본적으로 동일하며, 일부 구성에서는 약간 향상되었습니다.

마지막으로, 8개의 GPU 카드에서 배포 및 운영을 지원하는 Qwen-32B 기반의 프리필 및 디코딩 데모도 제공합니다. 실제 테스트 결과, 추론 가속 효과가 약 1.2배 향상되는 것으로 나타났습니다.

개방형 분산 컴파일 생태계 구축

현재 우리는 맞춤형 중복 시나리오라는 과제에 직면해 있는데, 과거에는 주로 수동 최적화에 의존해 해결해 왔고 이는 노동 집약적이고 비용이 많이 드는 작업이었습니다.우리는 Triton 분산 프레임워크를 제안하고 오픈 소스화했습니다.Triton 기반으로 구현되었지만, 각 회사에서 사용하는 컴파일러나 기반 통신 라이브러리가 무엇이든 통합하여 개방형 분산 생태계를 구축할 수 있습니다.

이 분야는 중국은 물론 전 세계적으로도 아직 상대적으로 미비한 상태입니다. 저희는 커뮤니티의 힘을 활용하여 구문 설계, 성능 최적화, 다양한 하드웨어 장치 지원 등 다양한 분야에서 더 많은 개발자의 참여를 유도하고, 기술 발전을 함께 촉진하고자 합니다. 마침내 좋은 성과를 거두었으며, 관련 사례는 모두 오픈 소스로 공개되어 있습니다. 적극적으로 소통하고, 더 나은 미래를 만들어갈 더 많은 파트너 여러분의 참여를 기대합니다!