feat: add --sft-assistant-turns option for multi-turn SFT

areno/api/data_utils.py

@@ -51,6 +51,110 @@ def response_to_tokens_and_mask(
     return prompt_ids + response_ids, [True] * len(prompt_ids) + [False] * len(response_ids)
 
 
+def _try_chat_template_encoding(
+    messages: list[dict[str, Any]], tokenizer, trainable_assistant_indices: set[int]
+) -> tuple[list[int], list[bool]] | None:
+    """Attempt incremental chat-template encoding; return None if not prefix-stable.
+
+    Tries to encode the conversation turn by turn using the tokenizer's
+    ``chat_template``.  If the re-encoded prefix ever differs from what was
+    already accumulated, the tokenizer is not prefix-stable and we return
+    ``None`` so the caller can fall back to plain-text concatenation.
+    """
+
+    if not getattr(tokenizer, "chat_template", None):
+        return None
+
+    tokens: list[int] = []
+    mask: list[bool] = []
+    for i in range(len(messages)):
+        partial_ids = normalize_token_ids(
+            apply_chat_template_with_options(
+                tokenizer, messages[: i + 1], tokenize=True, add_generation_prompt=False
+            )
+        )
+        # Guard against tokenizers whose chat_template is not prefix-stable.
+        if tokens and partial_ids[: len(tokens)] != tokens:
+            import warnings
+
+            warnings.warn(
+                "tokenizer chat_template is not prefix-stable; "
+                "falling back to plain-text encoding for multi-turn SFT",
+                RuntimeWarning,
+                stacklevel=3,
+            )
+            return None
+        # Only keep tokens added by the current turn.
+        new_tokens = partial_ids[len(tokens):]
+        role = messages[i].get("role", "user")
+        is_trainable = role == "assistant" and i in trainable_assistant_indices
+        tokens.extend(new_tokens)
+        mask.extend([not is_trainable] * len(new_tokens))
+    return tokens, mask
+
+
+def messages_to_tokens_and_mask(
+    messages: list[dict[str, Any]],
+    tokenizer,
+    eos_token_id: int,
+    *,
+    last_assistant_only: bool = False,
+) -> tuple[list[int], list[bool]]:
+    """Encode multi-turn chat messages into tokens with a training mask.
+
+    The mask follows the same convention as
+    :func:`prompt_response_to_tokens_and_mask`: ``True`` means "do not train"
+    (prompt context), ``False`` means "train" (assistant response).
+
+    * user / system / tool turns are always masked out (``True``).
+    * assistant turns are trainable (``False``) unless *last_assistant_only*
+      is set, in which case only the final assistant turn is trainable and
+      earlier assistant turns are treated as context (``True``).
+
+    The function uses the tokenizer chat template when available so turn
+    markers and special tokens match the model's expected format. For base
+    tokenizers without a chat template, a plain-text fallback concatenates
+    ``role: content`` per turn.
+
+    EOS is appended after the last message if not already present, so the
+    model learns to stop.
+    """
+
+    # Determine which assistant turns are trainable.
+    assistant_indices = [
+        i for i, msg in enumerate(messages) if msg.get("role") == "assistant"
+    ]
+    if last_assistant_only and assistant_indices:
+        trainable_assistant_indices = {assistant_indices[-1]}
+    else:
+        trainable_assistant_indices = set(assistant_indices)
+
+    chat_template_tokens = _try_chat_template_encoding(
+        messages, tokenizer, trainable_assistant_indices
+    )
+    if chat_template_tokens is not None:
+        tokens, mask = chat_template_tokens
+    else:
+        # Plain-text fallback: concatenate "role: content" per turn.
+        tokens = []
+        mask = []
+        for i, msg in enumerate(messages):
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            turn_text = f"{role}: {content}"
+            turn_ids = normalize_token_ids(tokenizer.encode(turn_text, add_special_tokens=False))
+            is_trainable = role == "assistant" and i in trainable_assistant_indices
+            tokens.extend(turn_ids)
+            mask.extend([not is_trainable] * len(turn_ids))
+
+    # Append EOS if not already present so the model learns to stop.
+    if eos_token_id is not None and (not tokens or tokens[-1] != eos_token_id):
+        tokens.append(eos_token_id)
+        mask.append(False)
+
+    return tokens, mask
+
+
 def has_any(record: dict[str, Any], keys: tuple[str, ...]) -> bool:
     """Return whether a record has any string field in keys."""
 

areno/api/trainer_config.py

@@ -79,6 +79,7 @@ class TrainerConfig:
     agent_timeout_s: float = 300.0
     train_tool_results: bool = False
     chat_template_enable_thinking: bool | None = None
+    sft_assistant_turns: str = "all"
     lora: LoraConfig | None = None
     reference_mode: Literal["independent", "reuse_actor_base"] = "independent"
 
@@ -95,6 +96,8 @@ def __post_init__(self) -> None:
             raise ValueError("attn_backend must be one of: flash, native")
         if self.model_hub not in {"hf", "modelscope"}:
             raise ValueError("model_hub must be one of: hf, modelscope")
+        if self.sft_assistant_turns not in {"all", "last"}:
+            raise ValueError("sft_assistant_turns must be one of: all, last")
         if isinstance(self.optimizer_state_offload, bool):
             self.optimizer_state_offload = "cpu" if self.optimizer_state_offload else "none"
         if self.optimizer_state_offload not in {"none", "cpu", "disk"}:

areno/api/trainers/sft.py

@@ -23,7 +23,7 @@
 
 import areno.api
 from areno.api.dashboard import record_dashboard_state
-from areno.api.data_utils import prompt_response_to_tokens_and_mask
+from areno.api.data_utils import messages_to_tokens_and_mask, prompt_response_to_tokens_and_mask
 from areno.api.multimodal import (
     encode_multimodal_prompt,
     expand_image_tokens,
@@ -72,6 +72,7 @@ def _fit_initialized(self) -> None:
                 processor,
                 max_prompt_tokens=self.config.max_prompt_tokens,
                 max_new_tokens=self.config.max_new_tokens,
+                sft_assistant_turns=getattr(self.config, "sft_assistant_turns", "all"),
             ):
                 if not train_batch:
                     continue
@@ -104,7 +105,7 @@ def _fit_initialized(self) -> None:
             self.logger.info("epoch=%d stage=epoch_end", epoch)
             record_dashboard_state(self.areno, stage="epoch_end", epoch=epoch, step=step, role="policy")
 
-    def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, max_new_tokens: int):
+    def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, max_new_tokens: int, sft_assistant_turns: str = "all"):
         # Dataset rows are converted lazily so large HF datasets do not need an
         # up-front tokenized copy. Rows that are empty, all-prompt, or exceed
         # the configured prompt or supervised-response budgets are dropped.
