256 lines
13 KiB
Python
256 lines
13 KiB
Python
"""think 태그 문법 강제 vLLM V1 로짓 프로세서 (2026-07-16).
|
||
|
||
근거 (temp0.75 1K×64 전수 실측): 생성의 40.1%가 think를 제대로 못 열고(34.3%는
|
||
</think>부터 시작), 정상 시작조차 태그를 평균 2.6개 사용(재개방·유사 멀티턴). 태그
|
||
방향 의미론이 학습되지 않아 성능과 무관하게 형식이 붕괴함 — RL이 이를 그대로 강화
|
||
하기 전에 문법을 생성 단계에서 강제한다.
|
||
|
||
규칙 (토큰 id 상태머신 — 디코드 불필요):
|
||
1) think가 열려 있으면 <think> 재호출 금지 (중첩/재개방 방지)
|
||
2) </think>가 1회 등장한 순간부터 <think>·</think> 모두 영구 금지
|
||
3) (옵션) <|im_start|> 금지 — 새 턴 환각 차단
|
||
<think>\n 프리필(ralo.custom_prompts.official_chat_think_prefill_prompt_fn)과 결합 시
|
||
문법이 완전 폐쇄된다: 열림 1회(프리필 보장) + 닫힘 정확히 1회 + 이후 답변부.
|
||
|
||
사용법 (boxed_eos/repetition_abort와 병행 등록 가능):
|
||
1) 엔진: vllm_kwargs.logits_processors: ["ralo.vllm_think_format:ThinkFormatLogitsProcessor"]
|
||
2) 요청: SamplingParams.extra_args = {"think_format": {
|
||
"think_open_id": <int>, # <think> 토큰 id
|
||
"think_close_id": <int>, # </think> 토큰 id
|
||
"prefilled_open": true, # 프롬프트가 <think>로 끝나는 경우 (프리필)
|
||
"ban_im_start": true, # <|im_start|> 재호출 금지 (옵션)
|
||
"im_start_id": <int>,
|
||
}}
|
||
extra_args에 think_format이 없는 요청은 완전히 무시된다.
|
||
|
||
주의 — async scheduling 비호환 (2026-07-16 실측): vLLM V1은 async scheduling을 기본
|
||
자동 활성화하는데, 이때 워커가 output_tok_ids에 실제 토큰 대신 -1 플레이스홀더를
|
||
채운다(gpu_model_runner의 use_async_scheduling 경로). 출력 토큰 값을 읽는 커스텀
|
||
프로세서(이 파일 + vllm_boxed_eos + vllm_repetition_abort)는 전부 무력화된다.
|
||
반드시 엔진에 `async_scheduling=False`를 함께 줘야 한다. 플레이스홀더가 감지되면
|
||
아래 상태머신이 1회 경고를 남긴다.
|
||
"""
|
||
|
||
import logging
|
||
|
||
from typing import Optional
|
||
|
||
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_OK = False
|
||
LogitsProcessor = object # type: ignore
|
||
BatchUpdate = None # type: ignore
|
||
|
||
logger = logging.getLogger(__name__)
|
||
_warned_placeholder = False
|
||
|
||
|
||
class ThinkFormatState:
|
||
"""토큰 id만으로 금지 목록을 결정하는 상태머신 (HF/vLLM 공용 코어)."""
|
||
|
||
__slots__ = ("open_id", "close_id", "im_start_id", "opened", "closed", "consumed",
|
||
"out", "max_think_tokens", "think_len", "force_prefill",
|
||
"force_whitelist", "_force_start")
|
||
|
||
def __init__(self, out, open_id, close_id, prefilled_open=False, im_start_id=None,
|
||
prefilled_closed=False, max_think_tokens=None,
|
||
force_prefill=None, force_whitelist=None):
|
||
self.out = out # 생성 토큰 리스트 (vLLM 라이브 참조 / HF에선 수동 feed)
|
||
self.open_id = int(open_id)
|
||
self.close_id = int(close_id)
|
||
self.im_start_id = int(im_start_id) if im_start_id is not None else None
|
||
# prefilled_closed: 프롬프트(프리픽스)에 이미 <think>…</think>가 완결되어 있는 경우
|
||
# (예: lrs/MCMC 청크 재개 — 누적 텍스트를 프롬프트로 넘기는 후속 요청). 이걸 안 주면
|
||
# 요청마다 상태가 리셋되어 닫힌 뒤에도 </think> 재방출이 허용된다 (2026-07-21 실측).
|
||
self.opened = bool(prefilled_open) or bool(prefilled_closed)
|
||
self.closed = bool(prefilled_closed)
|
||
self.consumed = 0
|
||
# max_think_tokens: <think> 안에서 이 토큰 수를 넘으면 강제 봉합
|
||
# (non-convergent 추론 루프 탈출 → 답변 단계로 밀어냄). None이면 비활성.
|
||
self.max_think_tokens = int(max_think_tokens) if max_think_tokens else None
|
||
self.think_len = 0 # <think> 열린 뒤 생성된 토큰 수
|
||
# 강제 봉합 프리필 시퀀스(예: [</think>, "\n\n"])를 순서대로 강제한 뒤,
|
||
# 첫 답변 토큰은 force_whitelist(학습데이터 top-N 답변시작 토큰)로만 허용 →
|
||
# 모델이 그중 최고를 스스로 고르게. 둘 다 없으면 </think> 하나만 강제(구 동작).
|
||
self.force_prefill = [int(x) for x in force_prefill] if force_prefill else [self.close_id]
|
||
self.force_whitelist = [int(x) for x in force_whitelist] if force_whitelist else None
|
||
self._force_start = None # 강제 봉합 시작 시점의 out 길이
|
||
|
||
def advance(self):
|
||
"""새 토큰 소비 후 현재 시점의 금지 토큰 id 리스트 반환."""
|
||
global _warned_placeholder
|
||
while self.consumed < len(self.out):
|
||
t = self.out[self.consumed]
|
||
self.consumed += 1
|
||
if t == -1 and not _warned_placeholder:
|
||
_warned_placeholder = True
|
||
logger.warning(
|
||
"[ThinkFormat] output_tok_ids에 -1 플레이스홀더 감지 — vLLM async "
|
||
"scheduling이 켜져 있어 닫힘 감지가 불가능합니다. 엔진에 "
|
||
"async_scheduling=False를 전달하세요.")
|
||
if self.opened and not self.closed:
|
||
self.think_len += 1
|
||
if t == self.open_id:
|
||
self.opened = True
|
||
elif t == self.close_id and self.opened:
|
||
self.closed = True
|
||
return self.banned_ids()
|
||
|
||
def force_allowed_ids(self):
|
||
"""강제 봉합 진행 중이면 이번 스텝 허용 토큰 id 리스트, 아니면 None.
|
||
|
||
프리필 시퀀스를 순서대로 1개씩 강제 → 끝나면 답변 첫 토큰을 화이트리스트로
|
||
1스텝 제한 → 그 뒤 해제(None). 진입 후 self.closed가 True가 돼도 _force_start
|
||
기준으로 계속 진행한다."""
|
||
if self.max_think_tokens is None:
|
||
return None
|
||
if self._force_start is None:
|
||
if (self.opened and not self.closed
|
||
and self.think_len >= self.max_think_tokens):
|
||
self._force_start = len(self.out) # 강제 봉합 개시
|
||
else:
|
||
return None
|
||
progress = len(self.out) - self._force_start
|
||
seq = self.force_prefill
|
||
if progress < len(seq):
|
||
return [seq[progress]] # 프리필: 그 자리 토큰만 허용
|
||
if self.force_whitelist and progress == len(seq):
|
||
return list(self.force_whitelist) # 답변 첫 토큰: top-N만 허용
|
||
return None # 프리필+화이트리스트 끝 → 해제
|
||
|
||
def banned_ids(self):
|
||
banned = []
|
||
if self.closed:
|
||
banned = [self.open_id, self.close_id] # 닫힌 후엔 둘 다 영구 금지
|
||
elif self.opened:
|
||
banned = [self.open_id] # 열려 있는 동안 재개방 금지
|
||
if self.im_start_id is not None:
|
||
banned.append(self.im_start_id)
|
||
return banned
|
||
|
||
|
||
class ThinkFormatLogitsProcessor(LogitsProcessor):
|
||
"""think 태그 문법 강제 — 금지 토큰 로짓만 -inf, 그 외 무변경."""
|
||
|
||
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, ThinkFormatState] = {}
|
||
self._rows: list[int] = []
|
||
self._cols: list[int] = []
|
||
self._rows_t = None
|
||
self._cols_t = None
|
||
# 강제 봉합: force_rows는 전체 -inf, (allow_rows, allow_cols)만 0으로 살림
|
||
self._force_rows_t = None
|
||
self._allow_rows_t = None
|
||
self._allow_cols_t = None
|
||
|
||
def is_argmax_invariant(self) -> bool:
|
||
return False
|
||
|
||
@staticmethod
|
||
def add_request(params: "SamplingParams", _prompt, output_tok_ids) -> Optional[ThinkFormatState]:
|
||
cfg = (getattr(params, "extra_args", None) or {}).get("think_format")
|
||
if not cfg or cfg.get("think_open_id") is None or cfg.get("think_close_id") is None:
|
||
return None
|
||
return ThinkFormatState(
|
||
out=output_tok_ids,
|
||
open_id=cfg["think_open_id"],
|
||
close_id=cfg["think_close_id"],
|
||
prefilled_open=bool(cfg.get("prefilled_open", False)),
|
||
prefilled_closed=bool(cfg.get("prefilled_closed", False)),
|
||
im_start_id=cfg.get("im_start_id") if cfg.get("ban_im_start", False) else None,
|
||
max_think_tokens=cfg.get("max_think_tokens"),
|
||
force_prefill=cfg.get("force_prefill"),
|
||
force_whitelist=cfg.get("force_whitelist"),
|
||
)
|
||
|
||
def update_state(self, batch_update: "BatchUpdate | None") -> None:
|
||
process_dict_updates(self.states, batch_update, self.add_request)
|
||
rows, cols = [], [] # 일반 금지(-inf)
|
||
force_rows, allow_rows, allow_cols = [], [], [] # 강제: 행 전체 -inf 후 allow만 0
|
||
for idx, st in self.states.items():
|
||
banned = st.advance()
|
||
allowed = st.force_allowed_ids()
|
||
if allowed is not None:
|
||
force_rows.append(idx)
|
||
for tid in allowed:
|
||
allow_rows.append(idx)
|
||
allow_cols.append(tid)
|
||
else:
|
||
for tid in banned:
|
||
rows.append(idx)
|
||
cols.append(tid)
|
||
|
||
def _t(vals):
|
||
return torch.tensor(vals, device="cpu", dtype=torch.int64,
|
||
pin_memory=self.pin_memory).to(self.device, non_blocking=True)
|
||
self._rows_t, self._cols_t = (_t(rows), _t(cols)) if rows else (None, None)
|
||
if force_rows:
|
||
self._force_rows_t = _t(force_rows)
|
||
self._allow_rows_t = _t(allow_rows)
|
||
self._allow_cols_t = _t(allow_cols)
|
||
else:
|
||
self._force_rows_t = self._allow_rows_t = self._allow_cols_t = None
|
||
|
||
def apply(self, logits: torch.Tensor) -> torch.Tensor:
|
||
if self._rows_t is not None:
|
||
logits[self._rows_t, self._cols_t] = float("-inf")
|
||
if self._force_rows_t is not None:
|
||
# 강제 봉합: 해당 행 전체 -inf 후 허용 토큰만 0 (프리필=1개, 화이트리스트=N개)
|
||
logits[self._force_rows_t] = float("-inf")
|
||
logits[self._allow_rows_t, self._allow_cols_t] = 0.0
|
||
return logits
|
||
|
||
|
||
def build_think_format_extra_args(algo_cfg: dict, tokenizer,
|
||
prefilled_open: bool = False) -> Optional[dict]:
|
||
"""dapo_kwargs/lrs_kwargs의 think_format 블록 → SamplingParams.extra_args."""
|
||
cfg = (algo_cfg or {}).get("think_format") or {}
|
||
if not cfg.get("enabled", False):
|
||
return None
|
||
open_id = tokenizer.convert_tokens_to_ids("<think>")
|
||
close_id = tokenizer.convert_tokens_to_ids("</think>")
|
||
im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
||
if open_id is None or close_id is None:
|
||
return None
|
||
|
||
# 강제 봉합 프리필: </think> + "\n\n" 시퀀스 후, 답변 첫 토큰을 학습데이터
|
||
# top-N 답변시작 토큰(화이트리스트)으로만 허용해 모델이 스스로 고르게 한다.
|
||
# (2026-07-21 amq3 clean 292K 실측: </think> 뒤 100% "\n\n", 첫 실토큰 top20이 97%)
|
||
force_prefill = force_whitelist = None
|
||
if cfg.get("force_close_prefill"):
|
||
nn = tokenizer("\n\n", add_special_tokens=False)["input_ids"]
|
||
force_prefill = [int(close_id)] + [int(x) for x in nn]
|
||
# config에서 직접 준 화이트리스트(id 리스트) 우선, 없으면 학습데이터 기반 기본값.
|
||
# 사람이름(John/Mary/James/Maria)·단일문자는 제외 — 답변 요약이 아니라 스토리
|
||
# 문제 재서술로 빠질 수 있어서. 일반 답변시작 토큰만 (커버리지 ~96%):
|
||
# To/We/Let/The/Given/###/(/First/In
|
||
force_whitelist = cfg.get("force_close_whitelist") or [
|
||
1249, 1654, 10061, 785, 22043, 14374, 7, 5338, 641,
|
||
]
|
||
force_whitelist = [int(x) for x in force_whitelist]
|
||
|
||
return {
|
||
"think_format": {
|
||
"think_open_id": int(open_id),
|
||
"think_close_id": int(close_id),
|
||
"prefilled_open": bool(cfg.get("prefilled_open", prefilled_open)),
|
||
"prefilled_closed": bool(cfg.get("prefilled_closed", False)),
|
||
"ban_im_start": bool(cfg.get("ban_im_start", True)),
|
||
"im_start_id": int(im_start_id) if im_start_id is not None else None,
|
||
"max_think_tokens": (int(cfg["max_think_tokens"])
|
||
if cfg.get("max_think_tokens") else None),
|
||
"force_prefill": force_prefill,
|
||
"force_whitelist": force_whitelist,
|
||
}
|
||
}
|