freesolo-flash-dev 1.0.1__tar.gz → 1.0.3__tar.gz

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/.github/workflows/bake-kernel-cache.yml

@@ -22,7 +22,7 @@ on:
     inputs:
       sms:
         description: "comma-separated sm list to bake; default is all validated arches"
-        default: "sm80,sm86,sm89,sm90,sm120"
+        default: "sm80,sm86,sm89,sm90,sm120,sm100"
         required: false
   # auto-rebake.yml calls this for the EXPENSIVE path (the cache is invalid and must be re-warmed on a
   # real GPU). It passes the arch subset, the fingerprints to stamp, and a digest-pinned base so the
@@ -105,10 +105,9 @@ jobs:
               gpu_type_id: "NVIDIA RTX PRO 6000 Blackwell Server Edition",
               allowed_cuda: "13.0",
             }
-          # Datacenter Blackwell (sm100, distinct from the sm120 RTX Pro 6000). B200 has 180 GB and
-          # good secure-cloud capacity. CUDA-13 host like the other Blackwell parts. NOT in the default
-          # `sms` list yet (B200 is unvalidated) -- bake it explicitly via `sms=sm100` until a smoke
-          # passes, then add it to the default.
+          # datacenter blackwell (sm100, distinct from the sm120 rtx pro 6000). b200 has 180 gb,
+          # good secure-cloud capacity, and is validated. promote its image only after a successful
+          # real-b200 bake, then include it in default and automatic rebakes.
           - {
               sm: sm100,
               arch: "10.0",
@@ -120,7 +119,7 @@ jobs:
       - name: Should bake ${{ matrix.sm }}?
         id: gate
         run: |
-          req="${{ inputs.sms || 'sm80,sm86,sm89,sm90,sm120' }}"
+          req="${{ inputs.sms || 'sm80,sm86,sm89,sm90,sm120,sm100' }}"
           req="${req// /}"   # tolerate spaces, e.g. "sm80, sm86"
           case ",$req," in
             *,${{ matrix.sm }},*) echo "run=true" >> "$GITHUB_OUTPUT" ;;

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.1
+Version: 1.0.3
 Summary: Flash — managed LoRA post-training (SFT/GRPO/OPD) 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-1.0.1 → freesolo_flash_dev-1.0.3}/flash/catalog.py

@@ -13,6 +13,22 @@ ALGORITHMS = ("sft", "grpo", "opd")
 # the cost model, VRAM sizing and worker allocator branch on — import ``samples_on_policy`` instead
 # of hand-rolling the ("grpo", "opd") tuple at each site.
 _ON_POLICY_ALGORITHMS = frozenset({"grpo", "opd"})
+_IMAGE_TRAINING_MODELS = frozenset(
+    {
+        "Qwen/Qwen3.5-0.8B",
+        "Qwen/Qwen3.5-2B",
+        "Qwen/Qwen3.5-4B",
+        "Qwen/Qwen3.5-9B",
+        "Qwen/Qwen3.6-27B",
+        "Qwen/Qwen3.6-35B-A3B",
+    }
+)
+
+
+def supports_image_training(model: str | ModelInfo | None) -> bool:
+    """Return whether a curated model supports image-bearing SFT, GRPO, and OPD."""
+    model_id = model.id if isinstance(model, ModelInfo) else model
+    return bool(model_id and model_id in _IMAGE_TRAINING_MODELS)
 
 
 def samples_on_policy(algorithm: str) -> bool:
@@ -146,10 +162,10 @@ class ModelInfo:
 
 DEFAULT_MODEL = "Qwen/Qwen3.5-4B"
 