@@ -120,6 +121,7 @@ def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, m
                 processor,
                 max_prompt_tokens=max_prompt_tokens,
                 max_new_tokens=max_new_tokens,
+                sft_assistant_turns=sft_assistant_turns,
             )
             if seq is None:
                 skipped += 1
@@ -152,17 +154,31 @@ def _maybe_save(self, epoch: int, step: int) -> None:
         record_dashboard_state(self.areno, stage="save_checkpoint_end", epoch=epoch, step=step, role="policy")
 
 
-def _record_to_train_sequence(record: Any, tokenizer, processor=None, *, max_prompt_tokens: int, max_new_tokens: int):
+def _record_to_train_sequence(
+    record: Any, tokenizer, processor=None, *, max_prompt_tokens: int, max_new_tokens: int, sft_assistant_turns: str = "all"
+):
     """Normalize one loader-produced SFT row into backend training format.
 
     `prompt_mask=True` means "do not train this source token"; the backend loss
     is next-token aligned, so the loss function later uses positions after the
     prompt prefix. RL-only fields are filled with zeros to satisfy the shared
     `TrainSequence` packing contract.
+
+    Two row schemas are accepted:
+
+    * ``{"prompt": str, "response": str}`` – single-turn (original format).
+    * ``{"messages": list[dict]}`` – multi-turn chat; each dict has ``role``
+      and ``content`` keys. ``sft_assistant_turns`` controls which assistant
+      turns are trainable: ``"all"`` (default) trains every assistant turn,
+      ``"last"`` trains only the final assistant turn.
     """
 
     record = dict(record)
     eos_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
