mirror of https://github.com/razor-ai/soup.git
97 lines
3.6 KiB
Python
97 lines
3.6 KiB
Python
"""GRPO long-context + memory-efficient RL — v0.50.0 Part B.
|
|
|
|
Schema helpers for two unsloth/axolotl features:
|
|
|
|
1. ``long_context_grpo`` — wires Tiled MLP (v0.56.0 Part A) when available,
|
|
otherwise accepts as-is so users can opt in early. Schema-only in v0.50.0;
|
|
live Tiled MLP integration deferred to v0.56.0.
|
|
2. ``vllm_sleep_mode`` — between-rollouts vLLM standby (memory savings during
|
|
the optimisation step). Requires vLLM ≥ 0.7 and is delegated to the vLLM
|
|
AsyncEngineArgs ``enable_sleep_mode`` flag.
|
|
|
|
Security:
|
|
- Both flags are strict bools (project bool-as-int policy enforced at the
|
|
schema layer via Pydantic).
|
|
- ``validate_long_context_grpo_compat`` rejects multi-stage combinations
|
|
that would silently no-op (e.g. ring-attention + long_context_grpo on a
|
|
non-Llama base).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
# Backends where vLLM is genuinely available (mirrors v0.30.0 vllm.py policy).
|
|
_VLLM_SUPPORTED_BACKENDS: frozenset[str] = frozenset({"transformers", "unsloth"})
|
|
|
|
|
|
def validate_long_context_grpo_compat(
|
|
*,
|
|
task: str,
|
|
backend: str,
|
|
use_ring_attention: bool,
|
|
) -> None:
|
|
"""Schema-time gate for ``long_context_grpo=True``.
|
|
|
|
Rejects on:
|
|
- non-GRPO task (long_context_grpo only meaningful for RL rollouts);
|
|
- mlx backend (no Tiled MLP support);
|
|
- ``use_ring_attention=True`` (both rewrite attention — pick one;
|
|
mirrors v0.49.0 LongLoRA + ring-attention exclusivity).
|
|
"""
|
|
if not isinstance(task, str) or not task:
|
|
raise ValueError("task must be a non-empty string")
|
|
if "\x00" in task:
|
|
raise ValueError("task must not contain null bytes")
|
|
if task != "grpo":
|
|
raise ValueError(
|
|
f"long_context_grpo requires task='grpo'; got task={task!r}"
|
|
)
|
|
if not isinstance(backend, str) or not backend:
|
|
raise ValueError("backend must be a non-empty string")
|
|
if "\x00" in backend:
|
|
raise ValueError("backend must not contain null bytes")
|
|
if backend == "mlx":
|
|
raise ValueError(
|
|
"long_context_grpo is not supported on backend=mlx in v0.50.0"
|
|
)
|
|
if not isinstance(use_ring_attention, bool):
|
|
raise ValueError(
|
|
"use_ring_attention must be a bool, got "
|
|
f"{type(use_ring_attention).__name__}"
|
|
)
|
|
if use_ring_attention:
|
|
raise ValueError(
|
|
"long_context_grpo and use_ring_attention are mutually "
|
|
"exclusive — pick one (both rewrite the attention kernel)"
|
|
)
|
|
|
|
|
|
def validate_vllm_sleep_mode_compat(*, backend: str) -> None:
|
|
"""Schema-time gate for ``vllm_sleep_mode=True``.
|
|
|
|
Sleep mode is a vLLM-only feature; rejected on backends that do not
|
|
use vLLM for rollouts.
|
|
"""
|
|
if not isinstance(backend, str) or not backend:
|
|
raise ValueError("backend must be a non-empty string")
|
|
if "\x00" in backend:
|
|
raise ValueError("backend must not contain null bytes")
|
|
if backend not in _VLLM_SUPPORTED_BACKENDS:
|
|
raise ValueError(
|
|
f"vllm_sleep_mode requires backend in {sorted(_VLLM_SUPPORTED_BACKENDS)}; "
|
|
f"got backend={backend!r}"
|
|
)
|
|
|
|
|
|
def apply_vllm_sleep_mode(engine_args: object) -> None: # noqa: ARG001
|
|
"""Live wiring for vLLM sleep mode — deferred to v0.50.1.
|
|
|
|
The schema flag is accepted in v0.50.0 but the actual
|
|
``AsyncEngineArgs.enable_sleep_mode=True`` plumbing is wired in
|
|
v0.50.1 (mirrors v0.27.0 MII stub-then-live pattern).
|
|
"""
|
|
raise NotImplementedError(
|
|
"vllm_sleep_mode live wiring deferred to v0.50.1; "
|
|
"schema flag accepts the value but rollout-time sleep/wake "
|
|
"is not yet wired into the vLLM engine factory."
|
|
)
|