표본 마스크

표본 마스크 (Sampling Mask, Distribution Replay)

RL 롤아웃을 위해 top-k/top-p 샘플링을 쓰다 보면(예: GRPO), 은근히 헷갈리는 불일치가 생겨요. 샘플러가 실제로 표본을 뽑은 분포는 잘린(truncated) 분포인데, 훈련 중 로그 확률을 계산할 때 쓰는 소프트맥스는 전체 어휘(vocabulary)를 대상으로 하거든요. 이 문서는 그 간격을 메워주는 sampling mask 기능을 설명해요.

출처: vLLM 공식 문서 — sampling_mask

배경 (Background)

이 기능은 DeepSeek-V3.2 기술 보고서의 3.3절에 설명된 Keep Sampling Mask 전략을 구현해요. 핵심 통찰은 이렇습니다. 롤아웃 샘플링 중 top-k/top-p 절단은 π_old와 π_θ의 행동 공간(action space) 사이에 불일치를 만들어, 중요도 샘플링(importance sampling)의 원리를 위반하고 훈련을 불안정하게 만든다는 거예요. π_old의 절단 마스크를 보존해 훈련 중 π_θ에 적용하면, 두 정책이 동일한 행동 부분공간을 공유하게 됩니다. DeepSeek은 top-p 샘플링과 Keep Sampling Mask 전략을 결합하면 RL 훈련 중 언어 일관성을 효과적으로 보존한다고 보고해요.

예를 들어 LLM이 롤아웃 단계에서 토큰을 생성할 때 top-k/top-p로 어휘 일부만 남기고 샘플링한다고 해볼게요. 이때 샘플러가 실제로 "고려했던" 토큰 집합이 무엇인지를 훈련 쪽에 그대로 알려주지 않으면, 훈련에서 로그 확률을 계산하는 기준과 생성에서 샘플링하는 기준이 어긋나요. sampling mask는 생성의 각 단계마다 top-k/top-p/min-p 필터링을 통과해 살아남은 토큰 ID의 정확한 집합을 돌려줘서, 훈련 쪽도 같은 지지 집합(support) 위에서 정규화할 수 있게 해줍니다.

빠른 시작 (Quick start)

vllm serve <model> \
    --return-sampling-mask \
    --logprobs-mode processed_logprobs

from vllm import LLM, SamplingParams

llm = LLM(model, return_sampling_mask=True,
           logprobs_mode="processed_logprobs")
output = llm.generate(
    "The capital of France is",
    SamplingParams(temperature=1.0, top_k=50, top_p=0.95, logprobs=1),
)
mask = output[0].outputs[0].sampling_mask
# mask.token_ids: [[187, 326, 512], [42, 88], ...]
#   mask.token_ids[i] = token IDs in the sampling support for generated token i

마스크는 /inference/v1/generate HTTP 엔드포인트를 통해서도 얻을 수 있어요.

{
  "choices": [{
    "token_ids": [187, 42, 303],
    "sampling_mask": [[187, 326, 512], [42, 88], [303, 11, 22]],
    "finish_reason": "stop"
  }]
}

요구 사항 (Requirements)

엔진은 지원하지 않는 조합을 시작 시점이나 요청 시점에 거부해요.

• 추측 디코딩 (Speculative decoding)
• Diffussion 모델
• 커스텀 로짓 프로세서 (엔진 수준 --logits-processors)

동작 방식 (How it works)

1. 샘플러가 모든 로짓 프로세서(penalties, logit bias, bad words, temperature, min-p)를 적용한 뒤 top-k/top-p 필터링을 적용하고, 제외된 로짓을 -inf로 설정해요.
2. 샘플링 후 torch.isfinite(processed_logits)로 살아남은 토큰 ID를 식별합니다. 이것이 sampling mask예요.
3. 마스크는 샘플링된 토큰과 함께 비동기적으로 GPU → CPU로 전송돼요.
4. 요청이 완료되면 단계별 마스크가 병합되고 응답을 위해 list[list[int]]로 변환됩니다.

RL 훈련에서의 사용 (RL training usage)

훈련 쪽은 중요도 비율 π_θ/π_old를 위해 두 가지가 필요해요.

π_old(a|s) — 이전 정책의 nucleus 정규화 로그 확률: --logprobs-mode processed_logprobs가 설정되면 vLLM이 이미 반환해요. log_softmax가 처리된 로짓(필터링된 토큰은 -inf) 위에서 계산되므로, 분모에는 nucleus만 포함됩니다.

π_θ(a|s) — 현재 정책의 nucleus 정규화 로그 확률: 훈련 프레임워크가 마스크를 이용해 계산해요.

# mask_ids: list[int], the sampling support for this token
# logits: the training model's raw logits for this position
keep = torch.zeros(vocab_size, dtype=torch.bool)
keep[mask_ids] = True
masked_logits = logits.masked_fill(~keep, float("-inf"))
log_prob = log_softmax(masked_logits)[sampled_token_id]

양쪽이 같은 토큰 집합 위에서 정규화되므로 중요도 비율이 일관됩니다.

제한 사항 (Limitations)

• 엔진 수준 플래그: --return-sampling-mask는 전역적으로 FlashInfer 융합 샘플러를 비활성화해요. 마스크가 필요 없는 요청도 모든 요청이 PyTorch 샘플링 경로의 비용을 부담합니다.
• 스트리밍 미지원: 마스크는 중간 스트리밍 청크가 아니라 최종 응답에서만 반환돼요.

더 알아보기 (Learn more)

• 비동기 강화학습 (Async RL)
• RLHF
• 가중치 전송 (Weight Transfer)