+    if sft_assistant_turns not in ("all", "last"):
+        raise ValueError(
+            f"sft_assistant_turns must be 'all' or 'last', got {sft_assistant_turns!r}"
+        )
     if record_has_image(record):
         if "response" not in record:
             raise ValueError("SFT image rows must contain `response`")
@@ -240,18 +256,28 @@ def _record_to_train_sequence(record: Any, tokenizer, processor=None, *, max_pro
             features=features,
             eos_token_id=int(record.get("eos_token_id", eos_token_id)),
         )
-    if "prompt" not in record or "response" not in record:
+    if "messages" in record:
+        messages = record["messages"]
+        if not isinstance(messages, list) or not messages:
+            return None
+        if not any(msg.get("role") == "assistant" for msg in messages):
+            return None
+        tokens, prompt_mask = messages_to_tokens_and_mask(
+            messages, tokenizer, eos_token_id, last_assistant_only=(sft_assistant_turns == "last")
+        )
+    elif "prompt" in record and "response" in record:
+        if record["prompt"] is None or record["response"] is None:
+            return None
+        prompt = str(record["prompt"])
+        response = str(record["response"])
+        if not response:
+            return None
+        tokens, prompt_mask = prompt_response_to_tokens_and_mask(prompt, response, tokenizer, eos_token_id)
+    else:
         raise ValueError(
-            "SFT dataset loader must return rows with `prompt` and `response`; "
-            "normalize raw dataset fields in --dataset-loader-fn"
+            "SFT dataset loader must return rows with `prompt` and `response`, "
+            "or `messages`; normalize raw dataset fields in --dataset-loader-fn"
         )
-    if record["prompt"] is None or record["response"] is None:
-        return None
-    prompt = str(record["prompt"])
-    response = str(record["response"])
-    if not response:
-        return None
-    tokens, prompt_mask = prompt_response_to_tokens_and_mask(prompt, response, tokenizer, eos_token_id)
 
     if len(tokens) < 2:
         return None

areno/cli/train.py

@@ -122,6 +122,7 @@ def flash_attention_unsupported_model_reason(model_config):
             "agent_fn",
             "agent_timeout_s",
             "train_tool_results",
+            "sft_assistant_turns",
             "reward_fn_path",
             "reward_ckpt",
         ),
@@ -976,6 +977,7 @@ def _trainer_config_from_args(args) -> TrainerConfig:
             agent_timeout_s=args.agent_timeout_s,
             train_tool_results=args.train_tool_results,
             chat_template_enable_thinking=chat_template_enable_thinking,
+            sft_assistant_turns=args.sft_assistant_turns,
             lora=lora,
             reference_mode=args.reference_mode,
         )
@@ -1254,6 +1256,7 @@ def section(title: str, names: list[str]) -> dict:
                 "agent_fn",
                 "agent_timeout_s",
                 "train_tool_results",
+                "sft_assistant_turns",
                 "reward_fn_path",
                 "reward_ckpt",
             ],
@@ -1792,6 +1795,13 @@ def _dataset_builder_for_suffix(suffix: str) -> str:
     "--agent-timeout-s", type=float, default=300.0, show_default=True, help="Agentic rollout proxy request timeout."
 )
 @click.option("--train-tool-results", is_flag=True, help="Include tool-result spans in agentic policy loss.")
