freesolo-flash-dev 0.2.40__tar.gz → 0.2.41__tar.gz

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.40
+Version: 0.2.41
 Summary: Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/catalog.py

@@ -12,7 +12,7 @@ ALGORITHMS = ("sft", "grpo")
 def normalize_algorithm(value: str) -> str:
     """Canonical (lowercased, validated) algorithm name."""
     if not value:
-        value = "grpo"
+        value = "sft"
     elif not isinstance(value, str):
         # A truthy non-string (e.g. a JSON number/bool/array) would AttributeError on .lower(), which
         # escapes the callers' ValueError/ConfigError guards -> uncaught 500. Raise ValueError instead.
@@ -95,6 +95,16 @@ class ModelInfo:
     num_layers: int = 0
     hidden_size: int = 0
 
+    @property
+    def is_moe(self) -> bool:
+        """True for a mixture-of-experts model — a token routes through only a subset of experts.
+
+        Keyed off ``active_params_b`` (0.0 == dense, "every token hits every param"). Used by the
+        GRPO worker to pick REENTRANT gradient checkpointing for MoE (its router re-dispatches tokens
+        on recompute, which the non-reentrant metadata-equality assert rejects).
+        """
+        return 0.0 < self.active_params_b < self.params_b
+
     def to_dict(self) -> dict[str, Any]:
         data = asdict(self)
         serving = data.get("serving")
