180 lines
7.9 KiB
Python
180 lines
7.9 KiB
Python
"""과도 반복(폭주) 롤아웃 조기 종료용 vLLM V1 커스텀 logits processor.
|
|
|
|
배경: long-CoT eval 실측에서 롤아웃의 ~35%가 20480 캡까지 도달하며, 상당수가
|
|
n-gram 반복 루프(어차피 reward에서 오답 처리됨)다. 이들이 가장 비싼(긴 컨텍스트)
|
|
디코드 토큰을 소모한다.
|
|
|
|
동작: 요청별로 생성 토큰 목록(vLLM이 라이브 참조로 제공)을 감시하다가
|
|
n-gram 반복 비율이 임계 이상이면 해당 요청의 로짓을 EOS만 남기고 -inf로 마스킹
|
|
→ 그 요청만 즉시 종료. **트리거 전에는 로짓을 전혀 건드리지 않으므로 정상
|
|
롤아웃의 출력은 비트 단위로 동일하다** (출력 불변 원칙 하에 폭주 커팅만 승인됨,
|
|
2026-07-09 결정). 트리거된 롤아웃은 짧게 잘린 채 수집되고 기존 repetition
|
|
필터/reward가 오답 처리한다.
|
|
|
|
사용법:
|
|
1) 엔진: vllm_kwargs.logits_processors: ["ralo.vllm_repetition_abort:RepetitionEosLogitsProcessor"]
|
|
2) 요청: SamplingParams.extra_args = {"repetition_abort": {
|
|
"eos_token_id": <int>, # 필수 — 강제할 EOS
|
|
"ngram": 7, # n-gram 크기 (dapo 필터와 동일 계열)
|
|
"threshold": 0.9, # 반복 비율 임계 (필터 0.8보다 보수적 기본값 —
|
|
# 출력이 실제로 잘리는 개입이므로)
|
|
"min_tokens": 2048, # 이 길이 전에는 검사하지 않음
|
|
"check_interval": 512, # 검사 주기 (토큰)
|
|
}}
|
|
extra_args에 repetition_abort가 없는 요청은 완전히 무시된다.
|
|
|
|
주의 (2026-07-16 실측): vLLM V1의 async scheduling(기본 자동 활성화)에서는 워커가
|
|
output_tok_ids에 실제 토큰 대신 -1 플레이스홀더를 채워 이 프로세서가 무력화된다.
|
|
반드시 엔진에 `async_scheduling=False`를 함께 전달할 것 (think_format·boxed_eos 공통).
|
|
"""
|
|
|
|
from typing import Optional, Sequence
|
|
|
|
import torch
|
|
|
|
try:
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor
|
|
from vllm.v1.sample.logits_processor.builtin import process_dict_updates
|
|
_VLLM_OK = True
|
|
except ImportError: # vLLM 미설치 환경(트레이너 전용 노드 등)에서도 임포트 가능하게
|
|
_VLLM_OK = False
|
|
LogitsProcessor = object # type: ignore
|
|
BatchUpdate = None # type: ignore
|
|
|
|
|
|
def ngram_repetition_ratio(token_ids: Sequence[int], n: int) -> float:
|
|
"""반복 n-gram에 속하는 "위치"의 비율 (0..1).
|
|
|
|
주의: dapo/sampler.py의 필터 함수와 공식이 다르다. 필터는 "2회 이상 나온 고유
|
|
n-gram 개수 ÷ 위치 수"인데, 이 정의는 순수 반복 루프에서 오히려 0에 수렴한다
|
|
(고유 n-gram이 몇 개 안 되므로) — 재보정 주석의 실측치(정상 max 0.58, 폭주
|
|
0.975)와 부합하는 것은 여기 구현한 "반복 n-gram이 차지하는 위치 비율" 쪽이다.
|
|
순수 루프 → ~1.0, 무반복 텍스트 → 0.0.
|
|
"""
|
|
if n <= 0 or len(token_ids) < n * 2:
|
|
return 0.0
|
|
counts: dict = {}
|
|
total = 0
|
|
for i in range(len(token_ids) - n + 1):
|
|
total += 1
|
|
key = tuple(token_ids[i : i + n])
|
|
counts[key] = counts.get(key, 0) + 1
|
|
if total == 0:
|
|
return 0.0
|
|
repeated_positions = sum(c for c in counts.values() if c > 1)
|
|
return repeated_positions / max(total, 1)
|
|
|
|
|
|
class _ReqState:
|
|
__slots__ = ("out", "eos", "ngram", "threshold", "interval", "next_check", "triggered")
|
|
|
|
def __init__(self, out, eos, ngram, threshold, min_tokens, interval):
|
|
self.out = out # vLLM이 넘겨주는 라이브 출력 토큰 리스트 참조
|
|
self.eos = int(eos)
|
|
self.ngram = int(ngram)
|
|
self.threshold = float(threshold)
|
|
self.interval = max(1, int(interval))
|
|
self.next_check = max(1, int(min_tokens))
|
|
self.triggered = False
|
|
|
|
def should_trigger(self) -> bool:
|
|
"""검사 시점 도달 시 반복 비율 평가. 트리거되면 True (한 번만)."""
|
|
if self.triggered:
|
|
return False
|
|
if len(self.out) < self.next_check:
|
|
return False
|
|
ratio = ngram_repetition_ratio(self.out, self.ngram)
|
|
if ratio >= self.threshold:
|
|
self.triggered = True
|
|
return True
|
|
self.next_check = len(self.out) + self.interval
|
|
return False
|
|
|
|
|
|
class RepetitionEosLogitsProcessor(LogitsProcessor):
|
|
"""반복 폭주 감지 시 해당 요청 로짓을 EOS 단일화 — 다른 요청/트리거 전 요청은 무변경."""
|
|
|
|
def __init__(self, vllm_config, device: torch.device, is_pin_memory: bool):
|
|
if not _VLLM_OK:
|
|
raise RuntimeError("vLLM V1 logits processor API를 찾을 수 없음")
|
|
self.device = device
|
|
self.pin_memory = is_pin_memory
|
|
self.states: dict[int, _ReqState] = {}
|
|
# 트리거된 (배치 인덱스, eos id) 마스크 텐서
|
|
self._rows = self._tensor([], torch.int64)
|
|
self._eos = self._tensor([], torch.int64)
|
|
self._have_triggered = False
|
|
|
|
def _tensor(self, data, dtype):
|
|
return torch.tensor(data, device="cpu", dtype=dtype, pin_memory=self.pin_memory).to(
|
|
device=self.device, non_blocking=True
|
|
)
|
|
|
|
def is_argmax_invariant(self) -> bool:
|
|
return False # 트리거 시 argmax를 EOS로 바꾼다
|
|
|
|
@staticmethod
|
|
def add_request(params: "SamplingParams", _prompt, output_tok_ids) -> Optional[_ReqState]:
|
|
cfg = (getattr(params, "extra_args", None) or {}).get("repetition_abort")
|
|
if not cfg or cfg.get("eos_token_id") is None:
|
|
return None
|
|
return _ReqState(
|
|
out=output_tok_ids,
|
|
eos=cfg["eos_token_id"],
|
|
ngram=cfg.get("ngram", 7),
|
|
threshold=cfg.get("threshold", 0.9),
|
|
min_tokens=cfg.get("min_tokens", 2048),
|
|
interval=cfg.get("check_interval", 512),
|
|
)
|
|
|
|
def update_state(self, batch_update: "BatchUpdate | None") -> None:
|
|
changed = process_dict_updates(self.states, batch_update, self.add_request)
|
|
|
|
newly = False
|
|
for state in self.states.values():
|
|
if state.should_trigger():
|
|
newly = True
|
|
|
|
if changed or newly:
|
|
rows, eos = [], []
|
|
for idx, state in self.states.items():
|
|
if state.triggered:
|
|
rows.append(idx)
|
|
eos.append(state.eos)
|
|
self._rows = self._tensor(rows, torch.int64)
|
|
self._eos = self._tensor(eos, torch.int64)
|
|
self._have_triggered = bool(rows)
|
|
if newly:
|
|
try:
|
|
print(f"[RepetitionAbort] {len(rows)} request(s) forced to EOS "
|
|
f"(runaway n-gram repetition)", flush=True)
|
|
except Exception:
|
|
pass
|
|
|
|
def apply(self, logits: torch.Tensor) -> torch.Tensor:
|
|
if self._have_triggered and self._rows.numel() > 0:
|
|
# 트리거 행 전체 -inf 후 EOS만 0 → softmax에서 EOS 확률 1
|
|
logits[self._rows] = float("-inf")
|
|
logits[self._rows, self._eos] = 0.0
|
|
return logits
|
|
|
|
|
|
def build_repetition_abort_extra_args(algo_cfg: dict, eos_token_id) -> Optional[dict]:
|
|
"""dapo_kwargs/lrs_kwargs의 repetition_abort 블록 → SamplingParams.extra_args.
|
|
|
|
블록이 없거나 enabled가 아니면 None (프로세서가 해당 요청을 완전히 무시).
|
|
"""
|
|
cfg = (algo_cfg or {}).get("repetition_abort") or {}
|
|
if not cfg.get("enabled", False) or eos_token_id is None:
|
|
return None
|
|
return {
|
|
"repetition_abort": {
|
|
"eos_token_id": int(eos_token_id),
|
|
"ngram": int(cfg.get("ngram", 7)),
|
|
"threshold": float(cfg.get("threshold", 0.9)),
|
|
"min_tokens": int(cfg.get("min_tokens", 2048)),
|
|
"check_interval": int(cfg.get("check_interval", 512)),
|
|
}
|
|
}
|