Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts
https://arxiv.org/pdf/2502.19811
1. 제안 기법
대규모 Mixture-of-Experts (MoE) 모델은 각 입력 토큰에 대해 소수의 전문가 네트워크만을 활성화함으로써, 방대한 수의 파라미터를 활용하면서도 연산 비용을 효율적으로 관리할 수 있다는 장점을 지닌다. 그러나 이러한 모델 구조를 분산 환경에 적용할 경우, GPU 간의 통신 빈도가 급증하여 전체 실행 시간 중 상당 부분을 통신에 할애하게 되는 문제가 발생한다. 실제로 일부 연구 결과에 따르면, 주요 MoE 모델의 forward 연산 과정에서 GPU 간 통신이 전체 모델 실행 시간의 약 47%를 차지하는 것으로 보고되었다. 이러한 성능 저하 문제를 해결하기 위해, MoE 레이어 내 통신과 연산을 파이프라인 방식으로 병행 처리하는 기존의 오버랩 기법들이 시도되었다. 하지만 기존의 coarse-grained 오버랩 방식들은 연산과 통신을 지나치게 큰 덩어리로 묶어 동시 수행함으로써, 연산 효율성을 저하시키거나 통신 지연을 완전히 은닉하지 못하는 등의 한계점을 드러냈다. 예를 들어 일부 프레임워크에서는 전문가 연산과 all-to-all 통신을 단순히 2단계 파이프라인으로 구성하거나, 제한적인 탐색 공간 내에서 수동으로 오버랩 단계를 설정하는 등의 방법론을 채택하였으나, 이는 최적의 성능을 달성하는 데 미흡한 접근 방식이었다.
Comet은 이러한 한계를 극복하기 위해 미세한(granularity) 수준의 통신-연산 오버랩을 도입하여 MoE 레이어의 효율을 극대화하였다.
- Fine-grained 오버랩 방식 제안: 통신과 연산 사이의 데이터 의존성을 정밀 분석하고 작업 재구성을 통해, 이전보다 훨씬 잘게 쪼갠 단위로 통신과 계산을 겹친다. 이를 통해 통신이 완전히 끝나기를 기다리지 않고도 연산을 부분적으로 수행할 수 있으며, 반대로 연산 중간중간 통신을 진행하여 통신 지연을 효과적으로 숨긴다. 이러한 미세 오버랩 전략은 shared tensor라 불리는 통신-연산 간 공유 버퍼의 사용 패턴을 분석하고, 데이터를 특정 차원으로 분할한 후 연산 순서를 재조직함으로써 가능했다.
- Adaptive workload assignment: 본 논문에서는 통신 및 연산 작업에 스레드 블록 단위의 전용 자원을 할당하고, 각 작업에 할당된 블록 수를 동적으로 조절함으로써 통신과 연산 소요 시간의 정밀한 균형을 도모한다. 이를 통해 통신 시간이 상대적으로 길 경우, 더 많은 자원을 통신에 투입하여 지연을 감소시키고, 반대로 연산이 병목 현상을 보이는 경우에는 연산 자원을 증대시키는 방식으로 오버랩 구간의 유휴 시간(bubble)을 최소화한다. Comet은 이와 같은 adaptive workload assignment을 통해 다양한 입력 크기 및 하드웨어 환경에서도 통신과 연산의 균형 있는 중첩을 가능하게 하였다.
- 성능 및 실용적 성과: 실제 Nvidia H800 GPU 기반 실험 결과, 단일 MoE 레이어에서 최대 1.96배의 가속 효과를 보였으며, 엔드투엔드 전체 모델 실행에서는 평균 1.71배의 속도 향상을 달성하였다. 이러한 성능을 바탕으로 Comet은 현재 수만 개 GPU 규모의 프로덕션 클러스터에 적용되어 수백만 GPU시간을 절감하는 실질적인 효과를 나타내고 있으며, Megatron-LM 등 기존 프레임워크와의 통합 가능성을 입증하였다. 이는 Comet의 프로그래밍 모델이 실제 사용 환경에서도 적용 가능함을 보여주는 것이며, 향후 Triton 또는 TVM과 같은 컴파일러를 활용하여 해당 모델을 구현하는 등의 확장 가능성 또한 제시하고 있다.
이 논문의 저자들은 MoE 레이어 내 연산들을 producer-consumer 파이프라인으로 모델링 하였고, MoE 아키텍처에는 두 가지 상이한 파이프라인이 존재한다는 것을 알게 되었다.
- 통신-연산 파이프라인은 forward 과정의 초기 단계에 해당하며, 통신 작업(예: 토큰 분배)이 데이터 생산자 역할을 수행하고, 해당 출력 데이터를 기반으로 연산 작업(GEMM 등)이 소비자 역할을 수행하는 흐름을 나타낸다. 구체적으로, 한 GPU에서 다른 GPU로 토큰을 전송하는 all-to-all 통신이 이루어진 후, 수신된 토큰들에 대해 해당 GPU의 전문가 네트워크 연산(GEMM)이 진행되는 단계가 이에 해당한다. 통신의 출력 버퍼는 곧바로 연산의 입력 버퍼로 활용되며, Comet에서는 이 공유 버퍼를 shared tensor라고 지칭한다.
- 연산-통신 파이프라인은 MoE 레이어의 두 번째 부분(예: forward에서 전문가 연산 결과를 다시 모으는 단계)에 해당하며, 각 전문가의 출력 값을 합치거나 선택(top-K)하여 최종 토큰 출력을 계산하고, 이를 원래 위치로 보내는 combine 통신이 이에 속한다. 이때도 앞 단계 연산의 출력이 다음 통신의 입력으로 shared tensor를 통해 연결된다.
Comet의 미세 조정된 오버랩은 두 파이프라인 각각에서 공유 텐서의 데이터 의존성을 정밀하게 해결함으로써 가능했다. 구체적으로, Comet은 (1) 공유 텐서를 적절한 차원으로 분해하고, (2) 그 조각들에 대한 연산 순서를 재조정하는 의존성 해결 절차를 도입했다.
그림 3에서 나타나듯이, 이러한 최적화를 통해 통신 및 연산 간의 자료 불일치 문제를 해소함으로써 보다 균일하고 밀도 높은 중첩 패턴을 획득할 수 있다. 핵심적인 원리는 생산자와 소비자가 공유 텐서 내 독립적인 데이터에 한하여 동시 작동이 가능하다는 것으로, Comet은 각 연산이 공유 텐서에 접근하는 양상을 분석하여, 소비자 측면에서 데이터 간 상호 의존성이 존재하지 않는 차원을 식별하고 해당 축을 따라 텐서를 세밀하게 분할한다.
이 전략의 각 파이프라인 적용 시 동작 원리를 구체적인 예시를 들어 설명하면 아래와 같다.
- Forward 1단계 (통신-연산 파이프라인, MoE layer0): 본 단계에서는 입력 토큰 라우팅을 위한 All-to-All 통신(토큰 분산)이 다수의 GPU에 걸쳐 이루어진다. 이어서 각 GPU에서 수신된 토큰들에 대해 Expert FFN 연산(GEMM)이 수행된다. Comet은 이 과정에서 shared tensor를 토큰 차원으로 분해하여, 일부 토큰 데이터가 도착하는 즉시 해당 부분의 연산을 개시할 수 있도록 한다. 특히 토큰들을 원래 소스 GPU별로 분류 및 정렬함으로써, 로컬 GPU에서 기 보유하고 있던 토큰들의 연산을 우선적으로 수행하고, 동시에 원격에서 전송되는 토큰들은 백그라운드로 수신한다. 이와 같이 원격 토큰의 전체 도착을 대기하지 않고 로컬 토큰을 먼저 처리함으로써 통신과 연산을 최대한 병행할 수 있다. Comet은 이러한 전략을 구현하기 위해 GroupGEMM 기법을 도입, 단일 커널 내에서 여러 expert의 GEMM 연산을 통합 처리하며, 로컬 데이터 우선으로 계산 타일 순서를 재배치한다. 결과적으로 각 expert 연산에 필요한 토큰 행렬이 완전히 준비될 때까지 기다리지 않고, 수신된 부분부터 즉시 처리하는 파이프라인 실행이 가능해진다.
- 2단계 순방향 연산(연산-통신 파이프라인, MoE 레이어1)에서는 각 GPU의 expert들이 두 번째 FFN 연산을 완료한 후, 다른 GPU에 위치한 동일 토큰의 나머지 부분 결과들과 결합/축소 통신을 수행하여 최종 출력을 생성한다. 이때, 소비자인 결합 통신은 일반적으로 top-K reduction(각 토큰별 상위 K개의 expert 출력만 취합)으로 구현되는데, 토큰 차원에서 높은 상호 의존성이 존재하여 모든 expert 연산이 완료되어야만 시작할 수 있다는 제약이 따른다. Comet은 이러한 레이어1 단계에서도 공유 텐서를 적절히 분해하고 expert 연산 순서를 재조정함으로써 문제 해결을 시도하였다. 구체적으로, expert들의 GroupGEMM 연산을 기존처럼 각 expert별 순차 실행하는 대신, 출력 행렬을 열 방향으로 점진적으로 번갈아가며 계산하는 방식(column-wise)으로 변경하였다. 이로써 초기 몇 개 열의 출력 결과가 도출되는 즉시 해당 부분에 대한 축소 및 통신을 개시할 수 있게 되었다. 즉, 일부 expert들의 부분적인 출력이 준비되는 대로 결합 연산을 부분적으로 진행하여 통신을 선행적으로 수행하는 동시에 나머지 expert 연산을 지속적으로 진행하는 방식으로 오버랩이 발생한다. 기존 방식에서는 모든 expert 출력이 계산 완료된 이후에야 토큰별 출력 결합이 가능하였으나, Comet의 재스케줄링을 통해 연산 도중에도 통신을 병행함으로써 총 지연 시간을 단축시킬 수 있었다.
2. CUDA 커널 구조 및 통신 알고리즘 세부 분석
Comet은 통신과 연산을 논리적으로 중첩시키는 것을 넘어, 실질적인 GPU 환경에서 단일 CUDA 커널 실행 내에 통신과 연산을 통합하는 융합 기법을 적용하였다. 이는 개별적인 연산/통신 커널을 순차적으로 호출하며 스트림 방식으로 중첩하는 기존 방식에 비해 훨씬 더 정밀한 동기화와 오버랩을 가능하게 한다. Comet 연구진은 다양한 커널 융합 방법 중 “스레드 블록 수준 분리” (수평적 융합) 방식을 채택하였으며, 이를 이해하기 위해서는 먼저 대비되는 “수직적 융합” 방식을 고찰할 필요가 있다.
- 수직적 융합(vertical fusion)은 단일 스레드 블록 내에서 통신 I/O와 연산(GEMM)을 통합하는 방식입니다. 이는 각 스레드 블록이 먼저 통신을 통해 필요한 데이터를 수신한 후 즉시 GEMM 연산을 수행하고, 완료되면 결과를 다시 통신으로 전송하는 프롤로그/에필로그 구조와 유사하다. 이러한 접근 방식을 통해 여러 블록이 병렬로 실행되면서 통신과 연산의 중첩이 가능하지만, 오버랩 타이밍이 불규칙하고 비결정적이므로 통신 대 연산 시간의 정밀한 조율이 어렵다. 특히 토큰 단위의 세밀한 I/O가 블록 내부에서 빈번하게 발생할 경우, GPU의 연산 자원을 충분히 활용하지 못하여 GEMM 성능 저하를 야기할 수 있으며, 이러한 문제는 최신 Hopper 아키텍처에서 더욱 두드러집니다. 수직적 융합은 개념적으로는 단순하지만, 연산 효율성 저하와 타이밍 조절의 제약이라는 한계점을 지닌다.
- 수평적 fusion (thread block 특수화): 통신과 연산 작업을 아예 다른 스레드 블록들에서 수행하도록 분리하는 것이다 . 즉 일부 스레드 블록은 오직 연산 전용 블록으로, 나머지 블록은 통신 전용 블록으로 역할을 고정한다. 이렇게 하면 한 종류의 작업은 한 블록들에서만 수행되므로 각자의 자원 할당을 격리할 수 있고, 통신과 연산 각각의 소요 시간을 별도로 측정 및 조율하기가 용이해진다 . Comet은 이 방법을 통해 하드웨어 자원 할당을 정밀 제어하고 통신/연산의 실행 시간을 맞춤으로써 latency hiding을 극대화하였다. 더욱이 통신과 연산이 분리되어 있으므로, 연산 블록에서 실행되는 GEMM 커널은 기존에 최적화된 구현(예: CUTLASS 라이브러리 기반 템플릿)을 거의 그대로 사용할 수 있어 계산 성능을 떨어뜨리지 않고 유지할 수 있다. 실제 Comet의 fused 커널에서 연산 부분은 기본 Megatron-LM의 CUTLASS GEMM과 동일한 방식으로 동작하므로, 오버랩으로 인한 연산 속도 감소가 없다 . 이는 thread block 특수화 기법의 큰 장점으로, 오버랩 달성 + 연산 효율 보존이라는 두 마리 토끼를 잡은 셈이다.
Comet의 스레드 블록 특화 커널 구조는 Hopper 아키텍처를 예시로 들어 살펴보았을 때, 그림 7에서 볼 수 있듯이 각 스트리밍 멀티프로세서(SM)에 하나의 블록만이 배치되는 형태로 구성된다. 연산 전용 블록 내부에서는 일부 워프(warp)들이 글로벌 메모리에서 데이터를 로드하는 생산자(Producer) 역할을 수행하고, 다른 워프들은 로드된 데이터를 즉시 텐서 코어 연산에 투입하는 소비자(Consumer) 역할을 담당한다. Hopper GPU에서는 비동기 텐서 메모리 가속기(async TMA) 명령어를 활용하여 한 워프가 입력 행렬을 글로벌 메모리에서 공유 메모리로 비동기적으로 적재한 후, 다른 워프가 해당 데이터를 받아 텐서 코어 MMA 연산을 수행하는 방식으로 GEMM을 진행한다. 이로써 메모리 로드와 연산을 워프 수준에서 파이프라인으로 중첩시켜 연산 유닛 활용도를 높일 수 있다. 한편, 통신 전용 블록들은 연산 블록들이 생성한 결과물을 받아 후속 처리를 진행한다. 구체적으로, 연산 블록의 소비자 워프가 계산 완료된 출력 결과를 글로벌 메모리에 기록해두면, 통신 블록의 워프들이 이를 읽어 Top-K 감소 등의 필요한 처리를 수행한다. 이후 각 토큰 결과를 로컬 메모리에 쓰거나 원격 GPU 메모리로 직접 전송함으로써 통신을 완료한다. 통신 블록 내부에서도 워프들이 역할을 분담하여 일부는 데이터 패킹/전송을, 일부는 로컬 기록을 담당하는 등 최적의 파이프라인 처리가 이루어진다. 이러한 구조 덕분에 Comet의 GEMM 연산 부분은 기존과 동일한 효율로 동작하면서도, 출력 결과가 생성되는 즉시 통신 블록들이 실시간으로 이를 다른 GPU에 전송하거나 로컬에 저장할 수 있어 통신 지연이 상당 부분 감소된다.
한 가지 고려할 점은, 왜 연산 warp와 통신 warp를 같은 스레드 블록에 두지 않았는가이다. 이론상으로는 하나의 블록 안에서 연산 warp들이 계산한 결과를 즉시 같은 블록 내 통신 warp들이 받아 원격으로 전송하면, 글로벌 메모리에 중간 결과를 쓸 필요 없이 공유 메모리에서 바로 보낼 수도 있을 것이다. 그러나 현재 GPU 하드웨어의 제약상, 한 스레드 블록에서 사용할 수 있는 warp 수가 한정되어 있어 통신까지 함께 넣으면 충분한 병렬 통신을 구현하기 어렵다. 예를 들어 Ampere/Ada 세대 GPU는 블록당 최대 32warps 정도이므로, 이 안에서 연산과 통신 warp를 모두 운용하면 통신이 이용할 수 있는 warp 수가 제한되어 통신 대역폭을 최대로 활용하지 못할 수 있다. 반대로 통신 warp가 연산 warp의 자원(레지스터, 공유메모리 등)을 잠식하여 연산 수행에도 간섭을 주게 된다. 이러한 이유로 Comet은 통신 warp와 연산 warp의 블록을 아예 분리함으로써 두 작업 간 간섭을 제거하고, 각자 필요한 만큼의 warp를 투입하여 GPU 자원을 최대로 활용하는 방향을 택한 것이다 . 결과적으로 Comet의 fused kernel에서는 통신 전용 블록들이 GPU의 남는 SM들을 활용하여 통신 작업을 수행하고, 연산 블록들은 전용 SM에서 full occupancy로 GEMM을 돌리게 되어 통신 대역폭과 연산 처리율이 모두 극대화된다.
Comet 커널의 통신 부분은 NVSHMEM 라이브러리를 활용하여 구현되었는데, NVSHMEM은 NVIDIA에서 제공하는 GPU 공유 메모리 라이브러리로서, 다수의 GPU 메모리를 포괄하는 글로벌 어드레스 공간을 형성하여 단일 GPU에서 다른 GPU의 메모리에 로컬 메모리처럼 접근할 수 있도록 지원한다. 이를 통해 GPU 커널 내부에서 CPU 개입이나 MPI 호출 없이 직접 원격 GPU 메모리에 데이터를 쓰거나 읽는 연산(put/get)을 수행할 수 있다. Comet은 NVSHMEM의 정밀한 GPU-initiated RMA(원격 메모리 조작) 기능을 활용하여, 앞서 설명된 통신 전용 블록의 워프들이 원격 GPU의 공유 텐서 영역에 직접 결과를 기록하거나 필요한 원격 데이터를 가져오도록 설계되었다. 이와 같이 NVSHMEM을 통한 통신은 기본적으로 NVIDIA의 GPUDirect 기술(동일 노드에서는 NVLink/NVSwitch, 서로 다른 노드 간에는 RDMA 네트워크)을 활용하여 하드웨어적으로 최적화된다. 따라서 Comet 커널 내의 각 통신 워프는 자신이 담당하는 토큰 데이터 조각을 원격 GPU의 지정 주소로 직접 기록하며, 해당 GPU에서는 기록된 데이터를 즉시 연산에 사용할 수 있다. NVSHMEM 통신은 GPU 메모리 간 직접 복사를 수행하므로 기존의 CPU 관여 MPI all-to-all 방식에 비해 지연 시간과 프로토콜 오버헤드가 감소하고, 세분화된 단위로 데이터를 전송할 수 있어 오버랩에 유리하다.
다만 NVSHMEM을 이용하려면 각 프로세스(GPU)에 일정 크기의 공유 버퍼 메모리를 미리 할당해야 한다. Comet에서는 모델 설정에 따라 이 버퍼 크기를 (입력 토큰 길이 M) × (토큰 임베딩 크기 H) 정도로 지정하며, 예를 들어 FP16/BF16 데이터형이면 이 크기의 2배 바이트를 차지한다 . 실험에 사용된 모델 기준으로 시퀀스 길이 4096일 때 GPU당 약 1632MB, 길이 8192일 때 3264MB 수준의 추가 메모리가 NVSHMEM 버퍼로 필요하다.
앞서 설명한 thread block 특수화 기법에서, 통신 전용 블록과 연산 전용 블록을 몇 대 몇 비율로 배정할지가 성능에 중요하다. Comet은 총 블록 수 B 중 b개를 연산(GEMM) 블록, 나머지 B - b개를 통신 블록으로 두는데, 이 분할점(b)을 동적으로 조율(adaptive assignment)하는 알고리즘을 제안했다.
다양한 입력 길이와 모델 구성(전문가 수 등)에서 실험한 결과, 환경별로 최적의 통신 대 연산 블록 비율이 존재하며, 이를 맞출 때 전체 지연이 최소화되었다고 한다. 예를 들어 입력 토큰 길이가 길어지면 통신과 연산 모두 처리 데이터가 비례해서 늘지만, GEMM 연산 시간은 메모리 대역폭보다 덜 증가하거나 통신 시간은 네트워크 병목으로 비선형 증가하는 등 서로 다른 스케일링을 보인다. Comet은 이러한 특성을 반영하여, 사전에 측정된 성능 모델 또는 런타임 힌트를 사용해 각 상황에 최적인 블록 할당을 결정한다. 그림 8의 결과에서 보이듯 입력 길이가 달라지면 최적 b 값도 달라지는데, Comet은 이러한 변화를 반영해 고정된 겹침 구조가 아닌 적응형 구조를 취함으로써 항상 최적에 가까운 통신-연산 겹침을 달성한다.
3. DeepEP와의 비교
DeepEP는 MoE 모델의 전문가 병렬(Expert Parallel, EP) 통신의 효율성을 제고하기 위해 개발된 최신 라이브러리로서, All-to-All 통신(토큰 분배 및 결과 집합)에 최적화된 고성능 커널들을 제공한다. DeepEP는 NVLink (노드 내) 및 RDMA (노드 간) 환경 모두에서 최대 대역폭에 근접하는 성능을 구현하는 것을 목표로 설계되었으며, FP8 저정밀도 전송 등을 통해 통신량 자체를 감소시키는 기능 또한 포함하고 있다. DeepEP의 주요 특징 중 하나는 통신-연산 오버랩을 위한 hook 기반 기법의 활용이며, 이는 GPU의 SM 자원을 점유하지 않고도 통신과 연산을 동시에 수행할 수 있는 방법이다.
공통점 및 유사한 점
- 목표와 문제 인식: 두 시스템 모두 대규모 MoE 환경에서 발생하는 통신 오버헤드가 성능 병목 현상의 주요 원인이라는 문제 인식에 기반하여 개발되었으며, 통신-연산 중첩을 통한 지연 시간 감소를 핵심 해결 방안으로 제시하고 있다.
- GPU 상의 통신 최적화 기술 활용: Comet과 DeepEP는 NVSHMEM/GPUDirect 등의 기술을 활용하여 GPU가 직접 통신을 수행하도록 구현하였다. Comet은 NVSHMEM 2.x를 커널 내에 통합하여 정밀한 RMA 통신을 수행하였으며, DeepEP 또한 자체 커널에서 NVSHMEM을 수정하여 사용하고 NVLink와 RDMA 모두를 지원한다. 이를 통해 CPU 개입 및 불필요한 메모리 복사를 최소화하고, NVLink 및 고속 네트워크의 이론적 대역폭에 근접하는 효율을 달성하였다.
- 연산 효율 유지: 두 기법 모두 통신을 중첩하면서도 모델 연산 성능 저하를 방지하도록 설계되었다는 공통점을 지닌다. Comet은 thread block 분리를 통해 GEMM 계산 커널이 원래의 성능을 유지하도록 보장하였으며, DeepEP는 통신을 GPU 연산에서 분리(Offload)함으로써 본래의 연산 커널(CUDA GEMM 등)이 영향을 받지 않도록 하였다. 다시 말해, Comet은 연산 경로 최적화를 유지하면서 통신 기능을 추가하였고, DeepEP는 통신을 별도의 경로로 처리하여 연산 경로를 간섭하지 않음으로써 전문가 계산 효율을 최대한 보존하였다.
차이점 및 고유한 접근 방식
- 오버랩 전략의 단위와 방식: Comet은 단일 미니배치 내에서 통신과 연산을 세분화하여 동시 실행하는 정밀한 동기식 오버랩을 구현하는 반면, DeepEP는 마이크로배치 수준의 비동기적 오버랩을 지원한다. Comet에서는 단일 forward 패스에서도 통신 및 연산을 수십 개의 타일로 분할하여 교차 수행하지만, DeepEP의 hook 기반 오버랩은 두 개의 마이크로배치를 파이프라인에 배치하여 하나가 통신하는 동안 다른 하나가 연산하는 방식으로 작동한다. DeepEP는 Dispatch 통신을 개시한 후 즉시 제어권을 반환하여 후속 연산을 수행할 수 있으며, 이때 통신 완료는 백그라운드 RDMA를 통해 진행되고 연산에는 GPU SM을 전혀 사용하지 않는다. 이러한 방식으로 DeepEP는 후속 배치의 연산과 이전 배치의 통신을 병행하는 파이프라인을 구성한다. 요약하면, Comet은 단일 작업 파이프라인을 미세 분할하여 동시 실행하는 반면, DeepEP는 복수의 작업을 중첩하는 방식으로 오버랩을 수행한다는 차이가 있다.
- 리소스 사용 및 커널 구조: Comet은 통신 또한 GPU 스레드가 담당하도록 설계되었으므로 일부 SM 자원을 통신에 할당한다(예: 전체 SM 중 일정 비율은 통신 스레드 블록이 점유). 반대로 DeepEP는 명시적으로 통신 자체에는 추가적인 SM을 사용하지 않도록 구현되었다. DeepEP의 훅 메커니즘은 GPU의 DMA 엔진/네트워크 카드가 데이터를 주고받는 동안, GPU 연산은 별개의 SM에서 독립적으로 진행될 수 있도록 하여 두 작업 간의 자원 경합을 최소화한다. 따라서 Comet은 통신을 위해 GPU 내부 자원 분배를 최적화한 설계이며, DeepEP는 통신을 GPU 외부 장치(네트워크)로 오프로딩하여 SM 간섭을 제거한 설계라고 할 수 있다. 이로 인해 Comet은 통신 오버랩을 위해 정교한 커널을 구현해야 하지만, DeepEP는 비교적 플러그인 형태로 작동하며 프레임워크의 기본 연산 커널과 병렬적으로 진행될 수 있다.
- 실현 방식 (통합 vs 모듈형): Comet은 Megatron-LM 프레임워크 내부에 최적화된 커널을 통합하는 형태로 구현되었다. 즉, MoE 레이어의 순방향/역방향 함수를 대체하는 CUDA 커널을 생성하여 프레임워크에 패치하는 방식이다. 반면 DeepEP는 독립적인 통신 라이브러리로서, 프레임워크에서는 DeepEP가 제공하는 API (dispatch 및 combine 함수 등)를 호출하여 통신을 수행하고, 그 반환된 훅을 적절히 사용하여 사용자가 연산 커널 호출을 조율하는 형태입니다. 따라서 Comet은 특정 프레임워크(Megatron) 전용으로 깊이 통합된 구조이며, DeepEP는 다양한 코드베이스에 비교적 용이하게 연결할 수 있는 모듈형 구조라고 할 수 있다. 예를 들어 PyTorch에서 MoE 구현 시, Comet을 사용하려면 Flux 커널을 직접 호출하거나 사용자 정의 Extension을 사용해야 하지만, DeepEP는 PyTorch 연산 훅으로 통신 텐서를 미리 받아오고 이후 연산에서 사용하는 방식으로 비교적 유연하게 통합할 수 있다.
Comet의 적응형 할당 방식은 내부적으로 최적 블록 분할 비율을 자동으로 산출하여 다양한 상황에 적용 가능한 자동 튜닝 기능을 제공한다. 반면, DeepEP는 사용자에게 훅 사용 여부, FP8 사용 여부, SM 커널 사용 여부 등 여러 옵션을 제공하여 상황에 맞게 설정을 조정하도록 한다. 예를 들어 DeepEP는 저지연 모드와 고처리량 모드 커널을 구분하여, 추론과 같이 레이턴시에 민감한 경우에는 FP8와 RDMA 순수 전송 모드를, 훈련 시에는 NVLink와 RDMA 병행 고처리량 모드를 사용자가 선택할 수 있도록 한다. 또한 DeepEP는 SM 개수 제어 기능을 통해 통신 커널이 사용할 SM 비율을 제한하거나 조정할 수 있게 한다. 이와 같이 DeepEP는 다양한 설정 프로파일을 통해 여러 시나리오에 대응하며, Comet은 내부 성능 모델을 기반으로 동적으로 최적화한다는 차이점이 있다.
오버랩 효과만 놓고 보면, Comet은 통신 시간의 평균 80% 이상을 숨겼다고 보고된 반면, DeepEP는 마이크로 배치 2개 오버랩 시 이론적으로 통신 지연의 50%까지 겹칠 수 있다. 그러나 DeepEP는 추가적으로 통신 자체를 빠르게 하고(대역폭 증가) 통신량도 줄였기에 결과적으로 총 지연 감소 효과는 상당할 것이다. 한편, Comet은 전문가 수가 매우 많아질 경우(예: 한 모델에서) 스케줄링 오버헤드가 늘어날 수 있다고 지적한 반면, DeepEP는 수백 이상의 전문가 구성(DeepSeek-V3에서 320 전문가)에서도 효율적으로 동작하도록 설계된 것으로 보이는데, Comet은 단일 레이어 최적 겹침에 초점을 맞추고, DeepEP는 전체 MoE 시스템 전반의 최적화(통신+압축+병렬 전략) 측면이 강하여 직접적인 비교보다는 상호 보완적인 관계로 이해할 수 있다.
4. Comet 연구의 한계
Comet에서 제안된 정밀한 연산-통신 오버랩 기술은 혼합 전문가 모델(MoE)의 훈련 및 추론 성능 향상에 상당한 기여를 하였으나, 여전히 몇 가지 제한점과 실용적인 제약이 존재한다.
- 특정 워크로드(MoE)에 특화된 설계: Comet의 기법은 MoE 구조의 통신-연산 패턴에 최적화되어 있으며, shared tensor를 중심으로 한 producer-consumer 구조라는 전제를 활용한다. 이러한 가정은 대표적인 MoE (토큰 dispatch + 전문가 FFN + 토큰 combine)에는 잘 맞지만, 일반적인 다른 분산 DNN 레이어에 바로 적용하기는 어렵다. 예를 들어 파이프라인 병렬화나 Tensor 병렬화 같은 경우 데이터 의존 패턴이 달라 Comet식 분해/재스케줄 기법을 쓰기 어려울 수 있다. 따라서 Comet의 접근법을 범용 통신-연산 겹침 프레임워크로 확장하려면 각기 다른 패턴에 대한 추가 연구와 알고리즘 개발이 필요하다. 저자들도 이를 인지하여, 향후 Comet의 프로그래밍 모델을 **컴파일러(Triton/TVM)**로 일반화하는 방향을 제시했으나 , 아직 MoE 이외의 범주에의 직접적 적용은 제한적이다.
- 구현 복잡성과 모듈화 난이도: Comet은 성능을 극대화하기 위해 저수준 CUDA 커널을 세심하게 튜닝하고 NVSHMEM 같은 특수 라이브러리를 활용했다. 이러한 접근은 성능 면에서는 유리하지만, 프레임워크 관점에서 모듈화하거나 유지보수하기 어렵다는 단점이 있다. 예를 들어 PyTorch, TensorFlow 등의 기본 런타임과 완전히 통합하려면 커스텀 C++/CUDA 연산으로 깊숙이 들어가야 하므로, 프레임워크 업그레이드나 타사 디버깅 도구와의 호환성 문제가 생길 수 있다. Comet은 현재 Megatron-LM에 통합된 형태로 소개되었는데, 다른 MoE 프레임워크(예: DeepSpeed-MoE, FastMoE 등)에 이식하려면 Flux 라이브러리를 해당 환경에 맞게 붙이는 작업이 필요하다. 이는 상당한 전문성/공수가 요구되므로 일반 사용자가 즉시 활용하기엔 장벽이 있다. 반면 DeepEP 같은 경우는 hook 인터페이스를 제공해 비교적 쉽게 붙일 수 있도록 했는데, Comet도 장기적으로는 이런 사용성 개선이 뒤따라야 할 것이다.
- NVSHMEM 및 하드웨어 의존성: Comet의 성능 비밀 중 하나는 NVSHMEM을 통한 GPU 간 저지연 통신인데, 이는 NVIDIA GPU 및 해당 HPC 소프트웨어 스택에 강하게 의존한다. NVSHMEM은 현재 NVIDIA에서만 제공되며 AMD GPU 등 다른 하드웨어에서는 대등한 기능을 사용할 수 없다. 또한 NVSHMEM 2.x는 비교적 최신 CUDA와 HPC 환경에서 지원되므로, 일반적인 클라우드 GPU 환경이나 싱글 노드 멀티카드 환경에서는 설치나 활용이 번거로울 수 있다. Infiniband 네트워크를 사용하는 대규모 클러스터(예: NVIDIA NCCL 대신 NVSHMEM 사용 가능)가 아니면 Comet의 이점을 살리기 어려울 수 있다는 뜻이다. 따라서 Comet을 광범위한 플랫폼에 적용하려면, 예를 들어 MPI + NCCL 기반 환경에서도 동작하도록 수정하거나, AMD GPUs를 위한 ROC-shmem 대응 구현 등이 필요할 수 있다. 이러한 이식성 문제는 현재 Comet 논문에서 직접 다루진 않았지만 실용적으로 고려해야 한다.
- 메모리 오버헤드 및 설정: 앞서 언급했듯 NVSHMEM을 쓰기 위해서는 GPU마다 추가 메모리 버퍼를 확보해야 하며, 이는 입력 길이와 모델 크기에 비례하여 증가한다 . 매우 대용량의 시퀀스나 거대한 hidden dimension을 갖는 모델의 경우 이 버퍼가 커져서 메모리 부족을 야기할 가능성이 있다. 또한 NVSHMEM 버퍼 크기는 실행 전에 결정해야 하기 때문에, 동적으로 입력 크기가 변하는(workload 가변) 환경에서는 최악의 경우에 맞춰 큰 버퍼를 잡아놓아야 하는 비효율이 생길 수 있다. 이러한 메모리 트레이드오프는 Comet 사용 시 고려해야 하는 실용적 제약이다. DeepEP의 구현도 NVSHMEM 기반이라 유사한 제약이 있지만, DeepEP는 통신 버퍼를 큐로 관리하여 메모리 절약을 시도했다고 언급된다(다만 복잡성과 데드락 이슈로 향후 고정 버퍼로 변경 권고).
- 성능 최적화의 유효 범위: Comet은 network I/O 지연을 연산으로 숨기는 것에 중점을 두었으며, 통신 자체의 절대적인 시간 단축을 다루지 않는다. 즉, 겹침을 통해 통신 대기 시간은 감소하지만, 통신 자체에 걸리는 총 시간은 여전히 데이터 양과 네트워크 대역폭에 의해 좌우된다. 따라서 만약 네트워크가 지나치게 느리거나 통신량이 매우 많다면, Comet으로 완전히 해소되지 않을 수 있다.
- 복잡성 증가와 디버깅: Comet의 fine-grained 겹침은 매우 복잡한 동시성 패턴을 GPU 내부에 만들어낸다. 이것은 디버깅과 검증의 어려움으로 이어질 수 있는데, 예를 들어 한 GPU에서 예상치 못한 통신 지연이 생기거나, NVSHMEM으로 주고받는 데이터의 정합성 문제가 생기면 원인을 추적하기가 까다롭다. DeepEP도 성능 극대화를 위해 PTX 수준의 트릭을 쓰다가(undefined behavior 이용 ) 이슈가 발견되기도 하는 등, 이런 저수준 최적화에는 리스크가 따른다. Comet은 ByteDance 내부에서 충분히 검증되어 프로덕션에 들어갔다고 하지만, 일반 연구자가 이를 응용할 때는 예상치 못한 동기화 버그나 데드락 등을 만날 가능성도 있다. 따라서 Comet을 활용하거나 자체 구현할 계획이 있다면 철저한 테스트와 검증이 필요하며, 이러한 개발 난이도도 넓은 보급에 장애 요인이 될 수 있다.