샘플링 마스크

샘플링 마스크 (Sampling Mask, 분포 재생)

top-k/top-p 샘플링으로 RL rollout을 생성할 때(예: GRPO), 샘플러가 실제로 추출한 잘린(truncated) 분포와 학습 중 log-probability 계산에 쓰이는 전체 어휘 softmax 사이에는 체계적 불일치가 있습니다. sampling mask 기능은 각 생성 단계에서 top-k/top-p/min-p 필터링을 통과한 정확한 토큰 ID 집합을 반환해 이 격차를 메우므로, 학습 쪽이 같은 지지집합(support) 위에서 정규화할 수 있습니다.

출처: 문서

본문

배경 (Background)

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

빠른 시작 (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)

요구사항 이유
--return-sampling-mask 엔진 레벨 옵트인 (FlashInfer 샘플러 비활성화)
--logprobs-mode processed_logprobs 반환된 logprobs가 전체 어휘가 아닌 nucleus 위에서 정규화됨
temperature > 0 Greedy에는 잘린 분포가 없음
top_k > 0 마스크 크기 제한; 순수 top-p는 어휘 크기 마스크를 만들 수 있음
Model Runner V2 비동기 D2H 복사 파이프라인에 필요

엔진은 지원하지 않는 조합을 시작 또는 요청 시점에 거부합니다:

  • Speculative decoding
  • Diffusion 모델
  • 커스텀 logits 프로세서 (엔진 레벨 --logits-processors)

동작 방식 (How it works)

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

RL 학습 사용법 (RL training usage)

학습 쪽은 importance ratio π_θ/π_old 계산에 두 가지가 필요합니다:

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

π_θ(a|s) — 현재 정책의 nucleus 정규화 logprob: 학습 프레임워크가 마스크를 사용해 계산합니다:

# 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]

양쪽이 같은 토큰 집합 위에서 정규화하므로 importance ratio가 일관됩니다.

제한 사항 (Limitations)

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

더 알아보기 (Learn more)