-# The pre-quantized FP8 checkpoint each base model's serving engine LOADS (mirrors serving's
-# ``src.prequant_config``). Every dense model now serves a Freesolo-OWNED FP8_DYNAMIC checkpoint (no
-# community-repo dependence); the 35B VL MoE serves the OFFICIAL Qwen FP8 (it preserves the full
-# vision-language model). Informational for the catalog mirror — deploy gating reads only max_lora_rank.
+# the pre-quantized fp8 checkpoint each base model's serving engine loads (mirrors serving's
+# ``src.prequant_config``). every dense model serves a freesolo-owned fp8 checkpoint; the qwen3.6
+# 35b-a3b moe serves the official qwen fp8 variant. informational for the catalog mirror; deploy
+# gating reads only max_lora_rank.
 SERVING_FP8_MODEL_REPOS: dict[str, str] = {
     "openbmb/MiniCPM5-1B": "Freesolo-Co/MiniCPM5-1B-FP8",
     "Qwen/Qwen3.5-0.8B": "Freesolo-Co/Qwen3.5-0.8B-FP8",
@@ -157,6 +173,7 @@ SERVING_FP8_MODEL_REPOS: dict[str, str] = {
     "Qwen/Qwen3.5-4B": "Freesolo-Co/Qwen3.5-4B-FP8",
     "Qwen/Qwen3.5-9B": "Freesolo-Co/Qwen3.5-9B-FP8",
     "Qwen/Qwen3.6-35B-A3B": "Qwen/Qwen3.6-35B-A3B-FP8",
+    "Qwen/Qwen3.6-27B": "Freesolo-Co/Qwen3.6-27B-FP8",
 }
 
 MODELS: dict[str, ModelInfo] = {
@@ -182,7 +199,7 @@ MODELS: dict[str, ModelInfo] = {
     "Qwen/Qwen3.5-0.8B": ModelInfo(
         id="Qwen/Qwen3.5-0.8B",
         display_name="Qwen3.5 0.8B",
-        params="0.9B (text-only fine-tune)",
+        params="0.9B",
         params_b=0.9,
         vocab_size=248_320,
         algos=ALGORITHMS,
@@ -201,7 +218,7 @@ MODELS: dict[str, ModelInfo] = {
     "Qwen/Qwen3.5-2B": ModelInfo(
         id="Qwen/Qwen3.5-2B",
         display_name="Qwen3.5 2B",
-        params="2.3B (text-only fine-tune)",
+        params="2.3B",
         params_b=2.3,
         vocab_size=248_320,
         algos=ALGORITHMS,
@@ -219,7 +236,7 @@ MODELS: dict[str, ModelInfo] = {
     "Qwen/Qwen3.5-4B": ModelInfo(
         id="Qwen/Qwen3.5-4B",
         display_name="Qwen3.5 4B",
-        params="4.7B (text-only fine-tune)",
+        params="4.7B",
         params_b=4.7,
         vocab_size=248_320,
         algos=ALGORITHMS,
@@ -241,7 +258,7 @@ MODELS: dict[str, ModelInfo] = {
     "Qwen/Qwen3.5-9B": ModelInfo(
         id="Qwen/Qwen3.5-9B",
         display_name="Qwen3.5 9B",
-        params="9.7B (text-only fine-tune)",
+        params="9.7B",
         params_b=9.7,
         vocab_size=248_320,
         algos=ALGORITHMS,
@@ -264,6 +281,35 @@ MODELS: dict[str, ModelInfo] = {
         "(two bf16 copies + KV + the 248k-vocab fp32 logits) needs an 80 GB-class card "
         "(grpo_min_vram_gb floor).",
     ),
+    "Qwen/Qwen3.6-27B": ModelInfo(
+        id="Qwen/Qwen3.6-27B",
+        display_name="Qwen3.6 27B",
+        params="27B dense (multimodal VL, hybrid GDN)",
+        params_b=27.0,
+        num_layers=64,
+        hidden_size=5120,
+        vocab_size=248_320,
+        algos=("sft", "grpo", "opd"),
+        min_vram_gb=80,
+        sft_min_vram_gb=80,
+        quant="bf16",
+        recommended_gpu="A100 PCIe",
+        min_disk_gb=160,
+        serving=ServingCapacity(
+            gpu="H100",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.6-27B"],
+            max_loras=16,
+            max_lora_rank=64,
+            max_model_len=32768,
+            max_num_seqs=8,
+            gpu_memory_utilization=0.98,
+        ),
+        thinking="hybrid",
+        notes="Dense 27B multimodal VL checkpoint with image-capable bf16 LoRA training. SFT fits "
+        "the 80GB A100 (~54GB weights); colocated GRPO needs the B200 (trainer + vLLM rollout = two "
+        "~54GB copies). Serves the owned VL-preserving FP8 on an H100 tier (dense, so no MoE expert "
+        "LoRA-buffer multiplier).",
+    ),
     "Qwen/Qwen3.6-35B-A3B": ModelInfo(
         id="Qwen/Qwen3.6-35B-A3B",
         display_name="Qwen3.6 35B-A3B (MoE)",

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/flash/cli/commands.py

@@ -313,9 +313,8 @@ def cmd_train(args) -> int:
         runtime_secrets_from_local_env(args.config, keys=spec.environment.secrets) or None
     )
     if args.dry_run:
-        # dry-run is a faithful server-side preview: it sends the same declared secrets and runs the
-        # same config, warm-start, serving, and cost preflights as a real submit, but allocates no gpu
-        # and charges nothing. a rejection surfaces as the server's error with exit status 1.
+        # dry-run runs submit-time server preflights without importing user code, allocating a gpu,
+        # or charging anything. a rejection surfaces as the server's error with exit status 1.
         try:
             status = client.create_run(
                 payload,
@@ -330,6 +329,14 @@ def cmd_train(args) -> int:
             raise ApiError(exc.status, detail) from exc
         compatibility = status.pop("train_schema_compatibility", None)
         _print_train_schema_compatibility(compatibility)
+        print(
+            "dry-run validated: config/schema, model+algorithm compatibility, lora rank, "
+            "runtime-secret presence, warm-start source, serving context cap, and cost. it did NOT "
+            "import or run your environment.py; dataset loading, start_episode/episode shapes, "
+            "reward/scorer, worker imports, model load, and gpu/training are first exercised on the "
+            "worker after cold-start.",
+            file=sys.stderr,
+        )
         if render.styled():
             print(
                 render.object_panel(
@@ -700,7 +707,7 @@ def cmd_deployments(args) -> int:
         return 0
     print(
         f"{'RUN ID':<30}  {'STEP':<6}  {'REVISION':<40}  {'STATE':<14}  "
-        f"{'VERIFIED AT':<18}  {'OPENAI MODEL':<30}  DETAIL"
+        f"{'VERIFIED AT':<20}  {'OPENAI MODEL':<30}  DETAIL"
     )
     for row in rows:
         deployment = row.get("deployment") or {}
@@ -708,14 +715,16 @@ def cmd_deployments(args) -> int:
         step = deployment.get("checkpoint_step")
         step_text = "final" if step is None else str(step)
         verified_at = deployment.get("verified_at")
-        verified_text = "-" if verified_at is None else str(verified_at)
+        verified_text = (
+            "-" if verified_at is None else (render._humanize_ts(verified_at) or str(verified_at))
+        )
         revision = str(deployment.get("adapter_revision") or "-")
         state = str(deployment.get("state") or "-")
         openai_model = str(deployment.get("openai_model") or run_id)
         detail = str(deployment.get("error") or deployment.get("detail") or "")[:160]
         print(
             f"{run_id:<30}  {step_text:<6}  {revision:<40}  {state:<14}  "
-            f"{verified_text:<18}  {openai_model:<30}  {detail}"
+            f"{verified_text:<20}  {openai_model:<30}  {detail}"
         )
     return 0
 

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/flash/cli/render.py

@@ -488,7 +488,10 @@ def deployments_table(rows: list[dict]) -> str:
                 ("final" if step is None else str(step), _TEAL),
                 (str(deployment.get("adapter_revision") or "-"), _ACCENT2),
                 (state, color),
-                ("-" if verified_at is None else str(verified_at), _GRAY),
+                (
+                    "-" if verified_at is None else (_humanize_ts(verified_at) or str(verified_at)),
+                    _GRAY,
+                ),
                 (str(deployment.get("openai_model") or run_id), _GREEN),
                 (detail, _GRAY),
             ]

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/flash/cli/training_doc.py

@@ -154,6 +154,47 @@ flash undeploy <run-id>          # tear the endpoint down
 flash export --adapter-id <run-id> --repository <you>/<repo>  # copy adapter weights to your HF repo
 ```
 
+### Loading an exported adapter locally (transformers + peft)
+
+Flash trains the Qwen3.5/3.6 family against the full multimodal module tree, so adapter
+weights saved by the worker carry a `language_model.` infix in their tensor keys
+(`base_model.model.model.language_model.layers.*`). During `flash export`, adapters that
+can be represented in the text-only namespace are normalized to
+`base_model.model.model.layers.*`. When an export retains non-LM tensors, it keeps the
+multimodal namespace instead. Do not infer the namespace solely from whether `flash export`
+was used: inspect the keys. For normalized keys, load with vanilla peft and
+`AutoModelForCausalLM`:
+
+```python
+from peft import PeftModel
+from transformers import AutoModelForCausalLM
+
+base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-0.8B")  # qwen3.5 causal-lm class
+model = PeftModel.from_pretrained(base, "<you>/<repo>")
+```
+
+**Before you trust any local eval, check the key namespace matches the model class you
+loaded.** peft does NOT error on mismatched keys — it emits a `UserWarning` about missing
+adapter keys and silently applies *nothing*, so you would benchmark the bare base model
+believing it is your adapter. If the inspected keys carry the `language_model.` infix,
+load the base with the multimodal class whose parameters live under that namespace instead:
+
+```python
+# infixed keys (base_model.model.model.language_model.layers.*) need the multimodal class:
+from transformers import AutoModelForImageTextToText
+
+base = AutoModelForImageTextToText.from_pretrained("Qwen/Qwen3.5-0.8B")
+# qwen3.5/3.6 dense models resolve to Qwen3_5ForConditionalGeneration
+# qwen3.6-35b-a3b (moe) resolves to Qwen3_5MoeForConditionalGeneration
+model = PeftModel.from_pretrained(base, "<you>/<repo>")
+```
+
+To check which form you have, read the safetensors key names (stdlib, no torch needed):
+the JSON header of `adapter_model.safetensors` lists every tensor key. Rule of thumb:
+`model.layers.*` keys pair with `AutoModelForCausalLM`; `model.language_model.layers.*`
+keys pair with the `*ForConditionalGeneration` class. After loading, confirm the adapter
+actually applied: outputs (or a probe metric) must differ from the bare base model.
+
 The rest of this file is about doing the above *well* — designing a reward that teaches,
 and deciding honestly whether a run improved.
 
@@ -167,9 +208,11 @@ 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 — `flash train configs/sft.toml --dry-run` (server-side preview: catches config
-   errors, serving rank/context caps, and warm-start rank mismatches for free — no GPU, no charge;
-   a paid run on a broken config or an all-zero reward is wasted budget).
+4. Validate with `flash train configs/sft.toml --dry-run`. This server-side preview checks
+   config/schema, model+algorithm compatibility, LoRA rank, runtime-secret presence, warm-start
+   source, serving context cap, and cost for free, with no GPU or charge. It does not import or run
+   `environment.py`; dataset loading, episode shapes, reward/scorer, worker imports, model load, and
+   GPU/training are first exercised on the worker after cold-start.
 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.
@@ -206,7 +249,7 @@ spending another GPU run:
 
 | Issue | Symptom | Mitigation |
 | --- | --- | --- |
-| Environment id is blank or stale | `flash train --dry-run` fails, or the worker uses old reward/data | Run `flash env push --name my-env .` after every environment/data edit and paste the returned id into every config you submit. |
+| Environment id is blank or stale | A blank id fails config validation. A stale published id can pass `flash train --dry-run` because dry-run does not import or run `environment.py`; the worker then uses old reward/data. | Run `flash env push --name my-env .` after every environment/data edit and paste the returned id into every config you submit. |
 | Local-only env path in config | Config validation says there is no local path mode | Publish first, then use the returned slug in `[environment] id`. `flash train` only runs published env ids, not local paths. |
 | Config knobs are in the wrong table | Validation rejects `[grpo]`, `[sft]`, or unknown `[train]` keys | Put `epochs`, `group_size`, `max_completion_tokens`, `temperature`, `max_context_tokens`, LoRA, and other training knobs under `[train]`. |
 | Trying to pin managed infrastructure | `gpu.type`, `train.hf_repo`, or `model_policy` changes do not do what you expected | Treat GPU choice, model policy, and the run artifact repo as managed. Tune the model, algorithm, environment, and `[train]` knobs instead. |
@@ -221,6 +264,7 @@ spending another GPU run:
 | Run looks stuck after disconnecting | Terminal stopped streaming but the job may still be alive | Ctrl-C detaches. Use `flash log <run-id> --follow` to reattach, `flash log <run-id>` for the console/error output, or `flash cancel <run-id>` if you intentionally want to stop it. |
 | Final checkpoint regresses | Last step is worse than an earlier checkpoint | Run `flash checkpoints <run-id>`, deploy a specific step with `flash deploy <run-id>/step-N`, and compare with held-out probes before exporting or relying on the final adapter. |
 | Export fails before upload | CLI says no HuggingFace token | Pass `flash export --api-key hf_...`, or set `HF_TOKEN` in your shell, `.env`, or `.env.local`. Exports are private unless you pass `--public`. |
+| Exported adapter is a silent no-op locally | peft warns about missing adapter keys and local eval matches the bare base model | The adapter's key namespace does not match the loaded model class. `model.layers.*` keys pair with `AutoModelForCausalLM`; `model.language_model.layers.*` keys pair with `Qwen3_5ForConditionalGeneration` / `Qwen3_5MoeForConditionalGeneration` (via `AutoModelForImageTextToText`). See "Loading an exported adapter locally". |
 | SFT loss improves but quality does not | Train loss falls while held-out behavior stalls or degrades | Keep a held-out split outside training. Deploy and score that split; if quality drops, reduce epochs or improve data instead of adding more passes. |
 | Cost surprises | A quick experiment uses more GPU time than intended | Start with `--dry-run` and `--cost`, keep `epochs` and `max_examples` small for smoke tests, and scale only after reward/data wiring is proven. Setup time is reported for observability; customer cost is based on training-loop GPU time. |
 

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/flash/engine/multiturn_rollout.py

@@ -326,6 +326,7 @@ class _RolloutState:
         "done",
         "env_mask",
         "example",
+        "images",
         "logprobs",
         "messages",
         "prompt_ids",
@@ -333,8 +334,9 @@ class _RolloutState:
         "turns",
     )
 
-    def __init__(self, example, messages, prompt_ids, state, budget):
+    def __init__(self, example, messages, prompt_ids, state, budget, images=None):
         self.example = example
+        self.images = images
         self.messages = messages
         self.prompt_ids = prompt_ids
         self.cur_ids = list(prompt_ids)  # invariant: cur_ids == prompt_ids + completion_ids so far
@@ -398,23 +400,56 @@ def _advance_after_turn(
     r.cur_ids.extend(glue)
 
 
+def _extract_prompt_images(messages: list[dict]) -> list[object]:
+    """Extract decoded PIL images from TRL-injected prompt content blocks."""
+    from PIL import Image
+
+    images = []
+    for message in messages:
+        content = message.get("content")
+        if not isinstance(content, list):
+            continue
+        for part in content:
+            if not isinstance(part, dict) or part.get("type") not in {
+                "image",
+                "image_url",
+                "input_image",
+            }:
+                continue
+            image = part.get("image")
+            if not isinstance(image, Image.Image):
+                raise ValueError(
+                    "multimodal rollout image blocks must carry a PIL image at part['image']"
+                )
+            images.append(image)
+    return images
+
+
 def _build_rollout_states(
     examples: list[dict],
     active_env,
     render: Callable[[list, bool], list[int]],
     engine_max_len: int | None,
+    *,
+    extract_images: bool = False,
+    rollout_prompts: list[list[dict]] | None = None,
 ) -> list[_RolloutState]:
     """Initialise one :class:`_RolloutState` per example for :func:`rollout_async`."""
+    if rollout_prompts is not None and len(rollout_prompts) != len(examples):
+        raise ValueError("rollout prompts must align one-to-one with examples")
     rollouts: list[_RolloutState] = []
-    for example in examples:
+    for index, example in enumerate(examples):
         state = active_env.new_rollout_state(example)
         initial_messages = state.get("prompt") or state.get("messages")
+        if rollout_prompts is not None:
+            initial_messages = rollout_prompts[index]
         if not isinstance(initial_messages, list):
             raise KeyError("multi-turn rollout state must include prompt or messages")
         messages = [dict(m) for m in initial_messages]
+        images = _extract_prompt_images(messages) if extract_images else None
         prompt_ids = render(messages, True)
         budget = (engine_max_len - len(prompt_ids) - 8) if engine_max_len else None
-        rollouts.append(_RolloutState(example, messages, prompt_ids, state, budget))
+        rollouts.append(_RolloutState(example, messages, prompt_ids, state, budget, images))
     return rollouts
 
 
@@ -481,6 +516,7 @@ class _LogicalRequest:
     prefix_ids: tuple[int, ...]
     max_tokens: int
     initial: bool
+    images: list[object] | None
     attempts: int = 0
     started_at: float = 0.0
 
@@ -490,7 +526,7 @@ def rollout_async(
     examples: list[dict],
     active_env,
     render: Callable[[list, bool], list[int]],
-    submit: Callable[[str, list[int], int, bool], None],
+    submit: Callable[[str, list[int], int, bool, list[object] | None], None],
     poll: Callable[[], list[RolloutCompletion]],
     busy: Callable[[], bool],
     abort: Callable[[list[str]], None],
@@ -498,6 +534,8 @@ def rollout_async(
     max_turns: int,
     per_turn_max_tokens: int,
     engine_max_len: int | None = None,
+    extract_images: bool = False,
+    rollout_prompts: list[list[dict]] | None = None,
     request_timeout_seconds: float | None = None,
     request_max_attempts: int = 2,
     monotonic: Callable[[], float] = time.monotonic,
@@ -516,7 +554,14 @@ def rollout_async(
     if request_timeout_seconds is not None and request_timeout_seconds <= 0:
         raise ValueError("request_timeout_seconds must be positive when set")
 
-    rollouts = _build_rollout_states(examples, active_env, render, engine_max_len)
+    rollouts = _build_rollout_states(
+        examples,
+        active_env,
+        render,
+        engine_max_len,
+        extract_images=extract_images,
+        rollout_prompts=rollout_prompts,
+    )
     by_id: dict[str, _LogicalRequest] = {}
     to_env: queue.Queue = queue.Queue()
     to_submit: queue.Queue = queue.Queue()
@@ -527,7 +572,13 @@ def rollout_async(
         logical.started_at = monotonic()
         by_id[req_id] = logical
         try:
-            submit(req_id, list(logical.prefix_ids), logical.max_tokens, logical.initial)
+            submit(
+                req_id,
+                list(logical.prefix_ids),
+                logical.max_tokens,
+                logical.initial,
+                logical.images,
+            )
         except Exception:
             by_id.pop(req_id, None)
             with contextlib.suppress(Exception):
@@ -541,6 +592,7 @@ def rollout_async(
                 prefix_ids=tuple(prefix),
                 max_tokens=int(max_new),
                 initial=bool(initial),
+                images=r.images,
             )
         )
 
@@ -718,6 +770,8 @@ def build_rollout_func(
     active_env,
     tok,
     examples_by_key: dict,
+    processor=None,
+    multimodal: bool = False,
     max_completion: int,
     max_turns: int,
     temperature: float,
@@ -751,18 +805,47 @@ def build_rollout_func(
 
         _SOParams = StructuredOutputsParams
 
-    _render_cache = _LRUCache(8192)
-
-    def render(messages: list, add_generation_prompt: bool) -> list[int]:
-        cache_key = (
-            f"{add_generation_prompt}\x00{json.dumps(messages, sort_keys=True, default=str)}"
+    image_pad_id = None
+    if multimodal:
+        if processor is None:
+            raise ValueError("multimodal rollout requires a processor")
+        from flash.multimodal import (
+            collapse_image_pad_runs,
+            multimodal_prompt_key,
+            resolve_image_pad_token_id,
         )
-        cached = _render_cache.get(cache_key)
-        if cached is not None:
-            return cached
-        ids = render_message_ids(tok, messages, add_generation_prompt, thinking=thinking)
-        _render_cache.put(cache_key, ids)
-        return ids
+
+        image_pad_id = resolve_image_pad_token_id(processor, tok)
+
+        def render(messages: list, add_generation_prompt: bool) -> list[int]:
+            rendered = processor.apply_chat_template(
+                messages,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=True,
+                return_dict=True,
+                enable_thinking=thinking,
+            )
+            input_ids = rendered["input_ids"]
+            if hasattr(input_ids, "tolist"):
+                input_ids = input_ids.tolist()
+            ids = list(input_ids)
+            if ids and isinstance(ids[0], (list, tuple)):
+                ids = list(ids[0])
+            return [int(token_id) for token_id in ids]
+
+    else:
+        _render_cache = _LRUCache(8192)
+
+        def render(messages: list, add_generation_prompt: bool) -> list[int]:
+            cache_key = (
+                f"{add_generation_prompt}\x00{json.dumps(messages, sort_keys=True, default=str)}"
+            )
+            cached = _render_cache.get(cache_key)
+            if cached is not None:
+                return cached
+            ids = render_message_ids(tok, messages, add_generation_prompt, thinking=thinking)
+            _render_cache.put(cache_key, ids)
+            return ids
 
     env_glue = make_env_glue(tok, thinking=thinking)
 
@@ -774,7 +857,13 @@ def build_rollout_func(
         vocab_size = _engine_vocab_size(engine)
         active_ids: set[str] = set()
 
-        def submit(req_id: str, prefix_ids: list[int], max_tokens: int, initial: bool) -> None:
+        def submit(
+            req_id: str,
+            prefix_ids: list[int],
+            max_tokens: int,
+            initial: bool,
+            images: list[object] | None,
+        ) -> None:
             """Enqueue one assistant-turn request."""
             if not prefix_ids:
                 raise ValueError(
@@ -801,9 +890,21 @@ def build_rollout_func(
                 # env contract has no per-turn schema channel; unconstrained env/tool turns are
                 # not generated here.
                 sp_kwargs["structured_outputs"] = _SOParams(**copy.deepcopy(structured_outputs))
-            llm_engine.add_request(
-                req_id, {"prompt_token_ids": list(prefix_ids)}, SamplingParams(**sp_kwargs)
-            )
+            if images:
+                if image_pad_id is None:
+                    raise ValueError("multimodal rollout is missing the image-pad token id")
+                collapsed = collapse_image_pad_runs(list(prefix_ids), image_pad_id, len(images))
+                prompt = {
+                    "prompt_token_ids": collapsed,
+                    "multi_modal_data": {"image": images if len(images) > 1 else images[0]},
+                }
+                # suppress the image-pad token during generation: a stray pad emitted into the
+                # assistant turn would corrupt the next turn's collapse_image_pad_runs count and
+                # misalign vision data, aborting the rollout. mirrors the single-turn opd guard.
+                sp_kwargs["logit_bias"] = {int(image_pad_id): -100.0}
+            else:
+                prompt = {"prompt_token_ids": list(prefix_ids)}
+            llm_engine.add_request(req_id, prompt, SamplingParams(**sp_kwargs))
             active_ids.add(req_id)
 
         def poll() -> list[RolloutCompletion]:
@@ -838,7 +939,12 @@ def build_rollout_func(
             if sleep_mode:
                 engine.wake_up(tags=["kv_cache"])
                 woke = True
-            examples = [examples_by_key.get(_prompt_key(p), {"prompt": p}) for p in prompts]
+            if multimodal:
+                examples = [
+                    examples_by_key.get(multimodal_prompt_key(p), {"prompt": p}) for p in prompts
+                ]
+            else:
+                examples = [examples_by_key.get(_prompt_key(p), {"prompt": p}) for p in prompts]
             rollouts = rollout_async(
                 examples=examples,
                 active_env=active_env,
@@ -850,6 +956,8 @@ def build_rollout_func(
                 max_turns=max_turns,
                 per_turn_max_tokens=max_completion,
                 engine_max_len=engine_max_len,
+                extract_images=multimodal,
+                rollout_prompts=prompts if multimodal else None,
                 abort=abort,
                 request_timeout_seconds=request_timeout_seconds,
                 request_max_attempts=request_max_attempts,

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/flash/engine/recipe.py

@@ -35,6 +35,7 @@ class TeacherModel:
     # serverless models — confirm against the live fireworks.ai pricing before relying on the
     # `flash train` cost quote.
     usd_per_1m: tuple[float, float]
+    supports_images: bool = False
 
 
 # The alias used when [train] omits teacher_model — the historical fixed teacher, unchanged.
@@ -64,6 +65,7 @@ TEACHER_MODELS: dict[str, TeacherModel] = {
         model_id="accounts/fireworks/models/kimi-k2p6",
         display_name="Kimi K2.6",
         usd_per_1m=(0.95, 4.00),
+        supports_images=True,
     ),
 }
 
@@ -75,6 +77,11 @@ def normalize_teacher_alias(value: str) -> str:
     return "-".join(str(value).strip().lower().replace("_", " ").replace("-", " ").split())
 
 
+def teacher_supports_images(value: str) -> bool:
+    """Return whether the resolved managed teacher accepts image-conditioned echo scoring."""
+    return resolve_teacher(value).supports_images
+
+
 def resolve_teacher(value: str) -> TeacherModel:
     """Resolve a ``[train] teacher_model`` value to its catalog entry.
 

{freesolo_flash_dev-1.0.1 → freesolo_flash_dev-1.0.3}/flash/engine/worker/grpo.py

@@ -45,10 +45,16 @@ def build_grpo_prompt_dataset(prompts: list[dict]) -> tuple[list[dict], list]:
 
     Dataset.from_list infers one type per field across ALL rows; mixed-type metadata
     (e.g. int vs str in the same field) causes ArrowInvalid at RL startup. Fix: keep only
-    trivially-typed columns (prompt + integer example_idx); reward_fn maps the index back.
+    trivially typed columns (prompt, integer example_idx, and optional image descriptor strings);
+    reward_fn maps the index back.
     """
     examples = [p["example"] for p in prompts]
-    rows = [{"prompt": p["prompt"], "example_idx": i} for i, p in enumerate(prompts)]
+    rows = []
+    for i, prompt in enumerate(prompts):
+        row = {"prompt": prompt["prompt"], "example_idx": i}
+        if "images" in prompt:
+            row["images"] = list(prompt["images"])
+        rows.append(row)
     return rows, examples