docs/muon-clip.md

QK-Clip for MuonClip Optimizer (MLA)

Reference: Kimi K2 Technical Report, Section 2.1, Algorithm 1

개요

QK-Clip은 Muon optimizer에서 발생하는 attention logit explosion을 방지하기 위한 weight rescaling 기법이다. forward/backward에는 개입하지 않고, optimizer step 이후에 weight를 rescale하여 logit 성장을 원천 차단한다.

Algorithm 1: MuonClip

for each training step t:
    // 1. Muon optimizer step
    for each weight W:
        Mt = µ·Mt-1 + Gt
        Ot = Newton-Schulz(Mt) · √max(n,m) · 0.2
        Wt = Wt-1 - η·(Ot + λ·Wt-1)

    // 2. QK-Clip
    for each attention head h:
        S^h_max ← forward에서 기록한 head h의 max pre-softmax logit
        if S^h_max > τ:
            γ ← τ / S^h_max
            W^h_qc ← W^h_qc · √γ      (query compressed, q_nope)
            W^h_kc ← W^h_kc · √γ      (key compressed, k_nope)
            W^h_qr ← W^h_qr · γ       (query rotary, q_pe)
            // k_R (shared rotary, k_pe): 안 건드림

기존 코드 → MLA 수도코드

현재 코드 구조 (MHA/GQA)

parse_qk_layer(name)          → wq/wk 여부 판별, layer index 추출
get_qk_clip_info(config, n)   → QKClipInfo (kind, indices, head_dim, threshold, logit)
compute_scales(p, info)        → per-head √γ scales 텐서 반환
qk_clip(p, scales, head_dim)  → W.view(-1, head_dim, in_dim).mul_(scales)

현재 코드는 head_dim이 균일하고, Q/K weight 전체에 동일한 √γ를 적용한다.

MLA에서 달라지는 점

항목	MHA/GQA (현재)	MLA
Q weight	wq / q_proj	wq_b (up-proj from LoRA)
K weight	wk / k_proj	wkv_b (k_nope + v 합쳐져 있음)
Q head stride	qk_head_dim (균일)	qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
K head stride	qk_head_dim (균일)	kv_stride = qk_nope_head_dim + v_head_dim
Q scaling	전체 √γ	nope → √γ, rope → γ (서로 다름)
K scaling	전체 √γ	k_nope → √γ, v → 1.0 (부분만)
shared k_pe	없음	wkv_a 뒷부분, 안 건드림

수도코드: parse_qk_layer (MLA 확장)

def parse_qk_layer(name: str) -> tuple[str | None, int]:
    parts = normalize_fqn(name).split('.')
    kind = parts[-2]

    layer_idx = -1
    for part in reversed(parts):
        if part.isdigit():
            layer_idx = int(part)
            break

    # MHA/GQA: wq, wk, q_proj, k_proj
    # MLA:     wq_b (Q up-proj), wkv_b (KV up-proj)
    if kind in ('wq', 'wk', 'q_proj', 'k_proj', 'wq_b', 'wkv_b'):
        return kind, layer_idx

    return None, -1

수도코드: QKClipInfo (MLA 확장)

@dataclass
class QKClipInfo:
    kind: str | None          # 'wq_b' or 'wkv_b' (MLA) / 'wq','wk' (MHA)
    indices: list[int]        # clipping 대상 head indices
    head_dim: int             # 기존 MHA용 (uniform stride)
    threshold: float
    logit: torch.Tensor | None

    # MLA 전용 필드
    is_mla: bool = False
    qk_nope_head_dim: int = 0
    qk_rope_head_dim: int = 0
    v_head_dim: int = 0

수도코드: get_qk_clip_info (MLA 확장)

