EP: Expert를 나누고 토큰 보내기

MoE에서는 토큰마다 사용할 MLP인 Expert를 고릅니다. Expert가 많아지면 그 가중치를 모두 한 GPU에 보관하기 어려워집니다. 전문가 병렬화(Expert Parallelism, EP)는 서로 다른 Expert를 여러 GPU에 나누어 배치하는 방법입니다.

모델의 모든 부분에 같은 병렬화를 적용할 필요는 없습니다. 이번 글에서는 가장 단순하게 어텐션 등 공통 부분은 DP로 복제하고, Expert 부분은 EP로 나눈 구성을 살펴보겠습니다. 같은 GPU들이 두 방식을 함께 사용한다는 점을 먼저 확인하고, 토큰이 선택된 Expert로 흩어졌다가 원래 GPU와 토큰 순서로 돌아오는 과정을 따라갑니다. 이어서 여러 Expert의 결과를 합치는 방법과, Expert 수가 같아도 처리할 토큰 수는 달라지는 이유를 살펴보겠습니다.

어텐션은 복제하고, Expert는 나누어 배치하기

GPU 0이 입력 시퀀스 A를, GPU 1이 입력 시퀀스 B를 맡는다고 하겠습니다. A1·A2는 시퀀스 A의 두 토큰이고, B1·B2는 시퀀스 B의 두 토큰입니다. 그림의 작은 칸은 토큰 하나를 나타내며, 실제 계산하고 전달하는 것은 그 토큰의 활성값 벡터입니다.

어텐션 등 공통 부분의 가중치는 두 GPU에 동일하게 복제합니다. GPU 0은 A를, GPU 1은 B를 같은 가중치로 계산합니다. 앞서 살펴본 데이터 병렬화의 방식입니다. 이 예시에서는 시퀀스 하나를 GPU 사이에 나누지 않습니다. 어텐션은 각 GPU에서 자신이 맡은 시퀀스 안의 토큰 관계를 처리합니다.

Expert 부분의 배치는 다릅니다. E0·E1은 GPU 0에, E2·E3는 GPU 1에 둡니다. 이 네 Expert는 같은 가중치의 복제본이 아니라 서로 다른 가중치를 가진 MLP입니다. 토큰마다 어느 Expert를 사용할지 고르는 라우터(router)의 가중치는 두 GPU에 복제하고, 각 GPU가 자기 토큰에 대해 네 Expert 중에서 선택합니다.

각 GPU는 일부 Expert 가중치만 보관하지만, 어텐션과 라우터의 가중치는 여전히 복제합니다. 따라서 Expert를 반씩 나누었다고 전체 메모리 사용량까지 반으로 줄어드는 것은 아닙니다.

그림에서 GPU 0과 GPU 1의 위치는 위에서 아래까지 같습니다. 두 GPU가 어텐션을 계산하다가, MoE 구간에서는 선택된 Expert 쪽으로 토큰을 보냅니다. 계산 결과를 돌려받으면 다시 각자 맡은 입력의 계산을 이어갑니다. DP용 GPU와 EP용 GPU가 따로 있는 것이 아니라, 같은 GPU에서 모델의 부분마다 가중치 배치와 데이터 흐름이 달라지는 것입니다.

공통 부분에 DP를 사용하는 것은 이 글의 예시 구성입니다. EP가 어텐션의 병렬화까지 하나로 정하지는 않습니다. vLLM의 EP 문서도 어텐션을 복제하는 구성과 TP로 나누는 구성을 구별합니다. 여기서는 토큰의 왕복을 보기 위해 공통 부분의 TP와 다른 분할 방식은 생략하겠습니다.

토큰을 Expert로 보내고 원래 위치로 돌려받기

이제 MoE 구간만 확대하겠습니다. 먼저 토큰마다 Expert 하나를 고르는 Top-1을 사용합니다. 전달과 회수에 집중하도록 이 예시의 결합 가중치는 1로 두겠습니다. 라우터가 A1에는 E0, A2에는 E2, B1에는 E1, B2에는 E2를 골랐다고 하겠습니다.

다음 그림의 ‘다음’을 누르면 입력, Expert 선택, 전달, 계산, 회수, 순서 복구로 이어집니다. 토큰의 색과 이름은 유지되고, GPU와 Expert의 위치는 바뀌지 않습니다.

라우터는 목적지를 결정합니다. GPU 0에서 선택한 A1의 E0은 같은 GPU에 있지만, A2의 E2는 GPU 1에 있습니다. GPU 1에서도 B2의 E2는 로컬이고, B1의 E1은 GPU 0에 있습니다. 중앙의 라우터 하나가 모든 입력을 모아서 결정하는 구조가 아닙니다. 각 GPU가 자기 토큰으로 Expert 선택 정보를 계산합니다.

Dispatch는 그 선택에 따라 토큰 벡터를 Expert별로 모읍니다. A2의 벡터는 GPU 1로, B1의 벡터는 GPU 0으로 갑니다. A1과 B2는 같은 GPU 안에서 해당 Expert의 입력으로 배치합니다. Expert 가중치를 입력 쪽으로 매번 옮기는 대신, 가중치는 그대로 두고 토큰의 벡터를 보냅니다. Expert에 전달하는 것은 라우터의 점수가 아니라, 선택할 때 사용했던 토큰의 입력 벡터입니다.

