freesolo-flash-dev 1.2.48__tar.gz → 1.2.49__tar.gz

{freesolo_flash_dev-1.2.48 → freesolo_flash_dev-1.2.49}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.48
+Version: 1.2.49
 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.2.48 → freesolo_flash_dev-1.2.49}/flash/cli/scaffold/TRAINING.md

@@ -844,6 +844,15 @@ Pick SFT when you already have good answers and want the model to imitate them.
   `max_context_tokens` that plausibly fits prompt + completion, and only raise it when you see
   truncation (outputs cut off mid-thought, degraded loss). A bigger context just costs
   more.
+- **Multi-turn targets train on the assistant turns only.** When `sft_completion` returns a
+  full target transcript (assistant turns interleaved with the environment's tool/user
+  observations), Flash supervises the assistant turns and masks the observations out of the
+  loss — the model is trained to produce its own replies, not to predict the environment's.
+  For every ChatML target, including a single assistant turn, control strings `<|im_start|>` and
+  `<|im_end|>` are reserved in every message field the active template renders into the body,
+  including text content, reasoning, and tool-call serialization. They are indistinguishable from
+  structural turn delimiters after tokenization. The run logs role-aware masked rows and any
+  completion-only fallback rows separately.
 - **For Qwen3.5 thinking multi-turn SFT, put reasoning only in the final assistant
   turn.** Qwen3.5's chat template strips literal `<think>` blocks from prior assistant
   history and pre-opens `<think>\n` in the next generation prompt. If every assistant

{freesolo_flash_dev-1.2.48 → freesolo_flash_dev-1.2.49}/flash/engine/profiling/sft_workload.py

@@ -50,6 +50,8 @@ class PreparedSftWorkload:
     sampled_texts: list[str]
     multiturn_targets: int
     coerced_singleturn_targets: int
+    role_aware_multiturn_targets: int
+    fallback_multiturn_targets: int
     authored_reasoning_turns: int
     rendered_reasoning_spans: int
     truncated_reasoning_spans: int
@@ -104,7 +106,8 @@ def _processor_tokenized_row(
     *,
     max_length: int,
     thinking: bool,
-) -> tuple[list[int], list[int], bytes, int]:
+) -> tuple[list[int], list[int], bytes, int, bool]:
+    from flash.engine.worker.model.chatml_mask import assistant_only_mask
     from flash.engine.worker.model.packing import completion_mask_from_ids
 
     prepared_prompt = _multimodal_messages_with_images(prompt_messages, images)