def get_qk_clip_info(clip_config, n, qk_logits):
    if clip_config is None:
        return None

    threshold = clip_config['threshold']
    kind, layer_idx = parse_qk_layer(n)
    is_mla = clip_config.get('is_mla', False)

    logit, indices = None, []
    if qk_logits is not None and kind is not None:
        logit = qk_logits[layer_idx]
        if isinstance(logit, DTensor):
            logit = logit.full_tensor()

        if kind in ('wq_b', 'wq', 'q_proj'):
            indices = clip_config.get('q_indices', []) or []
        elif kind in ('wkv_b', 'wk', 'k_proj'):
            indices = clip_config.get('k_indices', []) or []

    if is_mla:
        return QKClipInfo(
            kind=kind,
            indices=indices,
            head_dim=clip_config['head_dim'],          # qk_head_dim (for wq_b)
            threshold=threshold,
            logit=logit,
            is_mla=True,
            qk_nope_head_dim=clip_config['qk_nope_head_dim'],
            qk_rope_head_dim=clip_config['qk_rope_head_dim'],
            v_head_dim=clip_config['v_head_dim'],
        )
    else:
        # 기존 MHA/GQA 경로
        return QKClipInfo(
            kind=kind, indices=indices,
            head_dim=clip_config['head_dim'],
            threshold=threshold, logit=logit,
        )

수도코드: compute_scales (MLA 확장)

기존과 동일하게 per-head γ를 계산한다. (γ 결정은 MHA와 동일) 달라지는 건 qk_clip 적용 시 head 내부를 sub-region별로 나눠서 다른 변환을 쓰는 것이다.

def compute_scales(p, qk_clip_state):
    """기존 코드와 동일. per-head √γ 반환."""
    kind = qk_clip_state.kind
    indices = qk_clip_state.indices
    threshold = qk_clip_state.threshold
    logit = qk_clip_state.logit

    head_scales = {}
    for logit_idx, head_idx in enumerate(indices):
        v_ele = float(logit[logit_idx])
        if v_ele > threshold:
            new_scale = math.sqrt(threshold / v_ele)      # √γ
            if head_idx not in head_scales or new_scale < head_scales[head_idx]:
                head_scales[head_idx] = new_scale

    if not head_scales:
        return None

    H_global = p.shape[0] // qk_clip_state.head_dim      # MLA: head_dim = qk_head_dim or kv_stride
    scales_full = torch.ones(H_global, device=p.data.device)
    for head_idx, scale in head_scales.items():
        scales_full[head_idx] = scale                     # √γ_h

    return scales_full

수도코드: qk_clip (MLA 확장)

per-head scales(√γ)는 동일하게 받되, head 내부 sub-region에 다른 함수를 적용한다.

def qk_clip(p, scales, head_dim, is_mla=False, kind=None, info=None):
    """
    scales: [n_heads] 텐서, 각 원소 = √γ_h

    is_mla=False: 기존 MHA/GQA (head 내 uniform √γ)
    is_mla=True:  MLA (head 내 sub-region별 다른 변환)
    """
    W = p.data if isinstance(p, torch.nn.Parameter) else p

    if not is_mla:
        # 기존: 모든 행에 √γ 균일 적용
        W.view(-1, head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1))
        return

    # MLA: head별로 sub-region 분리 적용
    if kind == 'wq_b':
        qk_nope = info.qk_nope_head_dim
        qk_rope = info.qk_rope_head_dim
        qk_head_dim = qk_nope + qk_rope

        for h in range(len(scales)):
            sqrt_gamma = scales[h].item()
            if sqrt_gamma >= 1.0:
                continue
            gamma = sqrt_gamma * sqrt_gamma      # √γ → γ
            s = h * qk_head_dim

            W[s : s + qk_nope]              *= sqrt_gamma   # q_nope → √γ
            W[s + qk_nope : s + qk_head_dim] *= gamma       # q_pe   → γ

    elif kind == 'wkv_b':
        qk_nope = info.qk_nope_head_dim
        kv_stride = qk_nope + info.v_head_dim

        for h in range(len(scales)):
            sqrt_gamma = scales[h].item()
            if sqrt_gamma >= 1.0:
                continue
            s = h * kv_stride

            W[s : s + qk_nope] *= sqrt_gamma                # k_nope → √γ
            # v 행: 안 건드림

수도코드: GQA에서 wkv_b indices 처리