+@click.option(
+    "--sft-assistant-turns",
+    type=click.Choice(["all", "last"], case_sensitive=False),
+    default="all",
+    show_default=True,
+    help="SFT: train on every assistant turn (all) or only the final assistant turn (last) in multi-turn data.",
+)
 @click.option(
     "--gspo-clip-eps", type=float, default=3.0e-4, show_default=True, help="GSPO sequence-ratio clipping epsilon."
 )

docs/cli/dataset_loaders.rst

@@ -23,7 +23,8 @@ SFT
 ---
 
 SFT always requires ``--dataset-loader-fn``. The loader must return rows with
-``prompt`` and ``response`` keys:
+either ``prompt`` and ``response`` keys (single-turn) or a ``messages`` key
+(multi-turn chat):
 
 .. code-block:: python
 
@@ -40,6 +41,21 @@ SFT always requires ``--dataset-loader-fn``. The loader must return rows with
            )
        return records
 
+For multi-turn chat data, return ``messages`` instead:
+
+.. code-block:: python
+
+   {"messages": [
+       {"role": "user", "content": "What is 2+2?"},
+       {"role": "assistant", "content": "4"},
+       {"role": "user", "content": "And 3+3?"},
+       {"role": "assistant", "content": "6"},
+   ]}
+
+The trainer trains on all assistant turns by default. Use
+``--sft-assistant-turns last`` to train only on the final assistant response.
+User, system, and tool-result tokens are always excluded.
+
 For a concrete example, use ``--dataset-path yahma/alpaca-cleaned`` with
 ``examples/sft/alpaca/dataset_loader.py``.
 

docs/cli/training.rst

@@ -240,6 +240,14 @@ and CUDA graph state; MLX retains one in-process model.
    results are environment observations rather than policy actions. Assistant
    text and assistant tool-call spans are trainable by default.
 
+``--sft-assistant-turns [all|last]``
+   SFT only. Controls which assistant turns in multi-turn chat data are
+   trainable. ``all`` (default) trains on every assistant turn. ``last``
+   trains only the final assistant turn, treating earlier assistant responses
+   as context. User, system, and tool-result tokens are always excluded from
+   training. This option has no effect on single-turn ``prompt``/``response``
+   SFT rows.
+
 Agentic trajectories can contain multiple chat-completion turns for the same
 prompt/sample pair. The agent owns the OpenAI-style message list and returns
 trajectory turns with the model response; Areno converts those turns into token
@@ -531,6 +539,38 @@ SFT instruction tuning
 SFT loaders must normalize raw rows to ``prompt`` and ``response`` dictionaries.
 The trainer performs tokenization and trains on the response suffix.
 
+SFT also supports multi-turn chat data. Instead of ``prompt``/``response``,
+the loader can return a ``messages`` field containing a list of
+``{"role": "user"|"assistant"|"system"|"tool", "content": "..."}`` dicts.
+The trainer tokenizes the full conversation and trains on assistant turns.
+Use ``--sft-assistant-turns last`` to train only on the final assistant
+response, which is useful for focused evaluation of end-to-end multi-turn
+behavior:
+
+.. code-block:: bash
+
+   areno train \
+     --ckpt Qwen/Qwen3-0.6B \
+     --dataset-path /path/to/multiturn.jsonl \
+     --dataset-loader-fn /path/to/multiturn_loader.py \
+     --algo sft \
+     --sft-assistant-turns last \
+     --tp-size 1 \
+     --world-size 1 \
+     --batch-size 2 \
+     --mini-bs 1
+
+A multi-turn SFT loader should produce rows like:
+
+.. code-block:: python
+
+   {"messages": [
+       {"role": "user", "content": "What is 2+2?"},
+       {"role": "assistant", "content": "4"},
+       {"role": "user", "content": "And 3+3?"},
+       {"role": "assistant", "content": "6"},
+   ]}
+
 DPO preference training
 ~~~~~~~~~~~~~~~~~~~~~~~