[transformers] DeepGEMM Mega MoE를 EP 구현체로 편입하기: 라우터 스타일을 버리고 슬라이스+all-gather로 최대 9배 빨라진 MoE 추론

PR 링크: huggingface/transformers#49606 상태: Merged | 변경: +194 / -108

들어가며

DeepGEMM의 Mega MoE는 MoE 레이어의 세 단계, 즉 token dispatch → experts GEMM → combine을 하나의 kernel로 fuse합니다. 그 대가로 kernel은 다음 세 가지를 요구합니다.

1. global expert id (rank-local로 remap되지 않은 원래 id)
2. unmasked routing weight
3. EP process group (kernel 내부에서 symmetric memory로 토큰을 교환하기 위해)

지금까지 transformers에서는 이 조건을 megamoe_router, megamoe_experts라는 별도의 parallel style로만 만족시킬 수 있었습니다. 문제는 DeepSeek V4가 기본으로 제공하는 plan이 ep_dispatch_experts라는 점입니다. 이 plan으로 experts_implementation="deepgemm_megamoe"를 지정해 로드하면 첫 forward에서 바로 터졌습니다.

DeepGEMM Mega MoE requires a process_group

결국 main에서는 Mega MoE를 쓰지 못하고 all_reduce EP 구현체 + DeepGEMM grouped FP8×FP4 kernel 조합으로 돌아가야 했습니다. 이 PR은 Mega MoE를 EP implementation의 하나("deepgemm_megamoe")로 편입합니다. 그리고 두 legacy style을 제거하는 🚨 breaking change를 함께 담고 있습니다.

코드 분석

1. distributed/expert_parallel.py — 라우터 스타일 제거

Before: 라우터 출력을 건드리지 않기 위해 EpRouterParallel을 상속해 post-forward를 no-op으로 만든 클래스가 있었습니다.

class RouterParallelMegaMoe(EpRouterParallel):
    def transform_output_post_forward(self, module, output, mesh):
        return output

After: 이 클래스는 삭제됐습니다. ep_dispatch_experts plan은 원래 experts 쪽에서 토큰 분배를 처리하므로, global id와 unmasked weight는 EP 구현체가 그대로 받습니다. 그래서 "라우터를 비틀어 마스킹을 끄는" 우회가 필요 없어졌습니다.

2. MegaMoeEP — 슬라이스, 패딩, all-gather

핵심은 새 구현체입니다. Mega MoE kernel은 각 rank의 행을 그 rank 고유의 토큰으로 취급합니다. 반면 현재 transformers 추론에서는 EP 그룹의 모든 rank가 같은 batch를 들고 있습니다. 그대로 넘기면 같은 토큰이 ep_size번 계산됩니다. 그래서 batch를 균등하게 잘라 각 rank가 자기 몫만 kernel에 넣습니다.

def __call__(self, hidden_states, top_k_index, top_k_weights):
    ep_group, ep_size, ep_rank = self.ep_mesh.get_group(), self.ep_mesh.size(), self.ep_mesh.get_local_rank()
    num_tokens = hidden_states.size(0)
    rows_per_rank = max(1, -(-num_tokens // ep_size))  # ceil div
    rows = slice(ep_rank * rows_per_rank, (ep_rank + 1) * rows_per_rank)
    hidden_states, top_k_index, top_k_weights = hidden_states[rows], top_k_index[rows], top_k_weights[rows]
    padding = rows_per_rank - hidden_states.size(0)
    if padding:
        hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, padding))
        top_k_index = torch.nn.functional.pad(top_k_index, (0, 0, 0, padding), value=-1)
        top_k_weights = torch.nn.functional.pad(top_k_weights, (0, 0, 0, padding))

    local_output = self.experts_forward(hidden_states, top_k_index, top_k_weights, process_group=ep_group)
    output = local_output.new_empty(rows_per_rank * ep_size, local_output.size(-1))
    all_gather_into_tensor(output, local_output.contiguous(), group=ep_group)
    return output[:num_tokens].to(hidden_states.dtype)

포인트는 세 가지입니다.