Q head → KV head 매핑이 필요하다. 여러 Q head가 같은 KV head를 공유하므로, group 내 최소 gamma 기준으로 한 번만 적용해야 한다.

def build_k_indices_for_mla(clip_config, n_heads, n_kv_heads):
    """
    Q head 기준 logit으로부터 KV head indices를 생성한다.
    q_indices가 Q head index 기준이라면,
    k_indices는 대응되는 KV head index로 변환해야 한다.

    주의: 같은 KV head에 매핑되는 여러 Q head 중
          가장 큰 logit (= 가장 작은 gamma)을 사용해야 한다.
    """
    heads_per_kv = n_heads // n_kv_heads
    q_indices = clip_config.get('q_indices', list(range(n_heads)))

    # Q head → KV head 매핑
    # logit 텐서에서 같은 kv_head에 대응되는 Q head들 중 max를 취하는 것은
    # compute_scales_mla 내부에서 min(gamma) 로 처리됨

    k_indices = []
    seen = set()
    for q_idx in q_indices:
        kv_idx = q_idx // heads_per_kv
        if kv_idx not in seen:
            k_indices.append(kv_idx)
            seen.add(kv_idx)

    return k_indices

수도코드: 호출 흐름 (통합)

# optimizer step 이후 호출되는 부분 (기존 코드 구조 유지)

for name, param in model.named_parameters():
    info = get_qk_clip_info(clip_config, name, qk_logits)
    if info is None or info.kind is None:
        continue

    scales = compute_scales(param, info)    # per-head √γ (MHA/MLA 공통)
    if scales is not None:
        qk_clip(param, scales, info.head_dim,
                is_mla=info.is_mla, kind=info.kind, info=info)

수도코드: clip_config 예시

# MHA/GQA (기존)
clip_config = {
    'head_dim': 128,
    'threshold': 100.0,
    'q_indices': list(range(n_heads)),
    'k_indices': list(range(n_kv_heads)),
}

# MLA (확장)
clip_config = {
    'is_mla': True,
    'head_dim': 192,                  # qk_head_dim (= qk_nope + qk_rope)
    'qk_nope_head_dim': 128,
    'qk_rope_head_dim': 64,
    'v_head_dim': 128,
    'threshold': 100.0,
    'q_indices': list(range(n_heads)),
    'k_indices': list(range(n_kv_heads)),  # build_k_indices_for_mla로 생성
}

행 인덱스 매핑 테이블

알고리즘 기호	텐서	행 범위	scale
W^h_qc	wq_b.weight	[h*qk_head_dim : h*qk_head_dim + qk_nope_head_dim]	√γ
W^h_qr	wq_b.weight	[h*qk_head_dim + qk_nope_head_dim : (h+1)*qk_head_dim]	γ
W^h_kc	wkv_b.weight	[kv_h*kv_stride : kv_h*kv_stride + qk_nope_head_dim]	√γ
k_R	wkv_a output 뒷부분	-	안 건드림

• kv_stride = qk_nope_head_dim + v_head_dim
• kv_h = h // (n_heads // n_kv_heads) (GQA head 매핑)

하이퍼파라미터

파라미터	값	비고
τ (threshold)	100	K2 full-scale 학습
τ (aggressive)	30	소규모 ablation, 성능 저하 없음 확인

참고사항

• Self-deactivation: K2에서 초기 70k step 동안 12.7%의 head만 trigger됨. 이후 모든 head의 S_max가 τ 아래로 내려가면서 자연스럽게 비활성화.
• DP/TP 환경: S^h_max를 all-reduce로 모든 rank에서 max 수집 필요.
• GQA 중복 적용 방지: 같은 KV head를 공유하는 Q head group에서 가장 작은 gamma(= 가장 큰 logit)를 기준으로 KV weight를 한 번만 scaling. compute_scales_mla에서 min(gamma) 로직으로 처리.
• wq_b_gate: attention logit이 아닌 output gate에만 관여하므로 QK-Clip 대상 아님.
• 기존 logit soft-cap: forward-level safety net으로 남겨두되, optimizer-level QK-Clip을 추가하는 것이 논문의 접근법.