전문가 부하 균형과 전문가 붕괴

전문가 몇 명에게 일이 몰리면

전문가 붕괴는 MoE 학습의 대표적인 실패야. 제약 없이 라우팅하면 초기에 조금 앞선 전문가에게 토큰이 더 몰려. 그 전문가는 더 빨리 배우고, 라우터는 다시 더 많은 토큰을 보내. 끝에는 몇 명만 모든 일을 하고 나머지는 메모리만 차지하는 죽은 가중치가 돼.

스스로 강화되는 쏠림

라우터와 전문가는 함께 학습돼. 부하를 퍼뜨리는 압력이 없으면 옵티마이저는 덜 쓰인 전문가를 되살릴 이유가 없어. 작은 초기 차이가 계속 커지는 구조라서 놓치면 값비싼 학습 전체를 망칠 수 있어.

균형을 잡는 방법

보조 부하 균형 손실: 배치 안에서 전문가 사용량이 비슷해지도록 벌점을 더해. Switch Transformer 이후 초기 MoE의 표준 방식이야.
보조 손실 없는 균형화: DeepSeek-V3는 라우터 로짓에 전문가별 편향을 두고 실제 사용량에 맞춰 조정해. 주 목표와 부딪히는 별도 손실 항이 없어.
전문가 용량 제한: 전문가 한 명이 배치에서 받을 토큰 수에 상한을 두고, 넘친 토큰은 다음 전문가로 보내거나 버려.
잡음 주입: 학습 중 라우터 점수에 잡음을 더해 경로가 너무 일찍 굳는 걸 막아.

추론할 때도 부하는 고르지 않아

학습이 끝나도 워크로드에 따라 인기 전문가가 생겨. 특정 전문가를 가진 GPU는 과부하되고 다른 GPU는 놀 수 있어. vLLM과 TensorRT-LLM 같은 서빙 도구가 이 쏠림을 동적으로 다뤄야 하는 이유야.

모델 카드에서 찾을 말

‘auxiliary loss’나 균형 계수가 나오면 전통적인 방식을 쓴 거야. ‘auxiliary-loss-free’나 ‘bias-based balancing’이 보이면 DeepSeek 계열 설계를 읽고 있는 셈이지. MoE 모델인데 균형 전략이 전혀 없다면 중요한 설명이 빠졌는지 의심해 봐.

Code

보조 부하 균형 손실(Switch Transformer 방식)·python

import torch

def load_balancing_loss(router_logits, expert_indices, num_experts):
    # router_logits: (B*T, num_experts)
    # expert_indices: (B*T, k) — top-k chosen experts
    fraction_per_expert = torch.zeros(num_experts, device=router_logits.device)
    for k in range(expert_indices.shape[-1]):
        fraction_per_expert.scatter_add_(
            0,
            expert_indices[..., k].view(-1),
            torch.ones_like(expert_indices[..., k].view(-1), dtype=torch.float),
        )
    fraction_per_expert /= expert_indices.shape[0]

    avg_router_prob = torch.softmax(router_logits, dim=-1).mean(dim=0)
    return num_experts * (fraction_per_expert * avg_router_prob).sum()

DeepSeek-V3식 편향 조정 의사 코드·python

# Each step: nudge expert_bias toward balancing observed load.
def adjust_expert_bias(expert_bias, observed_load, target_load, lr=1e-3):
    # observed_load: tokens routed to each expert this step
    # target_load:   ideal balanced load
    delta = (target_load - observed_load) * lr
    expert_bias += delta
    return expert_bias
# No backprop, no aux loss term, just a slow-feedback correction.

전문가 몇 명에게 일이 몰리면

스스로 강화되는 쏠림

균형을 잡는 방법

추론할 때도 부하는 고르지 않아

모델 카드에서 찾을 말

Code

External links

Exercise

Progress

댓글 0