@@ -143,9 +146,22 @@ def _processor_tokenized_row(
     untruncated_length = len(untruncated_ids)
     input_ids = untruncated_ids[:max_length]
     prompt_ids = ids(prompt["input_ids"])[:max_length]
-    loss_mask = completion_mask_from_ids(prompt_ids, input_ids)
+    # the span reader needs the text tokenizer's vocabulary; a processor exposes it as .tokenizer.
+    mask, role_aware = assistant_only_mask(
+        completion_mask_from_ids(prompt_ids, input_ids),
+        input_ids,
+        getattr(processor, "tokenizer", processor),
+        completion_messages,
+        # this path appends nothing of its own, so there is no appended EOS to preserve.
+        appended_eos=False,
+        template_source=processor,
+        # the IMAGE-PREPARED messages, which is what the processor actually rendered: probing the
+        # unprepared list would validate a different transcript than the one being masked.
+        source_messages=full_messages,
+        template_kwargs={"enable_thinking": thinking},
+    )
     full.pop("attention_mask", None)
-    return input_ids, loss_mask, _serialize_multimodal_inputs(full), untruncated_length
+    return input_ids, mask, _serialize_multimodal_inputs(full), untruncated_length, role_aware
 
 
 def _materialize_verl_images(
@@ -269,6 +285,10 @@ class _TokenizedSftRows:
     sampled_texts: list[str]
     multiturn_targets: int
     coerced_singleturn_targets: int
+    # multi-turn rows only: present means the row has a multi-turn target, and the value says whether
+    # its mask was role-aware or fell back to the contiguous span. one map rather than a set plus a
+    # parallel dict, so a row can never be counted as multi-turn while its mask status is missing.
+    multiturn_mask_applied: dict[int, bool]
     reasoning_by_index: dict[int, _RowReasoning]
     dropped: int
 
@@ -280,6 +300,8 @@ class _RetainedSftRows:
     authored_reasoning_turns: int
     rendered_reasoning_spans: int
     truncated_reasoning_spans: int
+    role_aware_multiturn_targets: int
+    fallback_multiturn_targets: int
     dropped: int
     # kept PER ROW, in the retained order, so the warning can be bounded to the rows the optimizer
     # actually consumes. the totals above describe the whole retained dataset, which is what the
@@ -386,6 +408,48 @@ def _row_reasoning(
     return _RowReasoning(authored, rendered, truncated)
 
 
+def _text_sft_row_spec(
+    spec,
+    prompt_messages: list[dict],
+    completion_messages: list[dict],
+    *,
+    tokenizer,
+    render_transcript: Callable[[list[dict]], str],
+    max_length: int,
+    row_index: int,
+    multiturn: bool,
+) -> tuple[str, _RowReasoning, dict[str, Any]]:
+    text = render_transcript([*prompt_messages, *completion_messages])
+    prompt_text = tokenizer.apply_chat_template(
+        prompt_messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=spec.thinking,
+    )
+    reasoning = _row_reasoning(
+        prompt_messages,
+        completion_messages,
+        render=render_transcript,
+        tokenizer=tokenizer,
+        max_length=max_length,
+    )
+    return (
+        text,
+        reasoning,
+        {
+            "text": text,
+            "prompt_text": prompt_text,
+            "target_messages": completion_messages,
+            "source_messages": [*prompt_messages, *completion_messages],
+            "template_kwargs": {"enable_thinking": spec.thinking},
+            # measurement carried alongside the row rather than in it: the parquet row is built from
+            # an explicit key list below, so neither key reaches verl.
+            "row_index": row_index,
+            "multiturn": multiturn,
+        },
+    )
+
+
 def _tokenize_prompt_rows(
     spec,
     prompt_rows: list[tuple[Any, list[dict], list[dict], bool]],
@@ -409,6 +473,7 @@ def _tokenize_prompt_rows(
     sampled_texts: list[str] = []
     multiturn_targets = 0
     coerced_singleturn_targets = 0
+    multiturn_mask_applied: dict[int, bool] = {}
     # kept per row rather than summed here, for the same reason as `untruncated_by_index`: rows that
     # lose their whole completion to the cap are dropped below, and folding their reasoning into a
     # running total would report loss from rows the run never trains on against a retained-row
@@ -430,7 +495,9 @@ def _tokenize_prompt_rows(
         coerced_scalar_output,
     ) in enumerate(prompt_rows):
         _reject_image_completion(completion_messages)
-        if len(completion_messages) > 1:
+        # read before the image branch rewrites `completion_messages` to its text-only form.
+        multiturn = len(completion_messages) > 1
+        if multiturn:
             multiturn_targets += 1
         elif (
             coerced_scalar_output
@@ -444,7 +511,13 @@ def _tokenize_prompt_rows(
             normalized = normalize_prompt_images(example, prompt_messages, package_root)
             completion_messages = text_only_prompt_messages(completion_messages)
             decoded_images = decode_image_descriptors(normalized.descriptors, package_root)
-            input_ids, loss_mask, multimodal_inputs, untruncated_length = _processor_tokenized_row(
+            (
+                input_ids,
+                loss_mask,
+                multimodal_inputs,
+                untruncated_length,
+                assistant_mask_applied,
+            ) = _processor_tokenized_row(
                 processor,
                 normalized.messages,
                 completion_messages,
@@ -453,6 +526,8 @@ def _tokenize_prompt_rows(
                 thinking=bool(spec.thinking),
             )
             untruncated_by_index[row_index] = untruncated_length
+            if multiturn:
+                multiturn_mask_applied[row_index] = assistant_mask_applied
             row_by_index[row_index] = {
                 "input_ids": input_ids,
                 "loss_mask": loss_mask,
@@ -480,22 +555,19 @@ def _tokenize_prompt_rows(
                 max_length=max_length - visual_inflation,
             )
         else:
-            text = render_transcript([*prompt_messages, *completion_messages])
-            prompt_text = tokenizer.apply_chat_template(
-                prompt_messages,
-                tokenize=False,
-                add_generation_prompt=True,
-                enable_thinking=spec.thinking,
-            )
-            sampled_texts.append(text)
-            reasoning_by_index[row_index] = _row_reasoning(
+            text, reasoning, text_spec = _text_sft_row_spec(
+                spec,
                 prompt_messages,
                 completion_messages,
-                render=render_transcript,
                 tokenizer=tokenizer,
+                render_transcript=render_transcript,
                 max_length=max_length,
+                row_index=row_index,
+                multiturn=multiturn,
             )
-            text_specs.append({"text": text, "prompt_text": prompt_text, "row_index": row_index})
+            sampled_texts.append(text)
+            reasoning_by_index[row_index] = reasoning
+            text_specs.append(text_spec)
 
     dropped = 0
     if text_specs:
@@ -506,9 +578,12 @@ def _tokenize_prompt_rows(
         )
         dropped += text_dropped
         for spec_row, tokenized in zip(kept_specs, tokenized_rows, strict=True):
+            row_index = spec_row["row_index"]
             input_ids = tokenized["input_ids"]
-            untruncated_by_index[spec_row["row_index"]] = tokenized["untruncated_length"]
-            row_by_index[spec_row["row_index"]] = {
+            untruncated_by_index[row_index] = tokenized["untruncated_length"]
+            if spec_row["multiturn"]:
+                multiturn_mask_applied[row_index] = tokenized["assistant_mask_applied"]
+            row_by_index[row_index] = {
                 "input_ids": input_ids,
                 "loss_mask": tokenized["completion_mask"],
                 "images": [],
@@ -520,6 +595,7 @@ def _tokenize_prompt_rows(
         sampled_texts=sampled_texts,
         multiturn_targets=multiturn_targets,
         coerced_singleturn_targets=coerced_singleturn_targets,
+        multiturn_mask_applied=multiturn_mask_applied,
         reasoning_by_index=reasoning_by_index,
         dropped=dropped,
     )
@@ -533,6 +609,8 @@ def _filter_retained_rows(tokenized: _TokenizedSftRows, tokenizer) -> _RetainedS
     authored_reasoning = 0
     rendered_reasoning = 0
     truncated_reasoning = 0
+    role_aware_multiturn_targets = 0
+    fallback_multiturn_targets = 0
     per_row: list[_RowReasoning] = []
     dropped = tokenized.dropped
     for row_index in sorted(tokenized.row_by_index):
@@ -542,6 +620,11 @@ def _filter_retained_rows(tokenized: _TokenizedSftRows, tokenizer) -> _RetainedS
             # appended in lockstep with the row it measures, so the truncation counts below describe
             # the rows that are actually trained on rather than the ones that were dropped.
             retained_untruncated.append(tokenized.untruncated_by_index[row_index])
+            if row_index in tokenized.multiturn_mask_applied:
+                if tokenized.multiturn_mask_applied[row_index]:
+                    role_aware_multiturn_targets += 1
+                else:
+                    fallback_multiturn_targets += 1
             # summed here for the same reason, so the reasoning-loss warning describes the rows the
             # run trains on. a dropped row contributes neither its authored reasoning nor its
             # survivors, which would otherwise be reported against a retained-row denominator.
@@ -566,6 +649,8 @@ def _filter_retained_rows(tokenized: _TokenizedSftRows, tokenizer) -> _RetainedS
         authored_reasoning_turns=authored_reasoning,
         rendered_reasoning_spans=rendered_reasoning,
         truncated_reasoning_spans=truncated_reasoning,
+        role_aware_multiturn_targets=role_aware_multiturn_targets,
+        fallback_multiturn_targets=fallback_multiturn_targets,
         dropped=dropped,
         row_reasoning=per_row,
     )
@@ -895,6 +980,8 @@ def prepare_sft_workload(
         sampled_texts=tokenized.sampled_texts,
         multiturn_targets=tokenized.multiturn_targets,
         coerced_singleturn_targets=tokenized.coerced_singleturn_targets,
+        role_aware_multiturn_targets=retained.role_aware_multiturn_targets,
+        fallback_multiturn_targets=retained.fallback_multiturn_targets,
         authored_reasoning_turns=retained.authored_reasoning_turns,
         rendered_reasoning_spans=retained.rendered_reasoning_spans,
         truncated_reasoning_spans=retained.truncated_reasoning_spans,

{freesolo_flash_dev-1.2.48 → freesolo_flash_dev-1.2.49}/flash/engine/worker/entry/sft.py

@@ -10,8 +10,10 @@ from __future__ import annotations
 
 import random
 
+from flash.engine.worker.model.chatml_mask import assistant_only_mask
 from flash.engine.worker.model.packing import (
     completion_mask_from_ids,
+    packing_appends_eos,
     tokenize_for_packing,
     untruncated_lengths_for_packing,
 )
@@ -37,7 +39,8 @@ def _pretokenize_completion_only(texts, tokenizer, max_length):
 
     Returns ``(kept_texts, pretok, n_dropped)``. Each pretok row also carries
     ``untruncated_length``: its token count BEFORE the cap. The post-slice length can never exceed
-    max_length, so it cannot say whether the cap actually bound; this can.
+    max_length, so it cannot say whether the cap actually bound; this can. It also carries
+    ``assistant_mask_applied``: whether the render was readable enough to mask non-assistant turns.
     """
     full_ids = tokenize_for_packing([t["text"] for t in texts], tokenizer, max_length)
     # the same encode without the cap, so a truncated row reports its real size rather than the cap.
@@ -45,14 +48,27 @@ def _pretokenize_completion_only(texts, tokenizer, max_length):
     prompt_ids = tokenizer(
         [t["prompt_text"] for t in texts], truncation=True, max_length=max_length
     )["input_ids"]
-    pretok = [
-        {
-            "input_ids": ids,
-            "completion_mask": completion_mask_from_ids(pids, ids),
-            "untruncated_length": length,
-        }
-        for ids, pids, length in zip(full_ids, prompt_ids, untruncated, strict=True)
-    ]
+    pretok = []
+    for text, ids, pids, length in zip(texts, full_ids, prompt_ids, untruncated, strict=True):
+        mask, role_aware = assistant_only_mask(
+            completion_mask_from_ids(pids, ids),
+            ids,
+            tokenizer,
+            text["target_messages"],
+            # the appended EOS counts only when it SURVIVED the cap: an equal-length row proves the
+            # token that terminates the example is still there to supervise.
+            appended_eos=packing_appends_eos(text["text"], tokenizer) and len(ids) == length,
+            source_messages=text["source_messages"],
+            template_kwargs=text["template_kwargs"],
+        )
+        pretok.append(
+            {
+                "input_ids": ids,
+                "completion_mask": mask,
+                "assistant_mask_applied": role_aware,
+                "untruncated_length": length,
+            }
+        )
     special_ids = set(getattr(tokenizer, "all_special_ids", None) or [])
     kept = [
         (t, r)