初始化项目，由ModelHub XC社区提供模型

Model: jepetolee/Qwen3-4B-AMQ3-Math-SFT Source: Original Platform
2026-07-27 06:54:11 +08:00
commit 1a95d2533b
16 changed files with 152933 additions and 0 deletions
--- a/vllm_repetition_abort.py
+++ b/vllm_repetition_abort.py
@@ -0,0 +1,179 @@
+"""과도 반복(폭주) 롤아웃 조기 종료용 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)),
+        }
+    }