@@ -263,9 +273,15 @@ MODELS: dict[str, ModelInfo] = {
             # rank-64 at only 6 hot slots: the fused-MoE LoRA buffer scales with
             # max_loras x rank x num_experts, so the A100-80GB ceiling is ~max_loras x rank = 384
             # (6 x 64 fits at 99.3% util; 16 x 64 OOMs on every single/multi GPU). Serving-validated.
+            # The 6 x 64 ceiling is WEIGHT-bound, not context-bound: the FP8 checkpoint's MoE experts
+            # load as bf16 under the LoRA path (~76 GiB on the 80 GiB card), so weights + the 6 x 64
+            # buffer leave only ~0.44 GiB for KV. Canary-measured 2026-07-04: at 8192 ctx that KV pool
+            # gives 2.2x concurrency; dropping to 4096 does NOT free any LoRA-slot room (7 x 64 and
+            # 8 x 64 both OOM at 4096 exactly as at 8192) but ~doubles concurrency to ~4.4x. So context
+            # is 4096 (concurrency win, matches the training-context cap) while slots stay 6 x 64.
             max_loras=6,
             max_lora_rank=64,
-            max_model_len=8192,
+            max_model_len=4096,
             max_num_seqs=8,
             max_num_batched_tokens=4096,
             gpu_memory_utilization=0.98,
@@ -313,6 +329,26 @@ def serving_lora_rank_cap(model: str | ModelInfo | None) -> int | None:
     return int(info.serving.max_lora_rank)
 
 
+def serving_context_cap(model: str | ModelInfo | None) -> int | None:
+    """Return the model's serving ``max_model_len`` (the context it is actually served at), or None
+    when Flash has no local serving entry (open-policy / uncataloged).
+
+    A LoRA trained at a longer context than it is served wastes compute and learns positions that are
+    never used at inference, so the control plane caps a run's training context to this (see
+    ``flash.lora_rank.preflight_train_context_within_serving``). Resolution mirrors
+    ``serving_lora_rank_cap``: unknown/open-policy models return None rather than a global fallback.
+    """
+    if isinstance(model, ModelInfo):
+        info = model
+    elif isinstance(model, str) and model.strip():
+        info = MODELS.get(model.strip())
+    else:
+        info = None
+    if info is None or info.serving is None:
+        return None
+    return int(info.serving.max_model_len)
+
+
 def vocab_size_for(model_id: str) -> int:
     """Curated vocab_size for a model, or the safe default for open-model-policy entries."""
     info = MODELS.get(model_id)

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/cli/__init__.py

@@ -168,7 +168,7 @@ class _FlashParser(_ThemedParser):
         footers = [
             f"new here? run `{CLI_NAME} login`, then `{CLI_NAME} env setup`",
             f"train after publishing: `{CLI_NAME} env push --name my-env .`, "
-            f"then `{CLI_NAME} train configs/rl.toml`",
+            f"then `{CLI_NAME} train configs/sft.toml`",
             f"any command in depth: `{CLI_NAME} <command> --help`",
             "docs: https://freesolo.co/docs",
         ]
@@ -234,7 +234,22 @@ def _build_parser() -> argparse.ArgumentParser:
     env = sub.add_parser("env", help="manage Freesolo environments")
     env_sub = env.add_subparsers(dest="env_cmd", required=True)
     setup = env_sub.add_parser("setup", help="create a starter Freesolo environment scaffold")
-    setup.set_defaults(func=cmd_env_setup)
+    setup_mode = setup.add_mutually_exclusive_group()
+    setup_mode.add_argument(
+        "--single-turn",
+        dest="turn_mode",
+        action="store_const",
+        const="single",
+        help="scaffold a single-turn environment (prompt -> one response). This is the default.",
+    )
+    setup_mode.add_argument(
+        "--multi-turn",
+        dest="turn_mode",
+        action="store_const",
+        const="multi",
+        help="scaffold a multi-turn environment (bounded episode with step_episode / score_episode).",
+    )
+    setup.set_defaults(func=cmd_env_setup, turn_mode="single")
 
     env_list = env_sub.add_parser("list", help="list local environment sources")
     env_list.set_defaults(func=cmd_env_list)

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/cli/commands.py

@@ -193,15 +193,177 @@ _STARTER_DATASET_JSONL = """\
 """
 
 
+_STARTER_ENV_MULTITURN_PY = '''\
+"""Starter Freesolo multi-turn environment.
+
+A multi-turn environment runs a bounded episode: the model produces an assistant
+action, `step_episode` advances the world (optionally appending an observation
+message), and the loop repeats until `done` or `max_episode_turns`. The finished
+transcript is graded by `score_episode`.
+
+Edit dataset/train.jsonl and the episode logic, then upload with
+`flash env push --name my-env .`.
+
+A managed run should use the returned [environment] id from
+`flash env push --name my-env .`.
+
+This starter implements a tiny "guess the secret number" game so you can see the
+episode hooks wired end-to-end. Replace it with your real task before a real run.
+
+Both SFT and GRPO train off this file:
+- GRPO (configs/rl.toml) rolls out full episodes and optimizes `score_episode`.
+- SFT (configs/sft.toml) learns the gold trajectory. Provide it per row as
+  `output = {"messages": [...]}` (a full assistant/tool trajectory) or a scalar
+  `output` for a single gold assistant turn.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from freesolo.datasets.types import TaskExample
+from freesolo.environments import (
+    EnvironmentEpisode,
+    EnvironmentMultiTurn,
+    EnvironmentStepResult,
+    RewardResult,
+)
+
+
+DEFAULT_DATASET_PATH = Path(__file__).parent / "dataset" / "train.jsonl"
+
+# How many assistant guesses the model gets before the episode is forced terminal.
+MAX_TURNS = 5
+
+
+def load_jsonl(path: str | Path):
+    rows = []
+    with Path(path).open() as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                rows.append(json.loads(line))
+    return rows
+
+
+def _secret(example: TaskExample) -> int:
+    return int(str(example.output).strip())
+
+
+class StarterMultiTurnEnv(EnvironmentMultiTurn):
+    dataset = load_jsonl(DEFAULT_DATASET_PATH)
+
+    def start_episode(self, example: TaskExample, prompt_text: str):
+        # The opening prompt shown to the model. `example.input` describes the range.
+        return [
+            {
+                "role": "user",
+                "content": (
+                    f"{example.input}\\n"
+                    f"Reply with a single integer per turn. I will say 'higher', "
+                    f"'lower', or 'correct'. You have {MAX_TURNS} guesses."
+                ),
+            }
+        ]
+
+    def max_episode_turns(self, example: TaskExample) -> int:
+        return MAX_TURNS
+
+    def step_episode(
+        self,
+        example: TaskExample,
+        messages: list,
+        assistant_response: str,
+    ) -> EnvironmentStepResult:
+        # Advance the world after one assistant action. Return done=True to end the
+        # episode, or append an observation message and keep going.
+        try:
+            guess = int(assistant_response.strip().split()[0])
+        except (ValueError, IndexError):
+            return EnvironmentStepResult(
+                done=False,
+                messages=[{"role": "user", "content": "Please reply with a single integer."}],
+            )
+        secret = _secret(example)
+        if guess == secret:
+            return EnvironmentStepResult(done=True, final_response_text=str(guess))
+        hint = "higher" if guess < secret else "lower"
+        return EnvironmentStepResult(
+            done=False,
+            messages=[{"role": "user", "content": hint}],
+        )
+
+    def score_episode(
+        self,
+        example: TaskExample,
+        episode: EnvironmentEpisode,
+    ) -> RewardResult:
+        # Grade the finished transcript. Reward a correct final guess; give partial
+        # credit for getting close so GRPO has a usable gradient.
+        secret = _secret(example)
+        try:
+            final = int(str(episode.response_text).strip())
+        except ValueError:
+            return RewardResult(score=0.0, threshold=1.0)
+        if final == secret:
+            return RewardResult(score=1.0, threshold=1.0, success=True)
+        # Closeness in [0, 1): further guesses score lower.
+        closeness = max(0.0, 1.0 - abs(final - secret) / 100.0)
+        return RewardResult(score=closeness * 0.5, threshold=1.0, success=False)
+
+
+def load_environment(dataset_path: str | None = None, **kwargs) -> StarterMultiTurnEnv:
+    env = StarterMultiTurnEnv()
+    if dataset_path:
+        env.dataset = load_jsonl(dataset_path)
+    return env
+'''
+
+
+# Multi-turn rows: `input` sets up the episode, `output` is the secret number.
+# The scalar `output` also serves as the gold single-turn SFT target; swap in
+# `{"messages": [...]}` to teach a full gold trajectory.
+_STARTER_DATASET_MULTITURN_JSONL = """\
+{"input":"I picked a secret whole number between 1 and 100.","output":"42"}
+{"input":"I picked a secret whole number between 1 and 100.","output":"73"}
+"""
+
+
 def cmd_env_setup(args) -> int:
+    requested_multi = getattr(args, "turn_mode", "single") == "multi"
+    starter_env = Path("environment.py")
+    dataset = Path("dataset/train.jsonl")
+    # An existing environment.py is the authoritative signal for which turn mode this
+    # scaffold already uses (the dataset is plain JSONL with no reliable mode marker).
+    # Anchor to it so a re-run never leaves a single-turn env beside a multi-turn
+    # dataset (or vice versa); the flag only decides the mode when starting fresh.
+    existing_multi: bool | None = None
+    anchor = "environment.py"
+    if starter_env.exists():
+        existing_multi = "EnvironmentMultiTurn" in starter_env.read_text(encoding="utf-8")
+    elif dataset.exists():
+        # No env.py to anchor on, but the starter multi-turn dataset carries a
+        # distinctive prompt; use it so we don't drop a single-turn env beside it.
+        existing_multi = "secret whole number" in dataset.read_text(encoding="utf-8")
+        anchor = "dataset/train.jsonl"
+    if existing_multi is not None and existing_multi != requested_multi:
+        have = "multi-turn" if existing_multi else "single-turn"
+        want = "multi-turn" if requested_multi else "single-turn"
+        msg = (
+            f"existing {anchor} is {have}; keeping it and ignoring --{want}. "
+            f"Delete environment.py and dataset/train.jsonl first to re-scaffold as {want}."
+        )
+        print(render.warn(msg) if render.styled() else f"warning: {msg}", file=sys.stderr)
+    multi_turn = requested_multi if existing_multi is None else existing_multi
+    env_py = _STARTER_ENV_MULTITURN_PY if multi_turn else _STARTER_ENV_PY
+    dataset_jsonl = _STARTER_DATASET_MULTITURN_JSONL if multi_turn else _STARTER_DATASET_JSONL
     Path("configs").mkdir(exist_ok=True)
     Path("dataset").mkdir(exist_ok=True)
-    dataset = Path("dataset/train.jsonl")
     if not dataset.exists():
-        dataset.write_text(_STARTER_DATASET_JSONL)
-    starter_env = Path("environment.py")
+        dataset.write_text(dataset_jsonl)
     if not starter_env.exists():
-        starter_env.write_text(_STARTER_ENV_PY)
+        starter_env.write_text(env_py)
     env_comment = (
         "# Environment: upload this project folder with\n"
         "# `flash env push --name my-env .`, then paste the returned id below.\n"
@@ -231,6 +393,7 @@ def cmd_env_setup(args) -> int:
             f"{env_comment}"
             "[train]\n"
             "epochs = 1\n"
+            "max_examples = 2  # rows to train on; the starter dataset has 2 (raise as your dataset grows)\n"
             "lora_rank = 32\n"
             "# GPU and HF artifacts are managed automatically by the platform: the GPU is\n"
             "# the cheapest fitting managed class, and artifacts live in a private environment-scoped repo.\n"
@@ -242,8 +405,8 @@ def cmd_env_setup(args) -> int:
     scaffolded = [
         "environment.py",
         "dataset/train.jsonl",
-        "configs/rl.toml",
         "configs/sft.toml",
+        "configs/rl.toml",
         "TRAINING.md",
     ]
     if render.styled():

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/cli/render.py

@@ -391,7 +391,7 @@ def models_table(rows: list[dict]) -> str:
     """Supported base models — a clean themed list of ids (the CLI lists ids only)."""
     dot = _glyph("•", "-")
     ids = "\n".join(f"  {_paint(dot, _FAINT)} {_paint(r['id'], _ACCENT2)}" for r in rows)
-    foot = arrow("train one with: flash train configs/rl.toml")
+    foot = arrow("train one with: flash train configs/sft.toml")
     return _safe(f"{header('models', 'supported base models')}\n{ids}\n\n{foot}")
 
 
@@ -643,8 +643,8 @@ def env_setup(paths: list[str]) -> str:
     labels = {
         "environment.py": "env entrypoint — edit the reward + prompt",
         "dataset/train.jsonl": "starter training rows",
-        "configs/rl.toml": "GRPO run config",
         "configs/sft.toml": "SFT run config",
+        "configs/rl.toml": "GRPO run config",
         "TRAINING.md": "how to train well — read this first",
     }
     keyw = max(len(p) for p in paths)

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/cli/training_doc.py

@@ -40,8 +40,8 @@ flash gpus                           # managed GPU classes with estimated $/hr
 ```text
 environment.py          # the task: how to prompt the model and how to score it
 dataset/train.jsonl     # training rows, one JSON object per line: {"input": ..., "output": ...}
-configs/rl.toml         # a GRPO (RL) run config
 configs/sft.toml        # an SFT run config
+configs/rl.toml         # a GRPO (RL) run config
 TRAINING.md             # this file
 ```
 
@@ -92,7 +92,7 @@ edit to `environment.py` or `dataset/` so the managed run uses your change.
 
 ```toml
 model = "Qwen/Qwen3.5-4B"   # see `flash models`
-algorithm = "grpo"          # "grpo" (RL) or "sft"
+algorithm = "sft"           # "sft" (supervised) or "grpo" (RL)
 # thinking = true           # opt-in reasoning mode, for models that support it
 
 [environment]
@@ -102,10 +102,11 @@ id = "your-org/my-env"      # the id printed by `flash env push`
                                    # never stored in the spec
 
 [train]
-steps = 150                 # GRPO is step-driven; SFT is epoch-driven (epochs = N)
+epochs = 1                  # SFT is epoch-driven; GRPO is step-driven (steps = N)
+max_examples = 2            # rows to train on (the starter dataset has 2)
 lora_rank = 32
 lora_alpha = 64
-# All GRPO/SFT knobs live under [train]. Do not add [grpo] or [sft] tables.
+# All SFT/GRPO knobs live under [train]. Do not add [sft] or [grpo] tables.
 ```
 
 GPU and HF artifacts are **fully managed** — do not pick `gpu.type` or set
@@ -117,10 +118,10 @@ extra.toml` (deep-merge) and `--set key=value` (e.g. `--set train.steps=300`).
 ### 4. Submit
 
 ```bash
-flash train configs/rl.toml --dry-run   # validate the config locally — no GPU, no charge
-flash train configs/rl.toml --cost      # pre-flight USD estimate, then exit
-flash train configs/rl.toml             # submit and follow logs (Ctrl-C detaches)
-flash train configs/rl.toml --background  # submit and return immediately
+flash train configs/sft.toml --dry-run     # validate the config locally — no GPU, no charge
+flash train configs/sft.toml --cost        # pre-flight USD estimate, then exit
+flash train configs/sft.toml               # submit and follow logs (Ctrl-C detaches)
+flash train configs/sft.toml --background  # submit and return immediately
 ```
 
 ### 5. Monitor
@@ -159,9 +160,9 @@ Work in tight, attributable iterations. Each one is a hypothesis:
 1. Reconstruct state — what's the best run so far, and what have you already tried?
 2. Form a hypothesis — pick ONE lever and say WHY it will move the metric.
 3. Change that ONE lever.
-4. Validate locally — `flash train configs/rl.toml --dry-run` (catches config errors
+4. Validate locally — `flash train configs/sft.toml --dry-run` (catches config errors
    for free; a paid run on a broken config or an all-zero reward is wasted budget).
-5. Submit — `flash train configs/rl.toml`.
+5. Submit — `flash train configs/sft.toml`.
 6. Judge — read the metric trend AND a sample of real rollouts (see below).
 7. Keep the best run; revert the change if it didn't beat the noise band. Repeat.
 ```
@@ -181,7 +182,7 @@ different value.
 A run is only evidence of improvement when **all** of these hold:
 
 - [ ] The run reached `done` (confirmed via `flash status <run-id>`), not merely submitted.
-- [ ] The reward trend rose (GRPO `reward_mean`) or the SFT loss fell — **beyond the noise band**, not within it.
+- [ ] The SFT loss fell or the reward trend rose (GRPO `reward_mean`) — **beyond the noise band**, not within it.
 - [ ] You **probed the trained adapter on real inputs** (`flash deploy` + `flash chat`), including cases it should fail — not just the metrics.
 - [ ] The score is real behavior, not empty/truncated/templated outputs, skipped rows, leakage, a swallowed exception, or a format-only win.
 - [ ] If you track a clean success signal separately from the shaped reward (an explicit `RewardMetric`), *that* moved too.
@@ -219,7 +220,7 @@ spending another GPU run:
 ## Judge the run, don't just finish it
 
 - **Judge the trend, not a single number.** The proof of training is the curve:
-  `reward_mean` rising over steps (GRPO) or loss falling (SFT). Record the base/early
+  loss falling (SFT) or `reward_mean` rising over steps (GRPO). Record the base/early
   value and the final value. A flat or noisy trend with no improvement is not success.
 - **Read the model's outputs, not just the metrics.** A rising reward can come from
   reward-hacking or a degenerate output the reward still credits — metrics alone never
@@ -480,9 +481,9 @@ flash env setup                       # scaffold environment.py, dataset/, confi
 flash env push --name my-env .        # publish the environment; paste the returned id into [environment]
 flash env pull your-org/my-env        # download a published environment into the current folder
 flash env delete your-org/my-env -y   # delete a published environment
-flash train configs/rl.toml --dry-run # validate the config locally (no GPU, no charge)
-flash train configs/rl.toml --cost    # pre-flight USD estimate, then exit
-flash train configs/rl.toml           # submit and follow logs (Ctrl-C detaches; --background to skip following)
+flash train configs/sft.toml --dry-run # validate the config locally (no GPU, no charge)
+flash train configs/sft.toml --cost    # pre-flight USD estimate, then exit
+flash train configs/sft.toml           # submit and follow logs (Ctrl-C detaches; --background to skip following)
 flash status <run-id>                 # state + accrued cost
 flash log <run-id>                    # reward/loss trend + worker console/error logs
 flash log <run-id> --follow           # stream a live run to completion

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/engine/chalk_kernels.py

@@ -22,7 +22,6 @@ _KERNELS: list[tuple[str, bool]] = [
     ("trainable_attn_epilogue", True),
     ("fused_embedding", True),
     ("gdn", True),
-    ("fused_mlp", False),  # off: eval-only bf16 MLP forward, not the training activation path
     ("attn_epilogue", False),  # off (eval-only; needs q/k/v out of LoRA)
     ("fp8_frozen_base", False),  # off by default: speed/memory tradeoff, enable only after run A/B
 ]
@@ -54,15 +53,24 @@ def chalk_fused_ce_available(model_id: str | None = None) -> bool:
     return callable(_apply)
 
 
-def install_chalk_kernels(model=None) -> dict:
+def install_chalk_kernels(model=None, *, fused_ce: bool = True) -> dict:
     """Apply chalk standalone kernels to ``model``; call AFTER TRL builds the trainer.
 
+    ``fused_ce=False`` skips the fused-linear-CE kernel. flce computes the loss inside a chunked
+    LM-head+CE kernel and returns ``logits=None`` (the whole point — it never materialises the
+    ``[batch, seq, vocab]`` logits). trl's ``SFTTrainer.compute_loss`` reads ``outputs.logits`` and
+    only skips it under ``use_liger_kernel=True`` — which flash can't set, since that makes trl apply
+    Liger and clash with chalk. So the SFT (trl) path passes ``fused_ce=False`` and lets the model
+    materialise logits; the custom GRPO/opd loops read the fused loss directly and keep it on.
+
     Returns chalk's per-kernel report, or ``{}`` when freesolo-chalk isn't installed.
     """
     if model is None:
         return {}
 
     kwargs = dict(_KERNELS)
+    if not fused_ce:
+        kwargs["fused_linear_cross_entropy"] = False
     try:
         from chalk.transformers import apply_chalk_kernel_to_qwen35
     except ImportError:

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/engine/worker/__init__.py

@@ -114,6 +114,7 @@ from flash.engine.worker.perf import (
     gpu_diagnostics,
     grad_checkpointing_on,
     grpo_sleep_mode,
+    grpo_use_reentrant,
     is_cuda_oom,
     liger_on,
     loraplus_optimizer_cls,
@@ -377,6 +378,7 @@ __all__ = [
     "grpo_mask_truncated_completions",
     "grpo_overrides",
     "grpo_sleep_mode",
+    "grpo_use_reentrant",
     "heartbeat",
     "hf_api",
     "hf_prefix",

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/engine/worker/perf/__init__.py

@@ -54,6 +54,7 @@ from flash.engine.worker.perf.memory import (
     fused_optim_name,
     grad_checkpointing_on,
     grpo_sleep_mode,
+    grpo_use_reentrant,
 )
 
 
@@ -350,6 +351,7 @@ __all__ = [
     "gpu_diagnostics",
     "grad_checkpointing_on",
     "grpo_sleep_mode",
+    "grpo_use_reentrant",
     "is_cuda_oom",
     "liger_on",
     "loraplus_optimizer_cls",

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/engine/worker/perf/memory.py

@@ -71,6 +71,50 @@ def grad_checkpointing_on(
     return True
 
 
+def _is_gdn_hybrid_family(model_id: str) -> bool:
+    """Offline family check for a Qwen3.5/3.6 GatedDeltaNet hybrid (no network/config probe).
+
+    Every curated Qwen3.5/3.6 model is a GDN hybrid; the non-Qwen curated model (MiniCPM, plain
+    Llama) and uncataloged open models are not. Kept as a string check so ``grpo_use_reentrant``
+    stays pure and hermetic (callable at config-build time, unit-testable without HF access).
+    """
+    mid = (model_id or "").lower()
+    return any(token in mid for token in ("qwen3.5", "qwen3_5", "qwen3.6", "qwen3_6"))
+
+
+def grpo_use_reentrant(model_id: str) -> bool:
+    """Whether GRPO gradient checkpointing must use REENTRANT recompute for this model.
+
+    MoE models AND GatedDeltaNet (GDN) hybrids need it. Non-reentrant checkpointing
+    (``use_reentrant=False``) asserts that every recomputed activation's metadata matches the
+    forward pass and dies on the FIRST backward — before a single optimizer step — with
+    ``torch.utils.checkpoint: Recomputed values ... different metadata`` whenever a decoder layer
+    contains a custom, data-dependent kernel whose saved-for-backward tensors the recompute lays out
+    differently:
+
+    - MoE (Qwen3.6-35B-A3B): the router re-dispatches tokens on recompute, so the grouped
+      expert-buffer shapes differ (forward expert-dispatch tokens 28192 vs recompute 3524 ==
+      group_size x). This is what #429 fixed.
+    - GDN hybrids (Qwen3.5/3.6 dense): FlashAttention-2 varlen-unpad on the full-attention layers,
+      the fused GatedDeltaNet chunk-scan on the linear-attention layers, and chalk's fused Triton
+      kernels each save shape-/data-dependent tensors that the non-reentrant metadata-equality check
+      can't positionally reconcile (live-confirmed on Qwen3.5-0.8B GRPO / RTX 4090: forward packed
+      varlen ``[1636, ...]`` vs recompute padded ``[1024, ...]``). Same failure mode as MoE.
+
+    Reentrant checkpointing re-runs the forward inside the same autograd context over the same
+    closed-over inputs (mask/position_ids threaded via the ``partial``; ``use_cache=False``) and does
+    NOT assert metadata equality, so it tolerates these recomputes and produces correct gradients.
+    Non-GDN dense models (MiniCPM / plain-attention) keep the faster, lower-overhead non-reentrant
+    path.
+    """
+    from flash.catalog import MODELS
+
+    info = MODELS.get(model_id)
+    if info is not None and info.is_moe:
+        return True
+    return _is_gdn_hybrid_family(model_id)
+
+
 def grpo_sleep_mode(
     model_id: str,
     *,

{freesolo_flash_dev-0.2.40 → freesolo_flash_dev-0.2.41}/flash/engine/worker/rl.py

@@ -27,6 +27,7 @@ from flash.engine.worker.perf import (
     fused_optim_name,
     gpu_diagnostics,
     grad_checkpointing_on,
+    grpo_use_reentrant,
     optimal_attn_impl,
     setup_perf_backends,
     wait_for_gpu,
@@ -354,7 +355,10 @@ def run_rl():
         "run_name": _w.wandb_run_name(),
         "seed": _w.SEED,
         "gradient_checkpointing": grad_checkpointing_on(model_id, vllm_max_len),
-        "gradient_checkpointing_kwargs": {"use_reentrant": False},
+        # MoE needs REENTRANT recompute: its router re-dispatches tokens on the backward recompute,
+        # so non-reentrant's metadata-equality assert fires on the first backward and kills the run
+        # (Qwen3.6-35B-A3B). Dense models keep the faster non-reentrant path. See grpo_use_reentrant.
+        "gradient_checkpointing_kwargs": {"use_reentrant": grpo_use_reentrant(model_id)},
         # Pin a stable GRPO recipe instead of TRL's defaults (which suppress the lift on short runs):
         # constant LR, group-mean-centered advantages (no std scaling), no length-norm; beta = KL coef.
         "lr_scheduler_type": "constant",