E2 아래에는 서로 다른 시퀀스에서 온 A2와 B2가 모입니다. 이것이 가능한 이유는 Expert의 MLP가 각 토큰을 독립적으로 변환하기 때문입니다. 두 벡터를 행으로 쌓아 같은 E2 가중치로 행렬 연산을 수행할 수 있지만, A2의 출력에 B2의 값을 섞지는 않습니다. 같은 Expert 아래에 모은다는 것은 계산을 묶는다는 뜻이지, 어텐션처럼 토큰 간 정보를 결합한다는 뜻은 아닙니다.

결과는 출발 GPU의 원래 토큰으로 돌아갑니다. E2가 만든 A2의 결과는 GPU 0으로, E1이 만든 B1의 결과는 GPU 1로 보냅니다. 로컬에서 계산한 A1과 B2의 결과는 같은 GPU에 남습니다. GPU 0은 A1·A2 순서로, GPU 1은 B1·B2 순서로 출력을 배치하여 다음 연산에 연결합니다. 이 예시에서는 결과가 하나씩이므로 Top-2에서 필요한 여러 결과의 가중합은 없습니다.

이 왕복은 입력별 배치 → Expert별 배치 → 다시 입력별 배치로 바뀌는 과정입니다. 목적지별 교환이라는 점에서 All-to-All과 연결되지만, 실제로는 목적지마다 토큰 수가 다를 수 있습니다. 각 결과가 어느 토큰에 해당하는지도 함께 추적해야 합니다. Megatron의 MoE 문서는 이 전달·회수 관계와 여러 dispatcher를 설명합니다. EP의 핵심은 특정 통신 함수의 이름보다, 선택된 Expert가 입력을 받고 그 결과가 올바른 토큰으로 돌아오는 관계입니다.

여러 Expert의 결과도 같은 토큰으로 돌아오기

토큰마다 두 Expert를 선택하는 Top-2라면 입력 벡터 하나가 두 Expert의 계산에 사용됩니다. 이번에는 GPU 0의 A2에 집중하겠습니다. 라우터가 로컬 E0과 GPU 1의 E2를 선택하고, 결과를 합칠 가중치를 각각 0.7과 0.3으로 정했다고 하겠습니다. 다른 토큰의 흐름은 생략합니다.

E0과 E2에는 같은 A2 벡터를 입력하지만, 서로 다른 가중치를 사용하므로 결과도 둘 생깁니다. 그림의 u0는 E0의 출력, u2는 E2의 출력입니다. u0는 GPU 0에서 계산하고, u2는 GPU 1에서 계산한 뒤 GPU 0으로 돌려보냅니다.

출발지인 GPU 0에서 같은 A2에 대한 두 결과만 가중합합니다. 그림에서는 y(A2) = 0.7 × u0 + 0.3 × u2입니다. 두 출력이 모두 d차원이라면 결과도 d차원이고, 원래 A2의 자리에 연결할 수 있습니다. 라우터가 Expert 선택에 사용한 정보와, 결과 결합에 사용할 가중치를 구별해서 유지해야 하는 이유입니다.

선택하는 Expert 수가 늘면 계산할 토큰·Expert 조합과 회수할 결과도 늘어납니다. 다만 두 Expert가 같은 GPU에 있다면 입력 전달을 함께 처리할 수 있으므로, 네트워크 전송량이 선택 수에 항상 그대로 비례하는 것은 아닙니다. 실제 전달량은 Expert 배치와 구현에 따라 달라집니다.

Expert가 같은 개수여도 처리할 토큰 수는 다르다

각 GPU가 같은 수의 입력으로 시작해도, 라우터가 고른 목적지가 같을 필요는 없습니다. 토큰 수를 늘려 GPU 0에는 A1부터 A4까지, GPU 1에는 B1부터 B4까지 네 개씩 있다고 하겠습니다. 같은 크기의 Expert 네 개를 여전히 GPU마다 두 개씩 두고, Top-1으로 선택합니다.

그림에서는 E0에 A1, E1에 B1을 보내고 나머지 여섯 토큰은 E2가 맡습니다. 따라서 GPU 0은 두 토큰, GPU 1은 여섯 토큰을 처리합니다. 공통 부분에서 입력을 고르게 나누었다고, Expert 구간의 작업량까지 고르게 나뉘는 것은 아닙니다.

특정 Expert가 많이 선택되면 그 GPU에 계산뿐 아니라 입력 버퍼와 통신 부담도 집중될 수 있습니다. 다른 GPU에서 출발한 토큰도 이 Expert의 결과를 기다려야 합니다. 반대로 Expert당 토큰이 너무 적으면 작은 행렬 연산이 많아져 GPU를 효율적으로 사용하기 어려울 수 있습니다.

토큰 수가 세 배라고 실행 시간도 정확히 세 배라는 뜻은 아닙니다. 이 그림이 보여주는 것은 가중치를 같은 개수로 나누는 일과 실제 들어오는 작업량을 고르게 만드는 일이 서로 다르다는 점입니다. 따라서 Expert 배치와 함께 라우팅 결과의 분포도 확인해야 합니다.

EP에서는 같은 GPU들이 공통 부분은 DP로 처리하고, Expert 부분에서는 필요한 토큰을 주고받도록 구성할 수 있습니다. 여러 병렬화 방식을 선택하고 조합하는 기준은 다음 글에서 살펴보겠습니다.