• -(-n // k) ceil division으로 모든 rank의 slice 크기를 맞춥니다. 그래서 all_gather_into_tensor가 고정 크기 출력을 바로 채울 수 있습니다. 가변 크기 gather였다면 host sync가 필요했을 것입니다.
• 패딩 슬롯은 **expert id -1, weight 0**으로 채웁니다. kernel이 이 슬롯을 skip하므로 계산 낭비가 없습니다. max(1, ...) 덕분에 토큰 수가 rank 수보다 적어도 shape이 0이 되지 않습니다.
• 마지막에 [:num_tokens]로 패딩을 잘라냅니다.

PR 설명의 예시가 이 흐름을 잘 보여줍니다.

GPU                    0                      1
batch           [t0, t1, t2, t3]       [t0, t1, t2, t3]
kernel input        [t0, t1]               [t2, t3]
kernel output       [y0, y1]               [y2, y3]
after all-gather         [y0, y1, y2, y3] on both ranks

리뷰어 vasqu는 "older torch 버전의 collective, 특히 backward가 늘 문제였다"며 안정성을 우려했습니다. 작성자는 이렇게 답했습니다. forward에서는 2.5 이전부터 있던 all_gather_into_tensor만 쓰고, Mega MoE kernel에는 backward가 없어 collective backward 자체가 존재하지 않는다는 것입니다. 클래스의 TODO도 눈여겨볼 만합니다. Mega MoE는 원래 DP attention + EP처럼 rank마다 다른 batch를 갖는 구조를 위해 설계됐습니다. 그 구조에서는 slice와 gather가 둘 다 필요 없으므로, 추론에 DP가 들어오면 이 오버헤드도 사라질 수 있습니다.

3. AutoEP의 선택 로직

def _uses_megamoe(experts: nn.Module) -> bool:
    return getattr(getattr(experts, "config", None), "_experts_implementation", None) == "deepgemm_megamoe"

class AutoEP(EPImplementation):
    def __init__(self, experts, ep_mesh, tp_mesh):
        ...
        self.megamoe = MegaMoeEP(experts, ep_mesh, tp_mesh) if _uses_megamoe(experts) else None

    def __call__(self, hidden_states, top_k_index, top_k_weights):
        if self.megamoe is not None:
            return self.megamoe(hidden_states, top_k_index, top_k_weights)
        if self.all_reduce is not None and not (torch.is_grad_enabled() and hidden_states.requires_grad):
            return self.all_reduce(...)
        return self.all_to_all(...)

3outeille은 "auto가 추론 시 all_reduce와 deepgemm_megamoe 중 무엇을 고르는가"라고 물었습니다. 결론은 이것이 성능 기반 선택이 아니라 강제라는 것입니다. Mega MoE experts는 kernel 내부에서 토큰을 교환하므로 다른 EP 경로로는 애초에 실행될 수 없습니다. 그래서 auto에 넣어도 의미가 맞습니다. 반대로 명시적으로 다른 구현체를 지정하면 install_forward에서 조기에 에러를 냅니다.

if _uses_megamoe(module) and ep_implementation not in ("auto", "deepgemm_megamoe"):
    raise ValueError(...)

초기 버전에서는 quantizer가 config를 rewrite했습니다. 리뷰를 거치며 AutoEP가 직접 판단하는 구조로 정리되어 책임 위치가 더 자연스러워졌습니다. 이중 getattr이 과하게 방어적이라는 지적도 있었습니다. 작성자는 experts가 use_experts_implementation decorator가 붙지 않은 임의의 nn.Module일 수 있다는 이유로 유지했습니다.

4. integrations/deepgemm.py — symmetric buffer 공유

이전에는 experts 모듈마다 자기 buffer를 할당했습니다. 이제는 EP group(process_group.group_name)과 shape 단위로 캐시해 모든 레이어가 공유합니다. 메모리 해제는 DeepEP의 destroy_deepep_buffers()와 대칭인 destroy_megamoe_buffers()로 합니다.

이 설계가 중요한 이유는 리뷰 답변에 잘 드러납니다. 일반 CUDA buffer가 아니라 symmetric memory이기 때문에 할당이 collective rendezvous이고 느립니다(특히 decode 중에). 레이어마다 따로 할당하던 구조였다면 새로운 최대 batch가 들어올 때마다 레이어 수만큼 rendezvous가 일어났습니다. 최종 버전은 다음과 같이 동작합니다.

• 용량을 power-of-two가 아닌 실제 필요한 batch(384-token block 정렬)로 키웁니다.
• 키울 때 기존 buffer를 먼저 해제해 메모리를 재사용합니다.
• 따라서 가장 큰 batch로 warm-up하면 단 한 번만 할당됩니다.

다만 symmetric-memory pool은 high-water mark를 유지하므로, destroy는 buffer만 풀고 pool은 풀지 않습니다. DeepSeek-V4 크기 기준으로 rank당 토큰 용량 1개에 약 1 MiB가 듭니다.

왜 이게 좋은가

B200, EP 2, 64 experts, hidden 4096, top_k 6 환경에서 all_reduce 대비 측정한 결과입니다(cosine ≥ 0.9989).

Tokens	megamoe eager	all_reduce eager	가속	megamoe CUDA graph	all_reduce CUDA graph	가속
1	0.718 ms	1.871 ms	2.6x	0.102 ms	0.314 ms	3.1x
512	0.671 ms	1.886 ms	2.8x	0.221 ms	1.033 ms	4.7x
16384	2.018 ms	17.849 ms	8.8x	1.905 ms	17.402 ms	9.1x

all_reduce는 모든 rank가 batch 전체에 대해 자기 experts를 돌리므로 비용이 batch에 비례해 커집니다. Mega MoE 경로는 각 rank가 1/ep_size만 처리하고 토큰 교환은 fuse된 kernel이 담당합니다. 그래서 batch가 클수록 격차가 벌어집니다. kernel 자체는 1 token에서 0.076 ms에 불과합니다. 나머지 시간의 대부분은 PyTorch reference로 구현된 per_token_cast_to_fp8(약 12개 kernel)이 차지하므로, fused cast가 들어오면 개선 여지가 더 있습니다.

일반적 교훈

1. 특수 kernel은 별도 layer style이 아니라 dispatch 전략의 플러그인으로 넣어라. 기존 plan과 호환되면서 registry(EPImplementationInterface) 한 줄로 확장됩니다.
2. 균등 slice + 패딩으로 static shape를 유지하라. host sync 없는 고정 크기 collective는 CUDA graph capture와 궁합이 좋습니다.
3. 비싼 리소스는 소유 단위를 다시 생각하라. symmetric memory처럼 할당이 collective인 자원은 모듈 단위가 아니라 process group 단위로 공유해야 rendezvous 횟수가 줄어듭니다.
4. "선택"과 "강제"를 구분하라. auto 로직에 성능 휴리스틱이 아닌 기능적 제약을 넣을 때는, 다른 명시적 선택을 조기에 거부해 혼란을 막아야 합니다.

참고 자료

• torch.distributed.all_gather_into_tensor — MegaMoeEP에서 slice 출력을 모으는 collective
• torch.nn.functional.pad — 마지막 slice 패딩에 사용
• torch.distributed.device_mesh — ep_mesh의 group/rank 조회
• DeepGEMM — Mega MoE kernel을 제공하는 라이브러리

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.