From c3f22f94140c981ababe37c521317bff8f1c16ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=A5=9A=E8=B4=A2?= Date: Thu, 23 Jul 2026 20:10:33 +0800 Subject: [PATCH 1/8] feat: add Bailing V3 SWE SFT support Port the reviewed internal implementation to the public main branch while preserving newer upstream engine behavior. Key changes: - Add Bailing V3 KDA, gated MLA, and MoE model support - Add SWE SFT dataset loading and cache handling - Add focused model, loader, and dataset tests Refs: inclusionAI/AReaL#2188 Signed-off-by: chucai.dzq --- areal/api/cli_args.py | 16 +- areal/dataset/__init__.py | 18 + areal/dataset/swe_sft.py | 2692 +++++++++++++++++++++ areal/engine/megatron_engine.py | 193 +- areal/models/mcore/bailing_v3.py | 401 +++ areal/models/mcore/bailing_v3_bridge.py | 400 +++ areal/models/mcore/bailing_v3_mla.py | 119 + areal/models/mcore/hf_load.py | 43 +- areal/models/mcore/kda_attention.py | 1118 +++++++++ areal/models/mcore/lightning_attention.py | 82 +- areal/models/mcore/registry.py | 8 + areal/utils/saver.py | 6 + docs/en/cli_reference.md | 4 +- docs/zh/cli_reference.md | 4 +- examples/swe/config.py | 149 ++ examples/swe/train_sft.py | 114 + tests/models/test_zigzag_indices.py | 140 ++ tests/test_bailing_v3_hf_load.py | 138 ++ tests/test_bailing_v3_kda_cp_helpers.py | 224 ++ tests/test_dataset_swe_path_dispatch.py | 81 + tests/test_swe_sft_cache.py | 227 ++ 21 files changed, 6132 insertions(+), 45 deletions(-) create mode 100644 areal/dataset/swe_sft.py create mode 100644 areal/models/mcore/bailing_v3.py create mode 100644 areal/models/mcore/bailing_v3_bridge.py create mode 100644 areal/models/mcore/bailing_v3_mla.py create mode 100644 areal/models/mcore/kda_attention.py create mode 100644 examples/swe/config.py create mode 100644 examples/swe/train_sft.py create mode 100644 tests/models/test_zigzag_indices.py create mode 100644 tests/test_bailing_v3_hf_load.py create mode 100644 tests/test_bailing_v3_kda_cp_helpers.py create mode 100644 tests/test_dataset_swe_path_dispatch.py create mode 100644 tests/test_swe_sft_cache.py diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 117642a246..b5ad55b7d4 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -936,11 +936,12 @@ class MegatronEngineConfig: # MoE moe_router_dtype: str | None = "fp32" - moe_shared_expert_overlap: bool = field( - default=False, + moe_shared_expert_overlap: bool | None = field( + default=None, metadata={ "help": "Enable overlapping between shared expert computations and dispatcher communications. " - "Without this, the shared experts execute after the routed experts." + "Without this, the shared experts execute after the routed experts. " + "None keeps the model bridge's own default." }, ) moe_enable_deepep: bool = False @@ -961,13 +962,14 @@ class MegatronEngineConfig: "Requires TransformerEngine >= 2.7.0.", }, ) - moe_router_bias_update_rate: float = field( - default=0.0, + moe_router_bias_update_rate: float | None = field( + default=None, metadata={ "help": "Update rate for auxiliary-loss-free MoE load balancing " "(DeepSeek V3 style). Controls how fast expert_bias adjusts. " - "Default 0.0 disables bias updates; set a positive value such as " - "1e-3 to enable.", + "None keeps the model bridge's own default (AReaL bridges " + "disable it or derive it from the checkpoint). Set 0.0 to " + "disable explicitly; 1e-3 matches DeepSeek V3.", }, ) moe_z_loss_coeff: float | None = field( diff --git a/areal/dataset/__init__.py b/areal/dataset/__init__.py index 1ab23a4cb9..150676c400 100644 --- a/areal/dataset/__init__.py +++ b/areal/dataset/__init__.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import re from typing import TYPE_CHECKING, Optional from areal.api.cli_args import _DatasetConfig @@ -19,10 +20,17 @@ "virl39k", "hh-rlhf", "torl_data", + "swe_sft", ] logger = logging.getLogger("Dataset") +# Matches "swe" only as a path token delimited by /, _, -, or . (e.g. +# "swe_data/", "swe-bench", "my_swe.jsonl") so that paths merely containing +# the trigram (e.g. "answer_sft", "/home/swetha/") fall through to the +# generic load-from-disk fallback instead of the SWE trajectory pipeline. +_SWE_PATH_PATTERN = re.compile(r"(?:^|[/_\-.])swe(?:[/_\-.]|$)") + def _get_custom_dataset( path: str, @@ -133,6 +141,16 @@ def _get_custom_dataset( max_length=max_length, **kwargs, ) + elif _SWE_PATH_PATTERN.search(path.lower()) and type == "sft": + from .swe_sft import get_swe_sft_dataset + + return get_swe_sft_dataset( + path=path, + split=split, + tokenizer=tokenizer, + max_length=max_length, + **kwargs, + ) else: # Fallback: try loading as a generic HuggingFace dataset from disk. # This supports arbitrary datasets saved via dataset.save_to_disk(). diff --git a/areal/dataset/swe_sft.py b/areal/dataset/swe_sft.py new file mode 100644 index 0000000000..8548d7586b --- /dev/null +++ b/areal/dataset/swe_sft.py @@ -0,0 +1,2692 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""SWE SFT dataset loader. + +Loads SWE-bench trajectory data and converts it into progressive SFT +training pairs. Each trajectory is split at assistant-turn boundaries +so that every pair ends with an assistant segment (assistant message + +its subsequent tool responses). + +Example trajectory:: + + [system, user, asst1, tool1a, tool1b, asst2, tool2, asst3] + +Produces three pairs:: + + Pair 1: [system, user, asst1, tool1a, tool1b] + Pair 2: [system, user, asst1, tool1a, tool1b, asst2, tool2] + Pair 3: [system, user, asst1, tool1a, tool1b, asst2, tool2, asst3] + +In each pair, only the **last** assistant segment is trained (loss=1); +earlier assistant turns are treated as context (loss=0). + +By default, pairs whose current segment contains a tool result with +``is_error=True`` are discarded. Set ``filter_errors=False`` to keep them. + +The file is organized into the following sections: + +1. **Constants & Infrastructure** — shared constants, distributed sync +2. **Cleaning** — message content transforms (thinking tags, field cleanup) +3. **Filters** — keep/discard predicates (error, empty, bare-text, truncation) +4. **Splitting** — trajectory → progressive pairs (segment detection + split) +5. **Tokenization** — template detection, render→tokenize→loss_mask, dump +6. **Pipeline** — loading, processing, distributed cache, public API +7. **CLI** — ``python -m areal.dataset.swe_sft`` entry point +""" + +import json +import os +import random +import re +import shutil +import time + +from datasets import Dataset + +from areal.utils import logging + +logger = logging.getLogger("SWESFTDataset") + + +# ============================================================ +# 1. Constants & Infrastructure +# ============================================================ + +DATASET_NUM_PROC = 1 + +# Timeout (seconds) for non-rank-0 workers waiting for rank 0 to finish +# dataset processing. Progressive-pair tokenization of large trajectory +# corpora is single-process on rank 0 and can take hours; 10 h is the +# upper bound before workers give up. +_RANK0_CACHE_TIMEOUT = 36000 +_RANK0_CACHE_POLL_INTERVAL = 5 + + +def _extract_messages(record, record_idx): + """Extract messages and tools from a parsed JSONL record. + + Handles nested (``conversations`` wrapper) and flat formats. + Warns if multiple conversations are present. + + Returns: + Tuple of ``(messages, record_tools)``. *messages* may be empty. + """ + convs = record.get("conversations", []) + if convs: + if len(convs) > 1: + logger.warning( + "Record %d has %d conversations, using only the last one.", + record_idx, + len(convs), + ) + conv = convs[-1] + return conv.get("messages", []), conv.get("tools") + return record.get("messages", []), record.get("tools") + + +def _set_messages(record, messages): + """Write *messages* back into *record* (inverse of ``_extract_messages``). + + Used by the ``--save-trajectories`` CLI path to update truncated + messages in the original record structure before serialization. + """ + convs = record.get("conversations", []) + if convs: + convs[-1]["messages"] = messages + else: + record["messages"] = messages + + +def _iter_jsonl_records(path): + """Iterate trajectory JSONL records. + + Yields ``(record_idx, messages, record_tools)`` tuples. Handles + nested (``conversations`` wrapper) vs flat format auto-detection + via ``_extract_messages``. Records with empty messages are skipped. + + Warns about multi-user trajectories which break think-tag rendering + in templates with ``ns.last_query_index`` logic (e.g. Bailing). + """ + record_idx = 0 + n_multi_user = 0 + with open(path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + record = json.loads(line) + record_idx += 1 + messages, record_tools = _extract_messages(record, record_idx) + if not messages: + continue + n_user = sum(1 for m in messages if m.get("role") == "user") + if n_user > 1: + n_multi_user += 1 + if n_multi_user <= 3: + logger.warning( + "Record %d has %d user messages. Templates with " + "ns.last_query_index logic (e.g. Bailing) will NOT " + "render for assistant turns before the last " + "user message. Consider filtering!", + record_idx, + n_user, + ) + yield record_idx, messages, record_tools + if n_multi_user > 0: + logger.warning( + "Total %d/%d records have multiple user messages. " + "These may produce no-think training signal.", + n_multi_user, + record_idx, + ) + + +# ============================================================ +# 2. Cleaning — message content transforms +# ============================================================ + +# Match reasoning blocks with any common tag variant: +# ... (Qwen standard) +# ... (Claude) +# The opening and closing tag names need not match exactly — mixed pairs +# like ``...`` (seen in distillation data) are handled. +_THINK_OPEN_RE = re.compile(r"") +_THINK_CLOSE_RE = re.compile(r"") +_THINK_RE = re.compile(r"(.*?)", re.DOTALL) + + +def _normalize_thinking_tags(content): + """Normalise all thinking tag variants to ````/````. + + Distillation data from different models may use ```` (Claude) + vs ```` (Qwen). Non-standard variants are multi-token for the + Qwen tokenizer which breaks think/tool_call boundaries. + """ + if not content: + return content + content = _THINK_OPEN_RE.sub("", content) + content = _THINK_CLOSE_RE.sub("", content) + return content + + +def _extract_thinking(content): + """Strip thinking blocks from *content*. + + Callers must run ``_normalize_thinking_tags`` first so that all + tag variants have been converted to ````/````. + + Returns: + Cleaned content with thinking blocks removed, or the original + content unchanged if no thinking tags are found. + """ + if not content: + return content + cleaned = _THINK_RE.sub("", content).strip() + return cleaned if cleaned != content.strip() else content + + +def _clean_message(msg, strip_thinking=True, ensure_thinking=False): + """Remove non-standard fields before tokenization. + + Keeps only the fields expected by tokenizer chat templates: + role, content, reasoning_content (for assistant), tool_calls + (for assistant), tool_call_id (for tool). + + Handles thinking content in two representations: + + - Inline ``...`` tags in ``content`` + - Separate ``reasoning_content`` field (DeepSeek, Qwen3 API style) + + If both are present, inline tags take priority and + ``reasoning_content`` is dropped with a warning to avoid double + thinking blocks in the rendered template. + + Args: + msg: Raw message dict. + strip_thinking: If True, remove thinking from assistant messages + (both inline ```` tags and ``reasoning_content``). + Used for context turns. If False, preserve thinking as-is + (used for the training-target assistant turn). + ensure_thinking: If True, inject inline ``\n`` + on assistant turns that lack a thinking block (either + inline or in ``reasoning_content``). Requires the + patched Bailing template (via + ``_patch_chat_template_for_training``) which detects + ``had_think_tags`` and preserves empty think blocks. + """ + cleaned = {"role": msg["role"]} + + # Handle content — some assistant messages have content=None when + # they only contain tool_calls. Preserve None so chat templates + # that distinguish None vs "" render correctly. + content = msg.get("content") + # Some APIs (DeepSeek, Qwen3 with enable_thinking) return thinking + # in a separate ``reasoning_content`` field instead of inline + # ```` tags. Handle both representations. + raw_reasoning = msg.get("reasoning_content") if msg["role"] == "assistant" else None + has_thinking = False + if content is not None: + if msg["role"] == "assistant": + content = _normalize_thinking_tags(content) + has_inline_thinking = bool(_THINK_RE.search(content)) + if has_inline_thinking and raw_reasoning and raw_reasoning.strip(): + # Conflict: both reasoning_content and inline tags. + # Keep inline tags (they are already in the content the + # tokenizer will see) and drop reasoning_content to avoid + # double thinking blocks in the rendered template. + if not strip_thinking: + logger.warning( + "Message has both reasoning_content and inline " + " tags. Keeping inline tags, dropping " + "reasoning_content." + ) + raw_reasoning = None + elif not has_inline_thinking and raw_reasoning and raw_reasoning.strip(): + # Convert reasoning_content → inline in content. + # This ensures a single representation that templates + # render identically to the reasoning_content path, while + # being more transparent and debuggable. + if not strip_thinking: + content = ( + f"\n{raw_reasoning.strip(chr(10))}\n" + f"\n\n{content.lstrip(chr(10))}" + ) + has_inline_thinking = True + raw_reasoning = None + has_thinking = has_inline_thinking or bool( + raw_reasoning and raw_reasoning.strip() + ) + if strip_thinking: + content = _extract_thinking(content) + cleaned["content"] = content + elif msg["role"] == "assistant" and msg.get("tool_calls"): + # Assistant with tool_calls but content=None. + if raw_reasoning and raw_reasoning.strip(): + has_thinking = True + if not strip_thinking: + # Convert reasoning_content → inline in content. + cleaned["content"] = ( + f"\n{raw_reasoning.strip(chr(10))}\n" + ) + else: + cleaned["content"] = None + raw_reasoning = None + else: + cleaned["content"] = None + else: + # Non-assistant messages without content: default to empty string. + cleaned["content"] = "" + + # Preserve reasoning_content for target turns only when it was NOT + # already inlined above (i.e. only when raw_reasoning is still set). + if not strip_thinking and raw_reasoning is not None: + cleaned["reasoning_content"] = raw_reasoning + + # For the target assistant turn without a thinking block, inject + # inline ``\n`` so that the (patched) template detects + # think intent via ``had_think_tags`` and renders + # ``\n\n\n\n`` — identical token output to the old + # ``reasoning_content='\n'`` approach. + # + # Requires ``_patch_chat_template_for_training`` to have been called + # on the tokenizer, otherwise the stock Bailing template will extract + # and discard the empty ```` block. + if ensure_thinking and msg["role"] == "assistant" and not has_thinking: + cur_content = cleaned.get("content") + if cur_content is None or cur_content == "": + cleaned["content"] = "\n" + else: + cleaned["content"] = f"\n\n\n{cur_content.lstrip(chr(10))}" + + # Copy tool_calls for assistant messages + if msg["role"] == "assistant" and msg.get("tool_calls"): + cleaned_tool_calls = [] + for tc in msg["tool_calls"]: + cleaned_tc = { + "type": tc.get("type", "function"), + "function": { + "name": tc["function"]["name"], + "arguments": json.dumps(tc["function"]["arguments"]) + if isinstance(tc["function"]["arguments"], dict) + else tc["function"]["arguments"], + }, + } + if "id" in tc: + cleaned_tc["id"] = tc["id"] + cleaned_tool_calls.append(cleaned_tc) + cleaned["tool_calls"] = cleaned_tool_calls + + # Copy tool_call_id for tool messages + if msg["role"] == "tool" and msg.get("tool_call_id"): + cleaned["tool_call_id"] = msg["tool_call_id"] + + return cleaned + + +# ============================================================ +# 3. Filters — keep/discard predicates +# ============================================================ + + +def _segment_has_error(messages, start, end): + """Check if any tool message in ``messages[start:end]`` has ``is_error=True``.""" + for m in messages[start:end]: + if m.get("role") == "tool" and m.get("is_error") is True: + return True + return False + + +def _is_empty_tool_call(msg): + """True if assistant *msg* has no text content and no reasoning but has tool_calls.""" + content = msg.get("content") or "" + if content.strip() or not msg.get("tool_calls"): + return False + # If reasoning_content exists, the model did think — not a silent invocation. + reasoning = msg.get("reasoning_content") + if reasoning and reasoning.strip(): + return False + return True + + +def _is_bare_text_tool_call(msg): + """True if assistant *msg* has text without ```` tags and has tool_calls.""" + content = msg.get("content") or "" + if not content.strip() or not msg.get("tool_calls"): + return False + # If reasoning_content exists, thinking is in a separate field — not bare text. + reasoning = msg.get("reasoning_content") + if reasoning and reasoning.strip(): + return False + normalized = _THINK_OPEN_RE.sub("", content) + normalized = _THINK_CLOSE_RE.sub("", normalized) + match = _THINK_RE.search(normalized) + return not (match and match.group(1).strip()) + + +def _msg_has_thinking(msg): + """True if assistant *msg* has thinking content (inline or reasoning_content).""" + if msg.get("role") != "assistant": + return False + content = msg.get("content") or "" + normalized = _THINK_OPEN_RE.sub("", content) + normalized = _THINK_CLOSE_RE.sub("", normalized) + if _THINK_RE.search(normalized): + return True + rc = msg.get("reasoning_content") or "" + return bool(rc.strip()) + + +def _truncate_at_task_notification(messages): + """Truncate messages when a ```` follows a pure-text assistant. + + Claude Code emits ```` as a user message when a + background task (e.g. ``pip install``) completes. If the model has + already produced a text-only summary (no tool_calls), the notification + and all subsequent messages are noise — the model just replies + "nothing to do". Truncating here removes that noise. + + Only triggers when the pattern is: + assistant (text, no tool_calls) → user () + + Returns: + Truncated message list (or the original list if no truncation needed). + """ + for i, m in enumerate(messages): + if m.get("role") != "user": + continue + if "" not in (m.get("content") or ""): + continue + # Find preceding assistant + prev_asst = None + for j in range(i - 1, -1, -1): + if messages[j].get("role") == "assistant": + prev_asst = messages[j] + break + if prev_asst is None: + continue + content = prev_asst.get("content") or "" + if content.strip() and not prev_asst.get("tool_calls"): + # Truncate: keep everything up to (but not including) this user msg + return messages[:i] + return messages + + +# ============================================================ +# 3b. Balancing — downsample non-thinking pairs +# ============================================================ + + +def _classify_pair(pair): + """Classify a pair by its target assistant turn's content type. + + Returns one of: + ``"thinking"`` — target has actual ```` content or + non-empty ``reasoning_content``. + ``"no_thinking_tool_call"`` — target has no thinking but has + ``tool_calls`` (the dominant category that causes distribution + skew). + ``"pure_text"`` — target has no thinking and no tool_calls + (typically the final summary turn in a trajectory). + """ + target = pair[-1] + if target.get("role") != "assistant": + return "pure_text" + + content = target.get("content") or "" + rc = target.get("reasoning_content") or "" + # Require non-empty think content: pair cleaning runs BEFORE balancing + # and (with ensure_thinking) injects an empty \n into + # every no-think target, so a bare regex hit would classify everything + # as "thinking" and silently disable max_no_thinking_ratio. + _m = _THINK_RE.search(content) + has_thinking = bool(_m and _m.group(1).strip()) or bool(rc.strip()) + + if has_thinking: + return "thinking" + if target.get("tool_calls"): + return "no_thinking_tool_call" + return "pure_text" + + +def _balance_thinking_pairs(pairs, max_no_thinking_ratio, seed=42, tools_list=None): + """Downsample non-thinking **tool-call** pairs to control balance. + + Only ``no_thinking_tool_call`` pairs (no thinking but has tool_calls) + are subject to downsampling. ``thinking`` pairs and ``pure_text`` + pairs (the final summary turn, no thinking and no tool_calls) are + always kept — the latter are critical for the model to learn when + to stop calling tools and give a final answer. + + Args: + pairs: List of progressive SFT pairs. + max_no_thinking_ratio: Maximum ratio of non-thinking tool-call + pairs to thinking pairs. For example, ``1.0`` means at most + 1:1, ``2.0`` means at most 2 non-thinking per 1 thinking pair. + ``None`` disables downsampling. + seed: Random seed for reproducible downsampling. + + Returns: + Balanced list of pairs (order preserved, randomly sampled for + the downsampled category). + """ + if max_no_thinking_ratio is None: + return pairs, tools_list + + thinking = [] + no_think_tc = [] + pure_text = [] + for i, pair in enumerate(pairs): + cat = _classify_pair(pair) + if cat == "thinking": + thinking.append(i) + elif cat == "no_thinking_tool_call": + no_think_tc.append(i) + else: + pure_text.append(i) + + n_think = len(thinking) + n_no_think_tc = len(no_think_tc) + n_pure_text = len(pure_text) + + if n_think == 0: + logger.warning( + "No thinking pairs found; skipping balance " + "(all %d pairs have empty thinking).", + n_no_think_tc + n_pure_text, + ) + return pairs, tools_list + + max_no_think_tc = int(n_think * max_no_thinking_ratio) + if n_no_think_tc <= max_no_think_tc: + logger.info( + "Thinking balance OK: %d thinking + %d no-think-tc + %d pure-text " + "(ratio %.1f <= %.1f), no downsampling needed.", + n_think, + n_no_think_tc, + n_pure_text, + n_no_think_tc / n_think, + max_no_thinking_ratio, + ) + return pairs, tools_list + + rng = random.Random(seed) + sampled_tc = set(rng.sample(no_think_tc, max_no_think_tc)) + keep_indices = sorted(set(thinking) | sampled_tc | set(pure_text)) + balanced = [pairs[i] for i in keep_indices] + balanced_tools = ( + [tools_list[i] for i in keep_indices] if tools_list is not None else None + ) + + logger.info( + "Balanced thinking pairs: %d thinking + %d no-think-tc " + "(downsampled from %d, ratio %.1f → %.1f) + %d pure-text (kept all).", + n_think, + max_no_think_tc, + n_no_think_tc, + n_no_think_tc / n_think, + max_no_thinking_ratio, + n_pure_text, + ) + return balanced, balanced_tools + + +# ============================================================ +# 3c. Thinking augmentation stats +# ============================================================ + + +def _log_thinking_augmentation_stats( + n_variants, + prob, + n_total_trajs, + thinking_turns_per_traj, + total_asst_turns_per_traj, + patterns_per_traj, +): + """Log adaptive-thinking augmentation quality metrics. + + Called after the augmentation loop in loaders to report how well + the ``n_thinking_variants`` / ``random_strip_thinking_prob`` settings + produce diverse thinking-pattern variants. + + Args: + n_variants: ``n_thinking_variants`` setting (K). + prob: ``random_strip_thinking_prob`` setting. + n_total_trajs: Total number of source trajectories processed. + thinking_turns_per_traj: List of N_thinking per source trajectory. + total_asst_turns_per_traj: List of N_total_asst per source trajectory. + patterns_per_traj: List of ``set[frozenset]`` — the unique strip + patterns generated for each source trajectory (including the + empty frozenset for the original unstripped variant). + """ + n_eligible = sum(1 for n in thinking_turns_per_traj if n > 0) + total_thinking = sum(thinking_turns_per_traj) + total_asst = sum(total_asst_turns_per_traj) + + # 1. Thinking Turn Coverage + avg_thinking = total_thinking / max(n_total_trajs, 1) + thinking_ratio = total_thinking / max(total_asst, 1) + + # 2. Pattern Diversity + diversity_ratios = [] + for n_think, patterns in zip(thinking_turns_per_traj, patterns_per_traj): + if n_think == 0: + continue + theoretical_max = min(n_variants, 2**n_think) + actual_unique = len(patterns) + diversity_ratios.append(actual_unique / theoretical_max) + avg_diversity = sum(diversity_ratios) / max(len(diversity_ratios), 1) + + # 3. Augmentation Efficiency + n_non_trivial = 0 + for patterns in patterns_per_traj: + # Count variants that differ from the original (non-empty strip set) + n_non_trivial += sum(1 for p in patterns if p) + expected_aug = (n_variants - 1) * max(n_eligible, 1) + efficiency = n_non_trivial / max(expected_aug, 1) + + # 4. Total sample count + n_total_samples = sum(len(p) for p in patterns_per_traj) + + logger.info( + f"Thinking augmentation stats (K={n_variants}, p={prob:.2f}):\n" + f" Source trajectories: {n_total_trajs} " + f"({n_eligible} with thinking turns)\n" + f" Thinking coverage: {avg_thinking:.1f} thinking turns/traj, " + f"{thinking_ratio:.1%} of all assistant turns\n" + f" Pattern diversity: {avg_diversity:.2f} " + f"(1.0 = all variants unique)\n" + f" Augmentation efficiency: {efficiency:.2f} " + f"({n_non_trivial}/{expected_aug} non-trivial variants)\n" + f" Total samples after augmentation: {n_total_samples}" + ) + + +# ============================================================ +# 4. Splitting — trajectory → progressive pairs +# ============================================================ + + +def _find_segments(messages): + """Find assistant+tools segment boundaries. + + Returns: + List of ``(assistant_start_idx, segment_end_idx)`` tuples. + """ + segments = [] + i = 0 + while i < len(messages): + if messages[i].get("role") == "assistant": + j = i + 1 + while j < len(messages) and messages[j].get("role") == "tool": + j += 1 + segments.append((i, j)) + i = j + else: + i += 1 + return segments + + +def _split_and_filter( + messages, + filter_errors=True, + strip_all_thinking=False, + filter_empty_tool_calls=False, + filter_bare_text_tool_calls=False, + random_strip_thinking_prob=0.0, + rng=None, +): + """Split trajectory into progressive pairs and optionally filter. + + By default, thinking (``...``) is stripped from context + assistant turns only; the last assistant turn (training target) keeps + its content unchanged. Set *strip_all_thinking* to strip from every + assistant turn including the target. + + When *random_strip_thinking_prob* > 0, each target assistant turn that + has thinking content is independently stripped with that probability. + Stripped turns use the context-cleaned version (thinking fully removed, + no empty ```` injected). + + Args: + messages: Raw trajectory messages. + filter_errors: If True (default), discard pairs whose current segment + contains a tool result with ``is_error=True``. Set to False to + keep all pairs regardless of tool errors. + strip_all_thinking: If True, strip ```` blocks from every + assistant turn including the training target. + filter_empty_tool_calls: If True, discard pairs whose training-target + assistant turn has no text content but has tool_calls. + filter_bare_text_tool_calls: If True, discard pairs whose + training-target assistant turn has text content without + ```` tags and has tool_calls. + random_strip_thinking_prob: Probability of stripping thinking + from each target assistant turn. 0.0 = no stripping. + rng: ``random.Random`` instance for reproducible sampling. + + Returns: + Tuple of ``(pairs, n_filtered_errors, n_filtered_empty_tc, + n_filtered_bare_tc, n_stripped)``. + """ + segments = _find_segments(messages) + if not segments: + return [], 0, 0, 0, 0 + + pairs = [] + n_filtered_errors = 0 + n_filtered_empty_tc = 0 + n_filtered_bare_tc = 0 + n_stripped = 0 + + # Pre-clean all messages in context mode (thinking stripped). + # This avoids re-cleaning the same message for every progressive pair + # (O(N+K) instead of O(N*K) where K = number of segments). + context_cleaned = [_clean_message(m, strip_thinking=True) for m in messages] + + # For target assistant turns, clean with thinking preserved (unless + # strip_all_thinking is set, in which case context_cleaned is reusable). + # When stripping is active (augmented variant), use ensure_thinking=False + # so empty-thinking turns don't get \n injected. + stripping_active = random_strip_thinking_prob > 0.0 and rng is not None + target_ensure = not stripping_active + target_cleaned = {} + if not strip_all_thinking: + for asst_start, _ in segments: + target_cleaned[asst_start] = _clean_message( + messages[asst_start], + strip_thinking=False, + ensure_thinking=target_ensure, + ) + + for asst_start, seg_end in segments: + # Check if current segment has any tool errors + if filter_errors and _segment_has_error(messages, asst_start, seg_end): + n_filtered_errors += 1 + continue + + # Content-type filters operate on the raw assistant message. + asst_msg = messages[asst_start] + if filter_empty_tool_calls and _is_empty_tool_call(asst_msg): + n_filtered_empty_tc += 1 + continue + if filter_bare_text_tool_calls and _is_bare_text_tool_call(asst_msg): + n_filtered_bare_tc += 1 + continue + + # Build pair: include context up to the target assistant turn, + # truncating tool responses that follow it. This ensures the + # target assistant is always the *last* message so that chat + # templates with ``loop.last``-dependent rendering (e.g. Qwen3 + # ```` injection) behave consistently. The tool responses + # would have loss_mask=0 anyway and only add noise. + pair = list(context_cleaned[: asst_start + 1]) + if not strip_all_thinking: + # Randomly strip: leave context_cleaned version (thinking + # already removed) instead of replacing with target_cleaned. + should_strip = ( + rng is not None + and _msg_has_thinking(messages[asst_start]) + and rng.random() < random_strip_thinking_prob + ) + if not should_strip: + pair[asst_start] = target_cleaned[asst_start] + else: + n_stripped += 1 + pairs.append(pair) + + return pairs, n_filtered_errors, n_filtered_empty_tc, n_filtered_bare_tc, n_stripped + + +def _prepare_trajectory( + messages, + filter_errors=True, + filter_empty_tool_calls=False, + filter_bare_text_tool_calls=False, + random_strip_thinking_prob=0.0, + rng=None, +): + """Prepare a full trajectory for trajectory-level training. + + Cleans all messages preserving thinking for every assistant turn + (``strip_thinking=False``, ``ensure_thinking=True``). Identifies + which assistant segments should be masked (``loss_mask=0``) based + on error tool responses, empty tool calls, or bare-text tool calls. + + When *random_strip_thinking_prob* > 0, each assistant turn that has + thinking content is independently stripped with that probability. + Stripped turns have their ```` blocks and ``reasoning_content`` + completely removed (no empty ```` injected). + + Args: + messages: Raw trajectory messages. + filter_errors: If True (default), mask segments with error tool + responses. + filter_empty_tool_calls: If True, mask segments whose assistant + turn has no text content but has tool_calls. + filter_bare_text_tool_calls: If True, mask segments whose + assistant turn has text without ```` tags and has + tool_calls. + random_strip_thinking_prob: Probability of stripping thinking + from each assistant turn that has thinking content. + 0.0 (default) = no stripping, 1.0 = strip all. + rng: ``random.Random`` instance for reproducible sampling. + + Returns: + Tuple of ``(cleaned_messages, masked_segment_indices, + n_error, n_empty_tc, n_bare_tc, stripped_pattern)`` or ``None`` + if the trajectory has no assistant turns. *stripped_pattern* is + a ``frozenset`` of message indices whose thinking was stripped + (empty if no stripping occurred). + """ + segments = _find_segments(messages) + if not segments: + return None + + masked_indices = set() + n_error = 0 + n_empty_tc = 0 + n_bare_tc = 0 + for idx, (asst_start, seg_end) in enumerate(segments): + if filter_errors and _segment_has_error(messages, asst_start, seg_end): + masked_indices.add(idx) + n_error += 1 + continue + asst_msg = messages[asst_start] + if filter_empty_tool_calls and _is_empty_tool_call(asst_msg): + masked_indices.add(idx) + n_empty_tc += 1 + continue + if filter_bare_text_tool_calls and _is_bare_text_tool_call(asst_msg): + masked_indices.add(idx) + n_bare_tc += 1 + + # Determine which assistant turns to randomly strip thinking from. + strip_thinking_indices = set() + stripping_active = random_strip_thinking_prob > 0.0 and rng is not None + if stripping_active: + for asst_start, _seg_end in segments: + if _msg_has_thinking(messages[asst_start]): + if rng.random() < random_strip_thinking_prob: + strip_thinking_indices.add(asst_start) + + # When stripping is active (augmented variant), use ensure_thinking=False + # for ALL turns so that empty-thinking turns don't get \n + # injected. Only real thinking content is preserved. + # When stripping is inactive (variant 0 or no augmentation), keep + # ensure_thinking=True to match the standard training format. + default_ensure = not stripping_active + + cleaned = [] + for i, m in enumerate(messages): + if i in strip_thinking_indices: + cleaned.append( + _clean_message(m, strip_thinking=True, ensure_thinking=False) + ) + else: + cleaned.append( + _clean_message(m, strip_thinking=False, ensure_thinking=default_ensure) + ) + + return ( + cleaned, + sorted(masked_indices), + n_error, + n_empty_tc, + n_bare_tc, + frozenset(strip_thinking_indices), + ) + + +# ============================================================ +# 5. Tokenization — template detection, render, loss mask, dump +# ============================================================ + + +# -- Chat template patch (runtime, no file modification) -------- + +# Both Bailing and Qwen3 templates have ``ns.last_query_index`` logic +# that prevents ```` rendering for assistant turns BEFORE the +# last user message, AND discards inline empty ``\n`` +# extracted from content. +# +# This breaks trajectory-mode training: +# - Multi-user trajectories: turns before the last user msg lack +# - Empty ensure_thinking via inline gets stripped +# +# The patch below handles both Bailing (`ASSISTANT` style) +# and Qwen3 (`<|im_start|>assistant` style) templates: +# 1. Adds ``had_think_tags`` detection so empty ```` survives. +# 2. Removes the ``ns.last_query_index`` gate so all assistant turns +# render ```` uniformly when think intent is detected. +# +# Applied at runtime via ``tokenizer.chat_template = patched`` — the +# original template file on disk is never modified. + +_BAILING_OLD_BLOCK = ( + "{%- if loop.index0 > ns.last_query_index %}\n" + " {%- if reasoning_content != '' %}\n" + " {{- 'ASSISTANT\\n' + '\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- 'ASSISTANT\\n' + content }}\n" + " {%- endif %}\n" + " {%- else %}\n" + " {{- 'ASSISTANT\\n' + content }}\n" + " {%- endif %}" +) +_BAILING_NEW_BLOCK = ( + "{%- if reasoning_content != '' or had_think_tags %}\n" + " {{- 'ASSISTANT\\n' + '\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- 'ASSISTANT\\n' + content }}\n" + " {%- endif %}" +) + +# Qwen3 uses `loop.last or (not loop.last and reasoning_content)` so the +# last turn always renders even with empty reasoning. We +# preserve `loop.last` and add `had_think_tags` for inline-empty support. +_QWEN3_OLD_BLOCK = ( + "{%- if loop.index0 > ns.last_query_index %}\n" + " {%- if loop.last or (not loop.last and reasoning_content) %}\n" + " {{- '<|im_start|>' + message.role + '\\n\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" + " {%- endif %}\n" + " {%- else %}\n" + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" + " {%- endif %}" +) +_QWEN3_NEW_BLOCK = ( + "{%- if loop.last or reasoning_content != '' or had_think_tags %}\n" + " {{- '<|im_start|>' + message.role + '\\n\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" + " {%- endif %}" +) + +_OLD_DETECT = "{%- set reasoning_content = '' %}" +_NEW_DETECT = ( + "{%- set reasoning_content = '' %}\n" + " {%- set had_think_tags = ('' in content) %}" +) + + +def _patch_chat_template_for_training(tokenizer): + """Patch Bailing/Qwen3 chat templates to render ```` uniformly. + + Detects template family by matching known render blocks: + - Bailing: ``ASSISTANT`` markers + - Qwen3: ``<|im_start|>assistant`` markers + + Other templates (e.g. plain ChatML without ``last_query_index``) + are left unchanged. If the template has ``last_query_index`` but + neither known block matches, logs a warning. + """ + template = getattr(tokenizer, "chat_template", None) + if not template or "last_query_index" not in template: + return + + if _BAILING_OLD_BLOCK in template: + family = "Bailing" + patched = template.replace(_BAILING_OLD_BLOCK, _BAILING_NEW_BLOCK) + elif _QWEN3_OLD_BLOCK in template: + family = "Qwen3" + patched = template.replace(_QWEN3_OLD_BLOCK, _QWEN3_NEW_BLOCK) + else: + # Reaching here means the template family needs the training patch + # (it gates rendering on last_query_index) but the verbatim block no + # longer matches — most likely an upstream template revision. Failing + # loudly beats silently training on data whose blocks the + # stock template strips (see _clean_message / ensure_thinking). + # + # Escape hatch for uses that do not depend on think normalization + # (e.g. precision-alignment forward dumps): set + # AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE=1 to proceed with a warning. + if os.environ.get("AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE", ""): + logger.warning( + "Chat template has last_query_index but matches neither known " + "render block; proceeding UNPATCHED because " + "AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE is set. Empty " + "blocks may be discarded by the stock template." + ) + return + raise ValueError( + "Chat template has last_query_index but matches neither the known " + "Bailing nor Qwen3 render block; the training patch cannot be " + "applied. Without it, empty blocks are discarded and " + "multi-turn thinking renders inconsistently. Update " + "_BAILING_OLD_BLOCK/_QWEN3_OLD_BLOCK for this template revision, " + "or set AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE=1 if think " + "normalization is irrelevant for this run." + ) + + if _OLD_DETECT not in patched: + raise ValueError( + "Chat template render block matched but the reasoning_content " + "detect line did not; had_think_tags would be undefined and " + "empty blocks would silently vanish. Update " + "_OLD_DETECT/_NEW_DETECT for this template revision." + ) + patched = patched.replace(_OLD_DETECT, _NEW_DETECT) + + tokenizer.chat_template = patched + logger.info( + f"Patched {family} chat template for training: removed " + "last_query_index gate, added had_think_tags detection." + ) + + +_TEMPLATE_PATTERNS = [ + # ChatML (Qwen, etc.): <|im_start|>assistant\n ... <|im_end|> + (r"<\|im_start\|>assistant\n", r"<\|im_end\|>"), + # Llama 3: <|start_header_id|>assistant<|end_header_id|>\n\n ... <|eot_id|> + (r"<\|start_header_id\|>assistant<\|end_header_id\|>\n\n", r"<\|eot_id\|>"), + # GLM: <|assistant|> ... (ends at next <|user|>, <|observation|>, or end of string) + (r"<\|assistant\|>", r"(?=<\|user\|>|<\|observation\|>|\Z)"), +] + + +def _parse_tool_call_arguments(messages): + """Parse JSON-string arguments in tool_calls to dicts. + + OpenAI returns tool_call arguments as JSON strings, but some chat + templates (e.g. GLM-4.x / GLM-5.x) expect parsed dicts. Most other + templates (Qwen / ChatML, Llama 3, Bailing, ...) accept the standard + OpenAI string form, so this conversion must be opt-in. + """ + patched = [] + for m in messages: + tool_calls = m.get("tool_calls") + if not tool_calls: + patched.append(m) + continue + new_tcs = [] + for tc in tool_calls: + fn = tc.get("function", tc) + args = fn.get("arguments") + if isinstance(args, str): + try: + parsed = json.loads(args) + except (json.JSONDecodeError, TypeError): + parsed = args + fn = {**fn, "arguments": parsed} + tc = {**tc, "function": fn} if "function" in tc else fn + new_tcs.append(tc) + patched.append({**m, "tool_calls": new_tcs}) + return patched + + +def _render_tokenize_mask( + messages, + tokenizer, + assistant_pattern, + tools=None, + *, + split_mode="pair", + error_indices=None, + parse_tool_call_args=False, +): + """Render, tokenize, and build loss_mask for a message list. + + In **pair mode** (default), only the **last** assistant turn gets + ``loss_mask=1``. In **trajectory mode**, **all** assistant turns + get ``loss_mask=1`` except those at indices in *error_indices*. + + When *parse_tool_call_args* is True, JSON-string ``tool_calls`` arguments + are converted to dicts before rendering (required by GLM chat templates; + other templates such as Qwen / Llama / Bailing must keep the OpenAI + string form). + + Returns: + Tuple of ``(full_text, input_ids, loss_mask, offset_mapping)``, or + ``None`` if ``apply_chat_template`` fails. + """ + # 1) Render the full template text. + try: + kwargs = {"tokenize": False} + if tools is not None: + kwargs["tools"] = tools + if parse_tool_call_args: + messages = _parse_tool_call_arguments(messages) + full_text = tokenizer.apply_chat_template(messages, **kwargs) + except Exception as e: + logger.warning( + "apply_chat_template failed: %s. Skipping sample.", + e, + ) + return None + + # 2) Tokenize with offset mapping so we can map char→token. + encoding = tokenizer( + full_text, add_special_tokens=False, return_offsets_mapping=True + ) + input_ids = encoding["input_ids"] + offset_mapping = encoding["offset_mapping"] + + # 3) Build loss_mask. + loss_mask = [0] * len(input_ids) + + if split_mode == "trajectory": + # Trajectory mode: mask ALL assistant segments, skip error_indices. + skip = set(error_indices) if error_indices else set() + matches = list(assistant_pattern.finditer(full_text)) + + # Verify regex matches correspond 1:1 to assistant messages. + n_asst = sum(1 for m in messages if m.get("role") == "assistant") + if len(matches) != n_asst: + # Fail closed: a spurious match (e.g. a tool output quoting the + # chat-template header literal) would otherwise put loss on + # user/tool tokens and silently defeat error masking. + logger.warning( + "Segment count mismatch: %d assistant messages but %d regex " + "matches in rendered text. Dropping this sample.", + n_asst, + len(matches), + ) + return None + + for seg_idx, m in enumerate(matches): + if seg_idx in skip: + continue + rs, re_ = m.start(1), m.end(0) + for tok_idx, (cs, ce) in enumerate(offset_mapping): + if ce > rs and cs < re_: + loss_mask[tok_idx] = 1 + else: + # Pair mode: mask only the LAST assistant segment. + last_match = None + for m in assistant_pattern.finditer(full_text): + last_match = m + if last_match is not None: + rs, re_ = last_match.start(1), last_match.end(0) + for tok_idx, (cs, ce) in enumerate(offset_mapping): + if ce > rs and cs < re_: + loss_mask[tok_idx] = 1 + else: + # Loss lands nowhere; the SFT loss path tolerates all-zero masks + # (kept, not dropped, to preserve cache compatibility) but this + # always signals template/pattern drift worth investigating. + logger.warning( + "No assistant segment matched the template pattern; sample " + "keeps an all-zero loss_mask." + ) + + return full_text, input_ids, loss_mask, offset_mapping + + +class _TokenizeAndMask: + """Picklable callable for ``Dataset.map(num_proc=N)``.""" + + def __init__( + self, + tokenizer, + assistant_pattern, + max_length=None, + *, + split_mode="pair", + parse_tool_call_args=False, + ): + self.tokenizer = tokenizer + self.assistant_pattern = assistant_pattern + self.max_length = max_length + self.split_mode = split_mode + self.parse_tool_call_args = parse_tool_call_args + + def __call__(self, sample): + error_indices = ( + sample.get("error_indices", []) if self.split_mode == "trajectory" else None + ) + tools_json = sample.get("tools_json") + tools = json.loads(tools_json) if tools_json else None + result = _render_tokenize_mask( + sample["messages"], + self.tokenizer, + self.assistant_pattern, + tools, + split_mode=self.split_mode, + error_indices=error_indices, + parse_tool_call_args=self.parse_tool_call_args, + ) + if result is None: + return {"input_ids": [], "loss_mask": []} + + _full_text, input_ids, loss_mask, _offset_mapping = result + + # Early exit: overlength or empty → return empty so a single + # filter pass removes it together with template-failure empties. + if self.max_length is not None and len(input_ids) > self.max_length: + return {"input_ids": [], "loss_mask": []} + + return {"input_ids": input_ids, "loss_mask": loss_mask} + + +def _detect_template_pattern(tokenizer, tools=None): + """Detect the assistant role delimiter used by this tokenizer's template. + + When *tools* is provided the probe is rendered with ``tools=`` so that + the detected delimiters match the actual training text (some templates + alter the system block when tools are present). + + Strategy: + 1. Try known ``_TEMPLATE_PATTERNS`` (fast, battle-tested). + 2. Fall back to double-probe diff: render the template with a known + marker and with empty content, then diff the two strings to extract + the exact header and end-of-turn delimiters. + + Raises: + ValueError: If both strategies fail to detect a usable pattern. + """ + _PROBE_CONTENT = "PROBE_MARKER" + + extra_kwargs = {} + if tools is not None: + extra_kwargs["tools"] = tools + + probe_msgs = [ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": _PROBE_CONTENT}, + ] + probe_text = tokenizer.apply_chat_template( + probe_msgs, tokenize=False, **extra_kwargs + ) + + # --- Strategy 1: known patterns --- + for hdr_re, eot_re in _TEMPLATE_PATTERNS: + if re.search(hdr_re, probe_text): + pattern = re.compile(hdr_re + r"(.*?)" + eot_re, re.DOTALL) + logger.info( + f"Detected template style (known pattern): " + f"header_re={hdr_re!r}, eot_re={eot_re!r}" + ) + return pattern + + # --- Strategy 2: double-probe diff --- + try: + probe_empty = [ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": ""}, + ] + text_empty = tokenizer.apply_chat_template( + probe_empty, tokenize=False, **extra_kwargs + ) + + marker_idx = probe_text.index(_PROBE_CONTENT) + header = probe_text[:marker_idx] + tail = probe_text[marker_idx + len(_PROBE_CONTENT) :] + + if text_empty == header + tail: + # Extract the assistant-specific header by removing the shared + # user-only prefix. + user_only = tokenizer.apply_chat_template( + [{"role": "user", "content": "x"}], + tokenize=False, + **extra_kwargs, + ) + asst_header = header[len(user_only) :] + # end-of-turn delimiter: strip leading newlines, then take + # up to the first newline (or the full string if none). + eot_stripped = tail.lstrip("\n") + eot = eot_stripped.split("\n")[0] if "\n" in eot_stripped else eot_stripped + + if asst_header and eot: + hdr_re = re.escape(asst_header) + eot_re = re.escape(eot) + pattern = re.compile(hdr_re + r"(.*?)" + eot_re, re.DOTALL) + logger.info( + f"Detected template style (probe diff): " + f"header={asst_header!r}, eot={eot!r}" + ) + return pattern + except (ValueError, IndexError): + pass # PROBE_CONTENT not found in rendered text, skip + + raise ValueError( + "Could not detect chat template assistant delimiters. " + "Unable to build a reliable loss mask. " + f"Probe text: {probe_text[:200]!r}" + ) + + +def _dump_samples( + samples, + tokenizer, + assistant_pattern, + tools_list, + dump_dir, + n_samples, + *, + split_mode="pair", + error_indices_list=None, + parse_tool_call_args=False, +): + """Dump sampled message lists as ``.txt`` + ``.json`` for inspection. + + Args: + samples: List of message-list samples (pairs or full trajectories). + tokenizer: Tokenizer with ``apply_chat_template`` support. + assistant_pattern: Compiled regex from ``_detect_template_pattern``. + tools_list: Per-sample tool definitions (parallel to *samples*), + or ``None`` when no tools are available. + dump_dir: Directory to write files into (created if needed). + n_samples: Number of random samples to dump. ``-1`` dumps all. + split_mode: ``"trajectory"`` for trajectory-mode loss masking. + error_indices_list: Per-sample error segment indices (trajectory mode). + """ + import random as _random + + os.makedirs(dump_dir, exist_ok=True) + + if n_samples == -1 or n_samples >= len(samples): + indices = list(range(len(samples))) + else: + indices = sorted(_random.sample(range(len(samples)), n_samples)) + + n_written = 0 + for i in indices: + sample = samples[i] + sample_tools = tools_list[i] if tools_list else None + err_idxs = ( + error_indices_list[i] + if split_mode == "trajectory" and error_indices_list + else None + ) + + result = _render_tokenize_mask( + sample, + tokenizer, + assistant_pattern, + sample_tools, + split_mode=split_mode, + error_indices=err_idxs, + parse_tool_call_args=parse_tool_call_args, + ) + if result is None: + continue + + full_text, input_ids, loss_mask, offset_mapping = result + n_loss = sum(loss_mask) + base = os.path.join(dump_dir, f"sample_{i}") + + # --- .txt --- + with open(base + ".txt", "w", encoding="utf-8") as fout: + fout.write( + f"Sample {i}: {len(sample)} messages, " + f"{len(input_ids)} tokens, loss=1: {n_loss}\n" + ) + fout.write(f"Last msg role: {sample[-1]['role']}\n") + fout.write(f"{'=' * 72}\n\n") + + fout.write("--- Rendered Text ---\n") + fout.write(full_text) + fout.write("\n\n") + + fout.write("--- Token / Loss Mask ---\n") + fout.write(f"{'Idx':>6} | {'TokenID':>8} | Loss | Token Text\n") + fout.write(f"{'-' * 6}-+-{'-' * 8}-+------+{'-' * 40}\n") + for t in range(len(input_ids)): + cs, ce = offset_mapping[t] + tok_text = repr(full_text[cs:ce]) + fout.write( + f"{t:>6} | {input_ids[t]:>8} | {loss_mask[t]:>4} | {tok_text}\n" + ) + + # --- .json --- + tokens_list = [] + for t in range(len(input_ids)): + cs, ce = offset_mapping[t] + tokens_list.append( + { + "idx": t, + "token_id": input_ids[t], + "text": full_text[cs:ce], + "loss": loss_mask[t], + } + ) + record = { + "sample_index": i, + "n_messages": len(sample), + "n_tokens": len(input_ids), + "n_loss_tokens": n_loss, + "rendered_text": full_text, + "tokens": tokens_list, + } + with open(base + ".json", "w", encoding="utf-8") as fout: + json.dump(record, fout, ensure_ascii=False) + + n_written += 1 + + logger.info(f"Dumped {n_written} samples to {dump_dir}/") + + +# ============================================================ +# 6. Pipeline — loading, processing, distributed cache, public API +# ============================================================ + + +def _load_trajectory_pairs( + path: str, + filter_errors: bool = True, + strip_all_thinking: bool = False, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + max_no_thinking_ratio: float | None = None, + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, +): + """Load trajectory JSONL and split into progressive pairs. + + When *n_thinking_variants* > 1, each trajectory is split K times: + variant 0 preserves all thinking, variants 1~K-1 randomly strip. + + Supports nested (``conversations`` wrapper) and flat JSONL formats + (auto-detected per record via ``_iter_jsonl_records``). + + Returns: + Tuple of ``(all_pairs, tools)`` where *tools* is ``None`` when no + tool definitions are found. + """ + all_pairs = [] + all_tools = [] + records_in = 0 + total_filtered_errors = 0 + total_filtered_empty_tc = 0 + total_filtered_bare_tc = 0 + total_truncated = 0 + total_stripped_thinking = 0 + + augment = n_thinking_variants > 1 + rng = ( + random.Random(random_strip_thinking_seed) + if random_strip_thinking_prob > 0.0 + else None + ) + + if augment and random_strip_thinking_prob <= 0.0: + logger.warning( + "n_thinking_variants=%d but random_strip_thinking_prob=0; " + "all variants will be identical.", + n_thinking_variants, + ) + + # Stats collectors for augmentation logging. + thinking_turns_per_traj = [] + total_asst_turns_per_traj = [] + patterns_per_traj = [] + + for record_idx, messages, record_tools in _iter_jsonl_records(path): + records_in = record_idx + + if truncate_task_notifications: + truncated = _truncate_at_task_notification(messages) + if len(truncated) < len(messages): + total_truncated += 1 + messages = truncated + + shared_kwargs = dict( + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + ) + + if augment: + # Variant 0: preserve all thinking. + pairs_orig, n_err, n_empty_tc, n_bare_tc, _ = _split_and_filter( + messages, **shared_kwargs, random_strip_thinking_prob=0.0, rng=None + ) + total_filtered_errors += n_err + total_filtered_empty_tc += n_empty_tc + total_filtered_bare_tc += n_bare_tc + all_pairs.extend(pairs_orig) + all_tools.extend([record_tools] * len(pairs_orig)) + # Collect stats. + segments = _find_segments(messages) + n_think = sum(1 for s, _ in segments if _msg_has_thinking(messages[s])) + n_asst = len(segments) + thinking_turns_per_traj.append(n_think) + total_asst_turns_per_traj.append(n_asst) + + # Variants 1 ~ K-1: random strip. + variant_patterns = {frozenset()} # original = no strip + for _k in range(n_thinking_variants - 1): + pairs_aug, _, _, _, n_stripped = _split_and_filter( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + total_stripped_thinking += n_stripped + all_pairs.extend(pairs_aug) + all_tools.extend([record_tools] * len(pairs_aug)) + # Approximate pattern: record which pairs had their target stripped. + # For stats, use the count as a proxy since _split_and_filter + # doesn't return per-pair strip info. + variant_patterns.add(frozenset([n_stripped])) + patterns_per_traj.append(variant_patterns) + else: + # Single variant (original behavior). + pairs, n_err, n_empty_tc, n_bare_tc, n_stripped = _split_and_filter( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + total_filtered_errors += n_err + total_filtered_empty_tc += n_empty_tc + total_filtered_bare_tc += n_bare_tc + total_stripped_thinking += n_stripped + all_pairs.extend(pairs) + all_tools.extend([record_tools] * len(pairs)) + + # Log extracted tools summary. + n_with_tools = sum(1 for t in all_tools if t is not None) + if n_with_tools > 0: + all_tool_names = set() + for t_list in all_tools: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info( + f"Extracted tools from {n_with_tools}/{len(all_tools)} pairs: " + f"{sorted(all_tool_names)}" + ) + + filter_parts = [] + if total_truncated: + filter_parts.append( + f"{total_truncated} trajectories truncated at task-notification" + ) + if total_filtered_errors: + filter_parts.append(f"{total_filtered_errors} with tool errors") + if total_filtered_empty_tc: + filter_parts.append(f"{total_filtered_empty_tc} empty-content tool calls") + if total_filtered_bare_tc: + filter_parts.append(f"{total_filtered_bare_tc} bare-text tool calls") + if total_stripped_thinking: + filter_parts.append(f"{total_stripped_thinking} thinking blocks stripped") + filter_msg = ", ".join(filter_parts) if filter_parts else "none" + + logger.info( + f"Loaded {records_in} trajectories, " + f"generated {len(all_pairs)} pairs " + f"(filtered: {filter_msg})" + ) + + if augment and patterns_per_traj: + _log_thinking_augmentation_stats( + n_thinking_variants, + random_strip_thinking_prob, + records_in, + thinking_turns_per_traj, + total_asst_turns_per_traj, + patterns_per_traj, + ) + + # Balance thinking / no-thinking pair ratio. + all_pairs, all_tools = _balance_thinking_pairs( + all_pairs, max_no_thinking_ratio, tools_list=all_tools + ) + + return all_pairs, all_tools + + +def _load_presplit_pairs( + path: str, + strip_all_thinking: bool = False, + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, +): + """Load pre-split pair JSONL where each line is ``{"messages": [...]}``. + + Messages are cleaned but no splitting or error-filtering is performed. + By default, thinking is stripped from context assistant turns but + preserved for the last assistant turn (the training target). Set + *strip_all_thinking* to strip from every assistant turn. + + When *n_thinking_variants* > 1, each pair is augmented: variant 0 + preserves thinking, variants 1~K-1 randomly strip the target turn. + + Also extracts per-record ``tools`` definitions so that each pair + carries its own tools, same as ``_load_trajectory_pairs``. + + Returns: + Tuple of ``(all_pairs, all_tools)`` where *all_tools* is a + parallel list of per-sample tool definitions (may be ``None``). + """ + all_pairs = [] + all_tools = [] + n_stripped = 0 + augment = n_thinking_variants > 1 + + rng = ( + random.Random(random_strip_thinking_seed) + if random_strip_thinking_prob > 0.0 + else None + ) + + def _build_pair(messages, last_asst, strip_target): + pair = [] + for idx, m in enumerate(messages): + is_target = m.get("role") == "assistant" and idx == last_asst + strip = strip_all_thinking or not is_target or strip_target + pair.append(_clean_message(m, strip_thinking=strip)) + return pair + + with open(path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + record = json.loads(line) + messages = record.get("messages", []) + if not messages: + continue + + record_tools = record.get("tools") + + # Find the last assistant index so we can preserve its thinking. + last_asst = None + for i, m in enumerate(messages): + if m.get("role") == "assistant": + last_asst = i + + has_thinking = ( + last_asst is not None + and not strip_all_thinking + and _msg_has_thinking(messages[last_asst]) + ) + + if augment: + # Variant 0: preserve all thinking. + all_pairs.append(_build_pair(messages, last_asst, strip_target=False)) + all_tools.append(record_tools) + + # Variants 1 ~ K-1: random strip. + for _k in range(n_thinking_variants - 1): + do_strip = ( + has_thinking + and rng is not None + and rng.random() < random_strip_thinking_prob + ) + if do_strip: + n_stripped += 1 + all_pairs.append( + _build_pair(messages, last_asst, strip_target=do_strip) + ) + all_tools.append(record_tools) + else: + # Single variant (original behavior). + strip_target = ( + has_thinking + and rng is not None + and rng.random() < random_strip_thinking_prob + ) + if strip_target: + n_stripped += 1 + all_pairs.append( + _build_pair(messages, last_asst, strip_target=strip_target) + ) + all_tools.append(record_tools) + + # Log extracted tools summary. + n_with_tools = sum(1 for t in all_tools if t is not None) + if n_with_tools > 0: + all_tool_names = set() + for t_list in all_tools: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info( + f"Extracted tools from {n_with_tools}/{len(all_tools)} pairs: " + f"{sorted(all_tool_names)}" + ) + + strip_msg = f", {n_stripped} thinking blocks stripped" if n_stripped else "" + logger.info(f"Loaded {len(all_pairs)} pre-split pairs from {path}{strip_msg}") + return all_pairs, all_tools + + +def _load_full_trajectories( + path: str, + filter_errors: bool = True, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, +): + """Load trajectory JSONL for trajectory-level training. + + Each trajectory becomes a single training sample with all assistant + turns as targets (``loss_mask=1``). When *filter_errors* is True, + assistant segments with error tool responses are identified so + tokenization can mask them (``loss_mask=0``) instead of discarding + the entire trajectory. + + When *n_thinking_variants* > 1, each trajectory is augmented into + K variants: the first preserves all thinking, the remaining K-1 + randomly strip thinking turns with *random_strip_thinking_prob*. + + Supports nested (``conversations`` wrapper) and flat JSONL formats + (auto-detected per record via ``_iter_jsonl_records``). + + Returns: + Tuple of ``(trajectories, error_indices_list, all_tools)`` where + *trajectories* is a list of cleaned message lists, + *error_indices_list* is a list of error segment index lists, + and *all_tools* is a parallel list of per-sample tool definitions. + """ + trajectories = [] + error_indices_list = [] + all_tools = [] + records_in = 0 + total_truncated = 0 + total_masked_errors = 0 + total_masked_empty_tc = 0 + total_masked_bare_tc = 0 + total_stripped_thinking = 0 + + augment = n_thinking_variants > 1 + rng = ( + random.Random(random_strip_thinking_seed) + if random_strip_thinking_prob > 0.0 + else None + ) + + if augment and random_strip_thinking_prob <= 0.0: + logger.warning( + "n_thinking_variants=%d but random_strip_thinking_prob=0; " + "all variants will be identical.", + n_thinking_variants, + ) + + # Stats collectors for augmentation logging. + thinking_turns_per_traj = [] + total_asst_turns_per_traj = [] + patterns_per_traj = [] + + for record_idx, messages, record_tools in _iter_jsonl_records(path): + records_in = record_idx + + if truncate_task_notifications: + truncated = _truncate_at_task_notification(messages) + if len(truncated) < len(messages): + total_truncated += 1 + messages = truncated + + shared_kwargs = dict( + filter_errors=filter_errors, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + ) + + if augment: + # Variant 0: preserve all thinking (no stripping). + result_orig = _prepare_trajectory( + messages, **shared_kwargs, random_strip_thinking_prob=0.0, rng=None + ) + if result_orig is None: + continue + cleaned_orig, masked_idxs, n_err, n_empty_tc, n_bare_tc, _ = result_orig + trajectories.append(cleaned_orig) + error_indices_list.append(masked_idxs) + all_tools.append(record_tools) + total_masked_errors += n_err + total_masked_empty_tc += n_empty_tc + total_masked_bare_tc += n_bare_tc + + # Collect stats: count thinking turns in this trajectory. + segments = _find_segments(messages) + n_think = sum(1 for s, _ in segments if _msg_has_thinking(messages[s])) + n_asst = len(segments) + thinking_turns_per_traj.append(n_think) + total_asst_turns_per_traj.append(n_asst) + + # Variants 1 ~ K-1: random strip thinking. + variant_patterns = {frozenset()} # original = empty pattern + for _k in range(n_thinking_variants - 1): + result_aug = _prepare_trajectory( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + if result_aug is None: + continue + cleaned_aug, _, _, _, _, strip_pattern = result_aug + trajectories.append(cleaned_aug) + error_indices_list.append(masked_idxs) # reuse + all_tools.append(record_tools) + total_stripped_thinking += len(strip_pattern) + variant_patterns.add(strip_pattern) + patterns_per_traj.append(variant_patterns) + else: + # Single variant (original behavior). + result = _prepare_trajectory( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + if result is None: + continue + cleaned, masked_idxs, n_err, n_empty_tc, n_bare_tc, strip_pattern = result + trajectories.append(cleaned) + error_indices_list.append(masked_idxs) + all_tools.append(record_tools) + total_masked_errors += n_err + total_masked_empty_tc += n_empty_tc + total_masked_bare_tc += n_bare_tc + total_stripped_thinking += len(strip_pattern) + + # Log extracted tools summary. + n_with_tools = sum(1 for t in all_tools if t is not None) + if n_with_tools > 0: + all_tool_names = set() + for t_list in all_tools: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info( + f"Extracted tools from {n_with_tools}/{len(all_tools)} " + f"trajectories: {sorted(all_tool_names)}" + ) + + parts = [] + if total_truncated: + parts.append(f"{total_truncated} trajectories truncated at task-notification") + if total_masked_errors: + parts.append(f"{total_masked_errors} with tool errors") + if total_masked_empty_tc: + parts.append(f"{total_masked_empty_tc} empty-content tool calls") + if total_masked_bare_tc: + parts.append(f"{total_masked_bare_tc} bare-text tool calls") + if total_stripped_thinking: + parts.append(f"{total_stripped_thinking} thinking blocks stripped") + mask_msg = ", ".join(parts) if parts else "none" + + logger.info( + f"Loaded {records_in} trajectories, " + f"kept {len(trajectories)} for training " + f"(masked: {mask_msg})" + ) + + if augment and patterns_per_traj: + _log_thinking_augmentation_stats( + n_thinking_variants, + random_strip_thinking_prob, + records_in, + thinking_turns_per_traj, + total_asst_turns_per_traj, + patterns_per_traj, + ) + + return trajectories, error_indices_list, all_tools + + +def _tokenize_samples( + messages_list, + tools_list, + tokenizer, + *, + split_mode: str = "pair", + error_indices_list: list | None = None, + max_length: int | None = None, + num_proc: int | None = None, + no_tools: bool = False, + dump_dir: str | None = None, + dump_n_samples: int = 0, + parse_tool_call_args: bool = False, +): + """Tokenize message lists into a training-ready Dataset. + + Works for both progressive pairs (``split_mode="pair"``) and + full trajectories (``split_mode="trajectory"``). + + In pair mode, only the last assistant turn per sample gets + ``loss_mask=1``. In trajectory mode, all assistant turns get + ``loss_mask=1`` except those at error segment indices. + + Args: + tools_list: Per-sample tool definitions (parallel to + *messages_list*). Each element is either ``None`` or a + list of tool dicts. + """ + if num_proc is None: + num_proc = max(1, min(os.cpu_count() or 1, DATASET_NUM_PROC)) + + # Find representative tools for template detection. + first_tools = None + if tools_list: + first_tools = next((t for t in tools_list if t is not None), None) + + if no_tools: + tools_list = None + first_tools = None + logger.info("Tool definitions disabled (no_tools=True)") + elif first_tools is not None: + all_tool_names = set() + for t_list in tools_list: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info(f"Using tools for chat template: {sorted(all_tool_names)}") + + if not messages_list: + raise ValueError("No valid samples to tokenize") + + # Build dataset columns. + data = {"messages": messages_list} + # Serialize per-sample tools as JSON strings for the Dataset column. + data["tools_json"] = ( + [json.dumps(t) if t else "" for t in tools_list] + if tools_list + else [""] * len(messages_list) + ) + remove_cols = ["messages", "tools_json"] + if split_mode == "trajectory": + data["error_indices"] = error_indices_list or [[] for _ in messages_list] + remove_cols.append("error_indices") + + dataset = Dataset.from_dict(data) + _patch_chat_template_for_training(tokenizer) + assistant_pattern = _detect_template_pattern(tokenizer, tools=first_tools) + + # Dump samples for inspection before the heavy map() pass. + if dump_dir and dump_n_samples != 0: + _dump_samples( + messages_list, + tokenizer, + assistant_pattern, + tools_list, + dump_dir, + dump_n_samples, + split_mode=split_mode, + error_indices_list=error_indices_list, + parse_tool_call_args=parse_tool_call_args, + ) + + process_fn = _TokenizeAndMask( + tokenizer, + assistant_pattern, + max_length=max_length, + split_mode=split_mode, + parse_tool_call_args=parse_tool_call_args, + ) + + dataset = dataset.map(process_fn, num_proc=num_proc).remove_columns(remove_cols) + + # Single filter pass: removes both apply_chat_template-failure empties and + # overlength samples (which _TokenizeAndMask also marks as empty). + before_filter = len(dataset) + dataset = dataset.filter(lambda x: len(x["input_ids"]) > 0, num_proc=num_proc) + n_filtered = before_filter - len(dataset) + if n_filtered > 0: + logger.info( + f"Filtered {n_filtered} samples " + f"(empty from template failures or exceeding max_length={max_length})" + ) + + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + + +def _process_swe_sft( + path: str, + tokenizer, + *, + max_length: int | None = None, + num_proc: int | None = None, + pre_split: bool = False, + filter_errors: bool = True, + strip_all_thinking: bool = False, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + no_tools: bool = False, + max_no_thinking_ratio: float | None = None, + split_mode: str = "pair", + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, + dump_dir: str | None = None, + dump_n_samples: int = 0, + parse_tool_call_args: bool = False, +): + """Load JSONL, split into pairs, tokenize, and filter. + + Combines file loading with ``_tokenize_samples`` so that the rank-0-only + path and the single-process path share the same logic. + + When *split_mode* is ``"trajectory"``, the full trajectory is kept as a + single training sample with all assistant turns as targets. + """ + error_indices_list = None + + if split_mode == "trajectory": + messages_list, error_indices_list, tools_list = _load_full_trajectories( + path, + filter_errors=filter_errors, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + ) + elif pre_split: + messages_list, tools_list = _load_presplit_pairs( + path, + strip_all_thinking=strip_all_thinking, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + ) + else: + messages_list, tools_list = _load_trajectory_pairs( + path, + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + max_no_thinking_ratio=max_no_thinking_ratio, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + ) + + return _tokenize_samples( + messages_list, + tools_list, + tokenizer, + split_mode=split_mode, + error_indices_list=error_indices_list, + max_length=max_length, + num_proc=num_proc, + no_tools=no_tools, + dump_dir=dump_dir, + dump_n_samples=dump_n_samples, + parse_tool_call_args=parse_tool_call_args, + ) + + +def get_swe_sft_dataset( + path: str, + split: str | None = None, + tokenizer=None, + max_length: int | None = None, + num_proc: int | None = None, + pre_split: bool = False, + filter_errors: bool = True, + strip_all_thinking: bool = False, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + no_tools: bool = False, + skip_pretokenized_filter: bool = False, + max_no_thinking_ratio: float | None = None, + split_mode: str = "pair", + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, + cache_dir: str | None = None, + dump_dir: str | None = None, + dump_samples: int = 0, + parse_tool_call_args: bool = False, +): + """Load SWE trajectory data and convert to SFT training pairs. + + By default, tool definitions are auto-extracted from the training data's + ``conversations[].tools`` field and passed to ``apply_chat_template`` + so that the tokenizer renders tool definitions in the system prompt + (e.g. Qwen3 ``# Tools`` block), matching the eval-time format. + Set *no_tools* to skip this and render without tool definitions. + + When *split_mode* is ``"trajectory"``, the full trajectory is kept as a + single training sample with all assistant turns as targets + (``loss_mask=1``). Error segments are masked (``loss_mask=0``) + when *filter_errors* is True, instead of being discarded. + Thinking is preserved by default but can be randomly stripped + per-turn via *random_strip_thinking_prob* (both modes). + + In distributed (SPMD) mode, only rank 0 performs the heavy processing + (JSONL loading, pair splitting, tokenization) and saves the result as + an Arrow dataset to *cache_dir*. Other ranks wait for rank 0 to + finish and then load the cached dataset directly via memory-mapped I/O. + + Args: + path: Path to the JSONL file containing SWE trajectories, or a + directory containing a pre-tokenized Arrow dataset (saved by + ``python -m areal.dataset.swe_sft --save-tokenized``). + split: Unused, kept for API compatibility. + tokenizer: Tokenizer with ``apply_chat_template`` support. + Not required when loading a pre-tokenized dataset. + max_length: Max token length. Longer sequences are filtered out. + num_proc: Number of parallel workers for tokenization. + Defaults to ``min(os.cpu_count(), DATASET_NUM_PROC)``. + pre_split: If True, treat input as pre-split pairs (each line is + ``{"messages": [...]}``) instead of full trajectories. + filter_errors: If True (default), discard pairs whose current segment + contains a tool result with ``is_error=True``. In trajectory + mode, sets ``loss_mask=0`` for error segments instead. + Set to False to keep/train all regardless of tool errors. + strip_all_thinking: If True, strip ``...`` from every + assistant turn including the training target. + Ignored in trajectory mode (thinking is always preserved). + filter_empty_tool_calls: If True, discard pairs whose training-target + assistant turn has no text content but has tool_calls. + filter_bare_text_tool_calls: If True, discard pairs whose + training-target assistant turn has text without ```` + tags and has tool_calls. + truncate_task_notifications: If True, truncate trajectories at the + first ```` that follows a pure-text assistant + turn, removing noise from background task completions. + no_tools: If True, do not pass tool definitions to + ``apply_chat_template`` even if the data contains them. + skip_pretokenized_filter: If True, skip the ``max_length`` filter + when loading a pre-tokenized dataset. Useful when the dataset + was already filtered during pretokenization and you want to + avoid NFS cache conflicts from concurrent ``dataset.filter()`` + calls across ranks. + max_no_thinking_ratio: Maximum ratio of non-thinking pairs to thinking + pairs. For example, ``1.0`` gives 1:1, ``2.0`` gives 1:2. + ``None`` (default) disables balancing. + split_mode: ``"pair"`` (default) splits trajectories into + progressive pairs. ``"trajectory"`` keeps the full trajectory + as a single sample — all assistant turns are targets with + ``loss_mask=1``, error segments are masked instead of filtered. + random_strip_thinking_prob: Probability of stripping thinking from + each target assistant turn. 0.0 (default) = no stripping, + 1.0 = strip all. Works in both pair and trajectory mode. + random_strip_thinking_seed: Random seed for reproducible thinking + stripping decisions. + n_thinking_variants: Number of thinking-pattern variants per + trajectory. ``1`` (default) = no augmentation. ``K > 1`` + = augment each trajectory into K variants: the first + preserves all thinking, the rest randomly strip with + *random_strip_thinking_prob*. + cache_dir: Directory to save/load the processed Arrow dataset. + When set in distributed mode, rank 0 processes the data and + saves here; other ranks load from this directory. If the + directory already contains a completed cache (``.done`` marker), + all ranks load from it directly without reprocessing. + dump_dir: Directory to write sample dump files (``.txt`` + ``.json``). + Only rank 0 writes. Set to None to disable. + dump_samples: Number of random samples to dump. ``-1`` = all, + ``0`` = disabled. + parse_tool_call_args: If True, convert OpenAI JSON-string + ``tool_calls.arguments`` to dicts before ``apply_chat_template``. + Required by GLM-4.x / GLM-5.x templates; leave at the default + (False) for Qwen / Llama / Bailing. + + Returns: + A HuggingFace ``Dataset`` with ``input_ids`` and ``loss_mask`` columns. + """ + from datasets import load_from_disk + + # Pre-tokenized Arrow dataset: load directly, skip all processing. + if os.path.isdir(path): + logger.info(f"Loading pre-tokenized dataset from {path}") + dataset = load_from_disk(path) + + if max_length is not None and not skip_pretokenized_filter: + before_filter = len(dataset) + dataset = dataset.filter( + lambda x: len(x["input_ids"]) <= max_length, num_proc=num_proc + ) + logger.info( + f"Filtered {before_filter - len(dataset)} samples " + f"exceeding max_length={max_length}" + ) + + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + + # --- Shared kwargs for _process_swe_sft --- + process_kwargs = dict( + max_length=max_length, + num_proc=num_proc, + pre_split=pre_split, + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + no_tools=no_tools, + max_no_thinking_ratio=max_no_thinking_ratio, + split_mode=split_mode, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + dump_dir=dump_dir, + dump_n_samples=dump_samples, + parse_tool_call_args=parse_tool_call_args, + ) + + # --- Distributed rank-0-only processing --- + rank = int(os.getenv("RANK", "0")) + world_size = int(os.getenv("WORLD_SIZE", "1")) + + if cache_dir is not None and world_size > 1: + done_marker = os.path.join(cache_dir, ".done") + meta_path = os.path.join(cache_dir, ".meta.json") + cache_meta = { + "version": 1, + "path": path, + "tokenizer": getattr(tokenizer, "name_or_path", None), + "process_kwargs": { + k: v + for k, v in process_kwargs.items() + if k not in ("dump_dir", "dump_n_samples") + }, + } + + def _filter_by_max_length(ds): + if max_length is None: + return ds + before = len(ds) + # Length via arrow list offsets: avoids decoding every row to + # Python lists, which for long-context datasets costs minutes of + # startup per rank while (on a validated cache) removing nothing — + # build-time _TokenizeAndMask already filtered with this max_length. + import pyarrow.compute as pc + + # ds.data is the underlying arrow table; a freshly built dataset + # carries an indices mapping (from .filter views) whose row count + # differs. Materialize the view first (no-op for load_from_disk). + if getattr(ds, "_indices", None) is not None: + ds = ds.flatten_indices() + lengths = pc.list_value_length(ds.data.column("input_ids")).to_pylist() + keep = [i for i, n in enumerate(lengths) if n <= max_length] + ds = ds.select(keep) + if len(ds) < before: + logger.info( + f"Rank {rank}: filtered {before - len(ds)} samples " + f"exceeding max_length={max_length}" + ) + if len(ds) == 0: + raise ValueError( + f"processed dataset at {cache_dir} has 0 samples after " + f"max_length={max_length} filtering" + ) + return ds + + def _load_valid_cache(): + if not os.path.exists(meta_path): + raise ValueError(f"cached dataset metadata is missing: {meta_path}") + with open(meta_path) as f: + cached_meta = json.load(f) + if cached_meta != cache_meta: + raise ValueError( + f"cached dataset metadata does not match current SWE settings: " + f"{meta_path}" + ) + dataset = load_from_disk(cache_dir) + if len(dataset) == 0: + raise ValueError(f"cached dataset is empty: {cache_dir}") + return dataset + + def _wait_for_valid_cache(): + start = time.monotonic() + last_error = None + while True: + if os.path.exists(done_marker): + try: + return _load_valid_cache() + except Exception as e: + last_error = e + elapsed = time.monotonic() - start + if elapsed > _RANK0_CACHE_TIMEOUT: + raise TimeoutError( + f"Waited {_RANK0_CACHE_TIMEOUT}s for rank 0 to rebuild " + f"a valid dataset cache at {cache_dir}. Last error: {last_error}" + ) + time.sleep(_RANK0_CACHE_POLL_INTERVAL) + + # Fast path: cache from a previous run (or rank 0 already finished). + if os.path.exists(done_marker): + if rank == 0: + try: + logger.info( + f"Rank {rank}: loading cached processed dataset from {cache_dir}" + ) + dataset = _load_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + except Exception as e: + logger.warning( + "Rank 0: invalid processed dataset cache at %s (%s); " + "rebuilding it.", + cache_dir, + e, + ) + shutil.rmtree(cache_dir, ignore_errors=True) + else: + try: + logger.info( + f"Rank {rank}: loading cached processed dataset from {cache_dir}" + ) + dataset = _load_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + except Exception as e: + logger.warning( + "Rank %d: cached processed dataset at %s is not usable " + "(%s); waiting for rank 0 to rebuild it.", + rank, + cache_dir, + e, + ) + dataset = _wait_for_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info( + f"Rank {rank}: loaded rebuilt dataset ({len(dataset)} samples)" + ) + return dataset + + if rank == 0: + # Rank 0: do the heavy processing and save for other ranks. + dataset = _process_swe_sft(path, tokenizer, **process_kwargs) + if len(dataset) == 0: + raise RuntimeError( + "SWE SFT preprocessing produced 0 samples; refusing to cache " + "an empty processed_dataset." + ) + shutil.rmtree(cache_dir, ignore_errors=True) + os.makedirs(cache_dir, exist_ok=True) + dataset.save_to_disk(cache_dir) + with open(meta_path, "w") as f: + json.dump(cache_meta, f, sort_keys=True) + # Write marker AFTER save completes so readers see a consistent dir. + with open(done_marker, "w") as f: + f.write(str(len(dataset))) + logger.info( + f"Rank 0: saved processed dataset " + f"({len(dataset)} samples) to {cache_dir}" + ) + dataset = _filter_by_max_length(dataset) + return dataset + else: + # Other ranks: wait for rank 0, then load with meta validation so a + # cache rebuilt for different settings (or mid-rmtree) is never + # silently loaded as this rank's dataset. + logger.info(f"Rank {rank}: waiting for rank 0 to process dataset...") + dataset = _wait_for_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info(f"Rank {rank}: loaded cached dataset ({len(dataset)} samples)") + return dataset + + # --- Non-distributed or no cache_dir: process in current process --- + return _process_swe_sft(path, tokenizer, **process_kwargs) + + +# ============================================================ +# 7. CLI — ``python -m areal.dataset.swe_sft`` +# ============================================================ + +if __name__ == "__main__": + import argparse + import sys + + from transformers import AutoTokenizer + + parser = argparse.ArgumentParser( + description="Verify SWE SFT pair generation and loss masking.", + ) + parser.add_argument("path", help="Path to SWE trajectory JSONL file") + parser.add_argument( + "--tokenizer", + default="Qwen/Qwen3-8B", + help="HuggingFace tokenizer name or path (default: Qwen/Qwen3-8B)", + ) + parser.add_argument( + "--max-length", + type=int, + default=None, + help="Filter samples exceeding this token length", + ) + parser.add_argument( + "--num-samples", + "-n", + type=int, + default=None, + help="Number of pairs to process. Controls loading, tokenization," + " display, and export. Default: all pairs.", + ) + parser.add_argument( + "--num-proc", + type=int, + default=None, + help=f"Number of parallel workers (default: min(cpu_count, {DATASET_NUM_PROC}))", + ) + parser.add_argument( + "--save-pairs", + "-o", + default=None, + metavar="FILE", + help='Save cleaned pairs to FILE (JSONL, each line: {"messages": [...]}).', + ) + parser.add_argument( + "--pre-split", + action="store_true", + help='Input is already in pair format (each line: {"messages": [...]}).' + " Skip trajectory splitting and error filtering.", + ) + parser.add_argument( + "--no-filter-errors", + action="store_true", + help="Keep pairs whose current segment contains tool results with " + "is_error=True (by default these are discarded).", + ) + parser.add_argument( + "--save-tokenized", + default=None, + metavar="DIR", + help="Save the tokenized dataset to DIR (Arrow format). " + "The saved directory can be used directly as the dataset path " + "during training, skipping all processing.", + ) + parser.add_argument( + "--strip-all-thinking", + action="store_true", + help="Strip ... from ALL assistant turns including " + "the training target. By default only context turns are stripped.", + ) + parser.add_argument( + "--no-tools", + action="store_true", + help="Do not pass tool definitions to apply_chat_template. " + "By default, tools are auto-extracted from the data and rendered " + "in the system prompt (e.g. Qwen3 '# Tools' block).", + ) + parser.add_argument( + "--parse-tool-call-args", + action="store_true", + help="Convert OpenAI JSON-string tool_calls.arguments to dicts " + "before apply_chat_template. Required by GLM-4.x / GLM-5.x " + "templates; leave off for Qwen / Llama / Bailing (which expect " + "the standard string form).", + ) + parser.add_argument( + "--filter-empty-tool-calls", + action="store_true", + help="Discard pairs whose training-target assistant turn has no " + "text content but has tool_calls (silent tool invocations).", + ) + parser.add_argument( + "--filter-bare-text-tool-calls", + action="store_true", + help="Discard pairs whose training-target assistant turn has text " + "content without tags and has tool_calls.", + ) + parser.add_argument( + "--truncate-task-notifications", + action="store_true", + help="Truncate trajectories at the first that " + "follows a pure-text assistant turn. Removes noise from background " + "task completions (e.g. pip install finishing after the model's summary).", + ) + parser.add_argument( + "--max-no-thinking-ratio", + type=float, + default=None, + help="Maximum ratio of non-thinking pairs to thinking pairs. " + "E.g. 1.0 = 1:1 balance, 2.0 = at most 2x non-thinking per " + "thinking pair. Non-thinking pairs are randomly downsampled. " + "Default: no balancing.", + ) + parser.add_argument( + "--split-mode", + choices=["pair", "trajectory"], + default="pair", + help="Sample construction mode. 'pair' (default): split trajectories " + "into progressive pairs. 'trajectory': keep the full trajectory " + "as a single sample with all assistant turns as targets.", + ) + parser.add_argument( + "--random-strip-thinking-prob", + type=float, + default=0.0, + help="Probability of stripping thinking from each target assistant " + "turn. 0.0 = no stripping (default), 1.0 = strip all. " + "Works in both pair mode and trajectory mode.", + ) + parser.add_argument( + "--random-strip-thinking-seed", + type=int, + default=42, + help="Random seed for reproducible thinking stripping decisions (default: 42).", + ) + parser.add_argument( + "--n-thinking-variants", + type=int, + default=1, + help="Number of thinking-pattern variants per trajectory. " + "1 = no augmentation (default). K > 1 = augment each trajectory " + "into K variants: the first preserves all thinking, the rest " + "randomly strip with --random-strip-thinking-prob.", + ) + parser.add_argument( + "--save-trajectories", + default=None, + metavar="FILE", + help="Save preprocessed trajectories to FILE (JSONL, original format) " + "after applying trajectory-level operations (e.g. " + "--truncate-task-notifications) but before pair splitting. " + "Each line preserves the original record structure with the " + "messages field updated.", + ) + parser.add_argument( + "--dump-samples", + default=None, + metavar="DIR", + help="Save sampled pairs to DIR, one file per pair. Each file " + "contains the rendered text and a token-by-token table with " + "token id, decoded text, and loss_mask.", + ) + parser.add_argument( + "--dump-n", + type=int, + default=None, + help="Number of pairs to dump when --dump-samples is set. " + "Default: all pairs. -1 also means all.", + ) + args = parser.parse_args() + + filter_errors = not args.no_filter_errors + strip_all_thinking = args.strip_all_thinking + filter_empty_tool_calls = args.filter_empty_tool_calls + filter_bare_text_tool_calls = args.filter_bare_text_tool_calls + truncate_task_notifications = args.truncate_task_notifications + max_no_thinking_ratio = args.max_no_thinking_ratio + + # --- Fast path: save preprocessed trajectories --- + if args.save_trajectories: + records_in = 0 + records_out = 0 + n_truncated = 0 + with ( + open(args.path, encoding="utf-8") as fin, + open(args.save_trajectories, "w", encoding="utf-8") as fout, + ): + for line in fin: + line = line.strip() + if not line: + continue + record = json.loads(line) + records_in += 1 + + messages, _ = _extract_messages(record, records_in) + + if truncate_task_notifications and messages: + truncated = _truncate_at_task_notification(messages) + if len(truncated) < len(messages): + n_truncated += 1 + _set_messages(record, truncated) + + fout.write(json.dumps(record, ensure_ascii=False) + "\n") + records_out += 1 + + parts = [] + if n_truncated: + parts.append(f"{n_truncated} truncated at task-notification") + op_msg = ", ".join(parts) if parts else "no changes" + print( + f"Saved {records_out}/{records_in} trajectories " + f"to {args.save_trajectories} ({op_msg})" + ) + sys.exit(0) + + # --- Load --- + split_mode = args.split_mode + error_indices_list = None + + if split_mode == "trajectory": + samples, error_indices_list, tools_list = _load_full_trajectories( + args.path, + filter_errors=filter_errors, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + random_strip_thinking_prob=args.random_strip_thinking_prob, + random_strip_thinking_seed=args.random_strip_thinking_seed, + n_thinking_variants=args.n_thinking_variants, + ) + label = "trajectories" + elif args.pre_split: + samples, tools_list = _load_presplit_pairs( + args.path, + strip_all_thinking=strip_all_thinking, + random_strip_thinking_prob=args.random_strip_thinking_prob, + random_strip_thinking_seed=args.random_strip_thinking_seed, + n_thinking_variants=args.n_thinking_variants, + ) + label = "pairs" + else: + samples, tools_list = _load_trajectory_pairs( + args.path, + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + max_no_thinking_ratio=max_no_thinking_ratio, + random_strip_thinking_prob=args.random_strip_thinking_prob, + random_strip_thinking_seed=args.random_strip_thinking_seed, + n_thinking_variants=args.n_thinking_variants, + ) + label = "pairs" + + # --- Slice + stats --- + total = len(samples) + if args.num_samples is not None: + samples = samples[: args.num_samples] + tools_list = tools_list[: args.num_samples] if tools_list else tools_list + if error_indices_list is not None: + error_indices_list = error_indices_list[: args.num_samples] + + print(f"Total {label}: {total}") + if args.num_samples is not None: + print(f"Using: {len(samples)}") + + if samples: + lengths = [len(s) for s in samples] + print( + f"Messages/sample: min={min(lengths)}, " + f"max={max(lengths)}, avg={sum(lengths) / len(lengths):.1f}" + ) + if error_indices_list is not None: + n_masked = sum(len(e) for e in error_indices_list) + print(f"Masked segments: {n_masked} (loss=0)") + + # --- Save cleaned samples as JSONL --- + if args.save_pairs: + with open(args.save_pairs, "w", encoding="utf-8") as fout: + err_iter = error_indices_list or [None] * len(samples) + tl_iter = tools_list if tools_list else [None] * len(samples) + for sample, sample_tools, err_idxs in zip(samples, tl_iter, err_iter): + record = {"messages": sample} + if sample_tools is not None: + record["tools"] = sample_tools + if err_idxs: + record["error_indices"] = err_idxs + fout.write(json.dumps(record, ensure_ascii=False) + "\n") + print(f"Wrote {len(samples)} {label} to {args.save_pairs}") + + # --- Tokenize / Dump --- + dump_dir = args.dump_samples if args.dump_samples else None + need_tokenize = args.save_tokenized + + # When --save-tokenized is set, auto-dump 50 samples alongside it + # unless the user explicitly set --dump-samples or --dump-n 0. + if need_tokenize and not dump_dir and args.dump_n != 0: + dump_dir = os.path.join(args.save_tokenized, "dumped_samples") + + if not need_tokenize and not dump_dir: + sys.exit(0) + + tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + _patch_chat_template_for_training(tok) + if args.dump_n is not None: + dump_n = args.dump_n + elif args.dump_samples: + # Explicit --dump-samples without --dump-n: dump all + dump_n = -1 if args.num_samples is None else args.num_samples + elif need_tokenize: + # Auto-dump with --save-tokenized: default 50 + dump_n = 50 + else: + dump_n = -1 + + # Dump can run independently without full tokenization. + if dump_dir and dump_n != 0: + dump_tools = None if args.no_tools else tools_list + first_tools = None + if dump_tools: + first_tools = next((t for t in dump_tools if t is not None), None) + assistant_pattern = _detect_template_pattern(tok, tools=first_tools) + _dump_samples( + samples, + tok, + assistant_pattern, + dump_tools, + dump_dir, + dump_n, + split_mode=split_mode, + error_indices_list=error_indices_list, + parse_tool_call_args=args.parse_tool_call_args, + ) + + if not need_tokenize: + sys.exit(0) + + ds = _tokenize_samples( + samples, + tools_list, + tok, + split_mode=split_mode, + error_indices_list=error_indices_list, + max_length=args.max_length, + num_proc=args.num_proc, + no_tools=args.no_tools, + parse_tool_call_args=args.parse_tool_call_args, + ) + + print(f"\nTokenized: {len(ds)} samples") + if args.save_tokenized: + ds.save_to_disk(args.save_tokenized) + print(f"Saved tokenized dataset ({len(ds)} samples) to {args.save_tokenized}") diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 7beb16ef69..7cf1e5d6b6 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -36,6 +36,7 @@ from transformers import PretrainedConfig import areal.models.mcore.bailing_moe_bridge # noqa: F401 # register bridge +import areal.models.mcore.bailing_v3_bridge # noqa: F401 # register bridge from areal.api import ( FinetuneSpec, InferenceEngine, @@ -87,6 +88,7 @@ ) from areal.infra.dist_rollout import DistRolloutCoordinator from areal.infra.platforms import current_platform, is_npu_available +from areal.models.mcore.bailing_v3_bridge import BailingV3Bridge from areal.models.mcore.hf_load import load_weights_from_hf_with_mbridge_fast from areal.models.mcore.hf_save import ( save_critic_value_head, @@ -507,6 +509,14 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec, *args, **kwargs): if self.mcore_config.use_deterministic_algorithms: set_deterministic_algorithms(self.tf_config, prebuild=True) + # Precision-alignment dumps (AReaL-friend tools/precision-alignment): + # when AREAL_DUMP_ROUTING is set, enable megatron RouterReplay + # recording so MoE expert indices can be captured during forward. + if os.environ.get("AREAL_DUMP_ROUTING", "") and hasattr( + self.tf_config, "moe_enable_routing_replay" + ): + self.tf_config.moe_enable_routing_replay = True + self.is_vision_model = is_valid_vision_model(self.hf_config.model_type) # GDN/SSM models (e.g. Qwen3.5) reject packed THD input and must run # the padded BSHD forward. Derived from model type rather than a @@ -714,9 +724,19 @@ def _build_glu_fc1_names(self) -> set[str]: def _build_hf_mcore_bridge(self): if self.bridge_cls == "mbridge": - self.bridge = mbridge.AutoBridge.from_pretrained( + hf_config = PretrainedConfig.from_pretrained( self.config.path, trust_remote_code=True ) + architectures = getattr(hf_config, "architectures", None) or [] + if "BailingMoeV3ForCausalLM" in architectures: + # BailingMoeV3 flash checkpoints keep model_type="bailing_hybrid", + # which overlaps the v2.5 bridge registration. Dispatch by + # architecture so KDA + gated-MLA weights use the v3 bridge. + self.bridge = BailingV3Bridge(hf_config) + else: + self.bridge = mbridge.AutoBridge.from_pretrained( + self.config.path, trust_remote_code=True + ) self.bridge.dtype = self.dtype if self.config.gradient_checkpointing: self.bridge.set_extra_args( @@ -728,13 +748,23 @@ def _build_hf_mcore_bridge(self): ) # Set MoE configuration overrides (aux-loss-free balancing, z-loss). + # mbridge extra_args override per-model bridge kwargs, so fields + # whose cli default may disagree with a bridge's deliberate + # default are forwarded only when explicitly configured + # (None = keep the bridge default). moe_extra_args: dict = { "moe_token_dispatcher_type": self.mcore_config.moe_token_dispatcher_type, "moe_permute_fusion": self.mcore_config.moe_permute_fusion, "moe_router_fusion": self.mcore_config.moe_router_fusion, - "moe_shared_expert_overlap": self.mcore_config.moe_shared_expert_overlap, - "moe_router_bias_update_rate": self.mcore_config.moe_router_bias_update_rate, } + if self.mcore_config.moe_shared_expert_overlap is not None: + moe_extra_args["moe_shared_expert_overlap"] = ( + self.mcore_config.moe_shared_expert_overlap + ) + if self.mcore_config.moe_router_bias_update_rate is not None: + moe_extra_args["moe_router_bias_update_rate"] = ( + self.mcore_config.moe_router_bias_update_rate + ) if self.mcore_config.moe_router_dtype is not None: moe_extra_args["moe_router_dtype"] = self.mcore_config.moe_router_dtype if self.mcore_config.moe_z_loss_coeff is not None: @@ -1020,6 +1050,11 @@ def save(self, meta: SaveLoadMeta): raise ValueError( "HF format does not support optimizer state saving, please use DCP format instead." ) + # HF export all-gathers full tensors across TP; reclaim allocator + # headroom first. Kept out of the dcp/recover path, which is + # frequency-driven and should not pay a full-heap GC per save. + gc.collect() + current_platform.empty_cache() self._save_model_to_hf( meta.path, tokenizer=meta.tokenizer, @@ -1099,6 +1134,11 @@ def optimizer_step(self): ) def lr_scheduler_step(self): + if os.environ.get("AREAL_DUMP_ROUTING", "") or os.environ.get( + "AREAL_DUMP_LOGP", "" + ): + # Precision-alignment forward-only mode: no optimizer/scheduler. + return assert self.lr_scheduler is not None, "LR Scheduler is not initialized." self.lr_scheduler.step(1) @@ -1172,6 +1212,23 @@ def forward_step(batch_iter, model): "BSHD is supported only for text-only models such as Qwen3.5" ) + # Precision-alignment routing dump: record MoE expert indices for the + # first microbatch via megatron RouterReplay (enabled in initialize). + _routing_dump_path = os.environ.get("AREAL_DUMP_ROUTING", "") + if _routing_dump_path and not getattr(self, "_routing_dumped", False): + try: + from megatron.core.transformer.moe.router_replay import ( + RouterReplay, + RouterReplayAction, + ) + + RouterReplay.set_global_router_replay_action( + RouterReplayAction.RECORD + ) + except Exception as e: + self.logger.warning(f"[ROUTING-DUMP] RECORD setup failed: {e}") + _routing_dump_path = "" + output = packed_context_parallel_forward( model, mb_input.padded_mb, @@ -1255,6 +1312,76 @@ def forward_step(batch_iter, model): ), ) + if _routing_dump_path and not getattr(self, "_routing_dumped", False): + try: + from megatron.core.transformer.moe.router_replay import ( + RouterReplay, + ) + + recorded = RouterReplay.get_recorded_data() + if recorded and any(r is not None for r in recorded): + pp_rank = mpu.get_pipeline_model_parallel_rank() + cp_rank = mpu.get_context_parallel_rank() + if ( + mpu.get_data_parallel_rank() == 0 + and mpu.get_tensor_model_parallel_rank() == 0 + ): + expert_indices = [ + r.detach().cpu() for r in recorded if r is not None + ] + # The router sees the CP-local token shard (zigzag + # split by packed_context_parallel_forward) and, + # with TP>1, the SP-local sub-shard of that. + # Apply the same CP split here; tp_rank 0 then + # holds the first contiguous SP chunk, so the + # row-count prefix slice below stays valid. + ids_src = mb_input.padded_mb["input_ids"] + cu = mb_input.padded_mb.get("cu_seqlens") + pos_src = torch.arange( + ids_src.numel(), device=ids_src.device + ) + if cp_size > 1 and cu is not None: + ids_src = split_packed_seqs_for_context_parallel( + ids_src, cu + ) + pos_src = split_packed_seqs_for_context_parallel( + pos_src, cu + ) + ids = ids_src.detach().cpu().reshape(-1) + pos = pos_src.detach().cpu().reshape(-1) + n_rows = expert_indices[0].shape[0] + save_data = { + "expert_indices": expert_indices, + "input_ids": ids[:n_rows], + # Canonical (padded packed-sequence) position + # of each recorded row so the merge tool can + # restore zigzag CP order; without it CP>1 + # shards cannot be mapped back. + "positions": pos[:n_rows], + "cp_size": cp_size, + "pp_rank": pp_rank, + "cp_rank": cp_rank, + } + if cu is not None: + save_data["padded_cu_seqlens"] = cu.detach().cpu() + orig_cu = mb_input.orig_mb.get("cu_seqlens") + if orig_cu is not None: + save_data["orig_cu_seqlens"] = orig_cu.detach().cpu() + out_file = ( + f"{_routing_dump_path}.pp{pp_rank}.cp{cp_rank}.pt" + ) + torch.save(save_data, out_file) + self.logger.info( + f"[ROUTING-DUMP] saved " + f"{len(save_data['expert_indices'])} MoE layers, " + f"pp={pp_rank} cp={cp_rank} -> {out_file}" + ) + RouterReplay.clear_global_indices() + RouterReplay.clear_global_router_replay_action() + self._routing_dumped = True + except Exception as e: + self.logger.warning(f"[ROUTING-DUMP] failed: {e}") + # Release tree attention metadata after forward pass for key in tree_attn_keys: del mb_input.padded_mb[key] @@ -1329,7 +1456,15 @@ def train_batch( self._ensure_ready() if self._awex_adapter is not None: self._awex_adapter.ensure_grad_buffers() - self.optimizer_zero_grad() + + # Precision-alignment forward-only mode: no optimizer exists (see + # _create_optimizer), so skip zero_grad/step and run forward only. + _fwd_only = bool( + os.environ.get("AREAL_DUMP_ROUTING", "") + or os.environ.get("AREAL_DUMP_LOGP", "") + ) + if not _fwd_only: + self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -1355,9 +1490,13 @@ def train_batch( # that extra division would shrink every gradient (and thus grad_norm and the # effective optimizer step) by `num_microbatches`. loss_multiplier = ( - mpu.get_data_parallel_world_size() - * self.optimizer.get_loss_scale().item() - * len(mb_list) + float(mpu.get_data_parallel_world_size()) + if _fwd_only + else ( + mpu.get_data_parallel_world_size() + * self.optimizer.get_loss_scale().item() + * len(mb_list) + ) ) def process_output( @@ -1375,9 +1514,12 @@ def process_output( self.forward_backward_batch( mb_list, process_output, - forward_only=False, + forward_only=_fwd_only, ) + if _fwd_only: + return {"num_micro_batches": len(mb_list.mbs)} + # Step 4: Optimizer step stats = self.optimizer_step() stats["num_micro_batches"] = len(mb_list.mbs) @@ -1794,6 +1936,14 @@ def _init_context_and_model_parallel_group(self) -> None: def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: if self.optimizer_config is None: return + if os.environ.get("AREAL_DUMP_ROUTING", "") or os.environ.get( + "AREAL_DUMP_LOGP", "" + ): + self.logger.info( + "[MegatronEngine] AREAL_DUMP_ROUTING/LOGP set, skipping optimizer " + "creation to save GPU memory (precision-alignment forward-only mode)." + ) + return assert self.model is not None and len(self.model) > 0 use_distributed_optimizer = ( @@ -2892,6 +3042,33 @@ def _compute_logprobs_and_loss( inputs = { k: v for k, v in inputs.items() if not k.startswith("_cp_") } + + # Precision-alignment logp dump: save final per-token logprobs for + # the first microbatch (last PP stage only; this branch already is). + _logp_dump_path = os.environ.get("AREAL_DUMP_LOGP", "") + if _logp_dump_path and not getattr(self, "_logp_dumped", False): + pp_rank = mpu.get_pipeline_model_parallel_rank() + cp_rank = mpu.get_context_parallel_rank() + if ( + mpu.get_data_parallel_rank() == 0 + and mpu.get_tensor_model_parallel_rank() == 0 + ): + save_data = { + "logprobs": logprobs.detach().cpu(), + "input_ids": inputs["input_ids"].detach().cpu(), + "pp_rank": pp_rank, + "cp_rank": cp_rank, + } + if "loss_mask" in inputs: + save_data["loss_mask"] = inputs["loss_mask"].detach().cpu() + out_file = f"{_logp_dump_path}.pp{pp_rank}.cp{cp_rank}.pt" + torch.save(save_data, out_file) + self.logger.info( + f"[LOGP-DUMP] pp={pp_rank} cp={cp_rank} " + f"logprobs={list(logprobs.shape)} -> {out_file}" + ) + self._logp_dumped = True + loss = loss_fn( logprobs, entropy, diff --git a/areal/models/mcore/bailing_v3.py b/areal/models/mcore/bailing_v3.py new file mode 100644 index 0000000000..c1263bca20 --- /dev/null +++ b/areal/models/mcore/bailing_v3.py @@ -0,0 +1,401 @@ +"""BailingMoeV3ForCausalLM (Ling V3) support for megatron-core. + +This module provides: +1. HF config -> MLATransformerConfig conversion +2. Heterogeneous layer spec construction (KDA + gated MLA) + +BailingMoeV3 uses: +- Mixed attention: KDA / Kimi Delta Attention (most layers) + gated MLA (every + ``layer_group_size``-th layer). KDA replaces the Lightning Attention used in V2.5. +- MoE: sigmoid routing, grouped TopK (n_group=8, topk_group=4), shared experts. +- Dense MLP for the first ``first_k_dense_replace`` layers, MoE for the rest. + +Layer pattern (layer_group_size=4): layers 0,1,2 = KDA, layer 3 = MLA, repeating. + +NOTE (first bring-up scope): MTP and FP8 are intentionally NOT wired here; the tiny +config is brought up in bf16 with MTP disabled. See docs/bailing-moe-v3-adaptation-plan.md. +""" + +import copy + +import torch +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, +) +from megatron.core.transformer.enums import LayerType +from megatron.core.transformer.multi_latent_attention import MLATransformerConfig +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import ( + TransformerBlockSubmodules, + get_num_layers_to_build, +) +from megatron.core.transformer.transformer_layer import get_transformer_layer_offset +from transformers import PretrainedConfig + +from areal.models.mcore.bailing_v3_mla import BailingV3MLASelfAttention +from areal.models.mcore.common import check_and_construct_configs, hf_to_mcore_base_args +from areal.models.mcore.kda_attention import ( + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, +) +from areal.utils import logging + +logger = logging.getLogger("BailingV3") + + +def is_kda_layer( + layer_number: int, layer_group_size: int, num_layers: int | None = None +) -> bool: + """Determine if a layer uses KDA (linear attention) vs MLA. + + In BailingMoeV3, layers are grouped by ``layer_group_size``. Within each group, the + last layer uses MLA (softmax attention) and the others use KDA (linear attention). + + For layer_group_size=4: layers 0,1,2 are KDA, layer 3 is MLA, repeating. + + Args: + layer_number: 0-indexed layer number. + layer_group_size: Number of layers per group. + + Returns: + True if the layer should use KDA. + """ + if layer_group_size <= 1: + return False + if num_layers is not None: + # Treat any incomplete tail group as softmax attention. HybridEngine v3 + # normally asserts divisibility, but this keeps AReaL's mapping total. + full_group_layers = (num_layers // layer_group_size) * layer_group_size + if layer_number >= full_group_layers: + return False + return (layer_number + 1) % layer_group_size != 0 + + +def _kda_head_dim(hf_config: PretrainedConfig) -> int: + """KDA key/value head dim. In Ling V3 this equals ``kv_channels`` (128).""" + return ( + getattr(hf_config, "kv_channels", None) + or getattr(hf_config, "head_dim", None) + or getattr(hf_config, "v_head_dim", None) + or getattr(hf_config, "qk_nope_head_dim", 128) + ) + + +def hf_to_mcore_config_bailing_v3( + hf_config: PretrainedConfig, + dtype: torch.dtype, +) -> MLATransformerConfig: + """Convert a BailingMoeV3 HuggingFace config to megatron-core MLATransformerConfig. + + KDA-specific knobs (conv kernel, no_kda_lora, head_dim) are NOT stored on the config; + they are passed to the KDA module via ModuleSpec params in ``_build_kda_attn_spec``. + """ + num_layers = hf_config.num_hidden_layers + first_k_dense_replace = getattr(hf_config, "first_k_dense_replace", 0) + moe_layer_freq = [0 if i < first_k_dense_replace else 1 for i in range(num_layers)] + + # Shared-expert intermediate size (direct value, else num_shared * moe_intermediate). + shared_expert_intermediate_size = getattr( + hf_config, "moe_shared_expert_intermediate_size", None + ) + if shared_expert_intermediate_size is None: + num_shared_experts = getattr(hf_config, "num_shared_experts", 0) + intermediate_size = getattr( + hf_config, "moe_intermediate_size", hf_config.intermediate_size + ) + shared_expert_intermediate_size = ( + num_shared_experts * intermediate_size if num_shared_experts > 0 else None + ) + + base_args = hf_to_mcore_base_args( + hf_config=hf_config, + dtype=dtype, + use_cpu_initialization=False, + add_bias_linear=False, + add_qkv_bias=False, + qk_layernorm=True, + attention_softmax_in_fp32=True, + cross_entropy_loss_fusion=False, + ) + + # MLA-specific parameters (for MLA layers). + # + # CRITICAL (carried over from V2.5): AReaL's mcore MLA receives qk_pos_emb_head_dim + # directly as the pure RoPE slice, so rotary_percent MUST be 1.0 (NOT the 0.5 that + # ant-megatron uses, where it passes kv_channels=128 and takes half). Using 0.5 here + # would halve the RoPE frequency table and badly degrade MLA accuracy. + # + # MLATransformerConfig defaults are tuned for DeepSeek-V2 YaRN; we pin them to plain + # rope to avoid spurious mscale scaling (rotary_scaling_factor=40 by default!). + rope_scaling = getattr(hf_config, "rope_scaling", None) or {} + rotary_scaling_factor = rope_scaling.get("factor", 1.0) + mla_args = { + "multi_latent_attention": True, + "q_lora_rank": getattr(hf_config, "q_lora_rank", None), + "kv_lora_rank": getattr(hf_config, "kv_lora_rank", 512), + "qk_head_dim": getattr(hf_config, "qk_nope_head_dim", 128), + "qk_pos_emb_head_dim": getattr(hf_config, "qk_rope_head_dim", 64), + "v_head_dim": getattr(hf_config, "v_head_dim", 128), + "rope_type": "rope", + "rotary_base": getattr(hf_config, "rope_theta", 10000.0), + "rotary_percent": 1.0, + "rotary_scaling_factor": rotary_scaling_factor, + "apply_rope_fusion": False, + "mscale": 0.707, + "mscale_all_dim": 0.707, + "original_max_position_embeddings": ( + rope_scaling.get("original_max_position_embeddings") + or getattr(hf_config, "original_max_position_embeddings", None) + or getattr(hf_config, "max_position_embeddings", 4096) + ), + } + + # MoE-specific parameters (same router family as V2.5: sigmoid + grouped TopK). + moe_args = { + "num_moe_experts": getattr(hf_config, "num_experts", None), + "moe_router_topk": getattr(hf_config, "num_experts_per_tok", 8), + "moe_router_score_function": getattr(hf_config, "scoring_func", "sigmoid"), + "moe_router_num_groups": getattr(hf_config, "n_group", 8), + "moe_router_group_topk": getattr(hf_config, "topk_group", 4), + "moe_router_topk_scaling_factor": getattr( + hf_config, "routed_scaling_factor", None + ), + "moe_ffn_hidden_size": getattr(hf_config, "moe_intermediate_size", None), + "moe_shared_expert_intermediate_size": shared_expert_intermediate_size, + "moe_layer_freq": moe_layer_freq, + "moe_router_enable_expert_bias": True, + "moe_router_load_balancing_type": "none", + "moe_grouped_gemm": True, + "moe_router_dtype": "fp32", + # Bias update rate only affects expert-bias drift across steps (not the forward), + # so it does not influence single-step SFT-loss alignment. Default frozen here for + # determinism; production training may set the HF value (~1e-3). + "moe_router_bias_update_rate": getattr( + hf_config, "router_bias_update_speed", 0.0 + ), + "moe_z_loss_coeff": 3.5e-6, + } + + all_args = {**base_args, **mla_args, **moe_args} + return check_and_construct_configs(all_args, MLATransformerConfig) + + +def _te_linear_and_norm(): + """Return (ColumnParallel, RowParallel, Norm) classes, preferring TE variants.""" + try: + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TENorm, + TERowParallelLinear, + ) + + return TEColumnParallelLinear, TERowParallelLinear, TENorm + except ImportError: + from megatron.core.tensor_parallel import ( + ColumnParallelLinear as TEColumnParallelLinear, + ) + from megatron.core.tensor_parallel import ( + RowParallelLinear as TERowParallelLinear, + ) + from megatron.core.transformer.torch_norm import WrappedTorchNorm as TENorm + + return TEColumnParallelLinear, TERowParallelLinear, TENorm + + +def _build_kda_attn_spec(hf_config: PretrainedConfig) -> ModuleSpec: + """Build a ModuleSpec for KimiDeltaAttention with params from the HF config. + + KDA-specific knobs travel via ModuleSpec ``params`` (not the TransformerConfig) so we + do not have to extend cli_args / TransformerConfig for the bring-up. + """ + col, row, norm = _te_linear_and_norm() + return ModuleSpec( + module=KimiDeltaAttention, + submodules=KimiDeltaAttentionSubmodules( + in_proj=col, + beta_proj=col, + out_norm=norm, + out_proj=row, + ), + params={ + "head_dim": _kda_head_dim(hf_config), + # V3 uses ``short_conv_kernel_size``; fall back to the V2.5 name and then 4. + "conv_kernel_dim": getattr( + hf_config, + "short_conv_kernel_size", + getattr(hf_config, "linear_conv_kernel_dim", 4), + ), + "no_kda_lora": getattr(hf_config, "no_kda_lora", True), + "use_qk_l2norm": getattr(hf_config, "use_qk_l2norm", True), + # Clamped (safe) decay gate. BailingMoeV3 ckpts set kda_safe_gate=True / + # kda_lower_bound=-5.0; these must reach chunk_kda or the decay gate + # silently reverts to the unbounded softplus form (wrong loss). + "safe_gate": getattr(hf_config, "kda_safe_gate", False), + "lower_bound": getattr(hf_config, "kda_lower_bound", None), + }, + ) + + +def _build_gated_mla_spec(base_mla_spec: ModuleSpec, enable_gate: bool) -> ModuleSpec: + """Swap the MLA module for the v3 gated variant, keeping the same submodules/params.""" + spec = copy.deepcopy(base_mla_spec) + if enable_gate: + spec.submodules.self_attention.module = BailingV3MLASelfAttention + return spec + + +def make_mcore_layer_specs_bailing_v3( + tf_config: MLATransformerConfig, + hf_config: PretrainedConfig, + use_te: bool = True, + vp_stage: int | None = None, +) -> TransformerBlockSubmodules: + """Build heterogeneous layer specs for BailingMoeV3 (KDA + gated MLA). + + Creates 4 layer-spec variants (KDA/MLA x Dense/MoE). KDA layers use the custom + ``KimiDeltaAttention`` module; MLA layers use ``BailingV3MLASelfAttention`` (mcore MLA + plus head-wise gate). When PP>1, the full spec list is sliced for the current stage. + """ + assert tf_config.normalization == "RMSNorm", "only RMSNorm is supported" + + layer_group_size = getattr(hf_config, "layer_group_size", 1) + num_layers = tf_config.num_layers + first_k_dense_replace = getattr(hf_config, "first_k_dense_replace", 0) + gate_granularity = getattr(hf_config, "gated_attention_proj_granularity_type", None) + if gate_granularity is None: + enable_gate = bool(getattr(hf_config, "enable_gated_attention", False)) + else: + if gate_granularity != "head_wise": + raise ValueError( + "BailingMoeV3 currently supports only head-wise gated MLA; got " + f"gated_attention_proj_granularity_type={gate_granularity!r}." + ) + enable_gate = True + + # The KDA/MLA pattern is derived from layer_group_size (last layer of each group is + # MLA). Incomplete tail groups are kept as MLA as a conservative fallback; + # HybridEngine v3 normally asserts divisibility. + + _, _, te_norm = _te_linear_and_norm() + + # MLA layer specs (gated MLA via module swap on the standard MLA spec). + mla_dense_base = get_gpt_layer_with_transformer_engine_spec( + num_experts=None, + moe_grouped_gemm=False, + qk_layernorm=tf_config.qk_layernorm, + multi_latent_attention=True, + ) + mla_moe_base = get_gpt_layer_with_transformer_engine_spec( + num_experts=tf_config.num_moe_experts, + moe_grouped_gemm=tf_config.moe_grouped_gemm, + qk_layernorm=tf_config.qk_layernorm, + multi_latent_attention=True, + ) + mla_dense_spec = _build_gated_mla_spec(mla_dense_base, enable_gate) + mla_moe_spec = _build_gated_mla_spec(mla_moe_base, enable_gate) + + # KDA layer specs: start from standard (non-MLA) specs for correct MLP/layernorm, then + # replace self_attention with KDA. CRITICAL: restore a real input_layernorm (TENorm) + # because KDA's in_proj is a plain TEColumnParallelLinear (no fused layernorm), unlike + # the standard fused TELayerNormColumnParallelLinear QKV that the base spec assumes. + kda_attn_spec = _build_kda_attn_spec(hf_config) + kda_dense_base = get_gpt_layer_with_transformer_engine_spec( + num_experts=None, + moe_grouped_gemm=False, + qk_layernorm=True, + multi_latent_attention=False, + ) + kda_moe_base = get_gpt_layer_with_transformer_engine_spec( + num_experts=tf_config.num_moe_experts, + moe_grouped_gemm=tf_config.moe_grouped_gemm, + qk_layernorm=True, + multi_latent_attention=False, + ) + kda_dense_spec = copy.deepcopy(kda_dense_base) + kda_dense_spec.submodules.self_attention = kda_attn_spec + kda_dense_spec.submodules.input_layernorm = te_norm + kda_moe_spec = copy.deepcopy(kda_moe_base) + kda_moe_spec.submodules.self_attention = kda_attn_spec + kda_moe_spec.submodules.input_layernorm = te_norm + + # Per-layer assignment. + layer_specs = [] + for layer_idx in range(num_layers): + is_kda = is_kda_layer(layer_idx, layer_group_size, num_layers) + is_moe = layer_idx >= first_k_dense_replace + if is_kda: + spec = kda_moe_spec if is_moe else kda_dense_spec + else: + spec = mla_moe_spec if is_moe else mla_dense_spec + layer_specs.append(spec) + + n_kda = sum( + 1 for i in range(num_layers) if is_kda_layer(i, layer_group_size, num_layers) + ) + n_mla = num_layers - n_kda + n_moe = sum(1 for i in range(num_layers) if i >= first_k_dense_replace) + n_dense = num_layers - n_moe + logger.info( + f"Built BailingV3 layer specs: {num_layers} layers, " + f"layer_group_size={layer_group_size}, first_k_dense={first_k_dense_replace}, " + f"num_experts={tf_config.num_moe_experts}, gated_mla={enable_gate}" + ) + logger.info( + f"Layer composition: {n_kda} KDA + {n_mla} MLA, {n_dense} Dense + {n_moe} MoE" + ) + + # KDA supports HybridEngine-style all2all CP. Heads are moved from CP sequence + # shards into CP head shards before KDA, so each TP partition must have a head count + # divisible by CP. + if tf_config.context_parallel_size > 1 and n_kda > 0: + tp_size = tf_config.tensor_model_parallel_size + cp_size = tf_config.context_parallel_size + heads_per_tp = tf_config.num_attention_heads // tp_size + if heads_per_tp % cp_size != 0: + raise ValueError( + "For BailingMoeV3 KDA with CP, num_attention_heads / TP " + f"({heads_per_tp}) must be divisible by CP ({cp_size})." + ) + logger.info( + f"KDA all2all CP enabled: CP={cp_size}, " + f"heads_per_tp={heads_per_tp}, heads_per_cp={heads_per_tp // cp_size}" + ) + + # PP slicing: TransformerBlock._build_layers() builds ALL specs without slicing, so we + # must pre-slice for the current pipeline stage (mirrors get_gpt_decoder_block_spec). + num_layers_to_build = get_num_layers_to_build(tf_config, vp_stage=vp_stage) + if tf_config.pipeline_model_parallel_layout is not None: + local_layer_specs = [ + layer_specs[layer_id] + for layer_id in tf_config.pipeline_model_parallel_layout.get_layer_id_list( + layer_type=LayerType.decoder, vp_stage=vp_stage + ) + ] + elif num_layers_to_build < num_layers: + offset = get_transformer_layer_offset(tf_config, vp_stage=vp_stage) + local_layer_specs = layer_specs[offset : offset + num_layers_to_build] + else: + local_layer_specs = layer_specs + + if len(local_layer_specs) != num_layers: + logger.info( + f"PP slicing: building {len(local_layer_specs)}/{num_layers} layers " + f"for this pipeline stage" + ) + + if use_te: + layer_norm_impl = te_norm + else: + try: + from megatron.core.fusions.fused_layer_norm import FusedLayerNorm + + layer_norm_impl = FusedLayerNorm + except ImportError: + from megatron.core.transformer.torch_norm import WrappedTorchNorm + + layer_norm_impl = WrappedTorchNorm + + return TransformerBlockSubmodules( + layer_specs=local_layer_specs, + layer_norm=layer_norm_impl, + ) diff --git a/areal/models/mcore/bailing_v3_bridge.py b/areal/models/mcore/bailing_v3_bridge.py new file mode 100644 index 0000000000..a2ee2fec7e --- /dev/null +++ b/areal/models/mcore/bailing_v3_bridge.py @@ -0,0 +1,400 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mbridge Bridge for BailingMoeV3 (Ling V3, model_type bailing_moe_v3). + +Registers with mbridge so MegatronEngine.initialize() can use AutoBridge to load and +manage BailingMoeV3 models with heterogeneous attention (KDA + gated MLA). + +HF weight names follow the antllm BailingMoeV3 checkpoint layout (verified against +HybridEngine's bailing_moe_huggingface_ckpt_conversion.py: +``transform_kda_attention_and_layernorm_weights``): + +KDA layers (no_kda_lora=True): + self_attention.in_proj.weight -> attention.{q,k,v,f,g}_proj.weight (split q|k|v|g|gate) + self_attention.conv1d.weight -> attention.{q,k,v}_conv1d.weight (split q|k|v) + self_attention.beta_proj.weight -> attention.b_proj.weight + self_attention.out_norm.weight -> attention.o_norm.weight + self_attention.out_proj.weight -> attention.o_proj.weight + self_attention.dt_bias -> attention.dt_bias + self_attention.A_log -> attention.A_log + input_layernorm.weight -> input_layernorm.weight +MLA layers (gated): same as V2.5 MLA, plus + self_attention.linear_gate.weight -> attention.g_proj.weight + +The fused-in-proj / fused-conv1d split (multiple HF tensors -> one mcore tensor) is +handled on the load path by areal/models/mcore/hf_load.py (KDA branch in +_weight_to_mcore_tp). The reverse split for save is in _weight_to_hf_format below. +""" + +import torch +from mbridge.core import LLMBridge, register_model +from megatron.core.transformer import MLATransformerConfig +from megatron.core.transformer.enums import AttnBackend + +from areal.models.mcore.bailing_v3 import ( + is_kda_layer, + make_mcore_layer_specs_bailing_v3, +) +from areal.utils import logging + +logger = logging.getLogger("BailingV3Bridge") + +# KDA (linear-attention) mcore suffix -> HF name templates. +# NOTE: mcore in_proj split order is [query, key, value, g, gate]; HF names are +# [q_proj, k_proj, v_proj, f_proj, g_proj] (mcore "g" -> HF f_proj, mcore "gate" -> HF g_proj). +_KDA_ATTENTION_MAPPING = { + "input_layernorm.weight": ["model.layers.{layer_number}.input_layernorm.weight"], + "self_attention.in_proj.weight": [ + "model.layers.{layer_number}.attention.q_proj.weight", + "model.layers.{layer_number}.attention.k_proj.weight", + "model.layers.{layer_number}.attention.v_proj.weight", + "model.layers.{layer_number}.attention.f_proj.weight", + "model.layers.{layer_number}.attention.g_proj.weight", + ], + "self_attention.conv1d.weight": [ + "model.layers.{layer_number}.attention.q_conv1d.weight", + "model.layers.{layer_number}.attention.k_conv1d.weight", + "model.layers.{layer_number}.attention.v_conv1d.weight", + ], + "self_attention.beta_proj.weight": [ + "model.layers.{layer_number}.attention.b_proj.weight" + ], + "self_attention.out_norm.weight": [ + "model.layers.{layer_number}.attention.o_norm.weight" + ], + "self_attention.out_proj.weight": [ + "model.layers.{layer_number}.attention.o_proj.weight" + ], + "self_attention.dt_bias": ["model.layers.{layer_number}.attention.dt_bias"], + "self_attention.A_log": ["model.layers.{layer_number}.attention.A_log"], +} + +# Gated-MLA mcore suffix -> HF name templates. +_MLA_ATTENTION_MAPPING_Q_DIRECT = { + "self_attention.linear_q_proj.weight": [ + "model.layers.{layer_number}.attention.q_proj.weight" + ], +} +_MLA_ATTENTION_MAPPING_Q_LORA = { + "self_attention.linear_q_down_proj.weight": [ + "model.layers.{layer_number}.attention.q_a_proj.weight" + ], + "self_attention.linear_q_up_proj.layer_norm_weight": [ + "model.layers.{layer_number}.attention.q_a_layernorm.weight" + ], + "self_attention.linear_q_up_proj.weight": [ + "model.layers.{layer_number}.attention.q_b_proj.weight" + ], +} +_MLA_ATTENTION_MAPPING_COMMON = { + "input_layernorm.weight": ["model.layers.{layer_number}.input_layernorm.weight"], + "self_attention.linear_kv_down_proj.weight": [ + "model.layers.{layer_number}.attention.kv_a_proj_with_mqa.weight" + ], + "self_attention.linear_kv_up_proj.layer_norm_weight": [ + "model.layers.{layer_number}.attention.kv_a_layernorm.weight" + ], + "self_attention.linear_kv_up_proj.weight": [ + "model.layers.{layer_number}.attention.kv_b_proj.weight" + ], + "self_attention.linear_proj.weight": [ + "model.layers.{layer_number}.attention.dense.weight" + ], + # v3 head-wise gated attention + "self_attention.linear_gate.weight": [ + "model.layers.{layer_number}.attention.g_proj.weight" + ], +} + + +@register_model("bailing_moe_v3") +class BailingV3Bridge(LLMBridge): + """Bridge for BailingMoeV3 with heterogeneous KDA + gated MLA attention.""" + + TransformerConfigClass = MLATransformerConfig + + _DIRECT_MAPPING = { + "embedding.word_embeddings.weight": "model.word_embeddings.weight", + "decoder.final_layernorm.weight": "model.norm.weight", + "output_layer.weight": "lm_head.weight", + } + + _MLP_MAPPING = { + "mlp.linear_fc1.layer_norm_weight": [ + "model.layers.{layer_number}.post_attention_layernorm.weight" + ], + "mlp.linear_fc2.weight": ["model.layers.{layer_number}.mlp.down_proj.weight"], + "mlp.linear_fc1.weight": [ + "model.layers.{layer_number}.mlp.gate_proj.weight", + "model.layers.{layer_number}.mlp.up_proj.weight", + ], + "mlp.shared_experts.linear_fc2.weight": [ + "model.layers.{layer_number}.mlp.shared_experts.down_proj.weight" + ], + "mlp.shared_experts.linear_fc1.weight": [ + "model.layers.{layer_number}.mlp.shared_experts.gate_proj.weight", + "model.layers.{layer_number}.mlp.shared_experts.up_proj.weight", + ], + "pre_mlp_layernorm.weight": [ + "model.layers.{layer_number}.post_attention_layernorm.weight" + ], + "mlp.router.weight": ["model.layers.{layer_number}.mlp.gate.weight"], + "mlp.router.expert_bias": ["model.layers.{layer_number}.mlp.gate.expert_bias"], + "mlp.experts.linear_fc1.weight": [ + "model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight", + "model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight", + ], + "mlp.experts.linear_fc2.weight": [ + "model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight" + ], + } + + def _build_config(self): + hf_config = self.hf_config + num_layers = hf_config.num_hidden_layers + first_k_dense_replace = getattr(hf_config, "first_k_dense_replace", 0) + moe_layer_freq = [ + 0 if i < first_k_dense_replace else 1 for i in range(num_layers) + ] + shared_expert_intermediate_size = getattr( + hf_config, "moe_shared_expert_intermediate_size", None + ) + if shared_expert_intermediate_size is None: + num_shared_experts = getattr(hf_config, "num_shared_experts", 0) + if num_shared_experts > 0: + shared_expert_intermediate_size = num_shared_experts * getattr( + hf_config, "moe_intermediate_size", hf_config.intermediate_size + ) + + return self._build_base_config( + attention_backend=AttnBackend.fused, + layernorm_epsilon=hf_config.rms_norm_eps, + ffn_hidden_size=hf_config.intermediate_size, + qk_layernorm=True, + # MLA parameters (rotary_percent=1.0: qk_pos_emb_head_dim is the pure RoPE + # slice in AReaL's mcore MLA; see bailing_v3.hf_to_mcore_config_bailing_v3). + multi_latent_attention=True, + q_lora_rank=getattr(hf_config, "q_lora_rank", None), + kv_lora_rank=getattr(hf_config, "kv_lora_rank", 512), + qk_head_dim=getattr(hf_config, "qk_nope_head_dim", 128), + qk_pos_emb_head_dim=getattr(hf_config, "qk_rope_head_dim", 64), + v_head_dim=getattr(hf_config, "v_head_dim", 128), + rotary_base=getattr(hf_config, "rope_theta", 10000.0), + rope_type="rope", + rotary_percent=1.0, + rotary_scaling_factor=(getattr(hf_config, "rope_scaling", None) or {}).get( + "factor", 1.0 + ), + apply_rope_fusion=False, + # Keep in sync with bailing_v3.hf_to_mcore_config_bailing_v3: pin the + # YaRN mscale knobs so the softmax scale does not depend on the mcore + # (or HybridEngine-fork) MLATransformerConfig defaults. Inert while + # rotary_scaling_factor == 1.0, decisive for rope-scaled checkpoints. + mscale=0.707, + mscale_all_dim=0.707, + original_max_position_embeddings=( + (getattr(hf_config, "rope_scaling", None) or {}).get( + "original_max_position_embeddings" + ) + or getattr(hf_config, "original_max_position_embeddings", None) + or getattr(hf_config, "max_position_embeddings", 4096) + ), + # MoE parameters + moe_ffn_hidden_size=getattr(hf_config, "moe_intermediate_size", None), + moe_token_dispatcher_type="alltoall", + moe_router_enable_expert_bias=True, + moe_router_topk=getattr(hf_config, "num_experts_per_tok", 8), + num_moe_experts=getattr(hf_config, "num_experts", None), + moe_shared_expert_intermediate_size=shared_expert_intermediate_size, + moe_router_score_function=getattr(hf_config, "scoring_func", "sigmoid"), + moe_router_num_groups=getattr(hf_config, "n_group", 8), + moe_router_group_topk=getattr(hf_config, "topk_group", 4), + moe_router_topk_scaling_factor=getattr( + hf_config, "routed_scaling_factor", None + ), + moe_router_load_balancing_type="none", + moe_grouped_gemm=True, + moe_layer_freq=moe_layer_freq, + moe_router_dtype="fp32", + moe_router_bias_update_rate=getattr( + hf_config, "router_bias_update_speed", 0.0 + ), + moe_z_loss_coeff=3.5e-6, + persist_layer_norm=True, + bias_activation_fusion=True, + bias_dropout_fusion=True, + ) + + def _get_gptmodel_args(self) -> dict: + return dict( + vocab_size=self.hf_config.vocab_size, + max_sequence_length=self.hf_config.max_position_embeddings, + position_embedding_type="rope", + rotary_base=getattr(self.hf_config, "rope_theta", 10000.0), + ) + + def _get_transformer_layer_spec(self, vp_stage: int | None = None): + """Return heterogeneous layer specs (KDA + gated MLA). VPP is not supported.""" + assert self.config.normalization == "RMSNorm" + self.has_vp_stage = False + return make_mcore_layer_specs_bailing_v3( + self.config, self.hf_config, use_te=True, vp_stage=vp_stage + ) + + # ------------------------------------------------------------------ + # Weight name mapping + # ------------------------------------------------------------------ + def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]: + assert "_extra_state" not in mcore_weights_name + + if mcore_weights_name in self._DIRECT_MAPPING: + return [self._DIRECT_MAPPING[mcore_weights_name]] + + if ( + "self_attention" in mcore_weights_name + or "input_layernorm.weight" in mcore_weights_name + ): + return self._weight_name_mapping_attention(mcore_weights_name) + elif "mlp" in mcore_weights_name or "pre_mlp_layernorm" in mcore_weights_name: + return self._weight_name_mapping_mlp(mcore_weights_name) + else: + raise NotImplementedError( + f"Unsupported parameter name: {mcore_weights_name}" + ) + + def _weight_name_mapping_attention(self, name: str) -> list[str]: + """Dispatch to KDA or gated-MLA mapping based on the layer index.""" + layer_number_str = name.split(".")[2] + layer_number = int(layer_number_str) + layer_group_size = getattr(self.hf_config, "layer_group_size", 1) + + if is_kda_layer( + layer_number, + layer_group_size, + getattr(self.hf_config, "num_hidden_layers", None), + ): + mapping = _KDA_ATTENTION_MAPPING + else: + q_lora_rank = getattr(self.hf_config, "q_lora_rank", None) + q_mapping = ( + _MLA_ATTENTION_MAPPING_Q_LORA + if q_lora_rank is not None + else _MLA_ATTENTION_MAPPING_Q_DIRECT + ) + mapping = {**_MLA_ATTENTION_MAPPING_COMMON, **q_mapping} + + convert_names = [] + for keyword, mapping_names in mapping.items(): + if keyword in name: + convert_names.extend( + [x.format(layer_number=layer_number_str) for x in mapping_names] + ) + break + if not convert_names: + is_kda = is_kda_layer( + layer_number, + layer_group_size, + getattr(self.hf_config, "num_hidden_layers", None), + ) + raise NotImplementedError( + f"Unsupported attention parameter: {name} (kda={is_kda})" + ) + return convert_names + + def _weight_name_mapping_mlp(self, name: str) -> list[str]: + layer_number = name.split(".")[2] + convert_names = [] + for keyword, mapping_names in self._MLP_MAPPING.items(): + if keyword in name: + if "{expert_id}" in mapping_names[0]: + expert_id = name.split("weight")[-1] + convert_names.extend( + [ + x.format(layer_number=layer_number, expert_id=expert_id) + for x in mapping_names + ] + ) + else: + convert_names.extend( + [x.format(layer_number=layer_number) for x in mapping_names] + ) + break + if not convert_names: + raise NotImplementedError(f"Unsupported MLP parameter: {name}") + return convert_names + + def _weight_merge_across_tp( + self, + mcore_weights_name: str, + tp_shards: list[torch.Tensor], + param: torch.Tensor, + ) -> torch.Tensor: + """MLA's linear_q_down_proj / linear_kv_down_proj are replicated across TP.""" + if ( + "linear_q_down_proj." in mcore_weights_name + or "linear_kv_down_proj." in mcore_weights_name + ): + return tp_shards[0].clone() + return super()._weight_merge_across_tp(mcore_weights_name, tp_shards, param) + + # ------------------------------------------------------------------ + # KDA fused-weight split for the save path (mcore -> HF) + # ------------------------------------------------------------------ + def _kda_split_sections(self, kind: str) -> list[int]: + """Global split sections for KDA fused tensors (no_kda_lora layout).""" + from areal.models.mcore.bailing_v3 import _kda_head_dim + + head_dim = _kda_head_dim(self.hf_config) + num_heads = self.hf_config.num_attention_heads + qk_dim = head_dim * num_heads + v_dim = head_dim * num_heads + if kind == "in_proj": + # [query, key, value, g(->f_proj), gate(->g_proj)] + return [qk_dim, qk_dim, v_dim, qk_dim, v_dim] + elif kind == "conv1d": + return [qk_dim, qk_dim, v_dim] + raise ValueError(kind) + + def _deinterleave_and_split( + self, tensor: torch.Tensor, sections: list[int] + ) -> list[torch.Tensor]: + """Split a TP-merged (rank-interleaved) fused tensor into full components. + + After the default cross-TP merge, a fused ColumnParallel weight is laid out + rank-major: [c0_r0, c1_r0, ..., c0_r1, c1_r1, ...]. Deinterleave back into + contiguous full components [c0_full, c1_full, ...]. For tp_size==1 this is a + plain split. + """ + from megatron.core import parallel_state as mpu + + try: + tp_size = mpu.get_tensor_model_parallel_world_size() + except (RuntimeError, AssertionError): + tp_size = 1 + if tp_size <= 1: + return list(torch.split(tensor, sections, dim=0)) + + per_rank = [s // tp_size for s in sections] + components = [[] for _ in sections] + for chunk in torch.split(tensor, sum(per_rank), dim=0): + for i, part in enumerate(torch.split(chunk, per_rank, dim=0)): + components[i].append(part) + return [torch.cat(c, dim=0).contiguous() for c in components] + + def _weight_to_hf_format( + self, mcore_weights_name: str, mcore_weights: torch.Tensor + ) -> tuple[list[str], list[torch.Tensor]]: + """Convert mcore weights to HF format, splitting KDA fused in_proj / conv1d.""" + hf_names = self._weight_name_mapping_mcore_to_hf(mcore_weights_name) + if "self_attention.in_proj.weight" in mcore_weights_name: + comps = self._deinterleave_and_split( + mcore_weights, self._kda_split_sections("in_proj") + ) + return hf_names, comps + if "self_attention.conv1d.weight" in mcore_weights_name: + # conv1d weight is [conv_dim, 1, kernel]; split on dim 0. + comps = self._deinterleave_and_split( + mcore_weights, self._kda_split_sections("conv1d") + ) + return hf_names, comps + return super()._weight_to_hf_format(mcore_weights_name, mcore_weights) diff --git a/areal/models/mcore/bailing_v3_mla.py b/areal/models/mcore/bailing_v3_mla.py new file mode 100644 index 0000000000..efadc63dc2 --- /dev/null +++ b/areal/models/mcore/bailing_v3_mla.py @@ -0,0 +1,119 @@ +"""Gated MLA self-attention for BailingMoeV3 (Ling V3). + +v3 enables *head-wise gated attention* on top of standard MLA: a per-head sigmoid gate +(produced by a ``hidden_size -> num_attention_heads`` projection from the layer input) +scales the core attention output before the output projection. + + gate = sigmoid(linear_gate(hidden_states)) # [s, b, num_heads] + core_attn_out = core_attn_out.view(s, b, H, d) * gate[:, :, :, None] + output = linear_proj(core_attn_out) + +Reference: ant-megatron ``attention.py::apply_gated_attention_linear_gate`` with +``gated_attention_proj_granularity_type=head_wise`` and +``gated_attention_input_tensor_type=linear_qkv_input``. + +Implementation note: + Rather than copy mcore's (version-specific, monolithic) ``MLASelfAttention.forward``, + we inject the gate via a ``forward_pre_hook`` on ``linear_proj``. This only relies on + the MLA invariant that the forward ends with ``self.linear_proj(core_attn_out)`` where + ``core_attn_out`` is ``[s, b, num_heads_local * v_head_dim]`` (head-major), so it is + robust across megatron-core versions. +""" + +import torch +from megatron.core.transformer.multi_latent_attention import MLASelfAttention +from megatron.core.transformer.spec_utils import build_module + +from areal.utils import logging + +logger = logging.getLogger("BailingV3MLA") + +try: + from megatron.core.extensions.transformer_engine import TEColumnParallelLinear +except ImportError: # pragma: no cover + from megatron.core.tensor_parallel import ( + ColumnParallelLinear as TEColumnParallelLinear, + ) + + +class BailingV3MLASelfAttention(MLASelfAttention): + """MLA self-attention with v3 head-wise gated attention. + + Built exactly like mcore ``MLASelfAttention`` (same submodules / params) plus an + extra ``linear_gate`` projection. The gate is applied to the core attention output + just before ``linear_proj`` via a forward pre-hook. + """ + + def __init__( + self, + config, + submodules, + layer_number, + attn_mask_type=None, + **kwargs, + ): + if attn_mask_type is not None: + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + **kwargs, + ) + else: + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + **kwargs, + ) + + # Head-wise gate: hidden_size -> num_attention_heads (column-parallel; the head + # dimension is TP-sharded so the local output matches core_attn_out's local heads). + self.linear_gate = build_module( + TEColumnParallelLinear, + config.hidden_size, + config.num_attention_heads, + config=config, + init_method=config.init_method, + gather_output=False, + bias=False, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="gate", + ) + + self._gate_hidden_states = None + # Apply the gate to linear_proj's input. with_kwargs defaults to False, so the + # hook receives positional args and its return value replaces them. + self.linear_proj.register_forward_pre_hook(self._apply_gate_pre_hook) + + def _apply_gate_pre_hook(self, module, args): + """Scale ``core_attn_out`` (linear_proj input) head-wise by ``sigmoid(gate)``.""" + if self._gate_hidden_states is None or not args: + return None + core_attn_out = args[0] + gate, _ = self.linear_gate(self._gate_hidden_states) + gate = torch.sigmoid(gate.float()).type_as(core_attn_out) + seq_len, batch = core_attn_out.shape[:2] + num_heads_local = gate.shape[-1] + # core_attn_out must be head-major [s, b, num_heads_local * v_head_dim] for the + # per-head gate to align (the invariant this hook relies on). + assert core_attn_out.shape[-1] % num_heads_local == 0, ( + f"core_attn_out last dim {core_attn_out.shape[-1]} not divisible by " + f"num_heads_local {num_heads_local}; gated-MLA head-major assumption broken." + ) + gated = ( + core_attn_out.view(seq_len, batch, num_heads_local, -1) + * gate[:, :, :, None] + ) + gated = gated.reshape(core_attn_out.shape) + return (gated,) + tuple(args[1:]) + + def forward(self, hidden_states, *args, **kwargs): + # Stash the layer input so the linear_proj pre-hook can compute the gate. + self._gate_hidden_states = hidden_states + try: + return super().forward(hidden_states, *args, **kwargs) + finally: + self._gate_hidden_states = None diff --git a/areal/models/mcore/hf_load.py b/areal/models/mcore/hf_load.py index 6e44c0b382..72f5b85c71 100644 --- a/areal/models/mcore/hf_load.py +++ b/areal/models/mcore/hf_load.py @@ -269,6 +269,34 @@ def _slice_generic_weight( ] +def _merge_kda_fused_weight( + hf_weights_safe_slice: list, + tp_rank: int, + tp_size: int, +) -> torch.Tensor: + """Merge Bailing v3 KDA fused components into a single mcore tensor. + + KDA stores ``in_proj.weight`` as ``[q | k | v | f | g]`` and + ``conv1d.weight`` as ``[q | k | v]``. Each HF component is TP-sharded on dim 0 + before concatenation. + """ + comps = [] + for x in hf_weights_safe_slice: + shape = _get_shape(x) + if shape[0] % tp_size != 0: + raise ValueError( + f"KDA fused component dim0={shape[0]} is not divisible by TP {tp_size}." + ) + comps.append(x[_get_tp_slice(shape, dim=0, tp_rank=tp_rank, tp_size=tp_size)]) + return torch.cat(comps, dim=0).contiguous() + + +def _is_bailing_v3_config(hf_config) -> bool: + """Return whether an HF config selects the Bailing V3 architecture.""" + architectures = getattr(hf_config, "architectures", None) or () + return "BailingMoeV3ForCausalLM" in architectures + + def _convert_vision_qkv_hf_to_mcore( hf_config, mcore_weights_name: str, @@ -440,6 +468,11 @@ def _weight_to_mcore_tp( ) else: res = _slice_moe_expert_weight(hf_weights_safe_slice, tp_rank, tp_size) + elif _is_bailing_v3_config(hf_config) and ( + "self_attention.in_proj.weight" in mcore_weights_name + or "self_attention.conv1d.weight" in mcore_weights_name + ): + res = _merge_kda_fused_weight(hf_weights_safe_slice, tp_rank, tp_size) else: res = _slice_generic_weight( mcore_param_shape, hf_weights_safe_slice, tp_rank, tp_size @@ -534,6 +567,14 @@ def _load_weight_with_bridge_worker( if is_te_fp8_param and enable_fp8_param and hf_has_fp8 and not hf_all_fp8: raise RuntimeError("Expected all inputs to be FP8 for TE FP8 parameter") + target_dtype = bridge.dtype + if param.dtype == torch.float32 and ( + "mlp.router.expert_bias" in local_name + or "self_attention.dt_bias" in local_name + or "self_attention.A_log" in local_name + ): + target_dtype = torch.float32 + param_to_load = _weight_to_mcore_tp( hf_config=bridge.hf_config, mcore_weights_name=local_name, @@ -541,7 +582,7 @@ def _load_weight_with_bridge_worker( hf_weights_safe_slice=hf_weights_safe_slice, tp_rank=tp_rank, tp_size=tp_size, - dtype=bridge.dtype + dtype=target_dtype if not (is_te_fp8_param and hf_has_fp8 and hf_all_fp8) else None, ) diff --git a/areal/models/mcore/kda_attention.py b/areal/models/mcore/kda_attention.py new file mode 100644 index 0000000000..cd2cc380f4 --- /dev/null +++ b/areal/models/mcore/kda_attention.py @@ -0,0 +1,1118 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""KDA (Kimi Delta Attention) module for megatron-core, using fla Triton kernels. + +BailingMoeV3 (Ling V3) uses a heterogeneous architecture where most layers use KDA +(a gated delta-rule linear attention with a short causal convolution) and every +``layer_group_size``-th layer uses standard MLA. KDA replaces the Lightning Attention +used in BailingMoeV2.5. + +This module ports ant-megatron's ``megatron/core/ssm/kda.py`` (``KimiDeltaAttention``) +into AReaL's open-source megatron-core path. It is numerically faithful to that +reference but strips the ant-only FP8 / mxfp8 fast paths and the transpose/rms-norm +fusion optimizations (bf16 bring-up). + +Reference kernels (flash-linear-attention + causal-conv1d): + - fla.ops.kda.chunk_kda (chunked training kernel, gated delta rule) + - fla.ops.kda.gate.fused_kda_gate (decay gate from A_log / dt_bias) + - fla.modules.l2norm.l2_norm + - causal_conv1d.causal_conv1d_fn (depthwise causal short conv + silu) + +Architecture (no_kda_lora=True, the v3 setting): + in_proj(hidden) -> [qkv | g | gate] + qkv -> causal_conv1d (depthwise, silu) -> split [q, k, v] + beta = sigmoid(beta_proj(hidden)) + g (decay) via fused_kda_gate(A_log, dt_bias) or inside chunk_kda + core = chunk_kda(q, k, v, g, beta, A_log, dt_bias, qk_l2norm) + out = out_norm(core) * sigmoid(gate) # per-head RMSNorm over v_head_dim + out_proj(out) + +Notes: + - KDA does NOT use RoPE (it is a delta-rule SSM); ``rotary_pos_emb`` is ignored. + - Context parallelism uses HybridEngine's all2all strategy: CP sequence shards are + exchanged into head shards before KDA, then exchanged back before output norm. +""" + +import math +from dataclasses import dataclass, replace + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from megatron.core import parallel_state as mpu +from megatron.core.dist_checkpointing import ShardedTensor +from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.tensor_parallel.mappings import ( + gather_from_tensor_model_parallel_region, + scatter_to_sequence_parallel_region, +) +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.utils import ( + make_sharded_tensors_for_checkpoint, + sharded_state_dict_default, +) +from torch import Tensor + +from areal.utils import logging + +logger = logging.getLogger("KDAAttention") + +try: + from torch.distributed._functional_collectives import all_to_all_single_autograd +except ImportError: # pragma: no cover + all_to_all_single_autograd = None + +# Optional ant-megatron extension: packed-seq params that carry an explicit seq_idx. +try: + from megatron.core.packed_seq_params import PackedSeqParamsWithSeqidx +except ImportError: # pragma: no cover - not present in upstream megatron-core + PackedSeqParamsWithSeqidx = None + +# fla / causal-conv1d kernels (optional import; required at runtime for KDA). +try: + from fla.modules.l2norm import l2_norm as l2norm + from fla.ops.kda import chunk_kda, fused_recurrent_kda + + HAVE_FLA = True +except ImportError: # pragma: no cover + l2norm = None + chunk_kda = None + fused_recurrent_kda = None + HAVE_FLA = False + +try: + from fla.ops.kda.gate import fused_kda_gate +except ImportError: # pragma: no cover + fused_kda_gate = None + +try: + from causal_conv1d import causal_conv1d_fn +except ImportError: # pragma: no cover + causal_conv1d_fn = None + + +def _get_tp_world_size() -> int: + """Tensor-model-parallel world size, with a fallback for uninitialized mpu.""" + try: + if mpu.model_parallel_is_initialized(): + return mpu.get_tensor_model_parallel_world_size() + except (RuntimeError, AttributeError): + pass + return 1 + + +def _get_cp_world_size() -> int: + """Context-parallel world size, with a fallback for uninitialized mpu.""" + try: + if mpu.model_parallel_is_initialized(): + return mpu.get_context_parallel_world_size() + except (RuntimeError, AttributeError): + pass + return 1 + + +def _get_cp_rank() -> int: + """Context-parallel rank, with a fallback for uninitialized mpu.""" + try: + if mpu.model_parallel_is_initialized(): + return mpu.get_context_parallel_rank() + except (RuntimeError, AttributeError): + pass + return 0 + + +def _get_cp_group(): + """Context-parallel process group, with a fallback for uninitialized mpu.""" + try: + if mpu.model_parallel_is_initialized(): + return mpu.get_context_parallel_group() + except (RuntimeError, AttributeError): + pass + return None + + +class _AllToAll(torch.autograd.Function): + """Autograd fallback for equal-size all-to-all.""" + + @staticmethod + def forward(ctx, group, input_: Tensor): + ctx.group = group + world_size = dist.get_world_size(group=group) + if world_size == 1: + return input_ + input_ = input_.contiguous() + output = torch.empty_like(input_) + dist.all_to_all_single(output, input_, group=group) + return output + + @staticmethod + def backward(ctx, grad_output): + return None, _AllToAll.apply(ctx.group, grad_output) + + +def _all_to_all_equal(input_: Tensor, cp_group) -> Tensor: + """Equal-size all-to-all over the CP group with autograd support.""" + if cp_group is None: + return input_ + world_size = dist.get_world_size(group=cp_group) + if world_size == 1: + return input_ + input_shape = input_.shape + flat = input_.contiguous().reshape(-1) + if all_to_all_single_autograd is not None: + exchanged = all_to_all_single_autograd(flat, None, None, group=cp_group) + else: # pragma: no cover - old torch fallback + exchanged = _AllToAll.apply(cp_group, flat) + return exchanged.reshape(input_shape) + + +def _all_to_all_cp2hp( + input_: Tensor, + cp_group, + split_size_or_sections: list[int] | None = None, +) -> Tensor: + """CP sequence shard -> head shard. + + Shape: ``[S/CP, B, H] -> [S, B, H/CP]``. + """ + if cp_group is None: + return input_ + cp_size = dist.get_world_size(group=cp_group) + if cp_size == 1: + return input_ + if split_size_or_sections is not None: + chunks = torch.split(input_, split_size_or_sections, dim=-1) + return torch.cat( + [_all_to_all_cp2hp(chunk, cp_group) for chunk in chunks], + dim=-1, + ) + assert input_.dim() == 3, input_.shape + seq_len, batch, hidden = input_.shape + if hidden % cp_size != 0: + raise ValueError( + f"KDA all2all CP requires hidden dim {hidden} divisible by CP {cp_size}." + ) + hidden_per_cp = hidden // cp_size + flat = input_.reshape(seq_len * batch, hidden) + flat = torch.cat(torch.split(flat, hidden_per_cp, dim=-1), dim=0) + flat = _all_to_all_equal(flat, cp_group) + return flat.reshape(seq_len * cp_size, batch, hidden_per_cp) + + +def _all_to_all_hp2cp(input_: Tensor, cp_group) -> Tensor: + """Head shard -> CP sequence shard. + + Shape: ``[S, B, H/CP] -> [S/CP, B, H]``. + """ + if cp_group is None: + return input_ + cp_size = dist.get_world_size(group=cp_group) + if cp_size == 1: + return input_ + assert input_.dim() == 3, input_.shape + seq_len, batch, hidden_per_cp = input_.shape + if seq_len % cp_size != 0: + raise ValueError( + f"KDA all2all CP requires sequence dim {seq_len} divisible by CP {cp_size}." + ) + seq_per_cp = seq_len // cp_size + flat = input_.reshape(seq_len * batch, hidden_per_cp) + flat = _all_to_all_equal(flat, cp_group) + chunks = torch.split(flat, seq_per_cp * batch, dim=0) + return torch.cat(chunks, dim=-1).reshape(seq_per_cp, batch, hidden_per_cp * cp_size) + + +def _get_parameter_local_cp( + param: Tensor, + dim: int, + cp_rank: int, + cp_size: int, + split_size_or_sections: list[int] | None = None, +) -> Tensor: + """Slice a TP-local parameter for HybridEngine-style all2all CP.""" + if cp_size == 1: + return param + if split_size_or_sections is not None: + chunks = torch.split(param, split_size_or_sections, dim=dim) + return torch.cat( + [_get_parameter_local_cp(chunk, dim, cp_rank, cp_size) for chunk in chunks], + dim=dim, + ) + dim_size = param.size(dim) + if dim_size % cp_size != 0: + raise ValueError( + f"Cannot CP-slice parameter dim {dim} of size {dim_size} by CP {cp_size}." + ) + per_rank = dim_size // cp_size + slices = [slice(None)] * param.dim() + slices[dim] = slice(cp_rank * per_rank, (cp_rank + 1) * per_rank) + return param[tuple(slices)] + + +def _build_zigzag_undo_indices( + total_len: int, + cp_size: int, + cu_seqlens: Tensor | None, + device: torch.device, +) -> Tensor: + """Undo AReaL packed CP zigzag ordering after CP->HP all-to-all. + + Fully vectorized: for every canonical position the source index is + computed with tensor ops, so the only host sync is the single + divisibility check (and callers cache the result per microbatch via + ``_get_zigzag_undo_redo_indices``). + """ + indices = torch.arange(total_len, dtype=torch.long, device=device) + if cp_size <= 1: + return indices + + if cu_seqlens is None: + cu = torch.tensor([0, total_len], dtype=torch.long, device=device) + else: + cu = cu_seqlens.to(device=device, dtype=torch.long) + lens = cu[1:] - cu[:-1] + if bool(((lens % (2 * cp_size)) != 0).any()): + raise ValueError( + f"Packed sequence lengths {lens.tolist()} must be divisible by " + f"2*CP={2 * cp_size} for KDA CP zigzag reorder." + ) + + t_per_cp = total_len // cp_size + # For canonical (undone) position p in sequence i with offset o and + # zigzag chunk size c_i: chunk index k = o // c selects rank k for the + # front half (k < cp) and rank 2*cp-1-k for the mirrored back half; the + # source row lives in that rank's all-to-all block at the sequence's + # CP-local offset. + seq = torch.searchsorted(cu, indices, right=True) - 1 + o = indices - cu[seq] + c = (lens // (2 * cp_size))[seq] + cu_s = (cu[:-1] // cp_size)[seq] + k = o // c + j = o % c + front = k < cp_size + rank = torch.where(front, k, 2 * cp_size - 1 - k) + return rank * t_per_cp + cu_s + torch.where(front, j, j + c) + + +def _build_zigzag_redo_indices(undo_indices: Tensor) -> Tensor: + """Inverse permutation of ``_build_zigzag_undo_indices``.""" + redo = torch.empty_like(undo_indices) + redo[undo_indices] = torch.arange(undo_indices.numel(), device=undo_indices.device) + return redo + + +def _get_zigzag_undo_redo_indices( + total_len: int, + cp_size: int, + cu_seqlens: Tensor | None, + device: torch.device, +) -> tuple[Tensor, Tensor]: + """Build (or fetch cached) zigzag undo/redo index pairs. + + The permutation only depends on (total_len, cp_size, cu_seqlens), which + are identical for every KDA layer — and every recompute replay — of the + same microbatch, so cache on the cu_seqlens tensor object instead of + rebuilding (and syncing) per layer. + """ + key = (int(total_len), int(cp_size), str(device)) + cache = ( + getattr(cu_seqlens, "_kda_zigzag_cache", None) + if cu_seqlens is not None + else None + ) + if cache is not None and key in cache: + return cache[key] + undo = _build_zigzag_undo_indices(total_len, cp_size, cu_seqlens, device) + redo = _build_zigzag_redo_indices(undo) + if cu_seqlens is not None: + if cache is None: + cache = {} + cu_seqlens._kda_zigzag_cache = cache + cache[key] = (undo, redo) + return undo, redo + + +@dataclass +class KimiDeltaAttentionSubmodules: + """Module specs for the input/output projections and the output norm. + + For the v3 setting (``no_kda_lora=True``) ``f_b_proj`` and ``g_b_proj`` are + unused (Identity) because ``g`` and ``gate`` are produced directly by the fused + ``in_proj``. + """ + + in_proj: ModuleSpec | type = IdentityOp + beta_proj: ModuleSpec | type = IdentityOp + f_b_proj: ModuleSpec | type = IdentityOp + g_b_proj: ModuleSpec | type = IdentityOp + out_norm: ModuleSpec | type = IdentityOp + out_proj: ModuleSpec | type = IdentityOp + + +class KimiDeltaAttention(MegatronModule): + """KDA layer: input ``[s, b, h]`` -> output ``[s, b, h]``. + + Ported from ant-megatron ``megatron/core/ssm/kda.py``. FP8/mxfp8 fast paths and + the transpose / rms-norm fusion optimizations are intentionally omitted for the + bf16 bring-up; the numerical chain is preserved. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: KimiDeltaAttentionSubmodules, + layer_number: int = None, + attn_mask_type=None, + *, + head_dim: int | None = None, + conv_kernel_dim: int = 4, + no_kda_lora: bool = True, + use_qk_l2norm: bool = True, + safe_gate: bool = False, + lower_bound: float | None = None, + A_init_range: tuple[float, float] = (1, 16), + dt_min: float = 0.001, + dt_max: float = 0.1, + dt_init_floor: float = 1e-4, + bias: bool = False, + conv_bias: bool = False, + conv_init: float | None = None, + **kwargs, + ): + if not HAVE_FLA: # pragma: no cover + raise ImportError( + "flash-linear-attention (fla) with KDA kernels is required. " + "Install a version exposing fla.ops.kda.{chunk_kda,fused_recurrent_kda} " + "and fla.ops.kda.gate.fused_kda_gate." + ) + + super().__init__(config) + + # Attributes from arguments + self.layer_number = layer_number + self.bias = bias + self.conv_bias = conv_bias + self.conv_init = conv_init + assert A_init_range[0] > 0 and A_init_range[1] >= A_init_range[0] + self.A_init_range = A_init_range + self.use_qk_l2norm = use_qk_l2norm + self.no_kda_lora = no_kda_lora + # KDA decay-gate clamping. BailingMoeV3 HF checkpoints are trained with + # ``kda_safe_gate=True`` / ``kda_lower_bound=-5.0``: the decay gate is the + # bounded ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` rather than the + # default unbounded ``-exp(A_log) * softplus(g + dt_bias)``. These must be + # forwarded to the kernel or the loaded weights will not reproduce their forward. + self.safe_gate = safe_gate + self.lower_bound = lower_bound + self.tp_size = _get_tp_world_size() + self.cp_size = _get_cp_world_size() + self.cp_rank = _get_cp_rank() + self.sp_size = self.tp_size if config.sequence_parallel else 1 + + # Attributes from config + self.config = config + self.hidden_size = config.hidden_size + self.head_dim = head_dim if head_dim is not None else config.kv_channels + self.act_fn = config.activation_func + self.activation = getattr(self.act_fn, "__name__", "silu") + self.conv_kernel_dim = conv_kernel_dim + self.key_head_dim = self.head_dim + self.value_head_dim = self.head_dim + self.num_key_heads = config.num_attention_heads + self.num_value_heads = config.num_attention_heads + self.qk_dim = self.key_head_dim * self.num_key_heads + self.v_dim = self.value_head_dim * self.num_value_heads + assert self.num_value_heads % self.num_key_heads == 0 + if self.cp_size > 1: + heads_per_tp = self.num_key_heads // self.tp_size + if heads_per_tp % self.cp_size != 0: + raise ValueError( + "KDA all2all CP requires num_attention_heads / TP divisible " + f"by CP, got heads={self.num_key_heads}, TP={self.tp_size}, " + f"CP={self.cp_size}." + ) + logger.info( + f"KDA all2all CP enabled: cp_size={self.cp_size}, " + f"heads_per_tp={heads_per_tp}, heads_per_cp={heads_per_tp // self.cp_size}" + ) + + # nGPT value normalization is unused by released v3 configs. + self.use_nGPT = getattr(config, "use_nGPT", False) + self.value_norm = getattr(config, "value_norm", False) + + # Whether the decay gate is computed inside chunk_kda or via fused_kda_gate. + # If fla exposes kda_gate_ref we may apply the gate externally + # (use_gate_in_kernel=False). + try: + from fla.ops.kda.gate import kda_gate_ref # noqa: F401 + + self.use_gate_in_kernel = False + except ImportError: + self.use_gate_in_kernel = True + + # When the checkpoint requires a clamped (safe) gate, force the in-kernel path: + # it matches the HF BailingMoeV3 forward exactly (chunk_kda(use_gate_in_kernel= + # True, safe_gate=..., lower_bound=...)). The external fused_kda_gate path is + # avoided here because its clamping signature varies across fla forks. + # + # GUARD: chunk_kda only grew safe_gate/lower_bound in Arc fla >= v0.4.2. Older + # forks (e.g. the v1.5.0-pinned e131287, or v0.4.0) silently swallow them via + # **kwargs and revert to the unbounded softplus decay -> wrong (silent) loss. + # Fail loudly instead so the fla version mismatch is caught at build time. + if self.safe_gate or self.lower_bound is not None: + import inspect + + try: + _kda_params = inspect.signature(chunk_kda).parameters + except (TypeError, ValueError): # pragma: no cover - builtin/extension + _kda_params = {} + if "lower_bound" not in _kda_params or "safe_gate" not in _kda_params: + raise ImportError( + "This BailingMoeV3 checkpoint needs the clamped KDA gate " + f"(safe_gate={self.safe_gate}, lower_bound={self.lower_bound}), but " + "the installed flash-linear-attention's chunk_kda does not accept " + "these arguments (they would be silently ignored, giving the wrong " + "unbounded-softplus decay). Use Arc flash-linear-attention >= v0.4.2 " + "(e.g. branch bailing_fla_v0.4.2_env_autotune)." + ) + self.use_gate_in_kernel = True + logger.info( + f"KDA use_gate_in_kernel={self.use_gate_in_kernel} " + f"safe_gate={self.safe_gate} lower_bound={self.lower_bound}" + ) + + # Input projection (fused). no_kda_lora -> [qkv | g(qk_dim) | gate(v_dim)]. + if not self.no_kda_lora: + self.in_proj_dim = self.qk_dim * 2 + self.v_dim + self.value_head_dim * 2 + else: + self.in_proj_dim = self.qk_dim * 2 + self.v_dim + self.qk_dim + self.v_dim + + self.in_proj = build_module( + submodules.in_proj, + self.hidden_size, + self.in_proj_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="fc1", + ) + + self.beta_proj = build_module( + submodules.beta_proj, + self.hidden_size, + self.num_key_heads, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="fc1", + ) + + if not self.no_kda_lora: + self.f_b_proj = build_module( + submodules.f_b_proj, + self.value_head_dim, + self.qk_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="f_b_proj", + ) + else: + self.f_b_proj = nn.Identity() + + # Depthwise causal Conv1d over the (q, k, v) channels. + self.conv_dim = self.qk_dim * 2 + self.v_dim + self.conv_dim_local_tp = self.conv_dim // self.tp_size + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim_local_tp, + out_channels=self.conv_dim_local_tp, + bias=conv_bias, + kernel_size=self.conv_kernel_dim, + groups=self.conv_dim_local_tp, + padding=self.conv_kernel_dim - 1, + device=torch.cuda.current_device(), + dtype=config.params_dtype, + ) + # partition_dim is required alongside tensor_model_parallel: the mbridge + # HF-export fallback concatenates TP shards with + # ``torch.cat(shards, dim=param.partition_dim)``, and mbridge's get_model + # fills an unset partition_dim with -1 — which silently merges the 3-D + # conv1d weight along the kernel axis at TP>1. + setattr(self.conv1d.weight, "tensor_model_parallel", True) + setattr(self.conv1d.weight, "partition_dim", 0) + if conv_bias: + setattr(self.conv1d.bias, "tensor_model_parallel", True) + setattr(self.conv1d.bias, "partition_dim", 0) + + self.num_k_heads_local_tp = self.num_key_heads // self.tp_size + + with get_cuda_rng_tracker().fork(): + # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max. + # Keep dt_bias in fp32 (like A_log below): the HF ckpt stores it as fp32 and + # it feeds the decay gate exp(A_log)*(g+dt_bias), which is sensitive to + # precision. Also, hf_load's fp32-buffer exemption only fires when the target + # param is already fp32 — a bf16 dt_bias would be silently loaded in bf16. + dt = torch.exp( + torch.rand( + self.qk_dim // self.tp_size, + device=torch.cuda.current_device(), + dtype=torch.float32, + ) + * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ).clamp(min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + self.dt_bias._no_reinit = True + self.dt_bias._no_weight_decay = True + setattr(self.dt_bias, "tensor_model_parallel", True) + setattr(self.dt_bias, "partition_dim", 0) + + A = torch.empty( + self.num_k_heads_local_tp, + dtype=torch.float32, + device=torch.cuda.current_device(), + ).uniform_(*A_init_range) + A_log = torch.log(A) # keep A_log in fp32 + self.A_log = nn.Parameter(A_log) + self.A_log._no_weight_decay = True + setattr(self.A_log, "tensor_model_parallel", True) + setattr(self.A_log, "partition_dim", 0) + + if not self.no_kda_lora: + self.g_b_proj = build_module( + submodules.g_b_proj, + self.value_head_dim, + self.v_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="g_b_proj", + ) + else: + self.g_b_proj = nn.Identity() + + # Output norm (per-head RMSNorm over value_head_dim) applied before gating. + self.out_norm = build_module( + submodules.out_norm, + config=self.config, + hidden_size=self.value_head_dim, + eps=self.config.layernorm_epsilon, + ) + + self.out_proj = build_module( + submodules.out_proj, + self.v_dim, + self.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=bias, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="fc2", + ) + + self.reset_parameters() + + def reset_parameters(self): + """Reset the convolution parameters if a custom init range is provided.""" + if self.config.perform_initialization and self.conv_init is not None: + with get_cuda_rng_tracker().fork(): + nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) + + def _pad_packed_qkv(self, qkv: Tensor, cu_seqlens: Tensor) -> tuple[Tensor, int]: + """Pad a packed qkv tensor ``[Total_Seq, Dim]`` into ``[B, Dim, Max_Seq]``.""" + seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + batch_size = len(seqlens) + max_seqlen = seqlens.max().item() + + padded = torch.zeros( + batch_size, max_seqlen, qkv.shape[-1], dtype=qkv.dtype, device=qkv.device + ) + total_tokens = cu_seqlens[-1].item() + batch_indices = torch.arange( + batch_size, device=qkv.device, dtype=torch.long + ).repeat_interleave(seqlens) + offsets = cu_seqlens[:-1].repeat_interleave(seqlens) + seq_indices = ( + torch.arange(total_tokens, device=qkv.device, dtype=torch.long) - offsets + ) + padded[batch_indices, seq_indices] = qkv + return padded.transpose(1, 2).contiguous(), max_seqlen + + def _unpad_packed_qkv(self, padded_qkv: Tensor, cu_seqlens: Tensor) -> Tensor: + """Unpad ``[B, Dim, Max_Seq]`` convolution output back to ``[Total_Seq, Dim]``.""" + padded_qkv = padded_qkv.transpose(1, 2) + seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + batch_size = len(seqlens) + total_tokens = cu_seqlens[-1].item() + batch_indices = torch.arange( + batch_size, device=padded_qkv.device, dtype=torch.long + ).repeat_interleave(seqlens) + offsets = cu_seqlens[:-1].repeat_interleave(seqlens) + seq_indices = ( + torch.arange(total_tokens, device=padded_qkv.device, dtype=torch.long) + - offsets + ) + return padded_qkv[batch_indices, seq_indices] + + def forward( + self, + hidden_states: Tensor, + attention_mask: Tensor | None = None, + key_value_states: Tensor | None = None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos: Tensor | None = None, + rotary_pos_sin: Tensor | None = None, + attention_bias: Tensor | None = None, + packed_seq_params: PackedSeqParams | None = None, + sequence_len_offset: int | None = None, + *, + inference_params=None, + **kwargs, + ): + """Forward pass. ``hidden_states`` is ``[S, B, H]``; RoPE is ignored (KDA is an SSM).""" + if inference_context is not None or inference_params is not None: + raise NotImplementedError("KDA does not support inference for now.") + + seq_len_hidden_states, batch, _ = hidden_states.shape + seq_len = seq_len_hidden_states * self.sp_size * self.cp_size + + cu_seqlens_q = None + if packed_seq_params is not None: + cu_seqlens_q = packed_seq_params.cu_seqlens_q + + # Fused input projection -> [qkv | g | gate] + qkvfg, _ = self.in_proj(hidden_states) + qkv, g, gate = torch.split( + qkvfg, + [ + (self.qk_dim * 2 + self.v_dim) // self.tp_size, + (self.value_head_dim if not self.no_kda_lora else self.qk_dim) + // self.tp_size, + (self.value_head_dim if not self.no_kda_lora else self.v_dim) + // self.tp_size, + ], + dim=-1, + ) + + beta, _ = self.beta_proj(hidden_states) + + cp_group = _get_cp_group() if self.cp_size > 1 else None + undo_idx = None + redo_idx = None + qkv_split_sections = [ + self.qk_dim // self.tp_size, + self.qk_dim // self.tp_size, + self.v_dim // self.tp_size, + ] + + if self.cp_size > 1: + # Convert CP sequence shards into head shards before convolution/KDA. + # q/k/v must be exchanged independently to preserve fused channel layout: + # [S/CP, B, (Q|K|V)/TP] -> [S, B, (Q|K|V)/(TP*CP)]. + qkv = _all_to_all_cp2hp( + qkv, + cp_group, + split_size_or_sections=qkv_split_sections, + ) + beta = _all_to_all_cp2hp(beta, cp_group) + undo_idx, redo_idx = _get_zigzag_undo_redo_indices( + qkv.shape[0], self.cp_size, cu_seqlens_q, qkv.device + ) + qkv = qkv[undo_idx] + beta = beta[undo_idx] + + # seq-first [S, B, *] -> batch-first [B, S, *] + beta = beta.transpose(0, 1) + qkv = qkv.transpose(0, 1) + + qk_dim_local = self.qk_dim // (self.tp_size * self.cp_size) + v_dim_local = self.v_dim // (self.tp_size * self.cp_size) + conv1d_weight = self.conv1d.weight + conv1d_bias = self.conv1d.bias + if self.cp_size > 1: + conv1d_weight = _get_parameter_local_cp( + conv1d_weight, + dim=0, + cp_rank=self.cp_rank, + cp_size=self.cp_size, + split_size_or_sections=qkv_split_sections, + ) + if conv1d_bias is not None: + conv1d_bias = _get_parameter_local_cp( + conv1d_bias, + dim=0, + cp_rank=self.cp_rank, + cp_size=self.cp_size, + split_size_or_sections=qkv_split_sections, + ) + + # Depthwise causal short convolution (+ silu) over the q/k/v channels. + if packed_seq_params is not None: + b, s, d = qkv.shape + if causal_conv1d_fn is None: + qkv_flat = qkv.reshape(-1, d) + qkv_padded, max_seqlen_q = self._pad_packed_qkv(qkv_flat, cu_seqlens_q) + qkv_conv = self.act_fn( + F.conv1d( + qkv_padded, + conv1d_weight, + conv1d_bias, + padding=self.conv_kernel_dim - 1, + groups=conv1d_weight.shape[0], + ) + )[..., :max_seqlen_q] + qkv = self._unpad_packed_qkv(qkv_conv, cu_seqlens_q).reshape(b, s, d) + else: + assert self.activation in ["silu", "swish"] + if ( + PackedSeqParamsWithSeqidx is not None + and isinstance(packed_seq_params, PackedSeqParamsWithSeqidx) + and packed_seq_params.seq_idx is not None + ): + seq_idx = packed_seq_params.seq_idx + else: + seqlens = cu_seqlens_q[1:] - cu_seqlens_q[:-1] + seq_idx = ( + torch.repeat_interleave( + torch.arange( + len(seqlens), device=qkv.device, dtype=torch.int32 + ), + seqlens, + ) + .unsqueeze(0) + .contiguous() + ) + qkv_input = qkv.reshape(1, -1, d).contiguous().transpose(1, 2) + qkv_conv = causal_conv1d_fn( + x=qkv_input, + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + seq_idx=seq_idx, + ) + qkv = qkv_conv.transpose(1, 2).reshape(b, s, d) + else: + qkv = qkv.transpose(1, 2).contiguous() # [B, Dim, S] + if causal_conv1d_fn is None: + qkv = self.act_fn( + F.conv1d( + qkv, + conv1d_weight, + conv1d_bias, + padding=self.conv_kernel_dim - 1, + groups=conv1d_weight.shape[0], + ) + )[..., :seq_len] + else: + assert self.activation in ["silu", "swish"] + qkv = causal_conv1d_fn( + x=qkv, + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + ) + qkv = qkv.transpose(1, 2) # [B, S, Dim] + + query, key, value = torch.split( + qkv, + [ + qk_dim_local, + qk_dim_local, + v_dim_local, + ], + dim=-1, + ) + + # Decay-gate features `g` (no_kda_lora -> Identity, already produced by in_proj). + if not self.no_kda_lora and self.tp_size > 1: + g = gather_from_tensor_model_parallel_region(g) + if self.config.sequence_parallel: + g = scatter_to_sequence_parallel_region(g) + if not self.no_kda_lora: + g, _ = self.f_b_proj(g) + if self.cp_size > 1: + g = _all_to_all_cp2hp(g, cp_group) + g = g[undo_idx] + g = g.transpose(0, 1) # [B, S, *] + + seq_len_for_kda = qkv.shape[1] + if packed_seq_params is not None: + query = query.reshape( + 1, seq_len_for_kda * batch, -1, self.key_head_dim + ).contiguous() + key = key.reshape( + 1, seq_len_for_kda * batch, -1, self.key_head_dim + ).contiguous() + value = value.reshape( + 1, seq_len_for_kda * batch, -1, self.value_head_dim + ).contiguous() + beta = beta.reshape(1, seq_len_for_kda * batch, -1).contiguous() + if self.use_gate_in_kernel: + g = g.reshape( + 1, seq_len_for_kda * batch, -1, self.key_head_dim + ).contiguous() + else: + g = g.reshape(1, seq_len_for_kda * batch, -1).contiguous() + else: + query = query.reshape( + batch, seq_len_for_kda, -1, self.key_head_dim + ).contiguous() + key = key.reshape( + batch, seq_len_for_kda, -1, self.key_head_dim + ).contiguous() + value = value.reshape( + batch, seq_len_for_kda, -1, self.value_head_dim + ).contiguous() + beta = beta.reshape(batch, seq_len_for_kda, -1).contiguous() + if self.use_gate_in_kernel: + g = g.reshape( + batch, seq_len_for_kda, -1, self.key_head_dim + ).contiguous() + + if self.use_qk_l2norm and self.use_nGPT and self.value_norm: + value = l2norm(value.contiguous()) + + A_log = self.A_log + dt_bias = self.dt_bias + if self.cp_size > 1: + A_log = _get_parameter_local_cp( + self.A_log, dim=0, cp_rank=self.cp_rank, cp_size=self.cp_size + ) + dt_bias = _get_parameter_local_cp( + self.dt_bias, dim=0, cp_rank=self.cp_rank, cp_size=self.cp_size + ) + + if not self.use_gate_in_kernel: + g = fused_kda_gate( + g, A_log.view(1, 1, -1, 1), self.key_head_dim, g_bias=dt_bias + ) + beta = beta.float().sigmoid() + + core_attn_out, _ = chunk_kda( + q=query, + k=key, + v=value, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + use_gate_in_kernel=self.use_gate_in_kernel, + safe_gate=self.safe_gate, + lower_bound=self.lower_bound, + cu_seqlens=cu_seqlens_q, + ) + + # Output gate `gate` (no_kda_lora -> Identity, already produced by in_proj). + if not self.no_kda_lora and self.tp_size > 1: + gate = gather_from_tensor_model_parallel_region(gate) + if self.config.sequence_parallel: + gate = scatter_to_sequence_parallel_region(gate) + if not self.no_kda_lora: + gate, _ = self.g_b_proj(gate) + + if self.cp_size > 1: + core_attn_out = core_attn_out.reshape(batch, seq_len_for_kda, -1) + core_attn_out = core_attn_out.transpose(0, 1).contiguous() + core_attn_out = core_attn_out[redo_idx] + core_attn_out = _all_to_all_hp2cp(core_attn_out, cp_group) + + local_seq_len = core_attn_out.shape[0] + core_attn_out = core_attn_out.reshape( + local_seq_len, batch, -1, self.value_head_dim + ) + gate = gate.contiguous().reshape( + local_seq_len, batch, -1, self.value_head_dim + ) + norm_out = self._apply_gated_norm(core_attn_out, gate) + norm_out = norm_out.reshape(local_seq_len, batch, -1) + else: + gate = gate.transpose(0, 1) + gate = gate.contiguous().reshape( + batch, seq_len_for_kda, -1, self.value_head_dim + ) + norm_out = self._apply_gated_norm(core_attn_out, gate) + norm_out = norm_out.reshape(batch, seq_len_for_kda, -1) + norm_out = norm_out.transpose(0, 1).contiguous() # [S, B, v_dim_local] + + out, out_bias = self.out_proj(norm_out) + return out, out_bias + + def _apply_gated_norm(self, x: Tensor, gate: Tensor) -> Tensor: + """``out_norm(x) * sigmoid(gate)`` with per-head RMSNorm over ``value_head_dim``.""" + x_shape = x.shape + x_dtype = x.dtype + x = x.reshape(-1, self.value_head_dim) + y = self.out_norm(x) + gate = gate.reshape(-1, self.value_head_dim) + y = y * torch.sigmoid(gate.float()) + return y.to(x_dtype).reshape(x_shape) + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + """Sharded state dict for distributed checkpointing. + + ``A_log`` / ``dt_bias`` are TP-sharded on axis 0. The fused ``in_proj`` and the + depthwise ``conv1d`` weights are further split into named logical chunks + (query/key/value[/g/gate]) so they can be resharded across TP independently. + """ + sharded_state_dict = {} + self._save_to_state_dict(sharded_state_dict, "", keep_vars=True) + sharded_state_dict = make_sharded_tensors_for_checkpoint( + sharded_state_dict, + prefix, + tensor_parallel_layers_axis_map={"A_log": 0, "dt_bias": 0}, + sharded_offsets=sharded_offsets, + ) + for name, module in self.named_children(): + if name == "conv1d": + module_sd = module.state_dict(prefix="", keep_vars=True) + tp_sharding_map = {"weight": 0} + if self.conv_bias: # pragma: no cover + tp_sharding_map["bias"] = 0 + module_sharded_sd = make_sharded_tensors_for_checkpoint( + module_sd, f"{prefix}{name}.", tp_sharding_map, sharded_offsets + ) + elif name == "out_norm": + # The KDA output norm is shared across heads and replicated across + # tensor-parallel ranks. Some norm implementations expose their own + # sharded_state_dict but do not encode the TP rank in replica_id, + # which makes Megatron DCP see duplicate replicated shards at TP>1. + module_sd = module.state_dict(prefix="", keep_vars=True) + module_sharded_sd = make_sharded_tensors_for_checkpoint( + module_sd, f"{prefix}{name}.", {}, sharded_offsets + ) + else: + module_sharded_sd = sharded_state_dict_default( + module, f"{prefix}{name}.", sharded_offsets, metadata + ) + sharded_state_dict.update(module_sharded_sd) + + in_proj_dim_local_tp = self.in_proj_dim // self.tp_size + assert ( + sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) + == in_proj_dim_local_tp + ), (in_proj_dim_local_tp, sharded_state_dict[f"{prefix}in_proj.weight"]) + sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory( + sharded_state_dict[f"{prefix}in_proj.weight"], + [ + self.qk_dim // self.tp_size, + self.qk_dim // self.tp_size, + self.v_dim // self.tp_size, + (self.value_head_dim if not self.no_kda_lora else self.qk_dim) + // self.tp_size, + (self.value_head_dim if not self.no_kda_lora else self.v_dim) + // self.tp_size, + ], + ["query", "key", "value", "g", "gate"], + 0, + ) + + conv_layer_name_list = ["conv1d.weight"] + assert ( + sharded_state_dict[f"{prefix}conv1d.weight"].data.size(0) + == self.conv_dim_local_tp + ), (self.conv_dim_local_tp, sharded_state_dict[f"{prefix}conv1d.weight"]) + if self.conv_bias: # pragma: no cover + conv_layer_name_list.append("conv1d.bias") + assert ( + sharded_state_dict[f"{prefix}conv1d.bias"].data.size(0) + == self.conv_dim_local_tp + ) + for conv_layer_name in conv_layer_name_list: + sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory( + sharded_state_dict[f"{prefix}{conv_layer_name}"], + [ + self.qk_dim // self.tp_size, + self.qk_dim // self.tp_size, + self.v_dim // self.tp_size, + ], + ["query", "key", "value"], + 0, + ) + + return sharded_state_dict + + +def _split_tensor_factory( + orig_sh_ten: ShardedTensor, + split_sections: list[int], + split_names: list[str], + split_dim: int, +) -> ShardedTensorFactory: + """Build a factory that splits a ShardedTensor into named independent chunks.""" + assert isinstance(orig_sh_ten, ShardedTensor), type(orig_sh_ten) + orig_sh_ten_no_data = orig_sh_ten.without_data() + + if sum(split_sections) != orig_sh_ten_no_data.local_shape[split_dim]: + raise ValueError( + f"Split sections must cover the whole dimension size, " + f"got {split_sections=} vs dimension size " + f"{orig_sh_ten_no_data.local_shape[split_dim]}" + ) + assert not isinstance(split_sections, int) + assert len(split_sections) == len(split_names) + + @torch.no_grad() + def sh_ten_build_fn( + key: str, + t: torch.Tensor, + replica_id: ReplicaId, + flattened_range: slice | None, + ): + factory_sh_ten = replace( + orig_sh_ten_no_data, + key=key, + data=t, + dtype=t.dtype, + replica_id=replica_id, + flattened_range=flattened_range, + ) + chunk_sh_tens = [] + split_start = 0 + for split_size, split_name in zip(split_sections, split_names): + split_chunks = factory_sh_ten.narrow(split_dim, split_start, split_size) + for sh_ten in split_chunks: + sh_ten.key = f"{sh_ten.key}.{split_name}" + chunk_sh_tens.extend(split_chunks) + split_start += split_size + + assert split_start == orig_sh_ten_no_data.local_shape[split_dim] + assert sum(sh_ten.data.numel() for sh_ten in chunk_sh_tens) == t.numel() + return chunk_sh_tens + + @torch.no_grad() + def sh_ten_merge_fn(sub_state_dict): + return torch.cat(sub_state_dict) + + return ShardedTensorFactory( + orig_sh_ten.key, + orig_sh_ten.data, + sh_ten_build_fn, + sh_ten_merge_fn, + orig_sh_ten.replica_id, + ) diff --git a/areal/models/mcore/lightning_attention.py b/areal/models/mcore/lightning_attention.py index 14771a8111..ba48db9483 100644 --- a/areal/models/mcore/lightning_attention.py +++ b/areal/models/mcore/lightning_attention.py @@ -194,33 +194,35 @@ def _build_zigzag_undo_indices( Supports both packed sequences (per-sequence zigzag via cu_seqlens) and fixed-length BSHD format (cu_seqlens=None -> single global sequence). """ - indices = torch.empty(total_len, dtype=torch.long, device=device) - t_per_cp = total_len // cp_size + indices = torch.arange(total_len, dtype=torch.long, device=device) + if cp_size <= 1: + return indices if cu_seqlens is None: - seq_bounds = [(0, total_len)] + cu = torch.tensor([0, total_len], dtype=torch.long, device=device) else: - seq_bounds = [ - (cu_seqlens[i].item(), cu_seqlens[i + 1].item()) - for i in range(cu_seqlens.shape[0] - 1) - ] - - for cu_start, cu_end in seq_bounds: - seq_len = cu_end - cu_start - chunk = seq_len // (2 * cp_size) - cu_s = cu_start // cp_size - - for j in range(cp_size): - block_start = j * t_per_cp + cu_s - base = torch.arange(chunk, device=device) - - dst_front = cu_start + j * chunk - indices[dst_front : dst_front + chunk] = block_start + base - - dst_back = cu_start + seq_len - (j + 1) * chunk - indices[dst_back : dst_back + chunk] = block_start + chunk + base + cu = cu_seqlens.to(device=device, dtype=torch.long) + lens = cu[1:] - cu[:-1] + if bool(((lens % (2 * cp_size)) != 0).any()): + raise ValueError( + f"Packed sequence lengths {lens.tolist()} must be divisible by " + f"2*CP={2 * cp_size} for CP zigzag reorder." + ) - return indices + t_per_cp = total_len // cp_size + # Fully vectorized (no per-sequence GPU->CPU sync): canonical position p + # in sequence i with offset o and zigzag chunk size c falls in chunk + # k = o // c — rank k's front half for k < cp, rank 2*cp-1-k's mirrored + # back half otherwise. + seq = torch.searchsorted(cu, indices, right=True) - 1 + o = indices - cu[seq] + c = (lens // (2 * cp_size))[seq] + cu_s = (cu[:-1] // cp_size)[seq] + k = o // c + j = o % c + front = k < cp_size + rank = torch.where(front, k, 2 * cp_size - 1 - k) + return rank * t_per_cp + cu_s + torch.where(front, j, j + c) def _build_zigzag_redo_indices(undo_indices: torch.Tensor) -> torch.Tensor: @@ -230,6 +232,37 @@ def _build_zigzag_redo_indices(undo_indices: torch.Tensor) -> torch.Tensor: return redo +def _get_zigzag_undo_redo_indices( + total_len: int, + cp_size: int, + cu_seqlens: torch.Tensor | None, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build (or fetch cached) zigzag undo/redo index pairs. + + The permutation only depends on (total_len, cp_size, cu_seqlens), which + are identical for every attention layer — and every recompute replay — + of the same microbatch, so cache on the cu_seqlens tensor object instead + of rebuilding per layer. + """ + key = (int(total_len), int(cp_size), str(device)) + cache = ( + getattr(cu_seqlens, "_zigzag_idx_cache", None) + if cu_seqlens is not None + else None + ) + if cache is not None and key in cache: + return cache[key] + undo = _build_zigzag_undo_indices(total_len, cp_size, cu_seqlens, device) + redo = _build_zigzag_redo_indices(undo) + if cu_seqlens is not None: + if cache is None: + cache = {} + cu_seqlens._zigzag_idx_cache = cache + cache[key] = (undo, redo) + return undo, redo + + @dataclass class LightningAttentionSubmodules: """Submodule specs for Lightning Self-Attention layer.""" @@ -657,7 +690,7 @@ def forward( full_seq_len = query.shape[0] # Undo zigzag: restore sequential token order for linear attention - undo_idx = _build_zigzag_undo_indices( + undo_idx, redo_idx = _get_zigzag_undo_redo_indices( full_seq_len, cp_size, cu_seqlens, query.device ) query = query[undo_idx] @@ -673,7 +706,6 @@ def forward( ) # Redo zigzag: restore zigzag order for all-to-all - redo_idx = _build_zigzag_redo_indices(undo_idx) attn_output = attn_output[redo_idx] # All-to-all: [S, B, H_local/CP, D] -> [S/CP, B, H_local, D] diff --git a/areal/models/mcore/registry.py b/areal/models/mcore/registry.py index c0112e04f6..ec8a1118e1 100644 --- a/areal/models/mcore/registry.py +++ b/areal/models/mcore/registry.py @@ -19,6 +19,10 @@ hf_to_mcore_config_bailing_moe, make_mcore_layer_specs_bailing_moe, ) +from areal.models.mcore.bailing_v3 import ( + hf_to_mcore_config_bailing_v3, + make_mcore_layer_specs_bailing_v3, +) from areal.models.mcore.qwen3 import ( hf_to_mcore_config_qwen3_dense, make_mcore_layer_specs_qwen3_dense, @@ -359,6 +363,8 @@ def make_hf_and_mcore_config( "BailingHybridForCausalLM", ): return hf_config, hf_to_mcore_config_bailing_moe(hf_config, dtype) + elif architecture == "BailingMoeV3ForCausalLM": + return hf_config, hf_to_mcore_config_bailing_v3(hf_config, dtype) else: raise ValueError( f"Architecture not registered for config conversion: {architecture}." @@ -376,6 +382,8 @@ def make_mcore_layer_specs(hf_config: PretrainedConfig, tf_config: TransformerCo "BailingHybridForCausalLM", ): return make_mcore_layer_specs_bailing_moe(tf_config, hf_config, use_te=True) + elif architecture == "BailingMoeV3ForCausalLM": + return make_mcore_layer_specs_bailing_v3(tf_config, hf_config, use_te=True) else: raise ValueError( f"Architecture not registered for config conversion: {architecture}." diff --git a/areal/utils/saver.py b/areal/utils/saver.py index 989ffd77cb..d6d2b02d67 100644 --- a/areal/utils/saver.py +++ b/areal/utils/saver.py @@ -130,6 +130,12 @@ def save( processor: AutoProcessor | None = None, base_model_path: str | None = None, ): + if ( + self.config.freq_epochs is None + and self.config.freq_steps is None + and self.config.freq_secs is None + ): + return if not self.freq_ctl.check( epochs=int(step == self.ft_spec.steps_per_epoch - 1), steps=1 ): diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 2ab6520d6f..1e9ade8b30 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -1109,12 +1109,12 @@ Refer to Megatron-LM documentation for implementation details. | `distribute_saved_activations` | boolean \| None | `None` | - | | `recompute_modules` | list of string \| None | `None` | - | | `moe_router_dtype` | string \| None | `"fp32"` | - | -| `moe_shared_expert_overlap` | boolean | `False` | Enable overlapping between shared expert computations and dispatcher communications. Without this, the shared experts execute after the routed experts. | +| `moe_shared_expert_overlap` | boolean \| None | `None` | Enable overlapping between shared expert computations and dispatcher communications. Without this, the shared experts execute after the routed experts. None keeps the model bridge's own default. | | `moe_enable_deepep` | boolean | `False` | - | | `moe_token_dispatcher_type` | string | `"alltoall"` | Type of token dispatcher. Options: 'allgather','alltoall' and 'flex'. | | `moe_permute_fusion` | boolean | `False` | Fuse token rearrangement ops during token dispatching. | | `moe_router_fusion` | boolean | `False` | Enable fusion for MoE TopK routing and aux-loss computation. Requires TransformerEngine >= 2.7.0. | -| `moe_router_bias_update_rate` | float | `0.0` | Update rate for auxiliary-loss-free MoE load balancing (DeepSeek V3 style). Controls how fast expert_bias adjusts. Default 0.0 disables bias updates; set a positive value such as 1e-3 to enable. | +| `moe_router_bias_update_rate` | float \| None | `None` | Update rate for auxiliary-loss-free MoE load balancing (DeepSeek V3 style). Controls how fast expert_bias adjusts. None keeps the model bridge's own default (AReaL bridges disable it or derive it from the checkpoint). Set 0.0 to disable explicitly; 1e-3 matches DeepSeek V3. | | `moe_z_loss_coeff` | float \| None | `None` | Scaling coefficient for router z-loss. Complements auxiliary-loss-free load balancing for router stability. A starting value of 1e-3 is recommended. None disables z-loss. | | `enable_chunked_logits` | boolean | `False` | Enable AReaL's CUDA-only chunked-logits path by replacing Megatron's native output layer with the vocab-parallel LM Head. NPU and tree training are unsupported. | | `entropy_requires_grad` | boolean | `False` | Whether the training loss requires entropy gradients. Defaults to False. With AReaL LM Head enabled, False permits destructive logits-storage reuse, so entropy is non-differentiable. Set True to use the differentiable fallback. | diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 8214bce1b1..5447c9c0af 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -1107,12 +1107,12 @@ Refer to Megatron-LM documentation for implementation details. | `distribute_saved_activations` | boolean \| None | `None` | - | | `recompute_modules` | list of string \| None | `None` | - | | `moe_router_dtype` | string \| None | `"fp32"` | - | -| `moe_shared_expert_overlap` | boolean | `False` | Enable overlapping between shared expert computations and dispatcher communications. Without this, the shared experts execute after the routed experts. | +| `moe_shared_expert_overlap` | boolean \| None | `None` | Enable overlapping between shared expert computations and dispatcher communications. Without this, the shared experts execute after the routed experts. None keeps the model bridge's own default. | | `moe_enable_deepep` | boolean | `False` | - | | `moe_token_dispatcher_type` | string | `"alltoall"` | Type of token dispatcher. Options: 'allgather','alltoall' and 'flex'. | | `moe_permute_fusion` | boolean | `False` | Fuse token rearrangement ops during token dispatching. | | `moe_router_fusion` | boolean | `False` | Enable fusion for MoE TopK routing and aux-loss computation. Requires TransformerEngine >= 2.7.0. | -| `moe_router_bias_update_rate` | float | `0.0` | Update rate for auxiliary-loss-free MoE load balancing (DeepSeek V3 style). Controls how fast expert_bias adjusts. Default 0.0 disables bias updates; set a positive value such as 1e-3 to enable. | +| `moe_router_bias_update_rate` | float \| None | `None` | Update rate for auxiliary-loss-free MoE load balancing (DeepSeek V3 style). Controls how fast expert_bias adjusts. None keeps the model bridge's own default (AReaL bridges disable it or derive it from the checkpoint). Set 0.0 to disable explicitly; 1e-3 matches DeepSeek V3. | | `moe_z_loss_coeff` | float \| None | `None` | Scaling coefficient for router z-loss. Complements auxiliary-loss-free load balancing for router stability. A starting value of 1e-3 is recommended. None disables z-loss. | | `enable_chunked_logits` | boolean | `False` | Enable AReaL's CUDA-only chunked-logits path by replacing Megatron's native output layer with the vocab-parallel LM Head. NPU and tree training are unsupported. | | `entropy_requires_grad` | boolean | `False` | Whether the training loss requires entropy gradients. Defaults to False. With AReaL LM Head enabled, False permits destructive logits-storage reuse, so entropy is non-differentiable. Set True to use the differentiable fallback. | diff --git a/examples/swe/config.py b/examples/swe/config.py new file mode 100644 index 0000000000..4c23dafe7d --- /dev/null +++ b/examples/swe/config.py @@ -0,0 +1,149 @@ +"""Configuration for SWE SFT training with AReaL.""" + +from dataclasses import dataclass, field + +from areal.api.cli_args import SFTConfig + + +@dataclass +class SweDataConfig: + """SWE-specific data processing configuration.""" + + filter_errors: bool = field( + default=True, + metadata={ + "help": "Discard pairs whose current segment contains a tool result " + "with is_error=True. Set to false to keep all pairs." + }, + ) + pre_split: bool = field( + default=False, + metadata={ + "help": "Input JSONL is already in pair format " + '(each line: {"messages": [...]}). ' + "Skip trajectory splitting and error filtering." + }, + ) + num_proc: int = field( + default=4, + metadata={"help": "Number of parallel workers for tokenization."}, + ) + strip_all_thinking: bool = field( + default=False, + metadata={ + "help": "Strip ... from ALL assistant turns " + "including the training target. By default only context " + "turns are stripped." + }, + ) + no_tools: bool = field( + default=False, + metadata={ + "help": "Do not pass tool definitions to apply_chat_template. " + "By default, tools are auto-extracted from the data and " + "rendered in the system prompt (e.g. Qwen3 '# Tools' block)." + }, + ) + + skip_pretokenized_filter: bool = field( + default=False, + metadata={ + "help": "Skip max_length filtering when loading a pre-tokenized " + "dataset. Useful when the dataset was already filtered during " + "pretokenization to avoid NFS cache conflicts from concurrent " + "dataset.filter() calls across ranks." + }, + ) + filter_empty_tool_calls: bool = field( + default=False, + metadata={ + "help": "Discard pairs whose training-target assistant turn has " + "no text content but has tool_calls (silent tool invocations)." + }, + ) + filter_bare_text_tool_calls: bool = field( + default=False, + metadata={ + "help": "Discard pairs whose training-target assistant turn has " + "text content without tags and has tool_calls." + }, + ) + truncate_task_notifications: bool = field( + default=False, + metadata={ + "help": "Truncate trajectories at the first " + "that follows a pure-text assistant turn. Removes noise from " + "background task completions." + }, + ) + parse_tool_call_args: bool = field( + default=False, + metadata={ + "help": "Convert OpenAI JSON-string tool_calls.arguments to dicts " + "before apply_chat_template. Required by GLM-4.x / GLM-5.x " + "templates; leave at the default (False) for Qwen / Llama / " + "Bailing, which expect the standard string form." + }, + ) + max_no_thinking_ratio: float | None = field( + default=None, + metadata={ + "help": "Maximum ratio of non-thinking pairs to thinking pairs. " + "For example, 1.0 gives 1:1 balance, 2.0 allows up to 2x " + "non-thinking pairs per thinking pair. " + "None (default) disables balancing." + }, + ) + split_mode: str = field( + default="pair", + metadata={ + "help": "Sample construction mode: 'pair' (default) splits " + "trajectories into progressive pairs; 'trajectory' keeps " + "the full trajectory as a single training sample." + }, + ) + random_strip_thinking_prob: float = field( + default=0.0, + metadata={ + "help": "Probability of stripping thinking from each target assistant " + "turn. 0.0 = no stripping (default), 1.0 = strip all. " + "Works in both pair mode and trajectory mode." + }, + ) + random_strip_thinking_seed: int = field( + default=42, + metadata={"help": "Random seed for reproducible thinking stripping decisions."}, + ) + n_thinking_variants: int = field( + default=1, + metadata={ + "help": "Number of thinking-pattern variants per trajectory. " + "1 = no augmentation (default). K > 1 = augment each " + "trajectory into K variants: the first preserves all " + "thinking, the rest randomly strip with " + "random_strip_thinking_prob." + }, + ) + cleanup_processed_dataset: bool = field( + default=True, + metadata={ + "help": "Remove the processed dataset cache directory after training. " + "Set to false to keep it for faster restarts." + }, + ) + dump_samples: int = field( + default=50, + metadata={ + "help": "Number of random samples to dump for inspection after " + "dataset processing. Each sample is saved as .txt and .json " + "in a dumped_samples/ directory alongside logs. " + "Set to 0 to disable, -1 to dump all." + }, + ) + + +@dataclass +class SweSFTConfig(SFTConfig): + """SFT configuration with SWE-specific data processing settings.""" + + swe: SweDataConfig = field(default_factory=SweDataConfig) diff --git a/examples/swe/train_sft.py b/examples/swe/train_sft.py new file mode 100644 index 0000000000..fe494437e9 --- /dev/null +++ b/examples/swe/train_sft.py @@ -0,0 +1,114 @@ +import getpass +import os +import pathlib +import shutil +import sys + +sys.path.append(str(pathlib.Path(__file__).parent)) +from config import SweSFTConfig + +from areal import SFTTrainer +from areal.api.cli_args import load_expr_config +from areal.dataset import get_custom_dataset +from areal.utils.hf_utils import load_hf_tokenizer +from areal.utils.logging import getLogger + +logger = getLogger("SweSFTTrain") + + +def _get_cache_dir(config: SweSFTConfig) -> str: + """Build the processed-dataset cache path next to checkpoints. + + Layout: ``{fileroot}/checkpoints/{user}/{experiment}/{trial}/processed_dataset``. + Scoped to trial_name so concurrent trials never share (and race on) a + cache directory. + """ + return os.path.join( + config.cluster.fileroot, + "checkpoints", + getpass.getuser(), + config.experiment_name, + config.trial_name, + "processed_dataset", + ) + + +def main(args): + config, _ = load_expr_config(args, SweSFTConfig) + + tokenizer = load_hf_tokenizer(config.tokenizer_path) + + rank = int(os.getenv("RANK", "0")) + cache_dir = _get_cache_dir(config) + + dump_dir = None + if config.swe.dump_samples != 0: + dump_dir = os.path.join( + config.cluster.fileroot, + "logs", + getpass.getuser(), + config.experiment_name, + config.trial_name, + "dumped_samples", + ) + + swe_kwargs = { + "num_proc": config.swe.num_proc, + "pre_split": config.swe.pre_split, + "filter_errors": config.swe.filter_errors, + "strip_all_thinking": config.swe.strip_all_thinking, + "filter_empty_tool_calls": config.swe.filter_empty_tool_calls, + "filter_bare_text_tool_calls": config.swe.filter_bare_text_tool_calls, + "truncate_task_notifications": config.swe.truncate_task_notifications, + "no_tools": config.swe.no_tools, + "skip_pretokenized_filter": config.swe.skip_pretokenized_filter, + "max_no_thinking_ratio": config.swe.max_no_thinking_ratio, + "split_mode": config.swe.split_mode, + "random_strip_thinking_prob": config.swe.random_strip_thinking_prob, + "random_strip_thinking_seed": config.swe.random_strip_thinking_seed, + "n_thinking_variants": config.swe.n_thinking_variants, + "dump_dir": dump_dir, + "dump_samples": config.swe.dump_samples, + "parse_tool_call_args": config.swe.parse_tool_call_args, + } + + train_dataset = get_custom_dataset( + split="train", + dataset_config=config.train_dataset, + tokenizer=tokenizer, + cache_dir=cache_dir, + **swe_kwargs, + ) + valid_dataset = None + if config.valid_dataset is not None: + # The cache dir holds one processed dataset (path is part of the cache + # meta); sharing it between train and valid would make each call + # invalidate and rebuild the other's cache. + valid_kwargs = dict(swe_kwargs) + if dump_dir is not None: + # Separate dump dir: sharing one directory would let the valid + # build overwrite the train build's sample_{i}.txt/.json dumps. + valid_kwargs["dump_dir"] = f"{dump_dir}_valid" + valid_dataset = get_custom_dataset( + split="test", + dataset_config=config.valid_dataset, + tokenizer=tokenizer, + cache_dir=None if cache_dir is None else f"{cache_dir}_valid", + **valid_kwargs, + ) + + with SFTTrainer( + config, train_dataset=train_dataset, valid_dataset=valid_dataset + ) as trainer: + trainer.train() + + # Cleanup processed dataset cache after training. + if config.swe.cleanup_processed_dataset and rank == 0: + for d in (cache_dir, f"{cache_dir}_valid"): + if os.path.isdir(d): + shutil.rmtree(d) + logger.info(f"Cleaned up processed dataset cache: {d}") + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/tests/models/test_zigzag_indices.py b/tests/models/test_zigzag_indices.py new file mode 100644 index 0000000000..ac6ac47d08 --- /dev/null +++ b/tests/models/test_zigzag_indices.py @@ -0,0 +1,140 @@ +"""Tests for the vectorized CP zigzag undo/redo index construction. + +The vectorized builders in lightning_attention/kda_attention replaced a +per-sequence Python loop that called ``.item()`` on GPU tensors in the layer +forward hot path. These tests pin the vectorized output to the original loop +reference and cover the per-microbatch cache.""" + +import pytest +import torch + +from areal.models.mcore.kda_attention import ( + _build_zigzag_undo_indices as kda_undo, +) +from areal.models.mcore.lightning_attention import ( + _build_zigzag_redo_indices, + _build_zigzag_undo_indices, + _get_zigzag_undo_redo_indices, +) + + +def _reference_undo_indices( + total_len: int, + cp_size: int, + cu_seqlens: torch.Tensor | None, + device: torch.device, +) -> torch.Tensor: + """Original loop-based implementation (pre-vectorization).""" + indices = torch.empty(total_len, dtype=torch.long, device=device) + if cp_size <= 1: + return torch.arange(total_len, dtype=torch.long, device=device) + t_per_cp = total_len // cp_size + + if cu_seqlens is None: + seq_bounds = [(0, total_len)] + else: + seq_bounds = [ + (int(cu_seqlens[i].item()), int(cu_seqlens[i + 1].item())) + for i in range(cu_seqlens.shape[0] - 1) + ] + + for cu_start, cu_end in seq_bounds: + seq_len = cu_end - cu_start + chunk = seq_len // (2 * cp_size) + cu_s = cu_start // cp_size + for j in range(cp_size): + block_start = j * t_per_cp + cu_s + base = torch.arange(chunk, device=device) + dst_front = cu_start + j * chunk + indices[dst_front : dst_front + chunk] = block_start + base + dst_back = cu_start + seq_len - (j + 1) * chunk + indices[dst_back : dst_back + chunk] = block_start + chunk + base + return indices + + +CASES = [ + (2, [8, 12]), + (2, [4]), + (2, [16, 4, 8, 24]), + (4, [16, 32]), + (4, [8, 8, 8]), + (8, [32, 64, 16]), +] + + +@pytest.mark.parametrize("cp_size,seq_lens", CASES) +@pytest.mark.parametrize("builder", [_build_zigzag_undo_indices, kda_undo]) +def test_vectorized_matches_loop_reference(cp_size, seq_lens, builder): + cu = torch.tensor([0] + list(torch.tensor(seq_lens).cumsum(0))) + total = int(cu[-1]) + ref = _reference_undo_indices(total, cp_size, cu, torch.device("cpu")) + out = builder(total, cp_size, cu, torch.device("cpu")) + torch.testing.assert_close(out, ref) + + +@pytest.mark.parametrize("cp_size", [2, 4]) +@pytest.mark.parametrize("builder", [_build_zigzag_undo_indices, kda_undo]) +def test_no_cu_seqlens_single_sequence(cp_size, builder): + total = 16 * cp_size + ref = _reference_undo_indices(total, cp_size, None, torch.device("cpu")) + out = builder(total, cp_size, None, torch.device("cpu")) + torch.testing.assert_close(out, ref) + + +def test_undo_is_a_permutation_and_redo_inverts(): + cu = torch.tensor([0, 16, 40]) + undo = _build_zigzag_undo_indices(40, 4, cu, torch.device("cpu")) + assert torch.equal(torch.sort(undo).values, torch.arange(40)) + redo = _build_zigzag_redo_indices(undo) + x = torch.randn(40) + torch.testing.assert_close(x[undo][redo], x) + + +def test_roundtrip_restores_zigzag_layout(): + """undo applied to the concatenated per-rank zigzag shards must yield + the canonical sequence.""" + cp_size = 2 + cu = torch.tensor([0, 8, 20]) + total = 20 + canonical = torch.arange(total) + # Build each rank's zigzag shard the way AReaL packs them. + shards = [] + for rank in range(cp_size): + rows = [] + for i in range(len(cu) - 1): + seq = canonical[cu[i] : cu[i + 1]] + half = len(seq) // (2 * cp_size) + rows.append(seq[half * rank : half * (rank + 1)]) + rows.append(seq[len(seq) - half * (rank + 1) : len(seq) - half * rank]) + shards.append(torch.cat(rows)) + zigzag = torch.cat(shards) + undo = _build_zigzag_undo_indices(total, cp_size, cu, torch.device("cpu")) + torch.testing.assert_close(zigzag[undo], canonical) + + +def test_indivisible_sequence_raises(): + cu = torch.tensor([0, 6]) # 6 not divisible by 2*cp for cp=2 + with pytest.raises(ValueError, match="divisible"): + _build_zigzag_undo_indices(6, 2, cu, torch.device("cpu")) + + +class TestUndoRedoCache: + def test_cache_hits_for_same_microbatch(self): + cu = torch.tensor([0, 8, 20]) + a = _get_zigzag_undo_redo_indices(20, 2, cu, torch.device("cpu")) + b = _get_zigzag_undo_redo_indices(20, 2, cu, torch.device("cpu")) + assert a[0] is b[0] and a[1] is b[1] + + def test_cache_scoped_to_cu_tensor_object(self): + cu1 = torch.tensor([0, 8, 20]) + cu2 = torch.tensor([0, 8, 20]) + a = _get_zigzag_undo_redo_indices(20, 2, cu1, torch.device("cpu")) + b = _get_zigzag_undo_redo_indices(20, 2, cu2, torch.device("cpu")) + assert a[0] is not b[0] + torch.testing.assert_close(a[0], b[0]) + + def test_no_cu_seqlens_not_cached(self): + a = _get_zigzag_undo_redo_indices(16, 2, None, torch.device("cpu")) + b = _get_zigzag_undo_redo_indices(16, 2, None, torch.device("cpu")) + assert a[0] is not b[0] + torch.testing.assert_close(a[0], b[0]) diff --git a/tests/test_bailing_v3_hf_load.py b/tests/test_bailing_v3_hf_load.py new file mode 100644 index 0000000000..9e1f22b934 --- /dev/null +++ b/tests/test_bailing_v3_hf_load.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +def _stub_module(monkeypatch, name: str, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_hf_load_module(monkeypatch): + class _Platform: + device_type = "cpu" + + class _Logger: + def info(self, *args, **kwargs): + return None + + class _FP8BlockwiseTensorHelper: + pass + + _stub_module(monkeypatch, "mbridge") + _stub_module(monkeypatch, "mbridge.core") + _stub_module(monkeypatch, "mbridge.core.bridge", Bridge=type("Bridge", (), {})) + _stub_module(monkeypatch, "megatron") + core = _stub_module(monkeypatch, "megatron.core") + core.parallel_state = _stub_module(monkeypatch, "megatron.core.parallel_state") + _stub_module( + monkeypatch, + "megatron.core.fp8_utils", + is_float8tensor=lambda value: False, + ) + _stub_module(monkeypatch, "areal") + _stub_module(monkeypatch, "areal.engine") + _stub_module(monkeypatch, "areal.engine.core") + _stub_module( + monkeypatch, + "areal.engine.core.model", + lang_config=lambda config: config, + ) + _stub_module(monkeypatch, "areal.engine.megatron_utils") + _stub_module( + monkeypatch, + "areal.engine.megatron_utils.fp8", + FP8BlockwiseTensorHelper=_FP8BlockwiseTensorHelper, + dequantize_params=lambda *args, **kwargs: None, + get_block_size_from_config=lambda *args, **kwargs: None, + ) + _stub_module(monkeypatch, "areal.infra") + _stub_module( + monkeypatch, + "areal.infra.platforms", + current_platform=_Platform(), + ) + _stub_module(monkeypatch, "areal.models") + _stub_module(monkeypatch, "areal.models.mcore") + _stub_module( + monkeypatch, + "areal.models.mcore.registry", + unwrap_to_gpt_model=lambda model: model, + ) + _stub_module(monkeypatch, "areal.utils") + _stub_module( + monkeypatch, + "areal.utils.logging", + getLogger=lambda name: _Logger(), + ) + + path = Path(__file__).resolve().parents[1] / "areal/models/mcore/hf_load.py" + spec = importlib.util.spec_from_file_location("_test_bailing_v3_hf_load", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "weight_name", + [ + "decoder.layers.0.self_attention.in_proj.weight", + "decoder.layers.0.self_attention.conv1d.weight", + ], +) +def test_kda_fused_dispatch_is_scoped_to_bailing_v3(monkeypatch, weight_name): + hf_load = _load_hf_load_module(monkeypatch) + kda_result = torch.tensor([1.0]) + generic_result = torch.tensor([2.0]) + monkeypatch.setattr(hf_load, "_merge_kda_fused_weight", lambda *args: kda_result) + monkeypatch.setattr(hf_load, "_slice_generic_weight", lambda *args: generic_result) + + gdn_config = SimpleNamespace(architectures=["Qwen3_5ForCausalLM"]) + result = hf_load._weight_to_mcore_tp( + hf_config=gdn_config, + mcore_weights_name=weight_name, + mcore_param_shape=[1], + hf_weights_safe_slice=[torch.ones(1)], + tp_rank=0, + tp_size=1, + ) + + assert result is generic_result + + +@pytest.mark.parametrize( + "weight_name", + [ + "decoder.layers.0.self_attention.in_proj.weight", + "decoder.layers.0.self_attention.conv1d.weight", + ], +) +def test_kda_fused_dispatch_handles_bailing_v3(monkeypatch, weight_name): + hf_load = _load_hf_load_module(monkeypatch) + kda_result = torch.tensor([1.0]) + generic_result = torch.tensor([2.0]) + monkeypatch.setattr(hf_load, "_merge_kda_fused_weight", lambda *args: kda_result) + monkeypatch.setattr(hf_load, "_slice_generic_weight", lambda *args: generic_result) + + bailing_config = SimpleNamespace(architectures=["BailingMoeV3ForCausalLM"]) + result = hf_load._weight_to_mcore_tp( + hf_config=bailing_config, + mcore_weights_name=weight_name, + mcore_param_shape=[1], + hf_weights_safe_slice=[torch.ones(1)], + tp_rank=0, + tp_size=1, + ) + + assert result is kda_result diff --git a/tests/test_bailing_v3_kda_cp_helpers.py b/tests/test_bailing_v3_kda_cp_helpers.py new file mode 100644 index 0000000000..813f325ff5 --- /dev/null +++ b/tests/test_bailing_v3_kda_cp_helpers.py @@ -0,0 +1,224 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import torch +import torch.nn as nn + + +def _stub_module(monkeypatch, name: str, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_kda_attention_module(monkeypatch): + class _CudaRngTracker: + def fork(self): + return self + + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + class _MegatronModule(nn.Module): + def __init__(self, config=None): + super().__init__() + self.config = config + + class _Logger: + def info(self, *args, **kwargs): + return None + + _stub_module(monkeypatch, "megatron") + _stub_module(monkeypatch, "megatron.core") + _stub_module( + monkeypatch, + "megatron.core.parallel_state", + model_parallel_is_initialized=lambda: False, + ) + _stub_module( + monkeypatch, + "megatron.core.dist_checkpointing", + ShardedTensor=type("ShardedTensor", (), {}), + ) + _stub_module( + monkeypatch, + "megatron.core.dist_checkpointing.mapping", + ReplicaId=type("ReplicaId", (), {}), + ShardedTensorFactory=type("ShardedTensorFactory", (), {}), + ) + _stub_module( + monkeypatch, + "megatron.core.packed_seq_params", + PackedSeqParams=type("PackedSeqParams", (), {}), + ) + _stub_module( + monkeypatch, + "megatron.core.tensor_parallel", + get_cuda_rng_tracker=lambda: _CudaRngTracker(), + ) + _stub_module( + monkeypatch, + "megatron.core.tensor_parallel.mappings", + gather_from_tensor_model_parallel_region=lambda x: x, + scatter_to_sequence_parallel_region=lambda x: x, + ) + _stub_module( + monkeypatch, + "megatron.core.transformer", + TransformerConfig=type("TransformerConfig", (), {}), + ) + _stub_module( + monkeypatch, + "megatron.core.transformer.identity_op", + IdentityOp=type("IdentityOp", (nn.Module,), {}), + ) + _stub_module( + monkeypatch, + "megatron.core.transformer.module", + MegatronModule=_MegatronModule, + ) + _stub_module( + monkeypatch, + "megatron.core.transformer.spec_utils", + ModuleSpec=type("ModuleSpec", (), {}), + build_module=lambda *args, **kwargs: None, + ) + _stub_module( + monkeypatch, + "megatron.core.transformer.utils", + make_sharded_tensors_for_checkpoint=lambda *args, **kwargs: {}, + sharded_state_dict_default=lambda *args, **kwargs: {}, + ) + _stub_module(monkeypatch, "areal") + _stub_module(monkeypatch, "areal.utils") + _stub_module( + monkeypatch, + "areal.utils.logging", + getLogger=lambda name: _Logger(), + ) + + path = Path(__file__).resolve().parents[1] / "areal/models/mcore/kda_attention.py" + spec = importlib.util.spec_from_file_location("_test_kda_attention", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_get_parameter_local_cp_slices_each_qkv_section(monkeypatch): + kda = _load_kda_attention_module(monkeypatch) + param = torch.arange(12).reshape(12, 1) + + sliced = kda._get_parameter_local_cp( + param, + dim=0, + cp_rank=1, + cp_size=2, + split_size_or_sections=[4, 4, 4], + ) + + torch.testing.assert_close(sliced.squeeze(1), torch.tensor([2, 3, 6, 7, 10, 11])) + + +def test_all_to_all_cp2hp_preserves_qkv_section_boundaries(monkeypatch): + kda = _load_kda_attention_module(monkeypatch) + monkeypatch.setattr(kda.dist, "get_world_size", lambda group=None: 2) + monkeypatch.setattr(kda, "_all_to_all_equal", lambda input_, cp_group: input_) + + qkv = torch.arange(2 * 1 * 12).reshape(2, 1, 12) + split_out = kda._all_to_all_cp2hp( + qkv, + cp_group=object(), + split_size_or_sections=[4, 4, 4], + ) + expected = torch.cat( + [ + kda._all_to_all_cp2hp(chunk, cp_group=object()) + for chunk in torch.split(qkv, [4, 4, 4], dim=-1) + ], + dim=-1, + ) + unsplit_out = kda._all_to_all_cp2hp(qkv, cp_group=object()) + + torch.testing.assert_close(split_out, expected) + assert not torch.equal(split_out, unsplit_out) + + +def test_out_norm_sharded_state_dict_uses_replicated_wrapper(monkeypatch): + kda = _load_kda_attention_module(monkeypatch) + + class _FakeShard: + def __init__(self, data): + self.data = data + + class _OutNormWithOwnShard(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(2)) + + def sharded_state_dict(self, *args, **kwargs): + raise AssertionError("out_norm own sharded_state_dict should be bypassed") + + calls = [] + + def fake_make_sharded( + state_dict, + prefix, + tensor_parallel_layers_axis_map=None, + sharded_offsets=(), + ): + calls.append( + ( + "make", + prefix, + tuple(state_dict), + dict(tensor_parallel_layers_axis_map or {}), + ) + ) + return { + f"{prefix}{name}": _FakeShard(tensor) + for name, tensor in state_dict.items() + } + + def fake_default(module, prefix="", sharded_offsets=(), metadata=None): + calls.append(("default", prefix)) + return { + f"{prefix}{name}": _FakeShard(tensor) + for name, tensor in module.state_dict(prefix="", keep_vars=True).items() + } + + monkeypatch.setattr(kda, "make_sharded_tensors_for_checkpoint", fake_make_sharded) + monkeypatch.setattr(kda, "sharded_state_dict_default", fake_default) + monkeypatch.setattr(kda, "_split_tensor_factory", lambda shard, *args: shard) + + attn = kda.KimiDeltaAttention.__new__(kda.KimiDeltaAttention) + nn.Module.__init__(attn) + attn.tp_size = 2 + attn.in_proj_dim = 20 + attn.qk_dim = 4 + attn.v_dim = 4 + attn.value_head_dim = 2 + attn.no_kda_lora = True + attn.conv_dim_local_tp = 6 + attn.conv_bias = False + attn.add_module("in_proj", nn.Linear(1, 10, bias=False)) + attn.add_module("conv1d", nn.Conv1d(6, 6, 1, groups=6, bias=False)) + attn.add_module("out_norm", _OutNormWithOwnShard()) + + state_dict = attn.sharded_state_dict(prefix="decoder.layers.0.self_attention.") + + assert "decoder.layers.0.self_attention.out_norm.weight" in state_dict + assert ( + "make", + "decoder.layers.0.self_attention.out_norm.", + ("weight",), + {}, + ) in calls + assert ("default", "decoder.layers.0.self_attention.out_norm.") not in calls diff --git a/tests/test_dataset_swe_path_dispatch.py b/tests/test_dataset_swe_path_dispatch.py new file mode 100644 index 0000000000..c57250ffd6 --- /dev/null +++ b/tests/test_dataset_swe_path_dispatch.py @@ -0,0 +1,81 @@ +import importlib.util +import logging +import sys +import types +from pathlib import Path + +import pytest + + +def _load_dataset_module(): + """Load areal/dataset/__init__.py standalone, stubbing heavy dependencies.""" + saved_modules = { + name: module + for name, module in sys.modules.items() + if name == "areal" or name.startswith("areal.") + } + for name in list(sys.modules): + if name == "areal" or name.startswith("areal."): + del sys.modules[name] + + areal_module = types.ModuleType("areal") + api_module = types.ModuleType("areal.api") + cli_args_module = types.ModuleType("areal.api.cli_args") + cli_args_module._DatasetConfig = object + utils_module = types.ModuleType("areal.utils") + utils_module.logging = logging + areal_module.api = api_module + areal_module.utils = utils_module + api_module.cli_args = cli_args_module + sys.modules["areal"] = areal_module + sys.modules["areal.api"] = api_module + sys.modules["areal.api.cli_args"] = cli_args_module + sys.modules["areal.utils"] = utils_module + + path = Path(__file__).parents[1] / "areal" / "dataset" / "__init__.py" + spec = importlib.util.spec_from_file_location("areal.dataset", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules["areal.dataset"] = module + try: + spec.loader.exec_module(module) + finally: + for name in list(sys.modules): + if name == "areal" or name.startswith("areal."): + del sys.modules[name] + sys.modules.update(saved_modules) + return module + + +_SWE_PATH_PATTERN = _load_dataset_module()._SWE_PATH_PATTERN + + +@pytest.mark.parametrize( + "path", + [ + "/storage/datasets/swe_data/sft_test_dataset.jsonl", + "/data/swe-bench/train.jsonl", + "/data/swe_sft/processed", + "swe.jsonl", + "/exp/my_swe.jsonl", + "/Data/SWE_data/file.jsonl", + ], +) +def test_swe_path_pattern_matches_swe_token_paths(path): + """Test that _SWE_PATH_PATTERN matches when 'swe' is a delimited path token.""" + assert _SWE_PATH_PATTERN.search(path.lower()) + + +@pytest.mark.parametrize( + "path", + [ + "/data/answer_sft/dataset", + "/home/swetha/my_sft_data", + "/data/sweep_results/train.jsonl", + "/corpora/swedish_sft", + "/data/answers.jsonl", + ], +) +def test_swe_path_pattern_ignores_incidental_trigram(path): + """Test that _SWE_PATH_PATTERN does not fire on paths merely containing 'swe'.""" + assert not _SWE_PATH_PATTERN.search(path.lower()) diff --git a/tests/test_swe_sft_cache.py b/tests/test_swe_sft_cache.py new file mode 100644 index 0000000000..0ac582fc25 --- /dev/null +++ b/tests/test_swe_sft_cache.py @@ -0,0 +1,227 @@ +import importlib.util +import json +import logging +import sys +import threading +import types +from pathlib import Path + +import pytest +from datasets import Dataset + + +def _load_swe_sft_module(): + saved_modules = { + name: module + for name, module in sys.modules.items() + if name == "areal" or name.startswith("areal.") + } + for name in list(sys.modules): + if name == "areal" or name.startswith("areal."): + del sys.modules[name] + + areal_module = types.ModuleType("areal") + dataset_module = types.ModuleType("areal.dataset") + utils_module = types.ModuleType("areal.utils") + utils_module.logging = logging + areal_module.dataset = dataset_module + areal_module.utils = utils_module + sys.modules["areal"] = areal_module + sys.modules["areal.dataset"] = dataset_module + sys.modules["areal.utils"] = utils_module + + path = Path(__file__).parents[1] / "areal" / "dataset" / "swe_sft.py" + spec = importlib.util.spec_from_file_location("areal.dataset.swe_sft", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules["areal.dataset.swe_sft"] = module + try: + spec.loader.exec_module(module) + finally: + for name in list(sys.modules): + if name == "areal" or name.startswith("areal."): + del sys.modules[name] + sys.modules.update(saved_modules) + return module + + +swe_sft = _load_swe_sft_module() + + +def _write_cache(cache_dir, input_ids, max_length=2): + dataset = Dataset.from_dict( + { + "input_ids": input_ids, + "loss_mask": [[1] * len(ids) for ids in input_ids], + } + ) + dataset.save_to_disk(str(cache_dir)) + meta = { + "version": 1, + "path": "unused.jsonl", + "tokenizer": None, + "process_kwargs": { + "max_length": max_length, + "num_proc": None, + "pre_split": False, + "filter_errors": True, + "strip_all_thinking": False, + "filter_empty_tool_calls": False, + "filter_bare_text_tool_calls": False, + "truncate_task_notifications": False, + "no_tools": False, + "max_no_thinking_ratio": None, + "split_mode": "pair", + "random_strip_thinking_prob": 0.0, + "random_strip_thinking_seed": 42, + "n_thinking_variants": 1, + "parse_tool_call_args": False, + }, + } + (cache_dir / ".meta.json").write_text(json.dumps(meta, sort_keys=True)) + (cache_dir / ".done").write_text(str(len(dataset))) + + +def test_get_swe_sft_dataset_loads_distributed_cache(tmp_path, monkeypatch): + cache_dir = tmp_path / "processed_dataset" + _write_cache(cache_dir, [[1, 2], [1, 2, 3]]) + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "2") + + dataset = swe_sft.get_swe_sft_dataset( + "unused.jsonl", + tokenizer=object(), + cache_dir=str(cache_dir), + max_length=2, + ) + + assert len(dataset) == 1 + assert dataset[0]["input_ids"] == [1, 2] + + +def test_get_swe_sft_dataset_rebuilds_cache_filtered_to_empty(tmp_path, monkeypatch): + cache_dir = tmp_path / "processed_dataset" + _write_cache(cache_dir, [[1, 2, 3]]) + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "2") + + def fake_process_swe_sft(*args, **kwargs): + return Dataset.from_dict({"input_ids": [[1]], "loss_mask": [[1]]}) + + monkeypatch.setattr(swe_sft, "_process_swe_sft", fake_process_swe_sft) + + dataset = swe_sft.get_swe_sft_dataset( + "unused.jsonl", + tokenizer=object(), + cache_dir=str(cache_dir), + max_length=2, + ) + + assert len(dataset) == 1 + assert (cache_dir / ".done").read_text() == "1" + + +def test_get_swe_sft_dataset_filters_dataset_with_indices_mapping( + tmp_path, monkeypatch +): + """Test that the max-length filter handles .filter() views (indices mapping).""" + cache_dir = tmp_path / "processed_dataset" + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "2") + + def fake_process_swe_sft(*args, **kwargs): + # Mimic _tokenize_samples: a .filter() view whose underlying arrow + # table has more rows than the visible dataset. + ds = Dataset.from_dict( + { + "input_ids": [[1], [], [1, 2], [], [1, 2, 3]], + "loss_mask": [[1], [], [1, 1], [], [1, 1, 1]], + } + ) + return ds.filter(lambda x: len(x["input_ids"]) > 0) + + monkeypatch.setattr(swe_sft, "_process_swe_sft", fake_process_swe_sft) + + dataset = swe_sft.get_swe_sft_dataset( + "unused.jsonl", + tokenizer=object(), + cache_dir=str(cache_dir), + max_length=2, + ) + + # 3 non-empty rows built, the len-3 row is filtered by max_length=2. + assert len(dataset) == 2 + assert dataset[0]["input_ids"] == [1] + assert dataset[1]["input_ids"] == [1, 2] + + +def test_get_swe_sft_dataset_refuses_to_cache_empty_processed_dataset( + tmp_path, monkeypatch +): + cache_dir = tmp_path / "processed_dataset" + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "2") + + def fake_process_swe_sft(*args, **kwargs): + return Dataset.from_dict({"input_ids": [], "loss_mask": []}) + + monkeypatch.setattr(swe_sft, "_process_swe_sft", fake_process_swe_sft) + + with pytest.raises(RuntimeError, match="produced 0 samples"): + swe_sft.get_swe_sft_dataset( + "unused.jsonl", + tokenizer=object(), + cache_dir=str(cache_dir), + max_length=2, + ) + + assert not (cache_dir / ".done").exists() + + +def test_get_swe_sft_dataset_worker_loads_cache_written_by_rank0(tmp_path, monkeypatch): + """Test that a non-rank-0 worker loads the cache once rank 0 publishes it.""" + cache_dir = tmp_path / "processed_dataset" + monkeypatch.setenv("RANK", "1") + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setattr(swe_sft, "_RANK0_CACHE_TIMEOUT", 10) + monkeypatch.setattr(swe_sft, "_RANK0_CACHE_POLL_INTERVAL", 0.05) + + writer = threading.Timer(0.2, _write_cache, args=(cache_dir, [[1, 2]])) + writer.start() + try: + dataset = swe_sft.get_swe_sft_dataset( + "unused.jsonl", + tokenizer=object(), + cache_dir=str(cache_dir), + max_length=2, + ) + finally: + writer.join() + + assert len(dataset) == 1 + assert dataset[0]["input_ids"] == [1, 2] + + +def test_get_swe_sft_dataset_worker_rejects_mismatched_cache(tmp_path, monkeypatch): + """Test that a worker times out instead of loading a cache built with other settings.""" + cache_dir = tmp_path / "processed_dataset" + monkeypatch.setenv("RANK", "1") + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setattr(swe_sft, "_RANK0_CACHE_TIMEOUT", 0.5) + monkeypatch.setattr(swe_sft, "_RANK0_CACHE_POLL_INTERVAL", 0.05) + + # Cache published for max_length=99; this worker asks for max_length=2. + writer = threading.Timer( + 0.1, _write_cache, args=(cache_dir, [[1, 2]]), kwargs={"max_length": 99} + ) + writer.start() + try: + with pytest.raises(TimeoutError): + swe_sft.get_swe_sft_dataset( + "unused.jsonl", + tokenizer=object(), + cache_dir=str(cache_dir), + max_length=2, + ) + finally: + writer.join() From 1343abf8bb3ead9c44bdc3d5c812fd72dbc367b0 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Wed, 12 Aug 2026 11:18:25 +0800 Subject: [PATCH 2/8] refactor(dataset): split SWE SFT loader into modules Separate message processing, tokenization, pipeline orchestration, and CLI code so each concern can evolve without growing a single dataset module. Signed-off-by: chucai.dzq --- areal/dataset/swe_sft.py | 2692 ------------------------- areal/dataset/swe_sft/__init__.py | 11 + areal/dataset/swe_sft/__main__.py | 381 ++++ areal/dataset/swe_sft/messages.py | 786 ++++++++ areal/dataset/swe_sft/pipeline.py | 1001 +++++++++ areal/dataset/swe_sft/tokenization.py | 539 +++++ tests/test_swe_sft_cache.py | 20 +- 7 files changed, 2732 insertions(+), 2698 deletions(-) delete mode 100644 areal/dataset/swe_sft.py create mode 100644 areal/dataset/swe_sft/__init__.py create mode 100644 areal/dataset/swe_sft/__main__.py create mode 100644 areal/dataset/swe_sft/messages.py create mode 100644 areal/dataset/swe_sft/pipeline.py create mode 100644 areal/dataset/swe_sft/tokenization.py diff --git a/areal/dataset/swe_sft.py b/areal/dataset/swe_sft.py deleted file mode 100644 index 8548d7586b..0000000000 --- a/areal/dataset/swe_sft.py +++ /dev/null @@ -1,2692 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""SWE SFT dataset loader. - -Loads SWE-bench trajectory data and converts it into progressive SFT -training pairs. Each trajectory is split at assistant-turn boundaries -so that every pair ends with an assistant segment (assistant message + -its subsequent tool responses). - -Example trajectory:: - - [system, user, asst1, tool1a, tool1b, asst2, tool2, asst3] - -Produces three pairs:: - - Pair 1: [system, user, asst1, tool1a, tool1b] - Pair 2: [system, user, asst1, tool1a, tool1b, asst2, tool2] - Pair 3: [system, user, asst1, tool1a, tool1b, asst2, tool2, asst3] - -In each pair, only the **last** assistant segment is trained (loss=1); -earlier assistant turns are treated as context (loss=0). - -By default, pairs whose current segment contains a tool result with -``is_error=True`` are discarded. Set ``filter_errors=False`` to keep them. - -The file is organized into the following sections: - -1. **Constants & Infrastructure** — shared constants, distributed sync -2. **Cleaning** — message content transforms (thinking tags, field cleanup) -3. **Filters** — keep/discard predicates (error, empty, bare-text, truncation) -4. **Splitting** — trajectory → progressive pairs (segment detection + split) -5. **Tokenization** — template detection, render→tokenize→loss_mask, dump -6. **Pipeline** — loading, processing, distributed cache, public API -7. **CLI** — ``python -m areal.dataset.swe_sft`` entry point -""" - -import json -import os -import random -import re -import shutil -import time - -from datasets import Dataset - -from areal.utils import logging - -logger = logging.getLogger("SWESFTDataset") - - -# ============================================================ -# 1. Constants & Infrastructure -# ============================================================ - -DATASET_NUM_PROC = 1 - -# Timeout (seconds) for non-rank-0 workers waiting for rank 0 to finish -# dataset processing. Progressive-pair tokenization of large trajectory -# corpora is single-process on rank 0 and can take hours; 10 h is the -# upper bound before workers give up. -_RANK0_CACHE_TIMEOUT = 36000 -_RANK0_CACHE_POLL_INTERVAL = 5 - - -def _extract_messages(record, record_idx): - """Extract messages and tools from a parsed JSONL record. - - Handles nested (``conversations`` wrapper) and flat formats. - Warns if multiple conversations are present. - - Returns: - Tuple of ``(messages, record_tools)``. *messages* may be empty. - """ - convs = record.get("conversations", []) - if convs: - if len(convs) > 1: - logger.warning( - "Record %d has %d conversations, using only the last one.", - record_idx, - len(convs), - ) - conv = convs[-1] - return conv.get("messages", []), conv.get("tools") - return record.get("messages", []), record.get("tools") - - -def _set_messages(record, messages): - """Write *messages* back into *record* (inverse of ``_extract_messages``). - - Used by the ``--save-trajectories`` CLI path to update truncated - messages in the original record structure before serialization. - """ - convs = record.get("conversations", []) - if convs: - convs[-1]["messages"] = messages - else: - record["messages"] = messages - - -def _iter_jsonl_records(path): - """Iterate trajectory JSONL records. - - Yields ``(record_idx, messages, record_tools)`` tuples. Handles - nested (``conversations`` wrapper) vs flat format auto-detection - via ``_extract_messages``. Records with empty messages are skipped. - - Warns about multi-user trajectories which break think-tag rendering - in templates with ``ns.last_query_index`` logic (e.g. Bailing). - """ - record_idx = 0 - n_multi_user = 0 - with open(path, encoding="utf-8") as f: - for line in f: - line = line.strip() - if not line: - continue - record = json.loads(line) - record_idx += 1 - messages, record_tools = _extract_messages(record, record_idx) - if not messages: - continue - n_user = sum(1 for m in messages if m.get("role") == "user") - if n_user > 1: - n_multi_user += 1 - if n_multi_user <= 3: - logger.warning( - "Record %d has %d user messages. Templates with " - "ns.last_query_index logic (e.g. Bailing) will NOT " - "render for assistant turns before the last " - "user message. Consider filtering!", - record_idx, - n_user, - ) - yield record_idx, messages, record_tools - if n_multi_user > 0: - logger.warning( - "Total %d/%d records have multiple user messages. " - "These may produce no-think training signal.", - n_multi_user, - record_idx, - ) - - -# ============================================================ -# 2. Cleaning — message content transforms -# ============================================================ - -# Match reasoning blocks with any common tag variant: -# ... (Qwen standard) -# ... (Claude) -# The opening and closing tag names need not match exactly — mixed pairs -# like ``...`` (seen in distillation data) are handled. -_THINK_OPEN_RE = re.compile(r"") -_THINK_CLOSE_RE = re.compile(r"") -_THINK_RE = re.compile(r"(.*?)", re.DOTALL) - - -def _normalize_thinking_tags(content): - """Normalise all thinking tag variants to ````/````. - - Distillation data from different models may use ```` (Claude) - vs ```` (Qwen). Non-standard variants are multi-token for the - Qwen tokenizer which breaks think/tool_call boundaries. - """ - if not content: - return content - content = _THINK_OPEN_RE.sub("", content) - content = _THINK_CLOSE_RE.sub("", content) - return content - - -def _extract_thinking(content): - """Strip thinking blocks from *content*. - - Callers must run ``_normalize_thinking_tags`` first so that all - tag variants have been converted to ````/````. - - Returns: - Cleaned content with thinking blocks removed, or the original - content unchanged if no thinking tags are found. - """ - if not content: - return content - cleaned = _THINK_RE.sub("", content).strip() - return cleaned if cleaned != content.strip() else content - - -def _clean_message(msg, strip_thinking=True, ensure_thinking=False): - """Remove non-standard fields before tokenization. - - Keeps only the fields expected by tokenizer chat templates: - role, content, reasoning_content (for assistant), tool_calls - (for assistant), tool_call_id (for tool). - - Handles thinking content in two representations: - - - Inline ``...`` tags in ``content`` - - Separate ``reasoning_content`` field (DeepSeek, Qwen3 API style) - - If both are present, inline tags take priority and - ``reasoning_content`` is dropped with a warning to avoid double - thinking blocks in the rendered template. - - Args: - msg: Raw message dict. - strip_thinking: If True, remove thinking from assistant messages - (both inline ```` tags and ``reasoning_content``). - Used for context turns. If False, preserve thinking as-is - (used for the training-target assistant turn). - ensure_thinking: If True, inject inline ``\n`` - on assistant turns that lack a thinking block (either - inline or in ``reasoning_content``). Requires the - patched Bailing template (via - ``_patch_chat_template_for_training``) which detects - ``had_think_tags`` and preserves empty think blocks. - """ - cleaned = {"role": msg["role"]} - - # Handle content — some assistant messages have content=None when - # they only contain tool_calls. Preserve None so chat templates - # that distinguish None vs "" render correctly. - content = msg.get("content") - # Some APIs (DeepSeek, Qwen3 with enable_thinking) return thinking - # in a separate ``reasoning_content`` field instead of inline - # ```` tags. Handle both representations. - raw_reasoning = msg.get("reasoning_content") if msg["role"] == "assistant" else None - has_thinking = False - if content is not None: - if msg["role"] == "assistant": - content = _normalize_thinking_tags(content) - has_inline_thinking = bool(_THINK_RE.search(content)) - if has_inline_thinking and raw_reasoning and raw_reasoning.strip(): - # Conflict: both reasoning_content and inline tags. - # Keep inline tags (they are already in the content the - # tokenizer will see) and drop reasoning_content to avoid - # double thinking blocks in the rendered template. - if not strip_thinking: - logger.warning( - "Message has both reasoning_content and inline " - " tags. Keeping inline tags, dropping " - "reasoning_content." - ) - raw_reasoning = None - elif not has_inline_thinking and raw_reasoning and raw_reasoning.strip(): - # Convert reasoning_content → inline in content. - # This ensures a single representation that templates - # render identically to the reasoning_content path, while - # being more transparent and debuggable. - if not strip_thinking: - content = ( - f"\n{raw_reasoning.strip(chr(10))}\n" - f"\n\n{content.lstrip(chr(10))}" - ) - has_inline_thinking = True - raw_reasoning = None - has_thinking = has_inline_thinking or bool( - raw_reasoning and raw_reasoning.strip() - ) - if strip_thinking: - content = _extract_thinking(content) - cleaned["content"] = content - elif msg["role"] == "assistant" and msg.get("tool_calls"): - # Assistant with tool_calls but content=None. - if raw_reasoning and raw_reasoning.strip(): - has_thinking = True - if not strip_thinking: - # Convert reasoning_content → inline in content. - cleaned["content"] = ( - f"\n{raw_reasoning.strip(chr(10))}\n" - ) - else: - cleaned["content"] = None - raw_reasoning = None - else: - cleaned["content"] = None - else: - # Non-assistant messages without content: default to empty string. - cleaned["content"] = "" - - # Preserve reasoning_content for target turns only when it was NOT - # already inlined above (i.e. only when raw_reasoning is still set). - if not strip_thinking and raw_reasoning is not None: - cleaned["reasoning_content"] = raw_reasoning - - # For the target assistant turn without a thinking block, inject - # inline ``\n`` so that the (patched) template detects - # think intent via ``had_think_tags`` and renders - # ``\n\n\n\n`` — identical token output to the old - # ``reasoning_content='\n'`` approach. - # - # Requires ``_patch_chat_template_for_training`` to have been called - # on the tokenizer, otherwise the stock Bailing template will extract - # and discard the empty ```` block. - if ensure_thinking and msg["role"] == "assistant" and not has_thinking: - cur_content = cleaned.get("content") - if cur_content is None or cur_content == "": - cleaned["content"] = "\n" - else: - cleaned["content"] = f"\n\n\n{cur_content.lstrip(chr(10))}" - - # Copy tool_calls for assistant messages - if msg["role"] == "assistant" and msg.get("tool_calls"): - cleaned_tool_calls = [] - for tc in msg["tool_calls"]: - cleaned_tc = { - "type": tc.get("type", "function"), - "function": { - "name": tc["function"]["name"], - "arguments": json.dumps(tc["function"]["arguments"]) - if isinstance(tc["function"]["arguments"], dict) - else tc["function"]["arguments"], - }, - } - if "id" in tc: - cleaned_tc["id"] = tc["id"] - cleaned_tool_calls.append(cleaned_tc) - cleaned["tool_calls"] = cleaned_tool_calls - - # Copy tool_call_id for tool messages - if msg["role"] == "tool" and msg.get("tool_call_id"): - cleaned["tool_call_id"] = msg["tool_call_id"] - - return cleaned - - -# ============================================================ -# 3. Filters — keep/discard predicates -# ============================================================ - - -def _segment_has_error(messages, start, end): - """Check if any tool message in ``messages[start:end]`` has ``is_error=True``.""" - for m in messages[start:end]: - if m.get("role") == "tool" and m.get("is_error") is True: - return True - return False - - -def _is_empty_tool_call(msg): - """True if assistant *msg* has no text content and no reasoning but has tool_calls.""" - content = msg.get("content") or "" - if content.strip() or not msg.get("tool_calls"): - return False - # If reasoning_content exists, the model did think — not a silent invocation. - reasoning = msg.get("reasoning_content") - if reasoning and reasoning.strip(): - return False - return True - - -def _is_bare_text_tool_call(msg): - """True if assistant *msg* has text without ```` tags and has tool_calls.""" - content = msg.get("content") or "" - if not content.strip() or not msg.get("tool_calls"): - return False - # If reasoning_content exists, thinking is in a separate field — not bare text. - reasoning = msg.get("reasoning_content") - if reasoning and reasoning.strip(): - return False - normalized = _THINK_OPEN_RE.sub("", content) - normalized = _THINK_CLOSE_RE.sub("", normalized) - match = _THINK_RE.search(normalized) - return not (match and match.group(1).strip()) - - -def _msg_has_thinking(msg): - """True if assistant *msg* has thinking content (inline or reasoning_content).""" - if msg.get("role") != "assistant": - return False - content = msg.get("content") or "" - normalized = _THINK_OPEN_RE.sub("", content) - normalized = _THINK_CLOSE_RE.sub("", normalized) - if _THINK_RE.search(normalized): - return True - rc = msg.get("reasoning_content") or "" - return bool(rc.strip()) - - -def _truncate_at_task_notification(messages): - """Truncate messages when a ```` follows a pure-text assistant. - - Claude Code emits ```` as a user message when a - background task (e.g. ``pip install``) completes. If the model has - already produced a text-only summary (no tool_calls), the notification - and all subsequent messages are noise — the model just replies - "nothing to do". Truncating here removes that noise. - - Only triggers when the pattern is: - assistant (text, no tool_calls) → user () - - Returns: - Truncated message list (or the original list if no truncation needed). - """ - for i, m in enumerate(messages): - if m.get("role") != "user": - continue - if "" not in (m.get("content") or ""): - continue - # Find preceding assistant - prev_asst = None - for j in range(i - 1, -1, -1): - if messages[j].get("role") == "assistant": - prev_asst = messages[j] - break - if prev_asst is None: - continue - content = prev_asst.get("content") or "" - if content.strip() and not prev_asst.get("tool_calls"): - # Truncate: keep everything up to (but not including) this user msg - return messages[:i] - return messages - - -# ============================================================ -# 3b. Balancing — downsample non-thinking pairs -# ============================================================ - - -def _classify_pair(pair): - """Classify a pair by its target assistant turn's content type. - - Returns one of: - ``"thinking"`` — target has actual ```` content or - non-empty ``reasoning_content``. - ``"no_thinking_tool_call"`` — target has no thinking but has - ``tool_calls`` (the dominant category that causes distribution - skew). - ``"pure_text"`` — target has no thinking and no tool_calls - (typically the final summary turn in a trajectory). - """ - target = pair[-1] - if target.get("role") != "assistant": - return "pure_text" - - content = target.get("content") or "" - rc = target.get("reasoning_content") or "" - # Require non-empty think content: pair cleaning runs BEFORE balancing - # and (with ensure_thinking) injects an empty \n into - # every no-think target, so a bare regex hit would classify everything - # as "thinking" and silently disable max_no_thinking_ratio. - _m = _THINK_RE.search(content) - has_thinking = bool(_m and _m.group(1).strip()) or bool(rc.strip()) - - if has_thinking: - return "thinking" - if target.get("tool_calls"): - return "no_thinking_tool_call" - return "pure_text" - - -def _balance_thinking_pairs(pairs, max_no_thinking_ratio, seed=42, tools_list=None): - """Downsample non-thinking **tool-call** pairs to control balance. - - Only ``no_thinking_tool_call`` pairs (no thinking but has tool_calls) - are subject to downsampling. ``thinking`` pairs and ``pure_text`` - pairs (the final summary turn, no thinking and no tool_calls) are - always kept — the latter are critical for the model to learn when - to stop calling tools and give a final answer. - - Args: - pairs: List of progressive SFT pairs. - max_no_thinking_ratio: Maximum ratio of non-thinking tool-call - pairs to thinking pairs. For example, ``1.0`` means at most - 1:1, ``2.0`` means at most 2 non-thinking per 1 thinking pair. - ``None`` disables downsampling. - seed: Random seed for reproducible downsampling. - - Returns: - Balanced list of pairs (order preserved, randomly sampled for - the downsampled category). - """ - if max_no_thinking_ratio is None: - return pairs, tools_list - - thinking = [] - no_think_tc = [] - pure_text = [] - for i, pair in enumerate(pairs): - cat = _classify_pair(pair) - if cat == "thinking": - thinking.append(i) - elif cat == "no_thinking_tool_call": - no_think_tc.append(i) - else: - pure_text.append(i) - - n_think = len(thinking) - n_no_think_tc = len(no_think_tc) - n_pure_text = len(pure_text) - - if n_think == 0: - logger.warning( - "No thinking pairs found; skipping balance " - "(all %d pairs have empty thinking).", - n_no_think_tc + n_pure_text, - ) - return pairs, tools_list - - max_no_think_tc = int(n_think * max_no_thinking_ratio) - if n_no_think_tc <= max_no_think_tc: - logger.info( - "Thinking balance OK: %d thinking + %d no-think-tc + %d pure-text " - "(ratio %.1f <= %.1f), no downsampling needed.", - n_think, - n_no_think_tc, - n_pure_text, - n_no_think_tc / n_think, - max_no_thinking_ratio, - ) - return pairs, tools_list - - rng = random.Random(seed) - sampled_tc = set(rng.sample(no_think_tc, max_no_think_tc)) - keep_indices = sorted(set(thinking) | sampled_tc | set(pure_text)) - balanced = [pairs[i] for i in keep_indices] - balanced_tools = ( - [tools_list[i] for i in keep_indices] if tools_list is not None else None - ) - - logger.info( - "Balanced thinking pairs: %d thinking + %d no-think-tc " - "(downsampled from %d, ratio %.1f → %.1f) + %d pure-text (kept all).", - n_think, - max_no_think_tc, - n_no_think_tc, - n_no_think_tc / n_think, - max_no_thinking_ratio, - n_pure_text, - ) - return balanced, balanced_tools - - -# ============================================================ -# 3c. Thinking augmentation stats -# ============================================================ - - -def _log_thinking_augmentation_stats( - n_variants, - prob, - n_total_trajs, - thinking_turns_per_traj, - total_asst_turns_per_traj, - patterns_per_traj, -): - """Log adaptive-thinking augmentation quality metrics. - - Called after the augmentation loop in loaders to report how well - the ``n_thinking_variants`` / ``random_strip_thinking_prob`` settings - produce diverse thinking-pattern variants. - - Args: - n_variants: ``n_thinking_variants`` setting (K). - prob: ``random_strip_thinking_prob`` setting. - n_total_trajs: Total number of source trajectories processed. - thinking_turns_per_traj: List of N_thinking per source trajectory. - total_asst_turns_per_traj: List of N_total_asst per source trajectory. - patterns_per_traj: List of ``set[frozenset]`` — the unique strip - patterns generated for each source trajectory (including the - empty frozenset for the original unstripped variant). - """ - n_eligible = sum(1 for n in thinking_turns_per_traj if n > 0) - total_thinking = sum(thinking_turns_per_traj) - total_asst = sum(total_asst_turns_per_traj) - - # 1. Thinking Turn Coverage - avg_thinking = total_thinking / max(n_total_trajs, 1) - thinking_ratio = total_thinking / max(total_asst, 1) - - # 2. Pattern Diversity - diversity_ratios = [] - for n_think, patterns in zip(thinking_turns_per_traj, patterns_per_traj): - if n_think == 0: - continue - theoretical_max = min(n_variants, 2**n_think) - actual_unique = len(patterns) - diversity_ratios.append(actual_unique / theoretical_max) - avg_diversity = sum(diversity_ratios) / max(len(diversity_ratios), 1) - - # 3. Augmentation Efficiency - n_non_trivial = 0 - for patterns in patterns_per_traj: - # Count variants that differ from the original (non-empty strip set) - n_non_trivial += sum(1 for p in patterns if p) - expected_aug = (n_variants - 1) * max(n_eligible, 1) - efficiency = n_non_trivial / max(expected_aug, 1) - - # 4. Total sample count - n_total_samples = sum(len(p) for p in patterns_per_traj) - - logger.info( - f"Thinking augmentation stats (K={n_variants}, p={prob:.2f}):\n" - f" Source trajectories: {n_total_trajs} " - f"({n_eligible} with thinking turns)\n" - f" Thinking coverage: {avg_thinking:.1f} thinking turns/traj, " - f"{thinking_ratio:.1%} of all assistant turns\n" - f" Pattern diversity: {avg_diversity:.2f} " - f"(1.0 = all variants unique)\n" - f" Augmentation efficiency: {efficiency:.2f} " - f"({n_non_trivial}/{expected_aug} non-trivial variants)\n" - f" Total samples after augmentation: {n_total_samples}" - ) - - -# ============================================================ -# 4. Splitting — trajectory → progressive pairs -# ============================================================ - - -def _find_segments(messages): - """Find assistant+tools segment boundaries. - - Returns: - List of ``(assistant_start_idx, segment_end_idx)`` tuples. - """ - segments = [] - i = 0 - while i < len(messages): - if messages[i].get("role") == "assistant": - j = i + 1 - while j < len(messages) and messages[j].get("role") == "tool": - j += 1 - segments.append((i, j)) - i = j - else: - i += 1 - return segments - - -def _split_and_filter( - messages, - filter_errors=True, - strip_all_thinking=False, - filter_empty_tool_calls=False, - filter_bare_text_tool_calls=False, - random_strip_thinking_prob=0.0, - rng=None, -): - """Split trajectory into progressive pairs and optionally filter. - - By default, thinking (``...``) is stripped from context - assistant turns only; the last assistant turn (training target) keeps - its content unchanged. Set *strip_all_thinking* to strip from every - assistant turn including the target. - - When *random_strip_thinking_prob* > 0, each target assistant turn that - has thinking content is independently stripped with that probability. - Stripped turns use the context-cleaned version (thinking fully removed, - no empty ```` injected). - - Args: - messages: Raw trajectory messages. - filter_errors: If True (default), discard pairs whose current segment - contains a tool result with ``is_error=True``. Set to False to - keep all pairs regardless of tool errors. - strip_all_thinking: If True, strip ```` blocks from every - assistant turn including the training target. - filter_empty_tool_calls: If True, discard pairs whose training-target - assistant turn has no text content but has tool_calls. - filter_bare_text_tool_calls: If True, discard pairs whose - training-target assistant turn has text content without - ```` tags and has tool_calls. - random_strip_thinking_prob: Probability of stripping thinking - from each target assistant turn. 0.0 = no stripping. - rng: ``random.Random`` instance for reproducible sampling. - - Returns: - Tuple of ``(pairs, n_filtered_errors, n_filtered_empty_tc, - n_filtered_bare_tc, n_stripped)``. - """ - segments = _find_segments(messages) - if not segments: - return [], 0, 0, 0, 0 - - pairs = [] - n_filtered_errors = 0 - n_filtered_empty_tc = 0 - n_filtered_bare_tc = 0 - n_stripped = 0 - - # Pre-clean all messages in context mode (thinking stripped). - # This avoids re-cleaning the same message for every progressive pair - # (O(N+K) instead of O(N*K) where K = number of segments). - context_cleaned = [_clean_message(m, strip_thinking=True) for m in messages] - - # For target assistant turns, clean with thinking preserved (unless - # strip_all_thinking is set, in which case context_cleaned is reusable). - # When stripping is active (augmented variant), use ensure_thinking=False - # so empty-thinking turns don't get \n injected. - stripping_active = random_strip_thinking_prob > 0.0 and rng is not None - target_ensure = not stripping_active - target_cleaned = {} - if not strip_all_thinking: - for asst_start, _ in segments: - target_cleaned[asst_start] = _clean_message( - messages[asst_start], - strip_thinking=False, - ensure_thinking=target_ensure, - ) - - for asst_start, seg_end in segments: - # Check if current segment has any tool errors - if filter_errors and _segment_has_error(messages, asst_start, seg_end): - n_filtered_errors += 1 - continue - - # Content-type filters operate on the raw assistant message. - asst_msg = messages[asst_start] - if filter_empty_tool_calls and _is_empty_tool_call(asst_msg): - n_filtered_empty_tc += 1 - continue - if filter_bare_text_tool_calls and _is_bare_text_tool_call(asst_msg): - n_filtered_bare_tc += 1 - continue - - # Build pair: include context up to the target assistant turn, - # truncating tool responses that follow it. This ensures the - # target assistant is always the *last* message so that chat - # templates with ``loop.last``-dependent rendering (e.g. Qwen3 - # ```` injection) behave consistently. The tool responses - # would have loss_mask=0 anyway and only add noise. - pair = list(context_cleaned[: asst_start + 1]) - if not strip_all_thinking: - # Randomly strip: leave context_cleaned version (thinking - # already removed) instead of replacing with target_cleaned. - should_strip = ( - rng is not None - and _msg_has_thinking(messages[asst_start]) - and rng.random() < random_strip_thinking_prob - ) - if not should_strip: - pair[asst_start] = target_cleaned[asst_start] - else: - n_stripped += 1 - pairs.append(pair) - - return pairs, n_filtered_errors, n_filtered_empty_tc, n_filtered_bare_tc, n_stripped - - -def _prepare_trajectory( - messages, - filter_errors=True, - filter_empty_tool_calls=False, - filter_bare_text_tool_calls=False, - random_strip_thinking_prob=0.0, - rng=None, -): - """Prepare a full trajectory for trajectory-level training. - - Cleans all messages preserving thinking for every assistant turn - (``strip_thinking=False``, ``ensure_thinking=True``). Identifies - which assistant segments should be masked (``loss_mask=0``) based - on error tool responses, empty tool calls, or bare-text tool calls. - - When *random_strip_thinking_prob* > 0, each assistant turn that has - thinking content is independently stripped with that probability. - Stripped turns have their ```` blocks and ``reasoning_content`` - completely removed (no empty ```` injected). - - Args: - messages: Raw trajectory messages. - filter_errors: If True (default), mask segments with error tool - responses. - filter_empty_tool_calls: If True, mask segments whose assistant - turn has no text content but has tool_calls. - filter_bare_text_tool_calls: If True, mask segments whose - assistant turn has text without ```` tags and has - tool_calls. - random_strip_thinking_prob: Probability of stripping thinking - from each assistant turn that has thinking content. - 0.0 (default) = no stripping, 1.0 = strip all. - rng: ``random.Random`` instance for reproducible sampling. - - Returns: - Tuple of ``(cleaned_messages, masked_segment_indices, - n_error, n_empty_tc, n_bare_tc, stripped_pattern)`` or ``None`` - if the trajectory has no assistant turns. *stripped_pattern* is - a ``frozenset`` of message indices whose thinking was stripped - (empty if no stripping occurred). - """ - segments = _find_segments(messages) - if not segments: - return None - - masked_indices = set() - n_error = 0 - n_empty_tc = 0 - n_bare_tc = 0 - for idx, (asst_start, seg_end) in enumerate(segments): - if filter_errors and _segment_has_error(messages, asst_start, seg_end): - masked_indices.add(idx) - n_error += 1 - continue - asst_msg = messages[asst_start] - if filter_empty_tool_calls and _is_empty_tool_call(asst_msg): - masked_indices.add(idx) - n_empty_tc += 1 - continue - if filter_bare_text_tool_calls and _is_bare_text_tool_call(asst_msg): - masked_indices.add(idx) - n_bare_tc += 1 - - # Determine which assistant turns to randomly strip thinking from. - strip_thinking_indices = set() - stripping_active = random_strip_thinking_prob > 0.0 and rng is not None - if stripping_active: - for asst_start, _seg_end in segments: - if _msg_has_thinking(messages[asst_start]): - if rng.random() < random_strip_thinking_prob: - strip_thinking_indices.add(asst_start) - - # When stripping is active (augmented variant), use ensure_thinking=False - # for ALL turns so that empty-thinking turns don't get \n - # injected. Only real thinking content is preserved. - # When stripping is inactive (variant 0 or no augmentation), keep - # ensure_thinking=True to match the standard training format. - default_ensure = not stripping_active - - cleaned = [] - for i, m in enumerate(messages): - if i in strip_thinking_indices: - cleaned.append( - _clean_message(m, strip_thinking=True, ensure_thinking=False) - ) - else: - cleaned.append( - _clean_message(m, strip_thinking=False, ensure_thinking=default_ensure) - ) - - return ( - cleaned, - sorted(masked_indices), - n_error, - n_empty_tc, - n_bare_tc, - frozenset(strip_thinking_indices), - ) - - -# ============================================================ -# 5. Tokenization — template detection, render, loss mask, dump -# ============================================================ - - -# -- Chat template patch (runtime, no file modification) -------- - -# Both Bailing and Qwen3 templates have ``ns.last_query_index`` logic -# that prevents ```` rendering for assistant turns BEFORE the -# last user message, AND discards inline empty ``\n`` -# extracted from content. -# -# This breaks trajectory-mode training: -# - Multi-user trajectories: turns before the last user msg lack -# - Empty ensure_thinking via inline gets stripped -# -# The patch below handles both Bailing (`ASSISTANT` style) -# and Qwen3 (`<|im_start|>assistant` style) templates: -# 1. Adds ``had_think_tags`` detection so empty ```` survives. -# 2. Removes the ``ns.last_query_index`` gate so all assistant turns -# render ```` uniformly when think intent is detected. -# -# Applied at runtime via ``tokenizer.chat_template = patched`` — the -# original template file on disk is never modified. - -_BAILING_OLD_BLOCK = ( - "{%- if loop.index0 > ns.last_query_index %}\n" - " {%- if reasoning_content != '' %}\n" - " {{- 'ASSISTANT\\n' + '\\n'" - " + reasoning_content.strip('\\n') + '\\n\\n\\n'" - " + content.lstrip('\\n') }}\n" - " {%- else %}\n" - " {{- 'ASSISTANT\\n' + content }}\n" - " {%- endif %}\n" - " {%- else %}\n" - " {{- 'ASSISTANT\\n' + content }}\n" - " {%- endif %}" -) -_BAILING_NEW_BLOCK = ( - "{%- if reasoning_content != '' or had_think_tags %}\n" - " {{- 'ASSISTANT\\n' + '\\n'" - " + reasoning_content.strip('\\n') + '\\n\\n\\n'" - " + content.lstrip('\\n') }}\n" - " {%- else %}\n" - " {{- 'ASSISTANT\\n' + content }}\n" - " {%- endif %}" -) - -# Qwen3 uses `loop.last or (not loop.last and reasoning_content)` so the -# last turn always renders even with empty reasoning. We -# preserve `loop.last` and add `had_think_tags` for inline-empty support. -_QWEN3_OLD_BLOCK = ( - "{%- if loop.index0 > ns.last_query_index %}\n" - " {%- if loop.last or (not loop.last and reasoning_content) %}\n" - " {{- '<|im_start|>' + message.role + '\\n\\n'" - " + reasoning_content.strip('\\n') + '\\n\\n\\n'" - " + content.lstrip('\\n') }}\n" - " {%- else %}\n" - " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" - " {%- endif %}\n" - " {%- else %}\n" - " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" - " {%- endif %}" -) -_QWEN3_NEW_BLOCK = ( - "{%- if loop.last or reasoning_content != '' or had_think_tags %}\n" - " {{- '<|im_start|>' + message.role + '\\n\\n'" - " + reasoning_content.strip('\\n') + '\\n\\n\\n'" - " + content.lstrip('\\n') }}\n" - " {%- else %}\n" - " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" - " {%- endif %}" -) - -_OLD_DETECT = "{%- set reasoning_content = '' %}" -_NEW_DETECT = ( - "{%- set reasoning_content = '' %}\n" - " {%- set had_think_tags = ('' in content) %}" -) - - -def _patch_chat_template_for_training(tokenizer): - """Patch Bailing/Qwen3 chat templates to render ```` uniformly. - - Detects template family by matching known render blocks: - - Bailing: ``ASSISTANT`` markers - - Qwen3: ``<|im_start|>assistant`` markers - - Other templates (e.g. plain ChatML without ``last_query_index``) - are left unchanged. If the template has ``last_query_index`` but - neither known block matches, logs a warning. - """ - template = getattr(tokenizer, "chat_template", None) - if not template or "last_query_index" not in template: - return - - if _BAILING_OLD_BLOCK in template: - family = "Bailing" - patched = template.replace(_BAILING_OLD_BLOCK, _BAILING_NEW_BLOCK) - elif _QWEN3_OLD_BLOCK in template: - family = "Qwen3" - patched = template.replace(_QWEN3_OLD_BLOCK, _QWEN3_NEW_BLOCK) - else: - # Reaching here means the template family needs the training patch - # (it gates rendering on last_query_index) but the verbatim block no - # longer matches — most likely an upstream template revision. Failing - # loudly beats silently training on data whose blocks the - # stock template strips (see _clean_message / ensure_thinking). - # - # Escape hatch for uses that do not depend on think normalization - # (e.g. precision-alignment forward dumps): set - # AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE=1 to proceed with a warning. - if os.environ.get("AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE", ""): - logger.warning( - "Chat template has last_query_index but matches neither known " - "render block; proceeding UNPATCHED because " - "AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE is set. Empty " - "blocks may be discarded by the stock template." - ) - return - raise ValueError( - "Chat template has last_query_index but matches neither the known " - "Bailing nor Qwen3 render block; the training patch cannot be " - "applied. Without it, empty blocks are discarded and " - "multi-turn thinking renders inconsistently. Update " - "_BAILING_OLD_BLOCK/_QWEN3_OLD_BLOCK for this template revision, " - "or set AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE=1 if think " - "normalization is irrelevant for this run." - ) - - if _OLD_DETECT not in patched: - raise ValueError( - "Chat template render block matched but the reasoning_content " - "detect line did not; had_think_tags would be undefined and " - "empty blocks would silently vanish. Update " - "_OLD_DETECT/_NEW_DETECT for this template revision." - ) - patched = patched.replace(_OLD_DETECT, _NEW_DETECT) - - tokenizer.chat_template = patched - logger.info( - f"Patched {family} chat template for training: removed " - "last_query_index gate, added had_think_tags detection." - ) - - -_TEMPLATE_PATTERNS = [ - # ChatML (Qwen, etc.): <|im_start|>assistant\n ... <|im_end|> - (r"<\|im_start\|>assistant\n", r"<\|im_end\|>"), - # Llama 3: <|start_header_id|>assistant<|end_header_id|>\n\n ... <|eot_id|> - (r"<\|start_header_id\|>assistant<\|end_header_id\|>\n\n", r"<\|eot_id\|>"), - # GLM: <|assistant|> ... (ends at next <|user|>, <|observation|>, or end of string) - (r"<\|assistant\|>", r"(?=<\|user\|>|<\|observation\|>|\Z)"), -] - - -def _parse_tool_call_arguments(messages): - """Parse JSON-string arguments in tool_calls to dicts. - - OpenAI returns tool_call arguments as JSON strings, but some chat - templates (e.g. GLM-4.x / GLM-5.x) expect parsed dicts. Most other - templates (Qwen / ChatML, Llama 3, Bailing, ...) accept the standard - OpenAI string form, so this conversion must be opt-in. - """ - patched = [] - for m in messages: - tool_calls = m.get("tool_calls") - if not tool_calls: - patched.append(m) - continue - new_tcs = [] - for tc in tool_calls: - fn = tc.get("function", tc) - args = fn.get("arguments") - if isinstance(args, str): - try: - parsed = json.loads(args) - except (json.JSONDecodeError, TypeError): - parsed = args - fn = {**fn, "arguments": parsed} - tc = {**tc, "function": fn} if "function" in tc else fn - new_tcs.append(tc) - patched.append({**m, "tool_calls": new_tcs}) - return patched - - -def _render_tokenize_mask( - messages, - tokenizer, - assistant_pattern, - tools=None, - *, - split_mode="pair", - error_indices=None, - parse_tool_call_args=False, -): - """Render, tokenize, and build loss_mask for a message list. - - In **pair mode** (default), only the **last** assistant turn gets - ``loss_mask=1``. In **trajectory mode**, **all** assistant turns - get ``loss_mask=1`` except those at indices in *error_indices*. - - When *parse_tool_call_args* is True, JSON-string ``tool_calls`` arguments - are converted to dicts before rendering (required by GLM chat templates; - other templates such as Qwen / Llama / Bailing must keep the OpenAI - string form). - - Returns: - Tuple of ``(full_text, input_ids, loss_mask, offset_mapping)``, or - ``None`` if ``apply_chat_template`` fails. - """ - # 1) Render the full template text. - try: - kwargs = {"tokenize": False} - if tools is not None: - kwargs["tools"] = tools - if parse_tool_call_args: - messages = _parse_tool_call_arguments(messages) - full_text = tokenizer.apply_chat_template(messages, **kwargs) - except Exception as e: - logger.warning( - "apply_chat_template failed: %s. Skipping sample.", - e, - ) - return None - - # 2) Tokenize with offset mapping so we can map char→token. - encoding = tokenizer( - full_text, add_special_tokens=False, return_offsets_mapping=True - ) - input_ids = encoding["input_ids"] - offset_mapping = encoding["offset_mapping"] - - # 3) Build loss_mask. - loss_mask = [0] * len(input_ids) - - if split_mode == "trajectory": - # Trajectory mode: mask ALL assistant segments, skip error_indices. - skip = set(error_indices) if error_indices else set() - matches = list(assistant_pattern.finditer(full_text)) - - # Verify regex matches correspond 1:1 to assistant messages. - n_asst = sum(1 for m in messages if m.get("role") == "assistant") - if len(matches) != n_asst: - # Fail closed: a spurious match (e.g. a tool output quoting the - # chat-template header literal) would otherwise put loss on - # user/tool tokens and silently defeat error masking. - logger.warning( - "Segment count mismatch: %d assistant messages but %d regex " - "matches in rendered text. Dropping this sample.", - n_asst, - len(matches), - ) - return None - - for seg_idx, m in enumerate(matches): - if seg_idx in skip: - continue - rs, re_ = m.start(1), m.end(0) - for tok_idx, (cs, ce) in enumerate(offset_mapping): - if ce > rs and cs < re_: - loss_mask[tok_idx] = 1 - else: - # Pair mode: mask only the LAST assistant segment. - last_match = None - for m in assistant_pattern.finditer(full_text): - last_match = m - if last_match is not None: - rs, re_ = last_match.start(1), last_match.end(0) - for tok_idx, (cs, ce) in enumerate(offset_mapping): - if ce > rs and cs < re_: - loss_mask[tok_idx] = 1 - else: - # Loss lands nowhere; the SFT loss path tolerates all-zero masks - # (kept, not dropped, to preserve cache compatibility) but this - # always signals template/pattern drift worth investigating. - logger.warning( - "No assistant segment matched the template pattern; sample " - "keeps an all-zero loss_mask." - ) - - return full_text, input_ids, loss_mask, offset_mapping - - -class _TokenizeAndMask: - """Picklable callable for ``Dataset.map(num_proc=N)``.""" - - def __init__( - self, - tokenizer, - assistant_pattern, - max_length=None, - *, - split_mode="pair", - parse_tool_call_args=False, - ): - self.tokenizer = tokenizer - self.assistant_pattern = assistant_pattern - self.max_length = max_length - self.split_mode = split_mode - self.parse_tool_call_args = parse_tool_call_args - - def __call__(self, sample): - error_indices = ( - sample.get("error_indices", []) if self.split_mode == "trajectory" else None - ) - tools_json = sample.get("tools_json") - tools = json.loads(tools_json) if tools_json else None - result = _render_tokenize_mask( - sample["messages"], - self.tokenizer, - self.assistant_pattern, - tools, - split_mode=self.split_mode, - error_indices=error_indices, - parse_tool_call_args=self.parse_tool_call_args, - ) - if result is None: - return {"input_ids": [], "loss_mask": []} - - _full_text, input_ids, loss_mask, _offset_mapping = result - - # Early exit: overlength or empty → return empty so a single - # filter pass removes it together with template-failure empties. - if self.max_length is not None and len(input_ids) > self.max_length: - return {"input_ids": [], "loss_mask": []} - - return {"input_ids": input_ids, "loss_mask": loss_mask} - - -def _detect_template_pattern(tokenizer, tools=None): - """Detect the assistant role delimiter used by this tokenizer's template. - - When *tools* is provided the probe is rendered with ``tools=`` so that - the detected delimiters match the actual training text (some templates - alter the system block when tools are present). - - Strategy: - 1. Try known ``_TEMPLATE_PATTERNS`` (fast, battle-tested). - 2. Fall back to double-probe diff: render the template with a known - marker and with empty content, then diff the two strings to extract - the exact header and end-of-turn delimiters. - - Raises: - ValueError: If both strategies fail to detect a usable pattern. - """ - _PROBE_CONTENT = "PROBE_MARKER" - - extra_kwargs = {} - if tools is not None: - extra_kwargs["tools"] = tools - - probe_msgs = [ - {"role": "user", "content": "x"}, - {"role": "assistant", "content": _PROBE_CONTENT}, - ] - probe_text = tokenizer.apply_chat_template( - probe_msgs, tokenize=False, **extra_kwargs - ) - - # --- Strategy 1: known patterns --- - for hdr_re, eot_re in _TEMPLATE_PATTERNS: - if re.search(hdr_re, probe_text): - pattern = re.compile(hdr_re + r"(.*?)" + eot_re, re.DOTALL) - logger.info( - f"Detected template style (known pattern): " - f"header_re={hdr_re!r}, eot_re={eot_re!r}" - ) - return pattern - - # --- Strategy 2: double-probe diff --- - try: - probe_empty = [ - {"role": "user", "content": "x"}, - {"role": "assistant", "content": ""}, - ] - text_empty = tokenizer.apply_chat_template( - probe_empty, tokenize=False, **extra_kwargs - ) - - marker_idx = probe_text.index(_PROBE_CONTENT) - header = probe_text[:marker_idx] - tail = probe_text[marker_idx + len(_PROBE_CONTENT) :] - - if text_empty == header + tail: - # Extract the assistant-specific header by removing the shared - # user-only prefix. - user_only = tokenizer.apply_chat_template( - [{"role": "user", "content": "x"}], - tokenize=False, - **extra_kwargs, - ) - asst_header = header[len(user_only) :] - # end-of-turn delimiter: strip leading newlines, then take - # up to the first newline (or the full string if none). - eot_stripped = tail.lstrip("\n") - eot = eot_stripped.split("\n")[0] if "\n" in eot_stripped else eot_stripped - - if asst_header and eot: - hdr_re = re.escape(asst_header) - eot_re = re.escape(eot) - pattern = re.compile(hdr_re + r"(.*?)" + eot_re, re.DOTALL) - logger.info( - f"Detected template style (probe diff): " - f"header={asst_header!r}, eot={eot!r}" - ) - return pattern - except (ValueError, IndexError): - pass # PROBE_CONTENT not found in rendered text, skip - - raise ValueError( - "Could not detect chat template assistant delimiters. " - "Unable to build a reliable loss mask. " - f"Probe text: {probe_text[:200]!r}" - ) - - -def _dump_samples( - samples, - tokenizer, - assistant_pattern, - tools_list, - dump_dir, - n_samples, - *, - split_mode="pair", - error_indices_list=None, - parse_tool_call_args=False, -): - """Dump sampled message lists as ``.txt`` + ``.json`` for inspection. - - Args: - samples: List of message-list samples (pairs or full trajectories). - tokenizer: Tokenizer with ``apply_chat_template`` support. - assistant_pattern: Compiled regex from ``_detect_template_pattern``. - tools_list: Per-sample tool definitions (parallel to *samples*), - or ``None`` when no tools are available. - dump_dir: Directory to write files into (created if needed). - n_samples: Number of random samples to dump. ``-1`` dumps all. - split_mode: ``"trajectory"`` for trajectory-mode loss masking. - error_indices_list: Per-sample error segment indices (trajectory mode). - """ - import random as _random - - os.makedirs(dump_dir, exist_ok=True) - - if n_samples == -1 or n_samples >= len(samples): - indices = list(range(len(samples))) - else: - indices = sorted(_random.sample(range(len(samples)), n_samples)) - - n_written = 0 - for i in indices: - sample = samples[i] - sample_tools = tools_list[i] if tools_list else None - err_idxs = ( - error_indices_list[i] - if split_mode == "trajectory" and error_indices_list - else None - ) - - result = _render_tokenize_mask( - sample, - tokenizer, - assistant_pattern, - sample_tools, - split_mode=split_mode, - error_indices=err_idxs, - parse_tool_call_args=parse_tool_call_args, - ) - if result is None: - continue - - full_text, input_ids, loss_mask, offset_mapping = result - n_loss = sum(loss_mask) - base = os.path.join(dump_dir, f"sample_{i}") - - # --- .txt --- - with open(base + ".txt", "w", encoding="utf-8") as fout: - fout.write( - f"Sample {i}: {len(sample)} messages, " - f"{len(input_ids)} tokens, loss=1: {n_loss}\n" - ) - fout.write(f"Last msg role: {sample[-1]['role']}\n") - fout.write(f"{'=' * 72}\n\n") - - fout.write("--- Rendered Text ---\n") - fout.write(full_text) - fout.write("\n\n") - - fout.write("--- Token / Loss Mask ---\n") - fout.write(f"{'Idx':>6} | {'TokenID':>8} | Loss | Token Text\n") - fout.write(f"{'-' * 6}-+-{'-' * 8}-+------+{'-' * 40}\n") - for t in range(len(input_ids)): - cs, ce = offset_mapping[t] - tok_text = repr(full_text[cs:ce]) - fout.write( - f"{t:>6} | {input_ids[t]:>8} | {loss_mask[t]:>4} | {tok_text}\n" - ) - - # --- .json --- - tokens_list = [] - for t in range(len(input_ids)): - cs, ce = offset_mapping[t] - tokens_list.append( - { - "idx": t, - "token_id": input_ids[t], - "text": full_text[cs:ce], - "loss": loss_mask[t], - } - ) - record = { - "sample_index": i, - "n_messages": len(sample), - "n_tokens": len(input_ids), - "n_loss_tokens": n_loss, - "rendered_text": full_text, - "tokens": tokens_list, - } - with open(base + ".json", "w", encoding="utf-8") as fout: - json.dump(record, fout, ensure_ascii=False) - - n_written += 1 - - logger.info(f"Dumped {n_written} samples to {dump_dir}/") - - -# ============================================================ -# 6. Pipeline — loading, processing, distributed cache, public API -# ============================================================ - - -def _load_trajectory_pairs( - path: str, - filter_errors: bool = True, - strip_all_thinking: bool = False, - filter_empty_tool_calls: bool = False, - filter_bare_text_tool_calls: bool = False, - truncate_task_notifications: bool = False, - max_no_thinking_ratio: float | None = None, - random_strip_thinking_prob: float = 0.0, - random_strip_thinking_seed: int = 42, - n_thinking_variants: int = 1, -): - """Load trajectory JSONL and split into progressive pairs. - - When *n_thinking_variants* > 1, each trajectory is split K times: - variant 0 preserves all thinking, variants 1~K-1 randomly strip. - - Supports nested (``conversations`` wrapper) and flat JSONL formats - (auto-detected per record via ``_iter_jsonl_records``). - - Returns: - Tuple of ``(all_pairs, tools)`` where *tools* is ``None`` when no - tool definitions are found. - """ - all_pairs = [] - all_tools = [] - records_in = 0 - total_filtered_errors = 0 - total_filtered_empty_tc = 0 - total_filtered_bare_tc = 0 - total_truncated = 0 - total_stripped_thinking = 0 - - augment = n_thinking_variants > 1 - rng = ( - random.Random(random_strip_thinking_seed) - if random_strip_thinking_prob > 0.0 - else None - ) - - if augment and random_strip_thinking_prob <= 0.0: - logger.warning( - "n_thinking_variants=%d but random_strip_thinking_prob=0; " - "all variants will be identical.", - n_thinking_variants, - ) - - # Stats collectors for augmentation logging. - thinking_turns_per_traj = [] - total_asst_turns_per_traj = [] - patterns_per_traj = [] - - for record_idx, messages, record_tools in _iter_jsonl_records(path): - records_in = record_idx - - if truncate_task_notifications: - truncated = _truncate_at_task_notification(messages) - if len(truncated) < len(messages): - total_truncated += 1 - messages = truncated - - shared_kwargs = dict( - filter_errors=filter_errors, - strip_all_thinking=strip_all_thinking, - filter_empty_tool_calls=filter_empty_tool_calls, - filter_bare_text_tool_calls=filter_bare_text_tool_calls, - ) - - if augment: - # Variant 0: preserve all thinking. - pairs_orig, n_err, n_empty_tc, n_bare_tc, _ = _split_and_filter( - messages, **shared_kwargs, random_strip_thinking_prob=0.0, rng=None - ) - total_filtered_errors += n_err - total_filtered_empty_tc += n_empty_tc - total_filtered_bare_tc += n_bare_tc - all_pairs.extend(pairs_orig) - all_tools.extend([record_tools] * len(pairs_orig)) - # Collect stats. - segments = _find_segments(messages) - n_think = sum(1 for s, _ in segments if _msg_has_thinking(messages[s])) - n_asst = len(segments) - thinking_turns_per_traj.append(n_think) - total_asst_turns_per_traj.append(n_asst) - - # Variants 1 ~ K-1: random strip. - variant_patterns = {frozenset()} # original = no strip - for _k in range(n_thinking_variants - 1): - pairs_aug, _, _, _, n_stripped = _split_and_filter( - messages, - **shared_kwargs, - random_strip_thinking_prob=random_strip_thinking_prob, - rng=rng, - ) - total_stripped_thinking += n_stripped - all_pairs.extend(pairs_aug) - all_tools.extend([record_tools] * len(pairs_aug)) - # Approximate pattern: record which pairs had their target stripped. - # For stats, use the count as a proxy since _split_and_filter - # doesn't return per-pair strip info. - variant_patterns.add(frozenset([n_stripped])) - patterns_per_traj.append(variant_patterns) - else: - # Single variant (original behavior). - pairs, n_err, n_empty_tc, n_bare_tc, n_stripped = _split_and_filter( - messages, - **shared_kwargs, - random_strip_thinking_prob=random_strip_thinking_prob, - rng=rng, - ) - total_filtered_errors += n_err - total_filtered_empty_tc += n_empty_tc - total_filtered_bare_tc += n_bare_tc - total_stripped_thinking += n_stripped - all_pairs.extend(pairs) - all_tools.extend([record_tools] * len(pairs)) - - # Log extracted tools summary. - n_with_tools = sum(1 for t in all_tools if t is not None) - if n_with_tools > 0: - all_tool_names = set() - for t_list in all_tools: - if t_list is not None: - for t in t_list: - all_tool_names.add(t.get("function", {}).get("name", "?")) - logger.info( - f"Extracted tools from {n_with_tools}/{len(all_tools)} pairs: " - f"{sorted(all_tool_names)}" - ) - - filter_parts = [] - if total_truncated: - filter_parts.append( - f"{total_truncated} trajectories truncated at task-notification" - ) - if total_filtered_errors: - filter_parts.append(f"{total_filtered_errors} with tool errors") - if total_filtered_empty_tc: - filter_parts.append(f"{total_filtered_empty_tc} empty-content tool calls") - if total_filtered_bare_tc: - filter_parts.append(f"{total_filtered_bare_tc} bare-text tool calls") - if total_stripped_thinking: - filter_parts.append(f"{total_stripped_thinking} thinking blocks stripped") - filter_msg = ", ".join(filter_parts) if filter_parts else "none" - - logger.info( - f"Loaded {records_in} trajectories, " - f"generated {len(all_pairs)} pairs " - f"(filtered: {filter_msg})" - ) - - if augment and patterns_per_traj: - _log_thinking_augmentation_stats( - n_thinking_variants, - random_strip_thinking_prob, - records_in, - thinking_turns_per_traj, - total_asst_turns_per_traj, - patterns_per_traj, - ) - - # Balance thinking / no-thinking pair ratio. - all_pairs, all_tools = _balance_thinking_pairs( - all_pairs, max_no_thinking_ratio, tools_list=all_tools - ) - - return all_pairs, all_tools - - -def _load_presplit_pairs( - path: str, - strip_all_thinking: bool = False, - random_strip_thinking_prob: float = 0.0, - random_strip_thinking_seed: int = 42, - n_thinking_variants: int = 1, -): - """Load pre-split pair JSONL where each line is ``{"messages": [...]}``. - - Messages are cleaned but no splitting or error-filtering is performed. - By default, thinking is stripped from context assistant turns but - preserved for the last assistant turn (the training target). Set - *strip_all_thinking* to strip from every assistant turn. - - When *n_thinking_variants* > 1, each pair is augmented: variant 0 - preserves thinking, variants 1~K-1 randomly strip the target turn. - - Also extracts per-record ``tools`` definitions so that each pair - carries its own tools, same as ``_load_trajectory_pairs``. - - Returns: - Tuple of ``(all_pairs, all_tools)`` where *all_tools* is a - parallel list of per-sample tool definitions (may be ``None``). - """ - all_pairs = [] - all_tools = [] - n_stripped = 0 - augment = n_thinking_variants > 1 - - rng = ( - random.Random(random_strip_thinking_seed) - if random_strip_thinking_prob > 0.0 - else None - ) - - def _build_pair(messages, last_asst, strip_target): - pair = [] - for idx, m in enumerate(messages): - is_target = m.get("role") == "assistant" and idx == last_asst - strip = strip_all_thinking or not is_target or strip_target - pair.append(_clean_message(m, strip_thinking=strip)) - return pair - - with open(path, encoding="utf-8") as f: - for line in f: - line = line.strip() - if not line: - continue - record = json.loads(line) - messages = record.get("messages", []) - if not messages: - continue - - record_tools = record.get("tools") - - # Find the last assistant index so we can preserve its thinking. - last_asst = None - for i, m in enumerate(messages): - if m.get("role") == "assistant": - last_asst = i - - has_thinking = ( - last_asst is not None - and not strip_all_thinking - and _msg_has_thinking(messages[last_asst]) - ) - - if augment: - # Variant 0: preserve all thinking. - all_pairs.append(_build_pair(messages, last_asst, strip_target=False)) - all_tools.append(record_tools) - - # Variants 1 ~ K-1: random strip. - for _k in range(n_thinking_variants - 1): - do_strip = ( - has_thinking - and rng is not None - and rng.random() < random_strip_thinking_prob - ) - if do_strip: - n_stripped += 1 - all_pairs.append( - _build_pair(messages, last_asst, strip_target=do_strip) - ) - all_tools.append(record_tools) - else: - # Single variant (original behavior). - strip_target = ( - has_thinking - and rng is not None - and rng.random() < random_strip_thinking_prob - ) - if strip_target: - n_stripped += 1 - all_pairs.append( - _build_pair(messages, last_asst, strip_target=strip_target) - ) - all_tools.append(record_tools) - - # Log extracted tools summary. - n_with_tools = sum(1 for t in all_tools if t is not None) - if n_with_tools > 0: - all_tool_names = set() - for t_list in all_tools: - if t_list is not None: - for t in t_list: - all_tool_names.add(t.get("function", {}).get("name", "?")) - logger.info( - f"Extracted tools from {n_with_tools}/{len(all_tools)} pairs: " - f"{sorted(all_tool_names)}" - ) - - strip_msg = f", {n_stripped} thinking blocks stripped" if n_stripped else "" - logger.info(f"Loaded {len(all_pairs)} pre-split pairs from {path}{strip_msg}") - return all_pairs, all_tools - - -def _load_full_trajectories( - path: str, - filter_errors: bool = True, - filter_empty_tool_calls: bool = False, - filter_bare_text_tool_calls: bool = False, - truncate_task_notifications: bool = False, - random_strip_thinking_prob: float = 0.0, - random_strip_thinking_seed: int = 42, - n_thinking_variants: int = 1, -): - """Load trajectory JSONL for trajectory-level training. - - Each trajectory becomes a single training sample with all assistant - turns as targets (``loss_mask=1``). When *filter_errors* is True, - assistant segments with error tool responses are identified so - tokenization can mask them (``loss_mask=0``) instead of discarding - the entire trajectory. - - When *n_thinking_variants* > 1, each trajectory is augmented into - K variants: the first preserves all thinking, the remaining K-1 - randomly strip thinking turns with *random_strip_thinking_prob*. - - Supports nested (``conversations`` wrapper) and flat JSONL formats - (auto-detected per record via ``_iter_jsonl_records``). - - Returns: - Tuple of ``(trajectories, error_indices_list, all_tools)`` where - *trajectories* is a list of cleaned message lists, - *error_indices_list* is a list of error segment index lists, - and *all_tools* is a parallel list of per-sample tool definitions. - """ - trajectories = [] - error_indices_list = [] - all_tools = [] - records_in = 0 - total_truncated = 0 - total_masked_errors = 0 - total_masked_empty_tc = 0 - total_masked_bare_tc = 0 - total_stripped_thinking = 0 - - augment = n_thinking_variants > 1 - rng = ( - random.Random(random_strip_thinking_seed) - if random_strip_thinking_prob > 0.0 - else None - ) - - if augment and random_strip_thinking_prob <= 0.0: - logger.warning( - "n_thinking_variants=%d but random_strip_thinking_prob=0; " - "all variants will be identical.", - n_thinking_variants, - ) - - # Stats collectors for augmentation logging. - thinking_turns_per_traj = [] - total_asst_turns_per_traj = [] - patterns_per_traj = [] - - for record_idx, messages, record_tools in _iter_jsonl_records(path): - records_in = record_idx - - if truncate_task_notifications: - truncated = _truncate_at_task_notification(messages) - if len(truncated) < len(messages): - total_truncated += 1 - messages = truncated - - shared_kwargs = dict( - filter_errors=filter_errors, - filter_empty_tool_calls=filter_empty_tool_calls, - filter_bare_text_tool_calls=filter_bare_text_tool_calls, - ) - - if augment: - # Variant 0: preserve all thinking (no stripping). - result_orig = _prepare_trajectory( - messages, **shared_kwargs, random_strip_thinking_prob=0.0, rng=None - ) - if result_orig is None: - continue - cleaned_orig, masked_idxs, n_err, n_empty_tc, n_bare_tc, _ = result_orig - trajectories.append(cleaned_orig) - error_indices_list.append(masked_idxs) - all_tools.append(record_tools) - total_masked_errors += n_err - total_masked_empty_tc += n_empty_tc - total_masked_bare_tc += n_bare_tc - - # Collect stats: count thinking turns in this trajectory. - segments = _find_segments(messages) - n_think = sum(1 for s, _ in segments if _msg_has_thinking(messages[s])) - n_asst = len(segments) - thinking_turns_per_traj.append(n_think) - total_asst_turns_per_traj.append(n_asst) - - # Variants 1 ~ K-1: random strip thinking. - variant_patterns = {frozenset()} # original = empty pattern - for _k in range(n_thinking_variants - 1): - result_aug = _prepare_trajectory( - messages, - **shared_kwargs, - random_strip_thinking_prob=random_strip_thinking_prob, - rng=rng, - ) - if result_aug is None: - continue - cleaned_aug, _, _, _, _, strip_pattern = result_aug - trajectories.append(cleaned_aug) - error_indices_list.append(masked_idxs) # reuse - all_tools.append(record_tools) - total_stripped_thinking += len(strip_pattern) - variant_patterns.add(strip_pattern) - patterns_per_traj.append(variant_patterns) - else: - # Single variant (original behavior). - result = _prepare_trajectory( - messages, - **shared_kwargs, - random_strip_thinking_prob=random_strip_thinking_prob, - rng=rng, - ) - if result is None: - continue - cleaned, masked_idxs, n_err, n_empty_tc, n_bare_tc, strip_pattern = result - trajectories.append(cleaned) - error_indices_list.append(masked_idxs) - all_tools.append(record_tools) - total_masked_errors += n_err - total_masked_empty_tc += n_empty_tc - total_masked_bare_tc += n_bare_tc - total_stripped_thinking += len(strip_pattern) - - # Log extracted tools summary. - n_with_tools = sum(1 for t in all_tools if t is not None) - if n_with_tools > 0: - all_tool_names = set() - for t_list in all_tools: - if t_list is not None: - for t in t_list: - all_tool_names.add(t.get("function", {}).get("name", "?")) - logger.info( - f"Extracted tools from {n_with_tools}/{len(all_tools)} " - f"trajectories: {sorted(all_tool_names)}" - ) - - parts = [] - if total_truncated: - parts.append(f"{total_truncated} trajectories truncated at task-notification") - if total_masked_errors: - parts.append(f"{total_masked_errors} with tool errors") - if total_masked_empty_tc: - parts.append(f"{total_masked_empty_tc} empty-content tool calls") - if total_masked_bare_tc: - parts.append(f"{total_masked_bare_tc} bare-text tool calls") - if total_stripped_thinking: - parts.append(f"{total_stripped_thinking} thinking blocks stripped") - mask_msg = ", ".join(parts) if parts else "none" - - logger.info( - f"Loaded {records_in} trajectories, " - f"kept {len(trajectories)} for training " - f"(masked: {mask_msg})" - ) - - if augment and patterns_per_traj: - _log_thinking_augmentation_stats( - n_thinking_variants, - random_strip_thinking_prob, - records_in, - thinking_turns_per_traj, - total_asst_turns_per_traj, - patterns_per_traj, - ) - - return trajectories, error_indices_list, all_tools - - -def _tokenize_samples( - messages_list, - tools_list, - tokenizer, - *, - split_mode: str = "pair", - error_indices_list: list | None = None, - max_length: int | None = None, - num_proc: int | None = None, - no_tools: bool = False, - dump_dir: str | None = None, - dump_n_samples: int = 0, - parse_tool_call_args: bool = False, -): - """Tokenize message lists into a training-ready Dataset. - - Works for both progressive pairs (``split_mode="pair"``) and - full trajectories (``split_mode="trajectory"``). - - In pair mode, only the last assistant turn per sample gets - ``loss_mask=1``. In trajectory mode, all assistant turns get - ``loss_mask=1`` except those at error segment indices. - - Args: - tools_list: Per-sample tool definitions (parallel to - *messages_list*). Each element is either ``None`` or a - list of tool dicts. - """ - if num_proc is None: - num_proc = max(1, min(os.cpu_count() or 1, DATASET_NUM_PROC)) - - # Find representative tools for template detection. - first_tools = None - if tools_list: - first_tools = next((t for t in tools_list if t is not None), None) - - if no_tools: - tools_list = None - first_tools = None - logger.info("Tool definitions disabled (no_tools=True)") - elif first_tools is not None: - all_tool_names = set() - for t_list in tools_list: - if t_list is not None: - for t in t_list: - all_tool_names.add(t.get("function", {}).get("name", "?")) - logger.info(f"Using tools for chat template: {sorted(all_tool_names)}") - - if not messages_list: - raise ValueError("No valid samples to tokenize") - - # Build dataset columns. - data = {"messages": messages_list} - # Serialize per-sample tools as JSON strings for the Dataset column. - data["tools_json"] = ( - [json.dumps(t) if t else "" for t in tools_list] - if tools_list - else [""] * len(messages_list) - ) - remove_cols = ["messages", "tools_json"] - if split_mode == "trajectory": - data["error_indices"] = error_indices_list or [[] for _ in messages_list] - remove_cols.append("error_indices") - - dataset = Dataset.from_dict(data) - _patch_chat_template_for_training(tokenizer) - assistant_pattern = _detect_template_pattern(tokenizer, tools=first_tools) - - # Dump samples for inspection before the heavy map() pass. - if dump_dir and dump_n_samples != 0: - _dump_samples( - messages_list, - tokenizer, - assistant_pattern, - tools_list, - dump_dir, - dump_n_samples, - split_mode=split_mode, - error_indices_list=error_indices_list, - parse_tool_call_args=parse_tool_call_args, - ) - - process_fn = _TokenizeAndMask( - tokenizer, - assistant_pattern, - max_length=max_length, - split_mode=split_mode, - parse_tool_call_args=parse_tool_call_args, - ) - - dataset = dataset.map(process_fn, num_proc=num_proc).remove_columns(remove_cols) - - # Single filter pass: removes both apply_chat_template-failure empties and - # overlength samples (which _TokenizeAndMask also marks as empty). - before_filter = len(dataset) - dataset = dataset.filter(lambda x: len(x["input_ids"]) > 0, num_proc=num_proc) - n_filtered = before_filter - len(dataset) - if n_filtered > 0: - logger.info( - f"Filtered {n_filtered} samples " - f"(empty from template failures or exceeding max_length={max_length})" - ) - - logger.info(f"Final dataset: {len(dataset)} samples") - return dataset - - -def _process_swe_sft( - path: str, - tokenizer, - *, - max_length: int | None = None, - num_proc: int | None = None, - pre_split: bool = False, - filter_errors: bool = True, - strip_all_thinking: bool = False, - filter_empty_tool_calls: bool = False, - filter_bare_text_tool_calls: bool = False, - truncate_task_notifications: bool = False, - no_tools: bool = False, - max_no_thinking_ratio: float | None = None, - split_mode: str = "pair", - random_strip_thinking_prob: float = 0.0, - random_strip_thinking_seed: int = 42, - n_thinking_variants: int = 1, - dump_dir: str | None = None, - dump_n_samples: int = 0, - parse_tool_call_args: bool = False, -): - """Load JSONL, split into pairs, tokenize, and filter. - - Combines file loading with ``_tokenize_samples`` so that the rank-0-only - path and the single-process path share the same logic. - - When *split_mode* is ``"trajectory"``, the full trajectory is kept as a - single training sample with all assistant turns as targets. - """ - error_indices_list = None - - if split_mode == "trajectory": - messages_list, error_indices_list, tools_list = _load_full_trajectories( - path, - filter_errors=filter_errors, - filter_empty_tool_calls=filter_empty_tool_calls, - filter_bare_text_tool_calls=filter_bare_text_tool_calls, - truncate_task_notifications=truncate_task_notifications, - random_strip_thinking_prob=random_strip_thinking_prob, - random_strip_thinking_seed=random_strip_thinking_seed, - n_thinking_variants=n_thinking_variants, - ) - elif pre_split: - messages_list, tools_list = _load_presplit_pairs( - path, - strip_all_thinking=strip_all_thinking, - random_strip_thinking_prob=random_strip_thinking_prob, - random_strip_thinking_seed=random_strip_thinking_seed, - n_thinking_variants=n_thinking_variants, - ) - else: - messages_list, tools_list = _load_trajectory_pairs( - path, - filter_errors=filter_errors, - strip_all_thinking=strip_all_thinking, - filter_empty_tool_calls=filter_empty_tool_calls, - filter_bare_text_tool_calls=filter_bare_text_tool_calls, - truncate_task_notifications=truncate_task_notifications, - max_no_thinking_ratio=max_no_thinking_ratio, - random_strip_thinking_prob=random_strip_thinking_prob, - random_strip_thinking_seed=random_strip_thinking_seed, - n_thinking_variants=n_thinking_variants, - ) - - return _tokenize_samples( - messages_list, - tools_list, - tokenizer, - split_mode=split_mode, - error_indices_list=error_indices_list, - max_length=max_length, - num_proc=num_proc, - no_tools=no_tools, - dump_dir=dump_dir, - dump_n_samples=dump_n_samples, - parse_tool_call_args=parse_tool_call_args, - ) - - -def get_swe_sft_dataset( - path: str, - split: str | None = None, - tokenizer=None, - max_length: int | None = None, - num_proc: int | None = None, - pre_split: bool = False, - filter_errors: bool = True, - strip_all_thinking: bool = False, - filter_empty_tool_calls: bool = False, - filter_bare_text_tool_calls: bool = False, - truncate_task_notifications: bool = False, - no_tools: bool = False, - skip_pretokenized_filter: bool = False, - max_no_thinking_ratio: float | None = None, - split_mode: str = "pair", - random_strip_thinking_prob: float = 0.0, - random_strip_thinking_seed: int = 42, - n_thinking_variants: int = 1, - cache_dir: str | None = None, - dump_dir: str | None = None, - dump_samples: int = 0, - parse_tool_call_args: bool = False, -): - """Load SWE trajectory data and convert to SFT training pairs. - - By default, tool definitions are auto-extracted from the training data's - ``conversations[].tools`` field and passed to ``apply_chat_template`` - so that the tokenizer renders tool definitions in the system prompt - (e.g. Qwen3 ``# Tools`` block), matching the eval-time format. - Set *no_tools* to skip this and render without tool definitions. - - When *split_mode* is ``"trajectory"``, the full trajectory is kept as a - single training sample with all assistant turns as targets - (``loss_mask=1``). Error segments are masked (``loss_mask=0``) - when *filter_errors* is True, instead of being discarded. - Thinking is preserved by default but can be randomly stripped - per-turn via *random_strip_thinking_prob* (both modes). - - In distributed (SPMD) mode, only rank 0 performs the heavy processing - (JSONL loading, pair splitting, tokenization) and saves the result as - an Arrow dataset to *cache_dir*. Other ranks wait for rank 0 to - finish and then load the cached dataset directly via memory-mapped I/O. - - Args: - path: Path to the JSONL file containing SWE trajectories, or a - directory containing a pre-tokenized Arrow dataset (saved by - ``python -m areal.dataset.swe_sft --save-tokenized``). - split: Unused, kept for API compatibility. - tokenizer: Tokenizer with ``apply_chat_template`` support. - Not required when loading a pre-tokenized dataset. - max_length: Max token length. Longer sequences are filtered out. - num_proc: Number of parallel workers for tokenization. - Defaults to ``min(os.cpu_count(), DATASET_NUM_PROC)``. - pre_split: If True, treat input as pre-split pairs (each line is - ``{"messages": [...]}``) instead of full trajectories. - filter_errors: If True (default), discard pairs whose current segment - contains a tool result with ``is_error=True``. In trajectory - mode, sets ``loss_mask=0`` for error segments instead. - Set to False to keep/train all regardless of tool errors. - strip_all_thinking: If True, strip ``...`` from every - assistant turn including the training target. - Ignored in trajectory mode (thinking is always preserved). - filter_empty_tool_calls: If True, discard pairs whose training-target - assistant turn has no text content but has tool_calls. - filter_bare_text_tool_calls: If True, discard pairs whose - training-target assistant turn has text without ```` - tags and has tool_calls. - truncate_task_notifications: If True, truncate trajectories at the - first ```` that follows a pure-text assistant - turn, removing noise from background task completions. - no_tools: If True, do not pass tool definitions to - ``apply_chat_template`` even if the data contains them. - skip_pretokenized_filter: If True, skip the ``max_length`` filter - when loading a pre-tokenized dataset. Useful when the dataset - was already filtered during pretokenization and you want to - avoid NFS cache conflicts from concurrent ``dataset.filter()`` - calls across ranks. - max_no_thinking_ratio: Maximum ratio of non-thinking pairs to thinking - pairs. For example, ``1.0`` gives 1:1, ``2.0`` gives 1:2. - ``None`` (default) disables balancing. - split_mode: ``"pair"`` (default) splits trajectories into - progressive pairs. ``"trajectory"`` keeps the full trajectory - as a single sample — all assistant turns are targets with - ``loss_mask=1``, error segments are masked instead of filtered. - random_strip_thinking_prob: Probability of stripping thinking from - each target assistant turn. 0.0 (default) = no stripping, - 1.0 = strip all. Works in both pair and trajectory mode. - random_strip_thinking_seed: Random seed for reproducible thinking - stripping decisions. - n_thinking_variants: Number of thinking-pattern variants per - trajectory. ``1`` (default) = no augmentation. ``K > 1`` - = augment each trajectory into K variants: the first - preserves all thinking, the rest randomly strip with - *random_strip_thinking_prob*. - cache_dir: Directory to save/load the processed Arrow dataset. - When set in distributed mode, rank 0 processes the data and - saves here; other ranks load from this directory. If the - directory already contains a completed cache (``.done`` marker), - all ranks load from it directly without reprocessing. - dump_dir: Directory to write sample dump files (``.txt`` + ``.json``). - Only rank 0 writes. Set to None to disable. - dump_samples: Number of random samples to dump. ``-1`` = all, - ``0`` = disabled. - parse_tool_call_args: If True, convert OpenAI JSON-string - ``tool_calls.arguments`` to dicts before ``apply_chat_template``. - Required by GLM-4.x / GLM-5.x templates; leave at the default - (False) for Qwen / Llama / Bailing. - - Returns: - A HuggingFace ``Dataset`` with ``input_ids`` and ``loss_mask`` columns. - """ - from datasets import load_from_disk - - # Pre-tokenized Arrow dataset: load directly, skip all processing. - if os.path.isdir(path): - logger.info(f"Loading pre-tokenized dataset from {path}") - dataset = load_from_disk(path) - - if max_length is not None and not skip_pretokenized_filter: - before_filter = len(dataset) - dataset = dataset.filter( - lambda x: len(x["input_ids"]) <= max_length, num_proc=num_proc - ) - logger.info( - f"Filtered {before_filter - len(dataset)} samples " - f"exceeding max_length={max_length}" - ) - - logger.info(f"Final dataset: {len(dataset)} samples") - return dataset - - # --- Shared kwargs for _process_swe_sft --- - process_kwargs = dict( - max_length=max_length, - num_proc=num_proc, - pre_split=pre_split, - filter_errors=filter_errors, - strip_all_thinking=strip_all_thinking, - filter_empty_tool_calls=filter_empty_tool_calls, - filter_bare_text_tool_calls=filter_bare_text_tool_calls, - truncate_task_notifications=truncate_task_notifications, - no_tools=no_tools, - max_no_thinking_ratio=max_no_thinking_ratio, - split_mode=split_mode, - random_strip_thinking_prob=random_strip_thinking_prob, - random_strip_thinking_seed=random_strip_thinking_seed, - n_thinking_variants=n_thinking_variants, - dump_dir=dump_dir, - dump_n_samples=dump_samples, - parse_tool_call_args=parse_tool_call_args, - ) - - # --- Distributed rank-0-only processing --- - rank = int(os.getenv("RANK", "0")) - world_size = int(os.getenv("WORLD_SIZE", "1")) - - if cache_dir is not None and world_size > 1: - done_marker = os.path.join(cache_dir, ".done") - meta_path = os.path.join(cache_dir, ".meta.json") - cache_meta = { - "version": 1, - "path": path, - "tokenizer": getattr(tokenizer, "name_or_path", None), - "process_kwargs": { - k: v - for k, v in process_kwargs.items() - if k not in ("dump_dir", "dump_n_samples") - }, - } - - def _filter_by_max_length(ds): - if max_length is None: - return ds - before = len(ds) - # Length via arrow list offsets: avoids decoding every row to - # Python lists, which for long-context datasets costs minutes of - # startup per rank while (on a validated cache) removing nothing — - # build-time _TokenizeAndMask already filtered with this max_length. - import pyarrow.compute as pc - - # ds.data is the underlying arrow table; a freshly built dataset - # carries an indices mapping (from .filter views) whose row count - # differs. Materialize the view first (no-op for load_from_disk). - if getattr(ds, "_indices", None) is not None: - ds = ds.flatten_indices() - lengths = pc.list_value_length(ds.data.column("input_ids")).to_pylist() - keep = [i for i, n in enumerate(lengths) if n <= max_length] - ds = ds.select(keep) - if len(ds) < before: - logger.info( - f"Rank {rank}: filtered {before - len(ds)} samples " - f"exceeding max_length={max_length}" - ) - if len(ds) == 0: - raise ValueError( - f"processed dataset at {cache_dir} has 0 samples after " - f"max_length={max_length} filtering" - ) - return ds - - def _load_valid_cache(): - if not os.path.exists(meta_path): - raise ValueError(f"cached dataset metadata is missing: {meta_path}") - with open(meta_path) as f: - cached_meta = json.load(f) - if cached_meta != cache_meta: - raise ValueError( - f"cached dataset metadata does not match current SWE settings: " - f"{meta_path}" - ) - dataset = load_from_disk(cache_dir) - if len(dataset) == 0: - raise ValueError(f"cached dataset is empty: {cache_dir}") - return dataset - - def _wait_for_valid_cache(): - start = time.monotonic() - last_error = None - while True: - if os.path.exists(done_marker): - try: - return _load_valid_cache() - except Exception as e: - last_error = e - elapsed = time.monotonic() - start - if elapsed > _RANK0_CACHE_TIMEOUT: - raise TimeoutError( - f"Waited {_RANK0_CACHE_TIMEOUT}s for rank 0 to rebuild " - f"a valid dataset cache at {cache_dir}. Last error: {last_error}" - ) - time.sleep(_RANK0_CACHE_POLL_INTERVAL) - - # Fast path: cache from a previous run (or rank 0 already finished). - if os.path.exists(done_marker): - if rank == 0: - try: - logger.info( - f"Rank {rank}: loading cached processed dataset from {cache_dir}" - ) - dataset = _load_valid_cache() - dataset = _filter_by_max_length(dataset) - logger.info(f"Final dataset: {len(dataset)} samples") - return dataset - except Exception as e: - logger.warning( - "Rank 0: invalid processed dataset cache at %s (%s); " - "rebuilding it.", - cache_dir, - e, - ) - shutil.rmtree(cache_dir, ignore_errors=True) - else: - try: - logger.info( - f"Rank {rank}: loading cached processed dataset from {cache_dir}" - ) - dataset = _load_valid_cache() - dataset = _filter_by_max_length(dataset) - logger.info(f"Final dataset: {len(dataset)} samples") - return dataset - except Exception as e: - logger.warning( - "Rank %d: cached processed dataset at %s is not usable " - "(%s); waiting for rank 0 to rebuild it.", - rank, - cache_dir, - e, - ) - dataset = _wait_for_valid_cache() - dataset = _filter_by_max_length(dataset) - logger.info( - f"Rank {rank}: loaded rebuilt dataset ({len(dataset)} samples)" - ) - return dataset - - if rank == 0: - # Rank 0: do the heavy processing and save for other ranks. - dataset = _process_swe_sft(path, tokenizer, **process_kwargs) - if len(dataset) == 0: - raise RuntimeError( - "SWE SFT preprocessing produced 0 samples; refusing to cache " - "an empty processed_dataset." - ) - shutil.rmtree(cache_dir, ignore_errors=True) - os.makedirs(cache_dir, exist_ok=True) - dataset.save_to_disk(cache_dir) - with open(meta_path, "w") as f: - json.dump(cache_meta, f, sort_keys=True) - # Write marker AFTER save completes so readers see a consistent dir. - with open(done_marker, "w") as f: - f.write(str(len(dataset))) - logger.info( - f"Rank 0: saved processed dataset " - f"({len(dataset)} samples) to {cache_dir}" - ) - dataset = _filter_by_max_length(dataset) - return dataset - else: - # Other ranks: wait for rank 0, then load with meta validation so a - # cache rebuilt for different settings (or mid-rmtree) is never - # silently loaded as this rank's dataset. - logger.info(f"Rank {rank}: waiting for rank 0 to process dataset...") - dataset = _wait_for_valid_cache() - dataset = _filter_by_max_length(dataset) - logger.info(f"Rank {rank}: loaded cached dataset ({len(dataset)} samples)") - return dataset - - # --- Non-distributed or no cache_dir: process in current process --- - return _process_swe_sft(path, tokenizer, **process_kwargs) - - -# ============================================================ -# 7. CLI — ``python -m areal.dataset.swe_sft`` -# ============================================================ - -if __name__ == "__main__": - import argparse - import sys - - from transformers import AutoTokenizer - - parser = argparse.ArgumentParser( - description="Verify SWE SFT pair generation and loss masking.", - ) - parser.add_argument("path", help="Path to SWE trajectory JSONL file") - parser.add_argument( - "--tokenizer", - default="Qwen/Qwen3-8B", - help="HuggingFace tokenizer name or path (default: Qwen/Qwen3-8B)", - ) - parser.add_argument( - "--max-length", - type=int, - default=None, - help="Filter samples exceeding this token length", - ) - parser.add_argument( - "--num-samples", - "-n", - type=int, - default=None, - help="Number of pairs to process. Controls loading, tokenization," - " display, and export. Default: all pairs.", - ) - parser.add_argument( - "--num-proc", - type=int, - default=None, - help=f"Number of parallel workers (default: min(cpu_count, {DATASET_NUM_PROC}))", - ) - parser.add_argument( - "--save-pairs", - "-o", - default=None, - metavar="FILE", - help='Save cleaned pairs to FILE (JSONL, each line: {"messages": [...]}).', - ) - parser.add_argument( - "--pre-split", - action="store_true", - help='Input is already in pair format (each line: {"messages": [...]}).' - " Skip trajectory splitting and error filtering.", - ) - parser.add_argument( - "--no-filter-errors", - action="store_true", - help="Keep pairs whose current segment contains tool results with " - "is_error=True (by default these are discarded).", - ) - parser.add_argument( - "--save-tokenized", - default=None, - metavar="DIR", - help="Save the tokenized dataset to DIR (Arrow format). " - "The saved directory can be used directly as the dataset path " - "during training, skipping all processing.", - ) - parser.add_argument( - "--strip-all-thinking", - action="store_true", - help="Strip ... from ALL assistant turns including " - "the training target. By default only context turns are stripped.", - ) - parser.add_argument( - "--no-tools", - action="store_true", - help="Do not pass tool definitions to apply_chat_template. " - "By default, tools are auto-extracted from the data and rendered " - "in the system prompt (e.g. Qwen3 '# Tools' block).", - ) - parser.add_argument( - "--parse-tool-call-args", - action="store_true", - help="Convert OpenAI JSON-string tool_calls.arguments to dicts " - "before apply_chat_template. Required by GLM-4.x / GLM-5.x " - "templates; leave off for Qwen / Llama / Bailing (which expect " - "the standard string form).", - ) - parser.add_argument( - "--filter-empty-tool-calls", - action="store_true", - help="Discard pairs whose training-target assistant turn has no " - "text content but has tool_calls (silent tool invocations).", - ) - parser.add_argument( - "--filter-bare-text-tool-calls", - action="store_true", - help="Discard pairs whose training-target assistant turn has text " - "content without tags and has tool_calls.", - ) - parser.add_argument( - "--truncate-task-notifications", - action="store_true", - help="Truncate trajectories at the first that " - "follows a pure-text assistant turn. Removes noise from background " - "task completions (e.g. pip install finishing after the model's summary).", - ) - parser.add_argument( - "--max-no-thinking-ratio", - type=float, - default=None, - help="Maximum ratio of non-thinking pairs to thinking pairs. " - "E.g. 1.0 = 1:1 balance, 2.0 = at most 2x non-thinking per " - "thinking pair. Non-thinking pairs are randomly downsampled. " - "Default: no balancing.", - ) - parser.add_argument( - "--split-mode", - choices=["pair", "trajectory"], - default="pair", - help="Sample construction mode. 'pair' (default): split trajectories " - "into progressive pairs. 'trajectory': keep the full trajectory " - "as a single sample with all assistant turns as targets.", - ) - parser.add_argument( - "--random-strip-thinking-prob", - type=float, - default=0.0, - help="Probability of stripping thinking from each target assistant " - "turn. 0.0 = no stripping (default), 1.0 = strip all. " - "Works in both pair mode and trajectory mode.", - ) - parser.add_argument( - "--random-strip-thinking-seed", - type=int, - default=42, - help="Random seed for reproducible thinking stripping decisions (default: 42).", - ) - parser.add_argument( - "--n-thinking-variants", - type=int, - default=1, - help="Number of thinking-pattern variants per trajectory. " - "1 = no augmentation (default). K > 1 = augment each trajectory " - "into K variants: the first preserves all thinking, the rest " - "randomly strip with --random-strip-thinking-prob.", - ) - parser.add_argument( - "--save-trajectories", - default=None, - metavar="FILE", - help="Save preprocessed trajectories to FILE (JSONL, original format) " - "after applying trajectory-level operations (e.g. " - "--truncate-task-notifications) but before pair splitting. " - "Each line preserves the original record structure with the " - "messages field updated.", - ) - parser.add_argument( - "--dump-samples", - default=None, - metavar="DIR", - help="Save sampled pairs to DIR, one file per pair. Each file " - "contains the rendered text and a token-by-token table with " - "token id, decoded text, and loss_mask.", - ) - parser.add_argument( - "--dump-n", - type=int, - default=None, - help="Number of pairs to dump when --dump-samples is set. " - "Default: all pairs. -1 also means all.", - ) - args = parser.parse_args() - - filter_errors = not args.no_filter_errors - strip_all_thinking = args.strip_all_thinking - filter_empty_tool_calls = args.filter_empty_tool_calls - filter_bare_text_tool_calls = args.filter_bare_text_tool_calls - truncate_task_notifications = args.truncate_task_notifications - max_no_thinking_ratio = args.max_no_thinking_ratio - - # --- Fast path: save preprocessed trajectories --- - if args.save_trajectories: - records_in = 0 - records_out = 0 - n_truncated = 0 - with ( - open(args.path, encoding="utf-8") as fin, - open(args.save_trajectories, "w", encoding="utf-8") as fout, - ): - for line in fin: - line = line.strip() - if not line: - continue - record = json.loads(line) - records_in += 1 - - messages, _ = _extract_messages(record, records_in) - - if truncate_task_notifications and messages: - truncated = _truncate_at_task_notification(messages) - if len(truncated) < len(messages): - n_truncated += 1 - _set_messages(record, truncated) - - fout.write(json.dumps(record, ensure_ascii=False) + "\n") - records_out += 1 - - parts = [] - if n_truncated: - parts.append(f"{n_truncated} truncated at task-notification") - op_msg = ", ".join(parts) if parts else "no changes" - print( - f"Saved {records_out}/{records_in} trajectories " - f"to {args.save_trajectories} ({op_msg})" - ) - sys.exit(0) - - # --- Load --- - split_mode = args.split_mode - error_indices_list = None - - if split_mode == "trajectory": - samples, error_indices_list, tools_list = _load_full_trajectories( - args.path, - filter_errors=filter_errors, - filter_empty_tool_calls=filter_empty_tool_calls, - filter_bare_text_tool_calls=filter_bare_text_tool_calls, - truncate_task_notifications=truncate_task_notifications, - random_strip_thinking_prob=args.random_strip_thinking_prob, - random_strip_thinking_seed=args.random_strip_thinking_seed, - n_thinking_variants=args.n_thinking_variants, - ) - label = "trajectories" - elif args.pre_split: - samples, tools_list = _load_presplit_pairs( - args.path, - strip_all_thinking=strip_all_thinking, - random_strip_thinking_prob=args.random_strip_thinking_prob, - random_strip_thinking_seed=args.random_strip_thinking_seed, - n_thinking_variants=args.n_thinking_variants, - ) - label = "pairs" - else: - samples, tools_list = _load_trajectory_pairs( - args.path, - filter_errors=filter_errors, - strip_all_thinking=strip_all_thinking, - filter_empty_tool_calls=filter_empty_tool_calls, - filter_bare_text_tool_calls=filter_bare_text_tool_calls, - truncate_task_notifications=truncate_task_notifications, - max_no_thinking_ratio=max_no_thinking_ratio, - random_strip_thinking_prob=args.random_strip_thinking_prob, - random_strip_thinking_seed=args.random_strip_thinking_seed, - n_thinking_variants=args.n_thinking_variants, - ) - label = "pairs" - - # --- Slice + stats --- - total = len(samples) - if args.num_samples is not None: - samples = samples[: args.num_samples] - tools_list = tools_list[: args.num_samples] if tools_list else tools_list - if error_indices_list is not None: - error_indices_list = error_indices_list[: args.num_samples] - - print(f"Total {label}: {total}") - if args.num_samples is not None: - print(f"Using: {len(samples)}") - - if samples: - lengths = [len(s) for s in samples] - print( - f"Messages/sample: min={min(lengths)}, " - f"max={max(lengths)}, avg={sum(lengths) / len(lengths):.1f}" - ) - if error_indices_list is not None: - n_masked = sum(len(e) for e in error_indices_list) - print(f"Masked segments: {n_masked} (loss=0)") - - # --- Save cleaned samples as JSONL --- - if args.save_pairs: - with open(args.save_pairs, "w", encoding="utf-8") as fout: - err_iter = error_indices_list or [None] * len(samples) - tl_iter = tools_list if tools_list else [None] * len(samples) - for sample, sample_tools, err_idxs in zip(samples, tl_iter, err_iter): - record = {"messages": sample} - if sample_tools is not None: - record["tools"] = sample_tools - if err_idxs: - record["error_indices"] = err_idxs - fout.write(json.dumps(record, ensure_ascii=False) + "\n") - print(f"Wrote {len(samples)} {label} to {args.save_pairs}") - - # --- Tokenize / Dump --- - dump_dir = args.dump_samples if args.dump_samples else None - need_tokenize = args.save_tokenized - - # When --save-tokenized is set, auto-dump 50 samples alongside it - # unless the user explicitly set --dump-samples or --dump-n 0. - if need_tokenize and not dump_dir and args.dump_n != 0: - dump_dir = os.path.join(args.save_tokenized, "dumped_samples") - - if not need_tokenize and not dump_dir: - sys.exit(0) - - tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) - _patch_chat_template_for_training(tok) - if args.dump_n is not None: - dump_n = args.dump_n - elif args.dump_samples: - # Explicit --dump-samples without --dump-n: dump all - dump_n = -1 if args.num_samples is None else args.num_samples - elif need_tokenize: - # Auto-dump with --save-tokenized: default 50 - dump_n = 50 - else: - dump_n = -1 - - # Dump can run independently without full tokenization. - if dump_dir and dump_n != 0: - dump_tools = None if args.no_tools else tools_list - first_tools = None - if dump_tools: - first_tools = next((t for t in dump_tools if t is not None), None) - assistant_pattern = _detect_template_pattern(tok, tools=first_tools) - _dump_samples( - samples, - tok, - assistant_pattern, - dump_tools, - dump_dir, - dump_n, - split_mode=split_mode, - error_indices_list=error_indices_list, - parse_tool_call_args=args.parse_tool_call_args, - ) - - if not need_tokenize: - sys.exit(0) - - ds = _tokenize_samples( - samples, - tools_list, - tok, - split_mode=split_mode, - error_indices_list=error_indices_list, - max_length=args.max_length, - num_proc=args.num_proc, - no_tools=args.no_tools, - parse_tool_call_args=args.parse_tool_call_args, - ) - - print(f"\nTokenized: {len(ds)} samples") - if args.save_tokenized: - ds.save_to_disk(args.save_tokenized) - print(f"Saved tokenized dataset ({len(ds)} samples) to {args.save_tokenized}") diff --git a/areal/dataset/swe_sft/__init__.py b/areal/dataset/swe_sft/__init__.py new file mode 100644 index 0000000000..7a8c2fe688 --- /dev/null +++ b/areal/dataset/swe_sft/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""SWE trajectory preprocessing and SFT dataset loading. + +The public entry point remains :func:`get_swe_sft_dataset`. Implementation +details are split by responsibility to keep the loader maintainable. +""" + +from .pipeline import get_swe_sft_dataset + +__all__ = ["get_swe_sft_dataset"] diff --git a/areal/dataset/swe_sft/__main__.py b/areal/dataset/swe_sft/__main__.py new file mode 100644 index 0000000000..2b114a7db0 --- /dev/null +++ b/areal/dataset/swe_sft/__main__.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Command-line tools for inspecting and preprocessing SWE SFT datasets.""" + +import json +import os + +from .messages import ( + _extract_messages, + _set_messages, + _truncate_at_task_notification, +) +from .pipeline import ( + _load_full_trajectories, + _load_presplit_pairs, + _load_trajectory_pairs, +) +from .tokenization import ( + DATASET_NUM_PROC, + _detect_template_pattern, + _dump_samples, + _patch_chat_template_for_training, + _tokenize_samples, +) + +def main(): + import argparse + import sys + + from transformers import AutoTokenizer + + parser = argparse.ArgumentParser( + description="Verify SWE SFT pair generation and loss masking.", + ) + parser.add_argument("path", help="Path to SWE trajectory JSONL file") + parser.add_argument( + "--tokenizer", + default="Qwen/Qwen3-8B", + help="HuggingFace tokenizer name or path (default: Qwen/Qwen3-8B)", + ) + parser.add_argument( + "--max-length", + type=int, + default=None, + help="Filter samples exceeding this token length", + ) + parser.add_argument( + "--num-samples", + "-n", + type=int, + default=None, + help="Number of pairs to process. Controls loading, tokenization," + " display, and export. Default: all pairs.", + ) + parser.add_argument( + "--num-proc", + type=int, + default=None, + help=f"Number of parallel workers (default: min(cpu_count, {DATASET_NUM_PROC}))", + ) + parser.add_argument( + "--save-pairs", + "-o", + default=None, + metavar="FILE", + help='Save cleaned pairs to FILE (JSONL, each line: {"messages": [...]}).', + ) + parser.add_argument( + "--pre-split", + action="store_true", + help='Input is already in pair format (each line: {"messages": [...]}).' + " Skip trajectory splitting and error filtering.", + ) + parser.add_argument( + "--no-filter-errors", + action="store_true", + help="Keep pairs whose current segment contains tool results with " + "is_error=True (by default these are discarded).", + ) + parser.add_argument( + "--save-tokenized", + default=None, + metavar="DIR", + help="Save the tokenized dataset to DIR (Arrow format). " + "The saved directory can be used directly as the dataset path " + "during training, skipping all processing.", + ) + parser.add_argument( + "--strip-all-thinking", + action="store_true", + help="Strip ... from ALL assistant turns including " + "the training target. By default only context turns are stripped.", + ) + parser.add_argument( + "--no-tools", + action="store_true", + help="Do not pass tool definitions to apply_chat_template. " + "By default, tools are auto-extracted from the data and rendered " + "in the system prompt (e.g. Qwen3 '# Tools' block).", + ) + parser.add_argument( + "--parse-tool-call-args", + action="store_true", + help="Convert OpenAI JSON-string tool_calls.arguments to dicts " + "before apply_chat_template. Required by GLM-4.x / GLM-5.x " + "templates; leave off for Qwen / Llama / Bailing (which expect " + "the standard string form).", + ) + parser.add_argument( + "--filter-empty-tool-calls", + action="store_true", + help="Discard pairs whose training-target assistant turn has no " + "text content but has tool_calls (silent tool invocations).", + ) + parser.add_argument( + "--filter-bare-text-tool-calls", + action="store_true", + help="Discard pairs whose training-target assistant turn has text " + "content without tags and has tool_calls.", + ) + parser.add_argument( + "--truncate-task-notifications", + action="store_true", + help="Truncate trajectories at the first that " + "follows a pure-text assistant turn. Removes noise from background " + "task completions (e.g. pip install finishing after the model's summary).", + ) + parser.add_argument( + "--max-no-thinking-ratio", + type=float, + default=None, + help="Maximum ratio of non-thinking pairs to thinking pairs. " + "E.g. 1.0 = 1:1 balance, 2.0 = at most 2x non-thinking per " + "thinking pair. Non-thinking pairs are randomly downsampled. " + "Default: no balancing.", + ) + parser.add_argument( + "--split-mode", + choices=["pair", "trajectory"], + default="pair", + help="Sample construction mode. 'pair' (default): split trajectories " + "into progressive pairs. 'trajectory': keep the full trajectory " + "as a single sample with all assistant turns as targets.", + ) + parser.add_argument( + "--random-strip-thinking-prob", + type=float, + default=0.0, + help="Probability of stripping thinking from each target assistant " + "turn. 0.0 = no stripping (default), 1.0 = strip all. " + "Works in both pair mode and trajectory mode.", + ) + parser.add_argument( + "--random-strip-thinking-seed", + type=int, + default=42, + help="Random seed for reproducible thinking stripping decisions (default: 42).", + ) + parser.add_argument( + "--n-thinking-variants", + type=int, + default=1, + help="Number of thinking-pattern variants per trajectory. " + "1 = no augmentation (default). K > 1 = augment each trajectory " + "into K variants: the first preserves all thinking, the rest " + "randomly strip with --random-strip-thinking-prob.", + ) + parser.add_argument( + "--save-trajectories", + default=None, + metavar="FILE", + help="Save preprocessed trajectories to FILE (JSONL, original format) " + "after applying trajectory-level operations (e.g. " + "--truncate-task-notifications) but before pair splitting. " + "Each line preserves the original record structure with the " + "messages field updated.", + ) + parser.add_argument( + "--dump-samples", + default=None, + metavar="DIR", + help="Save sampled pairs to DIR, one file per pair. Each file " + "contains the rendered text and a token-by-token table with " + "token id, decoded text, and loss_mask.", + ) + parser.add_argument( + "--dump-n", + type=int, + default=None, + help="Number of pairs to dump when --dump-samples is set. " + "Default: all pairs. -1 also means all.", + ) + args = parser.parse_args() + + filter_errors = not args.no_filter_errors + strip_all_thinking = args.strip_all_thinking + filter_empty_tool_calls = args.filter_empty_tool_calls + filter_bare_text_tool_calls = args.filter_bare_text_tool_calls + truncate_task_notifications = args.truncate_task_notifications + max_no_thinking_ratio = args.max_no_thinking_ratio + + # --- Fast path: save preprocessed trajectories --- + if args.save_trajectories: + records_in = 0 + records_out = 0 + n_truncated = 0 + with ( + open(args.path, encoding="utf-8") as fin, + open(args.save_trajectories, "w", encoding="utf-8") as fout, + ): + for line in fin: + line = line.strip() + if not line: + continue + record = json.loads(line) + records_in += 1 + + messages, _ = _extract_messages(record, records_in) + + if truncate_task_notifications and messages: + truncated = _truncate_at_task_notification(messages) + if len(truncated) < len(messages): + n_truncated += 1 + _set_messages(record, truncated) + + fout.write(json.dumps(record, ensure_ascii=False) + "\n") + records_out += 1 + + parts = [] + if n_truncated: + parts.append(f"{n_truncated} truncated at task-notification") + op_msg = ", ".join(parts) if parts else "no changes" + print( + f"Saved {records_out}/{records_in} trajectories " + f"to {args.save_trajectories} ({op_msg})" + ) + sys.exit(0) + + # --- Load --- + split_mode = args.split_mode + error_indices_list = None + + if split_mode == "trajectory": + samples, error_indices_list, tools_list = _load_full_trajectories( + args.path, + filter_errors=filter_errors, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + random_strip_thinking_prob=args.random_strip_thinking_prob, + random_strip_thinking_seed=args.random_strip_thinking_seed, + n_thinking_variants=args.n_thinking_variants, + ) + label = "trajectories" + elif args.pre_split: + samples, tools_list = _load_presplit_pairs( + args.path, + strip_all_thinking=strip_all_thinking, + random_strip_thinking_prob=args.random_strip_thinking_prob, + random_strip_thinking_seed=args.random_strip_thinking_seed, + n_thinking_variants=args.n_thinking_variants, + ) + label = "pairs" + else: + samples, tools_list = _load_trajectory_pairs( + args.path, + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + max_no_thinking_ratio=max_no_thinking_ratio, + random_strip_thinking_prob=args.random_strip_thinking_prob, + random_strip_thinking_seed=args.random_strip_thinking_seed, + n_thinking_variants=args.n_thinking_variants, + ) + label = "pairs" + + # --- Slice + stats --- + total = len(samples) + if args.num_samples is not None: + samples = samples[: args.num_samples] + tools_list = tools_list[: args.num_samples] if tools_list else tools_list + if error_indices_list is not None: + error_indices_list = error_indices_list[: args.num_samples] + + print(f"Total {label}: {total}") + if args.num_samples is not None: + print(f"Using: {len(samples)}") + + if samples: + lengths = [len(s) for s in samples] + print( + f"Messages/sample: min={min(lengths)}, " + f"max={max(lengths)}, avg={sum(lengths) / len(lengths):.1f}" + ) + if error_indices_list is not None: + n_masked = sum(len(e) for e in error_indices_list) + print(f"Masked segments: {n_masked} (loss=0)") + + # --- Save cleaned samples as JSONL --- + if args.save_pairs: + with open(args.save_pairs, "w", encoding="utf-8") as fout: + err_iter = error_indices_list or [None] * len(samples) + tl_iter = tools_list if tools_list else [None] * len(samples) + for sample, sample_tools, err_idxs in zip(samples, tl_iter, err_iter): + record = {"messages": sample} + if sample_tools is not None: + record["tools"] = sample_tools + if err_idxs: + record["error_indices"] = err_idxs + fout.write(json.dumps(record, ensure_ascii=False) + "\n") + print(f"Wrote {len(samples)} {label} to {args.save_pairs}") + + # --- Tokenize / Dump --- + dump_dir = args.dump_samples if args.dump_samples else None + need_tokenize = args.save_tokenized + + # When --save-tokenized is set, auto-dump 50 samples alongside it + # unless the user explicitly set --dump-samples or --dump-n 0. + if need_tokenize and not dump_dir and args.dump_n != 0: + dump_dir = os.path.join(args.save_tokenized, "dumped_samples") + + if not need_tokenize and not dump_dir: + sys.exit(0) + + tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + _patch_chat_template_for_training(tok) + if args.dump_n is not None: + dump_n = args.dump_n + elif args.dump_samples: + # Explicit --dump-samples without --dump-n: dump all + dump_n = -1 if args.num_samples is None else args.num_samples + elif need_tokenize: + # Auto-dump with --save-tokenized: default 50 + dump_n = 50 + else: + dump_n = -1 + + # Dump can run independently without full tokenization. + if dump_dir and dump_n != 0: + dump_tools = None if args.no_tools else tools_list + first_tools = None + if dump_tools: + first_tools = next((t for t in dump_tools if t is not None), None) + assistant_pattern = _detect_template_pattern(tok, tools=first_tools) + _dump_samples( + samples, + tok, + assistant_pattern, + dump_tools, + dump_dir, + dump_n, + split_mode=split_mode, + error_indices_list=error_indices_list, + parse_tool_call_args=args.parse_tool_call_args, + ) + + if not need_tokenize: + sys.exit(0) + + ds = _tokenize_samples( + samples, + tools_list, + tok, + split_mode=split_mode, + error_indices_list=error_indices_list, + max_length=args.max_length, + num_proc=args.num_proc, + no_tools=args.no_tools, + parse_tool_call_args=args.parse_tool_call_args, + ) + + print(f"\nTokenized: {len(ds)} samples") + if args.save_tokenized: + ds.save_to_disk(args.save_tokenized) + print(f"Saved tokenized dataset ({len(ds)} samples) to {args.save_tokenized}") + + +if __name__ == "__main__": + main() diff --git a/areal/dataset/swe_sft/messages.py b/areal/dataset/swe_sft/messages.py new file mode 100644 index 0000000000..6b24b767e0 --- /dev/null +++ b/areal/dataset/swe_sft/messages.py @@ -0,0 +1,786 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Message normalization, filtering, and trajectory splitting for SWE SFT.""" + +import json +import random +import re + +from areal.utils import logging + +logger = logging.getLogger("SWESFTDataset") + + +def _extract_messages(record, record_idx): + """Extract messages and tools from a parsed JSONL record. + + Handles nested (``conversations`` wrapper) and flat formats. + Warns if multiple conversations are present. + + Returns: + Tuple of ``(messages, record_tools)``. *messages* may be empty. + """ + convs = record.get("conversations", []) + if convs: + if len(convs) > 1: + logger.warning( + "Record %d has %d conversations, using only the last one.", + record_idx, + len(convs), + ) + conv = convs[-1] + return conv.get("messages", []), conv.get("tools") + return record.get("messages", []), record.get("tools") + + +def _set_messages(record, messages): + """Write *messages* back into *record* (inverse of ``_extract_messages``). + + Used by the ``--save-trajectories`` CLI path to update truncated + messages in the original record structure before serialization. + """ + convs = record.get("conversations", []) + if convs: + convs[-1]["messages"] = messages + else: + record["messages"] = messages + + +def _iter_jsonl_records(path): + """Iterate trajectory JSONL records. + + Yields ``(record_idx, messages, record_tools)`` tuples. Handles + nested (``conversations`` wrapper) vs flat format auto-detection + via ``_extract_messages``. Records with empty messages are skipped. + + Warns about multi-user trajectories which break think-tag rendering + in templates with ``ns.last_query_index`` logic (e.g. Bailing). + """ + record_idx = 0 + n_multi_user = 0 + with open(path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + record = json.loads(line) + record_idx += 1 + messages, record_tools = _extract_messages(record, record_idx) + if not messages: + continue + n_user = sum(1 for m in messages if m.get("role") == "user") + if n_user > 1: + n_multi_user += 1 + if n_multi_user <= 3: + logger.warning( + "Record %d has %d user messages. Templates with " + "ns.last_query_index logic (e.g. Bailing) will NOT " + "render for assistant turns before the last " + "user message. Consider filtering!", + record_idx, + n_user, + ) + yield record_idx, messages, record_tools + if n_multi_user > 0: + logger.warning( + "Total %d/%d records have multiple user messages. " + "These may produce no-think training signal.", + n_multi_user, + record_idx, + ) + + +# ============================================================ +# 2. Cleaning — message content transforms +# ============================================================ + +# Match reasoning blocks with any common tag variant: +# ... (Qwen standard) +# ... (Claude) +# The opening and closing tag names need not match exactly — mixed pairs +# like ``...`` (seen in distillation data) are handled. +_THINK_OPEN_RE = re.compile(r"") +_THINK_CLOSE_RE = re.compile(r"") +_THINK_RE = re.compile(r"(.*?)", re.DOTALL) + + +def _normalize_thinking_tags(content): + """Normalise all thinking tag variants to ````/````. + + Distillation data from different models may use ```` (Claude) + vs ```` (Qwen). Non-standard variants are multi-token for the + Qwen tokenizer which breaks think/tool_call boundaries. + """ + if not content: + return content + content = _THINK_OPEN_RE.sub("", content) + content = _THINK_CLOSE_RE.sub("", content) + return content + + +def _extract_thinking(content): + """Strip thinking blocks from *content*. + + Callers must run ``_normalize_thinking_tags`` first so that all + tag variants have been converted to ````/````. + + Returns: + Cleaned content with thinking blocks removed, or the original + content unchanged if no thinking tags are found. + """ + if not content: + return content + cleaned = _THINK_RE.sub("", content).strip() + return cleaned if cleaned != content.strip() else content + + +def _clean_message(msg, strip_thinking=True, ensure_thinking=False): + """Remove non-standard fields before tokenization. + + Keeps only the fields expected by tokenizer chat templates: + role, content, reasoning_content (for assistant), tool_calls + (for assistant), tool_call_id (for tool). + + Handles thinking content in two representations: + + - Inline ``...`` tags in ``content`` + - Separate ``reasoning_content`` field (DeepSeek, Qwen3 API style) + + If both are present, inline tags take priority and + ``reasoning_content`` is dropped with a warning to avoid double + thinking blocks in the rendered template. + + Args: + msg: Raw message dict. + strip_thinking: If True, remove thinking from assistant messages + (both inline ```` tags and ``reasoning_content``). + Used for context turns. If False, preserve thinking as-is + (used for the training-target assistant turn). + ensure_thinking: If True, inject inline ``\n`` + on assistant turns that lack a thinking block (either + inline or in ``reasoning_content``). Requires the + patched Bailing template (via + ``_patch_chat_template_for_training``) which detects + ``had_think_tags`` and preserves empty think blocks. + """ + cleaned = {"role": msg["role"]} + + # Handle content — some assistant messages have content=None when + # they only contain tool_calls. Preserve None so chat templates + # that distinguish None vs "" render correctly. + content = msg.get("content") + # Some APIs (DeepSeek, Qwen3 with enable_thinking) return thinking + # in a separate ``reasoning_content`` field instead of inline + # ```` tags. Handle both representations. + raw_reasoning = msg.get("reasoning_content") if msg["role"] == "assistant" else None + has_thinking = False + if content is not None: + if msg["role"] == "assistant": + content = _normalize_thinking_tags(content) + has_inline_thinking = bool(_THINK_RE.search(content)) + if has_inline_thinking and raw_reasoning and raw_reasoning.strip(): + # Conflict: both reasoning_content and inline tags. + # Keep inline tags (they are already in the content the + # tokenizer will see) and drop reasoning_content to avoid + # double thinking blocks in the rendered template. + if not strip_thinking: + logger.warning( + "Message has both reasoning_content and inline " + " tags. Keeping inline tags, dropping " + "reasoning_content." + ) + raw_reasoning = None + elif not has_inline_thinking and raw_reasoning and raw_reasoning.strip(): + # Convert reasoning_content → inline in content. + # This ensures a single representation that templates + # render identically to the reasoning_content path, while + # being more transparent and debuggable. + if not strip_thinking: + content = ( + f"\n{raw_reasoning.strip(chr(10))}\n" + f"\n\n{content.lstrip(chr(10))}" + ) + has_inline_thinking = True + raw_reasoning = None + has_thinking = has_inline_thinking or bool( + raw_reasoning and raw_reasoning.strip() + ) + if strip_thinking: + content = _extract_thinking(content) + cleaned["content"] = content + elif msg["role"] == "assistant" and msg.get("tool_calls"): + # Assistant with tool_calls but content=None. + if raw_reasoning and raw_reasoning.strip(): + has_thinking = True + if not strip_thinking: + # Convert reasoning_content → inline in content. + cleaned["content"] = ( + f"\n{raw_reasoning.strip(chr(10))}\n" + ) + else: + cleaned["content"] = None + raw_reasoning = None + else: + cleaned["content"] = None + else: + # Non-assistant messages without content: default to empty string. + cleaned["content"] = "" + + # Preserve reasoning_content for target turns only when it was NOT + # already inlined above (i.e. only when raw_reasoning is still set). + if not strip_thinking and raw_reasoning is not None: + cleaned["reasoning_content"] = raw_reasoning + + # For the target assistant turn without a thinking block, inject + # inline ``\n`` so that the (patched) template detects + # think intent via ``had_think_tags`` and renders + # ``\n\n\n\n`` — identical token output to the old + # ``reasoning_content='\n'`` approach. + # + # Requires ``_patch_chat_template_for_training`` to have been called + # on the tokenizer, otherwise the stock Bailing template will extract + # and discard the empty ```` block. + if ensure_thinking and msg["role"] == "assistant" and not has_thinking: + cur_content = cleaned.get("content") + if cur_content is None or cur_content == "": + cleaned["content"] = "\n" + else: + cleaned["content"] = f"\n\n\n{cur_content.lstrip(chr(10))}" + + # Copy tool_calls for assistant messages + if msg["role"] == "assistant" and msg.get("tool_calls"): + cleaned_tool_calls = [] + for tc in msg["tool_calls"]: + cleaned_tc = { + "type": tc.get("type", "function"), + "function": { + "name": tc["function"]["name"], + "arguments": json.dumps(tc["function"]["arguments"]) + if isinstance(tc["function"]["arguments"], dict) + else tc["function"]["arguments"], + }, + } + if "id" in tc: + cleaned_tc["id"] = tc["id"] + cleaned_tool_calls.append(cleaned_tc) + cleaned["tool_calls"] = cleaned_tool_calls + + # Copy tool_call_id for tool messages + if msg["role"] == "tool" and msg.get("tool_call_id"): + cleaned["tool_call_id"] = msg["tool_call_id"] + + return cleaned + + +# ============================================================ +# 3. Filters — keep/discard predicates +# ============================================================ + + +def _segment_has_error(messages, start, end): + """Check if any tool message in ``messages[start:end]`` has ``is_error=True``.""" + for m in messages[start:end]: + if m.get("role") == "tool" and m.get("is_error") is True: + return True + return False + + +def _is_empty_tool_call(msg): + """True if assistant *msg* has no text content and no reasoning but has tool_calls.""" + content = msg.get("content") or "" + if content.strip() or not msg.get("tool_calls"): + return False + # If reasoning_content exists, the model did think — not a silent invocation. + reasoning = msg.get("reasoning_content") + if reasoning and reasoning.strip(): + return False + return True + + +def _is_bare_text_tool_call(msg): + """True if assistant *msg* has text without ```` tags and has tool_calls.""" + content = msg.get("content") or "" + if not content.strip() or not msg.get("tool_calls"): + return False + # If reasoning_content exists, thinking is in a separate field — not bare text. + reasoning = msg.get("reasoning_content") + if reasoning and reasoning.strip(): + return False + normalized = _THINK_OPEN_RE.sub("", content) + normalized = _THINK_CLOSE_RE.sub("", normalized) + match = _THINK_RE.search(normalized) + return not (match and match.group(1).strip()) + + +def _msg_has_thinking(msg): + """True if assistant *msg* has thinking content (inline or reasoning_content).""" + if msg.get("role") != "assistant": + return False + content = msg.get("content") or "" + normalized = _THINK_OPEN_RE.sub("", content) + normalized = _THINK_CLOSE_RE.sub("", normalized) + if _THINK_RE.search(normalized): + return True + rc = msg.get("reasoning_content") or "" + return bool(rc.strip()) + + +def _truncate_at_task_notification(messages): + """Truncate messages when a ```` follows a pure-text assistant. + + Claude Code emits ```` as a user message when a + background task (e.g. ``pip install``) completes. If the model has + already produced a text-only summary (no tool_calls), the notification + and all subsequent messages are noise — the model just replies + "nothing to do". Truncating here removes that noise. + + Only triggers when the pattern is: + assistant (text, no tool_calls) → user () + + Returns: + Truncated message list (or the original list if no truncation needed). + """ + for i, m in enumerate(messages): + if m.get("role") != "user": + continue + if "" not in (m.get("content") or ""): + continue + # Find preceding assistant + prev_asst = None + for j in range(i - 1, -1, -1): + if messages[j].get("role") == "assistant": + prev_asst = messages[j] + break + if prev_asst is None: + continue + content = prev_asst.get("content") or "" + if content.strip() and not prev_asst.get("tool_calls"): + # Truncate: keep everything up to (but not including) this user msg + return messages[:i] + return messages + + +# ============================================================ +# 3b. Balancing — downsample non-thinking pairs +# ============================================================ + + +def _classify_pair(pair): + """Classify a pair by its target assistant turn's content type. + + Returns one of: + ``"thinking"`` — target has actual ```` content or + non-empty ``reasoning_content``. + ``"no_thinking_tool_call"`` — target has no thinking but has + ``tool_calls`` (the dominant category that causes distribution + skew). + ``"pure_text"`` — target has no thinking and no tool_calls + (typically the final summary turn in a trajectory). + """ + target = pair[-1] + if target.get("role") != "assistant": + return "pure_text" + + content = target.get("content") or "" + rc = target.get("reasoning_content") or "" + # Require non-empty think content: pair cleaning runs BEFORE balancing + # and (with ensure_thinking) injects an empty \n into + # every no-think target, so a bare regex hit would classify everything + # as "thinking" and silently disable max_no_thinking_ratio. + _m = _THINK_RE.search(content) + has_thinking = bool(_m and _m.group(1).strip()) or bool(rc.strip()) + + if has_thinking: + return "thinking" + if target.get("tool_calls"): + return "no_thinking_tool_call" + return "pure_text" + + +def _balance_thinking_pairs(pairs, max_no_thinking_ratio, seed=42, tools_list=None): + """Downsample non-thinking **tool-call** pairs to control balance. + + Only ``no_thinking_tool_call`` pairs (no thinking but has tool_calls) + are subject to downsampling. ``thinking`` pairs and ``pure_text`` + pairs (the final summary turn, no thinking and no tool_calls) are + always kept — the latter are critical for the model to learn when + to stop calling tools and give a final answer. + + Args: + pairs: List of progressive SFT pairs. + max_no_thinking_ratio: Maximum ratio of non-thinking tool-call + pairs to thinking pairs. For example, ``1.0`` means at most + 1:1, ``2.0`` means at most 2 non-thinking per 1 thinking pair. + ``None`` disables downsampling. + seed: Random seed for reproducible downsampling. + + Returns: + Balanced list of pairs (order preserved, randomly sampled for + the downsampled category). + """ + if max_no_thinking_ratio is None: + return pairs, tools_list + + thinking = [] + no_think_tc = [] + pure_text = [] + for i, pair in enumerate(pairs): + cat = _classify_pair(pair) + if cat == "thinking": + thinking.append(i) + elif cat == "no_thinking_tool_call": + no_think_tc.append(i) + else: + pure_text.append(i) + + n_think = len(thinking) + n_no_think_tc = len(no_think_tc) + n_pure_text = len(pure_text) + + if n_think == 0: + logger.warning( + "No thinking pairs found; skipping balance " + "(all %d pairs have empty thinking).", + n_no_think_tc + n_pure_text, + ) + return pairs, tools_list + + max_no_think_tc = int(n_think * max_no_thinking_ratio) + if n_no_think_tc <= max_no_think_tc: + logger.info( + "Thinking balance OK: %d thinking + %d no-think-tc + %d pure-text " + "(ratio %.1f <= %.1f), no downsampling needed.", + n_think, + n_no_think_tc, + n_pure_text, + n_no_think_tc / n_think, + max_no_thinking_ratio, + ) + return pairs, tools_list + + rng = random.Random(seed) + sampled_tc = set(rng.sample(no_think_tc, max_no_think_tc)) + keep_indices = sorted(set(thinking) | sampled_tc | set(pure_text)) + balanced = [pairs[i] for i in keep_indices] + balanced_tools = ( + [tools_list[i] for i in keep_indices] if tools_list is not None else None + ) + + logger.info( + "Balanced thinking pairs: %d thinking + %d no-think-tc " + "(downsampled from %d, ratio %.1f → %.1f) + %d pure-text (kept all).", + n_think, + max_no_think_tc, + n_no_think_tc, + n_no_think_tc / n_think, + max_no_thinking_ratio, + n_pure_text, + ) + return balanced, balanced_tools + + +# ============================================================ +# 3c. Thinking augmentation stats +# ============================================================ + + +def _log_thinking_augmentation_stats( + n_variants, + prob, + n_total_trajs, + thinking_turns_per_traj, + total_asst_turns_per_traj, + patterns_per_traj, +): + """Log adaptive-thinking augmentation quality metrics. + + Called after the augmentation loop in loaders to report how well + the ``n_thinking_variants`` / ``random_strip_thinking_prob`` settings + produce diverse thinking-pattern variants. + + Args: + n_variants: ``n_thinking_variants`` setting (K). + prob: ``random_strip_thinking_prob`` setting. + n_total_trajs: Total number of source trajectories processed. + thinking_turns_per_traj: List of N_thinking per source trajectory. + total_asst_turns_per_traj: List of N_total_asst per source trajectory. + patterns_per_traj: List of ``set[frozenset]`` — the unique strip + patterns generated for each source trajectory (including the + empty frozenset for the original unstripped variant). + """ + n_eligible = sum(1 for n in thinking_turns_per_traj if n > 0) + total_thinking = sum(thinking_turns_per_traj) + total_asst = sum(total_asst_turns_per_traj) + + # 1. Thinking Turn Coverage + avg_thinking = total_thinking / max(n_total_trajs, 1) + thinking_ratio = total_thinking / max(total_asst, 1) + + # 2. Pattern Diversity + diversity_ratios = [] + for n_think, patterns in zip(thinking_turns_per_traj, patterns_per_traj): + if n_think == 0: + continue + theoretical_max = min(n_variants, 2**n_think) + actual_unique = len(patterns) + diversity_ratios.append(actual_unique / theoretical_max) + avg_diversity = sum(diversity_ratios) / max(len(diversity_ratios), 1) + + # 3. Augmentation Efficiency + n_non_trivial = 0 + for patterns in patterns_per_traj: + # Count variants that differ from the original (non-empty strip set) + n_non_trivial += sum(1 for p in patterns if p) + expected_aug = (n_variants - 1) * max(n_eligible, 1) + efficiency = n_non_trivial / max(expected_aug, 1) + + # 4. Total sample count + n_total_samples = sum(len(p) for p in patterns_per_traj) + + logger.info( + f"Thinking augmentation stats (K={n_variants}, p={prob:.2f}):\n" + f" Source trajectories: {n_total_trajs} " + f"({n_eligible} with thinking turns)\n" + f" Thinking coverage: {avg_thinking:.1f} thinking turns/traj, " + f"{thinking_ratio:.1%} of all assistant turns\n" + f" Pattern diversity: {avg_diversity:.2f} " + f"(1.0 = all variants unique)\n" + f" Augmentation efficiency: {efficiency:.2f} " + f"({n_non_trivial}/{expected_aug} non-trivial variants)\n" + f" Total samples after augmentation: {n_total_samples}" + ) + + +# ============================================================ +# 4. Splitting — trajectory → progressive pairs +# ============================================================ + + +def _find_segments(messages): + """Find assistant+tools segment boundaries. + + Returns: + List of ``(assistant_start_idx, segment_end_idx)`` tuples. + """ + segments = [] + i = 0 + while i < len(messages): + if messages[i].get("role") == "assistant": + j = i + 1 + while j < len(messages) and messages[j].get("role") == "tool": + j += 1 + segments.append((i, j)) + i = j + else: + i += 1 + return segments + + +def _split_and_filter( + messages, + filter_errors=True, + strip_all_thinking=False, + filter_empty_tool_calls=False, + filter_bare_text_tool_calls=False, + random_strip_thinking_prob=0.0, + rng=None, +): + """Split trajectory into progressive pairs and optionally filter. + + By default, thinking (``...``) is stripped from context + assistant turns only; the last assistant turn (training target) keeps + its content unchanged. Set *strip_all_thinking* to strip from every + assistant turn including the target. + + When *random_strip_thinking_prob* > 0, each target assistant turn that + has thinking content is independently stripped with that probability. + Stripped turns use the context-cleaned version (thinking fully removed, + no empty ```` injected). + + Args: + messages: Raw trajectory messages. + filter_errors: If True (default), discard pairs whose current segment + contains a tool result with ``is_error=True``. Set to False to + keep all pairs regardless of tool errors. + strip_all_thinking: If True, strip ```` blocks from every + assistant turn including the training target. + filter_empty_tool_calls: If True, discard pairs whose training-target + assistant turn has no text content but has tool_calls. + filter_bare_text_tool_calls: If True, discard pairs whose + training-target assistant turn has text content without + ```` tags and has tool_calls. + random_strip_thinking_prob: Probability of stripping thinking + from each target assistant turn. 0.0 = no stripping. + rng: ``random.Random`` instance for reproducible sampling. + + Returns: + Tuple of ``(pairs, n_filtered_errors, n_filtered_empty_tc, + n_filtered_bare_tc, n_stripped)``. + """ + segments = _find_segments(messages) + if not segments: + return [], 0, 0, 0, 0 + + pairs = [] + n_filtered_errors = 0 + n_filtered_empty_tc = 0 + n_filtered_bare_tc = 0 + n_stripped = 0 + + # Pre-clean all messages in context mode (thinking stripped). + # This avoids re-cleaning the same message for every progressive pair + # (O(N+K) instead of O(N*K) where K = number of segments). + context_cleaned = [_clean_message(m, strip_thinking=True) for m in messages] + + # For target assistant turns, clean with thinking preserved (unless + # strip_all_thinking is set, in which case context_cleaned is reusable). + # When stripping is active (augmented variant), use ensure_thinking=False + # so empty-thinking turns don't get \n injected. + stripping_active = random_strip_thinking_prob > 0.0 and rng is not None + target_ensure = not stripping_active + target_cleaned = {} + if not strip_all_thinking: + for asst_start, _ in segments: + target_cleaned[asst_start] = _clean_message( + messages[asst_start], + strip_thinking=False, + ensure_thinking=target_ensure, + ) + + for asst_start, seg_end in segments: + # Check if current segment has any tool errors + if filter_errors and _segment_has_error(messages, asst_start, seg_end): + n_filtered_errors += 1 + continue + + # Content-type filters operate on the raw assistant message. + asst_msg = messages[asst_start] + if filter_empty_tool_calls and _is_empty_tool_call(asst_msg): + n_filtered_empty_tc += 1 + continue + if filter_bare_text_tool_calls and _is_bare_text_tool_call(asst_msg): + n_filtered_bare_tc += 1 + continue + + # Build pair: include context up to the target assistant turn, + # truncating tool responses that follow it. This ensures the + # target assistant is always the *last* message so that chat + # templates with ``loop.last``-dependent rendering (e.g. Qwen3 + # ```` injection) behave consistently. The tool responses + # would have loss_mask=0 anyway and only add noise. + pair = list(context_cleaned[: asst_start + 1]) + if not strip_all_thinking: + # Randomly strip: leave context_cleaned version (thinking + # already removed) instead of replacing with target_cleaned. + should_strip = ( + rng is not None + and _msg_has_thinking(messages[asst_start]) + and rng.random() < random_strip_thinking_prob + ) + if not should_strip: + pair[asst_start] = target_cleaned[asst_start] + else: + n_stripped += 1 + pairs.append(pair) + + return pairs, n_filtered_errors, n_filtered_empty_tc, n_filtered_bare_tc, n_stripped + + +def _prepare_trajectory( + messages, + filter_errors=True, + filter_empty_tool_calls=False, + filter_bare_text_tool_calls=False, + random_strip_thinking_prob=0.0, + rng=None, +): + """Prepare a full trajectory for trajectory-level training. + + Cleans all messages preserving thinking for every assistant turn + (``strip_thinking=False``, ``ensure_thinking=True``). Identifies + which assistant segments should be masked (``loss_mask=0``) based + on error tool responses, empty tool calls, or bare-text tool calls. + + When *random_strip_thinking_prob* > 0, each assistant turn that has + thinking content is independently stripped with that probability. + Stripped turns have their ```` blocks and ``reasoning_content`` + completely removed (no empty ```` injected). + + Args: + messages: Raw trajectory messages. + filter_errors: If True (default), mask segments with error tool + responses. + filter_empty_tool_calls: If True, mask segments whose assistant + turn has no text content but has tool_calls. + filter_bare_text_tool_calls: If True, mask segments whose + assistant turn has text without ```` tags and has + tool_calls. + random_strip_thinking_prob: Probability of stripping thinking + from each assistant turn that has thinking content. + 0.0 (default) = no stripping, 1.0 = strip all. + rng: ``random.Random`` instance for reproducible sampling. + + Returns: + Tuple of ``(cleaned_messages, masked_segment_indices, + n_error, n_empty_tc, n_bare_tc, stripped_pattern)`` or ``None`` + if the trajectory has no assistant turns. *stripped_pattern* is + a ``frozenset`` of message indices whose thinking was stripped + (empty if no stripping occurred). + """ + segments = _find_segments(messages) + if not segments: + return None + + masked_indices = set() + n_error = 0 + n_empty_tc = 0 + n_bare_tc = 0 + for idx, (asst_start, seg_end) in enumerate(segments): + if filter_errors and _segment_has_error(messages, asst_start, seg_end): + masked_indices.add(idx) + n_error += 1 + continue + asst_msg = messages[asst_start] + if filter_empty_tool_calls and _is_empty_tool_call(asst_msg): + masked_indices.add(idx) + n_empty_tc += 1 + continue + if filter_bare_text_tool_calls and _is_bare_text_tool_call(asst_msg): + masked_indices.add(idx) + n_bare_tc += 1 + + # Determine which assistant turns to randomly strip thinking from. + strip_thinking_indices = set() + stripping_active = random_strip_thinking_prob > 0.0 and rng is not None + if stripping_active: + for asst_start, _seg_end in segments: + if _msg_has_thinking(messages[asst_start]): + if rng.random() < random_strip_thinking_prob: + strip_thinking_indices.add(asst_start) + + # When stripping is active (augmented variant), use ensure_thinking=False + # for ALL turns so that empty-thinking turns don't get \n + # injected. Only real thinking content is preserved. + # When stripping is inactive (variant 0 or no augmentation), keep + # ensure_thinking=True to match the standard training format. + default_ensure = not stripping_active + + cleaned = [] + for i, m in enumerate(messages): + if i in strip_thinking_indices: + cleaned.append( + _clean_message(m, strip_thinking=True, ensure_thinking=False) + ) + else: + cleaned.append( + _clean_message(m, strip_thinking=False, ensure_thinking=default_ensure) + ) + + return ( + cleaned, + sorted(masked_indices), + n_error, + n_empty_tc, + n_bare_tc, + frozenset(strip_thinking_indices), + ) diff --git a/areal/dataset/swe_sft/pipeline.py b/areal/dataset/swe_sft/pipeline.py new file mode 100644 index 0000000000..144dd780f5 --- /dev/null +++ b/areal/dataset/swe_sft/pipeline.py @@ -0,0 +1,1001 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""SWE SFT loading, processing, distributed caching, and public dataset API.""" + +import json +import os +import random +import shutil +import time + +from datasets import Dataset + +from areal.utils import logging + +from .messages import ( + _balance_thinking_pairs, + _clean_message, + _find_segments, + _iter_jsonl_records, + _log_thinking_augmentation_stats, + _msg_has_thinking, + _prepare_trajectory, + _split_and_filter, + _truncate_at_task_notification, +) +from .tokenization import ( + DATASET_NUM_PROC, + _TokenizeAndMask, + _detect_template_pattern, + _dump_samples, + _patch_chat_template_for_training, +) + +logger = logging.getLogger("SWESFTDataset") + +_RANK0_CACHE_TIMEOUT = 36000 +_RANK0_CACHE_POLL_INTERVAL = 5 + + +def _load_trajectory_pairs( + path: str, + filter_errors: bool = True, + strip_all_thinking: bool = False, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + max_no_thinking_ratio: float | None = None, + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, +): + """Load trajectory JSONL and split into progressive pairs. + + When *n_thinking_variants* > 1, each trajectory is split K times: + variant 0 preserves all thinking, variants 1~K-1 randomly strip. + + Supports nested (``conversations`` wrapper) and flat JSONL formats + (auto-detected per record via ``_iter_jsonl_records``). + + Returns: + Tuple of ``(all_pairs, tools)`` where *tools* is ``None`` when no + tool definitions are found. + """ + all_pairs = [] + all_tools = [] + records_in = 0 + total_filtered_errors = 0 + total_filtered_empty_tc = 0 + total_filtered_bare_tc = 0 + total_truncated = 0 + total_stripped_thinking = 0 + + augment = n_thinking_variants > 1 + rng = ( + random.Random(random_strip_thinking_seed) + if random_strip_thinking_prob > 0.0 + else None + ) + + if augment and random_strip_thinking_prob <= 0.0: + logger.warning( + "n_thinking_variants=%d but random_strip_thinking_prob=0; " + "all variants will be identical.", + n_thinking_variants, + ) + + # Stats collectors for augmentation logging. + thinking_turns_per_traj = [] + total_asst_turns_per_traj = [] + patterns_per_traj = [] + + for record_idx, messages, record_tools in _iter_jsonl_records(path): + records_in = record_idx + + if truncate_task_notifications: + truncated = _truncate_at_task_notification(messages) + if len(truncated) < len(messages): + total_truncated += 1 + messages = truncated + + shared_kwargs = dict( + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + ) + + if augment: + # Variant 0: preserve all thinking. + pairs_orig, n_err, n_empty_tc, n_bare_tc, _ = _split_and_filter( + messages, **shared_kwargs, random_strip_thinking_prob=0.0, rng=None + ) + total_filtered_errors += n_err + total_filtered_empty_tc += n_empty_tc + total_filtered_bare_tc += n_bare_tc + all_pairs.extend(pairs_orig) + all_tools.extend([record_tools] * len(pairs_orig)) + # Collect stats. + segments = _find_segments(messages) + n_think = sum(1 for s, _ in segments if _msg_has_thinking(messages[s])) + n_asst = len(segments) + thinking_turns_per_traj.append(n_think) + total_asst_turns_per_traj.append(n_asst) + + # Variants 1 ~ K-1: random strip. + variant_patterns = {frozenset()} # original = no strip + for _k in range(n_thinking_variants - 1): + pairs_aug, _, _, _, n_stripped = _split_and_filter( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + total_stripped_thinking += n_stripped + all_pairs.extend(pairs_aug) + all_tools.extend([record_tools] * len(pairs_aug)) + # Approximate pattern: record which pairs had their target stripped. + # For stats, use the count as a proxy since _split_and_filter + # doesn't return per-pair strip info. + variant_patterns.add(frozenset([n_stripped])) + patterns_per_traj.append(variant_patterns) + else: + # Single variant (original behavior). + pairs, n_err, n_empty_tc, n_bare_tc, n_stripped = _split_and_filter( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + total_filtered_errors += n_err + total_filtered_empty_tc += n_empty_tc + total_filtered_bare_tc += n_bare_tc + total_stripped_thinking += n_stripped + all_pairs.extend(pairs) + all_tools.extend([record_tools] * len(pairs)) + + # Log extracted tools summary. + n_with_tools = sum(1 for t in all_tools if t is not None) + if n_with_tools > 0: + all_tool_names = set() + for t_list in all_tools: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info( + f"Extracted tools from {n_with_tools}/{len(all_tools)} pairs: " + f"{sorted(all_tool_names)}" + ) + + filter_parts = [] + if total_truncated: + filter_parts.append( + f"{total_truncated} trajectories truncated at task-notification" + ) + if total_filtered_errors: + filter_parts.append(f"{total_filtered_errors} with tool errors") + if total_filtered_empty_tc: + filter_parts.append(f"{total_filtered_empty_tc} empty-content tool calls") + if total_filtered_bare_tc: + filter_parts.append(f"{total_filtered_bare_tc} bare-text tool calls") + if total_stripped_thinking: + filter_parts.append(f"{total_stripped_thinking} thinking blocks stripped") + filter_msg = ", ".join(filter_parts) if filter_parts else "none" + + logger.info( + f"Loaded {records_in} trajectories, " + f"generated {len(all_pairs)} pairs " + f"(filtered: {filter_msg})" + ) + + if augment and patterns_per_traj: + _log_thinking_augmentation_stats( + n_thinking_variants, + random_strip_thinking_prob, + records_in, + thinking_turns_per_traj, + total_asst_turns_per_traj, + patterns_per_traj, + ) + + # Balance thinking / no-thinking pair ratio. + all_pairs, all_tools = _balance_thinking_pairs( + all_pairs, max_no_thinking_ratio, tools_list=all_tools + ) + + return all_pairs, all_tools + + +def _load_presplit_pairs( + path: str, + strip_all_thinking: bool = False, + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, +): + """Load pre-split pair JSONL where each line is ``{"messages": [...]}``. + + Messages are cleaned but no splitting or error-filtering is performed. + By default, thinking is stripped from context assistant turns but + preserved for the last assistant turn (the training target). Set + *strip_all_thinking* to strip from every assistant turn. + + When *n_thinking_variants* > 1, each pair is augmented: variant 0 + preserves thinking, variants 1~K-1 randomly strip the target turn. + + Also extracts per-record ``tools`` definitions so that each pair + carries its own tools, same as ``_load_trajectory_pairs``. + + Returns: + Tuple of ``(all_pairs, all_tools)`` where *all_tools* is a + parallel list of per-sample tool definitions (may be ``None``). + """ + all_pairs = [] + all_tools = [] + n_stripped = 0 + augment = n_thinking_variants > 1 + + rng = ( + random.Random(random_strip_thinking_seed) + if random_strip_thinking_prob > 0.0 + else None + ) + + def _build_pair(messages, last_asst, strip_target): + pair = [] + for idx, m in enumerate(messages): + is_target = m.get("role") == "assistant" and idx == last_asst + strip = strip_all_thinking or not is_target or strip_target + pair.append(_clean_message(m, strip_thinking=strip)) + return pair + + with open(path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + record = json.loads(line) + messages = record.get("messages", []) + if not messages: + continue + + record_tools = record.get("tools") + + # Find the last assistant index so we can preserve its thinking. + last_asst = None + for i, m in enumerate(messages): + if m.get("role") == "assistant": + last_asst = i + + has_thinking = ( + last_asst is not None + and not strip_all_thinking + and _msg_has_thinking(messages[last_asst]) + ) + + if augment: + # Variant 0: preserve all thinking. + all_pairs.append(_build_pair(messages, last_asst, strip_target=False)) + all_tools.append(record_tools) + + # Variants 1 ~ K-1: random strip. + for _k in range(n_thinking_variants - 1): + do_strip = ( + has_thinking + and rng is not None + and rng.random() < random_strip_thinking_prob + ) + if do_strip: + n_stripped += 1 + all_pairs.append( + _build_pair(messages, last_asst, strip_target=do_strip) + ) + all_tools.append(record_tools) + else: + # Single variant (original behavior). + strip_target = ( + has_thinking + and rng is not None + and rng.random() < random_strip_thinking_prob + ) + if strip_target: + n_stripped += 1 + all_pairs.append( + _build_pair(messages, last_asst, strip_target=strip_target) + ) + all_tools.append(record_tools) + + # Log extracted tools summary. + n_with_tools = sum(1 for t in all_tools if t is not None) + if n_with_tools > 0: + all_tool_names = set() + for t_list in all_tools: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info( + f"Extracted tools from {n_with_tools}/{len(all_tools)} pairs: " + f"{sorted(all_tool_names)}" + ) + + strip_msg = f", {n_stripped} thinking blocks stripped" if n_stripped else "" + logger.info(f"Loaded {len(all_pairs)} pre-split pairs from {path}{strip_msg}") + return all_pairs, all_tools + + +def _load_full_trajectories( + path: str, + filter_errors: bool = True, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, +): + """Load trajectory JSONL for trajectory-level training. + + Each trajectory becomes a single training sample with all assistant + turns as targets (``loss_mask=1``). When *filter_errors* is True, + assistant segments with error tool responses are identified so + tokenization can mask them (``loss_mask=0``) instead of discarding + the entire trajectory. + + When *n_thinking_variants* > 1, each trajectory is augmented into + K variants: the first preserves all thinking, the remaining K-1 + randomly strip thinking turns with *random_strip_thinking_prob*. + + Supports nested (``conversations`` wrapper) and flat JSONL formats + (auto-detected per record via ``_iter_jsonl_records``). + + Returns: + Tuple of ``(trajectories, error_indices_list, all_tools)`` where + *trajectories* is a list of cleaned message lists, + *error_indices_list* is a list of error segment index lists, + and *all_tools* is a parallel list of per-sample tool definitions. + """ + trajectories = [] + error_indices_list = [] + all_tools = [] + records_in = 0 + total_truncated = 0 + total_masked_errors = 0 + total_masked_empty_tc = 0 + total_masked_bare_tc = 0 + total_stripped_thinking = 0 + + augment = n_thinking_variants > 1 + rng = ( + random.Random(random_strip_thinking_seed) + if random_strip_thinking_prob > 0.0 + else None + ) + + if augment and random_strip_thinking_prob <= 0.0: + logger.warning( + "n_thinking_variants=%d but random_strip_thinking_prob=0; " + "all variants will be identical.", + n_thinking_variants, + ) + + # Stats collectors for augmentation logging. + thinking_turns_per_traj = [] + total_asst_turns_per_traj = [] + patterns_per_traj = [] + + for record_idx, messages, record_tools in _iter_jsonl_records(path): + records_in = record_idx + + if truncate_task_notifications: + truncated = _truncate_at_task_notification(messages) + if len(truncated) < len(messages): + total_truncated += 1 + messages = truncated + + shared_kwargs = dict( + filter_errors=filter_errors, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + ) + + if augment: + # Variant 0: preserve all thinking (no stripping). + result_orig = _prepare_trajectory( + messages, **shared_kwargs, random_strip_thinking_prob=0.0, rng=None + ) + if result_orig is None: + continue + cleaned_orig, masked_idxs, n_err, n_empty_tc, n_bare_tc, _ = result_orig + trajectories.append(cleaned_orig) + error_indices_list.append(masked_idxs) + all_tools.append(record_tools) + total_masked_errors += n_err + total_masked_empty_tc += n_empty_tc + total_masked_bare_tc += n_bare_tc + + # Collect stats: count thinking turns in this trajectory. + segments = _find_segments(messages) + n_think = sum(1 for s, _ in segments if _msg_has_thinking(messages[s])) + n_asst = len(segments) + thinking_turns_per_traj.append(n_think) + total_asst_turns_per_traj.append(n_asst) + + # Variants 1 ~ K-1: random strip thinking. + variant_patterns = {frozenset()} # original = empty pattern + for _k in range(n_thinking_variants - 1): + result_aug = _prepare_trajectory( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + if result_aug is None: + continue + cleaned_aug, _, _, _, _, strip_pattern = result_aug + trajectories.append(cleaned_aug) + error_indices_list.append(masked_idxs) # reuse + all_tools.append(record_tools) + total_stripped_thinking += len(strip_pattern) + variant_patterns.add(strip_pattern) + patterns_per_traj.append(variant_patterns) + else: + # Single variant (original behavior). + result = _prepare_trajectory( + messages, + **shared_kwargs, + random_strip_thinking_prob=random_strip_thinking_prob, + rng=rng, + ) + if result is None: + continue + cleaned, masked_idxs, n_err, n_empty_tc, n_bare_tc, strip_pattern = result + trajectories.append(cleaned) + error_indices_list.append(masked_idxs) + all_tools.append(record_tools) + total_masked_errors += n_err + total_masked_empty_tc += n_empty_tc + total_masked_bare_tc += n_bare_tc + total_stripped_thinking += len(strip_pattern) + + # Log extracted tools summary. + n_with_tools = sum(1 for t in all_tools if t is not None) + if n_with_tools > 0: + all_tool_names = set() + for t_list in all_tools: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info( + f"Extracted tools from {n_with_tools}/{len(all_tools)} " + f"trajectories: {sorted(all_tool_names)}" + ) + + parts = [] + if total_truncated: + parts.append(f"{total_truncated} trajectories truncated at task-notification") + if total_masked_errors: + parts.append(f"{total_masked_errors} with tool errors") + if total_masked_empty_tc: + parts.append(f"{total_masked_empty_tc} empty-content tool calls") + if total_masked_bare_tc: + parts.append(f"{total_masked_bare_tc} bare-text tool calls") + if total_stripped_thinking: + parts.append(f"{total_stripped_thinking} thinking blocks stripped") + mask_msg = ", ".join(parts) if parts else "none" + + logger.info( + f"Loaded {records_in} trajectories, " + f"kept {len(trajectories)} for training " + f"(masked: {mask_msg})" + ) + + if augment and patterns_per_traj: + _log_thinking_augmentation_stats( + n_thinking_variants, + random_strip_thinking_prob, + records_in, + thinking_turns_per_traj, + total_asst_turns_per_traj, + patterns_per_traj, + ) + + return trajectories, error_indices_list, all_tools + + +def _tokenize_samples( + messages_list, + tools_list, + tokenizer, + *, + split_mode: str = "pair", + error_indices_list: list | None = None, + max_length: int | None = None, + num_proc: int | None = None, + no_tools: bool = False, + dump_dir: str | None = None, + dump_n_samples: int = 0, + parse_tool_call_args: bool = False, +): + """Tokenize message lists into a training-ready Dataset. + + Works for both progressive pairs (``split_mode="pair"``) and + full trajectories (``split_mode="trajectory"``). + + In pair mode, only the last assistant turn per sample gets + ``loss_mask=1``. In trajectory mode, all assistant turns get + ``loss_mask=1`` except those at error segment indices. + + Args: + tools_list: Per-sample tool definitions (parallel to + *messages_list*). Each element is either ``None`` or a + list of tool dicts. + """ + if num_proc is None: + num_proc = max(1, min(os.cpu_count() or 1, DATASET_NUM_PROC)) + + # Find representative tools for template detection. + first_tools = None + if tools_list: + first_tools = next((t for t in tools_list if t is not None), None) + + if no_tools: + tools_list = None + first_tools = None + logger.info("Tool definitions disabled (no_tools=True)") + elif first_tools is not None: + all_tool_names = set() + for t_list in tools_list: + if t_list is not None: + for t in t_list: + all_tool_names.add(t.get("function", {}).get("name", "?")) + logger.info(f"Using tools for chat template: {sorted(all_tool_names)}") + + if not messages_list: + raise ValueError("No valid samples to tokenize") + + # Build dataset columns. + data = {"messages": messages_list} + # Serialize per-sample tools as JSON strings for the Dataset column. + data["tools_json"] = ( + [json.dumps(t) if t else "" for t in tools_list] + if tools_list + else [""] * len(messages_list) + ) + remove_cols = ["messages", "tools_json"] + if split_mode == "trajectory": + data["error_indices"] = error_indices_list or [[] for _ in messages_list] + remove_cols.append("error_indices") + + dataset = Dataset.from_dict(data) + _patch_chat_template_for_training(tokenizer) + assistant_pattern = _detect_template_pattern(tokenizer, tools=first_tools) + + # Dump samples for inspection before the heavy map() pass. + if dump_dir and dump_n_samples != 0: + _dump_samples( + messages_list, + tokenizer, + assistant_pattern, + tools_list, + dump_dir, + dump_n_samples, + split_mode=split_mode, + error_indices_list=error_indices_list, + parse_tool_call_args=parse_tool_call_args, + ) + + process_fn = _TokenizeAndMask( + tokenizer, + assistant_pattern, + max_length=max_length, + split_mode=split_mode, + parse_tool_call_args=parse_tool_call_args, + ) + + dataset = dataset.map(process_fn, num_proc=num_proc).remove_columns(remove_cols) + + # Single filter pass: removes both apply_chat_template-failure empties and + # overlength samples (which _TokenizeAndMask also marks as empty). + before_filter = len(dataset) + dataset = dataset.filter(lambda x: len(x["input_ids"]) > 0, num_proc=num_proc) + n_filtered = before_filter - len(dataset) + if n_filtered > 0: + logger.info( + f"Filtered {n_filtered} samples " + f"(empty from template failures or exceeding max_length={max_length})" + ) + + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + + +def _process_swe_sft( + path: str, + tokenizer, + *, + max_length: int | None = None, + num_proc: int | None = None, + pre_split: bool = False, + filter_errors: bool = True, + strip_all_thinking: bool = False, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + no_tools: bool = False, + max_no_thinking_ratio: float | None = None, + split_mode: str = "pair", + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, + dump_dir: str | None = None, + dump_n_samples: int = 0, + parse_tool_call_args: bool = False, +): + """Load JSONL, split into pairs, tokenize, and filter. + + Combines file loading with ``_tokenize_samples`` so that the rank-0-only + path and the single-process path share the same logic. + + When *split_mode* is ``"trajectory"``, the full trajectory is kept as a + single training sample with all assistant turns as targets. + """ + error_indices_list = None + + if split_mode == "trajectory": + messages_list, error_indices_list, tools_list = _load_full_trajectories( + path, + filter_errors=filter_errors, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + ) + elif pre_split: + messages_list, tools_list = _load_presplit_pairs( + path, + strip_all_thinking=strip_all_thinking, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + ) + else: + messages_list, tools_list = _load_trajectory_pairs( + path, + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + max_no_thinking_ratio=max_no_thinking_ratio, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + ) + + return _tokenize_samples( + messages_list, + tools_list, + tokenizer, + split_mode=split_mode, + error_indices_list=error_indices_list, + max_length=max_length, + num_proc=num_proc, + no_tools=no_tools, + dump_dir=dump_dir, + dump_n_samples=dump_n_samples, + parse_tool_call_args=parse_tool_call_args, + ) + + +def get_swe_sft_dataset( + path: str, + split: str | None = None, + tokenizer=None, + max_length: int | None = None, + num_proc: int | None = None, + pre_split: bool = False, + filter_errors: bool = True, + strip_all_thinking: bool = False, + filter_empty_tool_calls: bool = False, + filter_bare_text_tool_calls: bool = False, + truncate_task_notifications: bool = False, + no_tools: bool = False, + skip_pretokenized_filter: bool = False, + max_no_thinking_ratio: float | None = None, + split_mode: str = "pair", + random_strip_thinking_prob: float = 0.0, + random_strip_thinking_seed: int = 42, + n_thinking_variants: int = 1, + cache_dir: str | None = None, + dump_dir: str | None = None, + dump_samples: int = 0, + parse_tool_call_args: bool = False, +): + """Load SWE trajectory data and convert to SFT training pairs. + + By default, tool definitions are auto-extracted from the training data's + ``conversations[].tools`` field and passed to ``apply_chat_template`` + so that the tokenizer renders tool definitions in the system prompt + (e.g. Qwen3 ``# Tools`` block), matching the eval-time format. + Set *no_tools* to skip this and render without tool definitions. + + When *split_mode* is ``"trajectory"``, the full trajectory is kept as a + single training sample with all assistant turns as targets + (``loss_mask=1``). Error segments are masked (``loss_mask=0``) + when *filter_errors* is True, instead of being discarded. + Thinking is preserved by default but can be randomly stripped + per-turn via *random_strip_thinking_prob* (both modes). + + In distributed (SPMD) mode, only rank 0 performs the heavy processing + (JSONL loading, pair splitting, tokenization) and saves the result as + an Arrow dataset to *cache_dir*. Other ranks wait for rank 0 to + finish and then load the cached dataset directly via memory-mapped I/O. + + Args: + path: Path to the JSONL file containing SWE trajectories, or a + directory containing a pre-tokenized Arrow dataset (saved by + ``python -m areal.dataset.swe_sft --save-tokenized``). + split: Unused, kept for API compatibility. + tokenizer: Tokenizer with ``apply_chat_template`` support. + Not required when loading a pre-tokenized dataset. + max_length: Max token length. Longer sequences are filtered out. + num_proc: Number of parallel workers for tokenization. + Defaults to ``min(os.cpu_count(), DATASET_NUM_PROC)``. + pre_split: If True, treat input as pre-split pairs (each line is + ``{"messages": [...]}``) instead of full trajectories. + filter_errors: If True (default), discard pairs whose current segment + contains a tool result with ``is_error=True``. In trajectory + mode, sets ``loss_mask=0`` for error segments instead. + Set to False to keep/train all regardless of tool errors. + strip_all_thinking: If True, strip ``...`` from every + assistant turn including the training target. + Ignored in trajectory mode (thinking is always preserved). + filter_empty_tool_calls: If True, discard pairs whose training-target + assistant turn has no text content but has tool_calls. + filter_bare_text_tool_calls: If True, discard pairs whose + training-target assistant turn has text without ```` + tags and has tool_calls. + truncate_task_notifications: If True, truncate trajectories at the + first ```` that follows a pure-text assistant + turn, removing noise from background task completions. + no_tools: If True, do not pass tool definitions to + ``apply_chat_template`` even if the data contains them. + skip_pretokenized_filter: If True, skip the ``max_length`` filter + when loading a pre-tokenized dataset. Useful when the dataset + was already filtered during pretokenization and you want to + avoid NFS cache conflicts from concurrent ``dataset.filter()`` + calls across ranks. + max_no_thinking_ratio: Maximum ratio of non-thinking pairs to thinking + pairs. For example, ``1.0`` gives 1:1, ``2.0`` gives 1:2. + ``None`` (default) disables balancing. + split_mode: ``"pair"`` (default) splits trajectories into + progressive pairs. ``"trajectory"`` keeps the full trajectory + as a single sample — all assistant turns are targets with + ``loss_mask=1``, error segments are masked instead of filtered. + random_strip_thinking_prob: Probability of stripping thinking from + each target assistant turn. 0.0 (default) = no stripping, + 1.0 = strip all. Works in both pair and trajectory mode. + random_strip_thinking_seed: Random seed for reproducible thinking + stripping decisions. + n_thinking_variants: Number of thinking-pattern variants per + trajectory. ``1`` (default) = no augmentation. ``K > 1`` + = augment each trajectory into K variants: the first + preserves all thinking, the rest randomly strip with + *random_strip_thinking_prob*. + cache_dir: Directory to save/load the processed Arrow dataset. + When set in distributed mode, rank 0 processes the data and + saves here; other ranks load from this directory. If the + directory already contains a completed cache (``.done`` marker), + all ranks load from it directly without reprocessing. + dump_dir: Directory to write sample dump files (``.txt`` + ``.json``). + Only rank 0 writes. Set to None to disable. + dump_samples: Number of random samples to dump. ``-1`` = all, + ``0`` = disabled. + parse_tool_call_args: If True, convert OpenAI JSON-string + ``tool_calls.arguments`` to dicts before ``apply_chat_template``. + Required by GLM-4.x / GLM-5.x templates; leave at the default + (False) for Qwen / Llama / Bailing. + + Returns: + A HuggingFace ``Dataset`` with ``input_ids`` and ``loss_mask`` columns. + """ + from datasets import load_from_disk + + # Pre-tokenized Arrow dataset: load directly, skip all processing. + if os.path.isdir(path): + logger.info(f"Loading pre-tokenized dataset from {path}") + dataset = load_from_disk(path) + + if max_length is not None and not skip_pretokenized_filter: + before_filter = len(dataset) + dataset = dataset.filter( + lambda x: len(x["input_ids"]) <= max_length, num_proc=num_proc + ) + logger.info( + f"Filtered {before_filter - len(dataset)} samples " + f"exceeding max_length={max_length}" + ) + + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + + # --- Shared kwargs for _process_swe_sft --- + process_kwargs = dict( + max_length=max_length, + num_proc=num_proc, + pre_split=pre_split, + filter_errors=filter_errors, + strip_all_thinking=strip_all_thinking, + filter_empty_tool_calls=filter_empty_tool_calls, + filter_bare_text_tool_calls=filter_bare_text_tool_calls, + truncate_task_notifications=truncate_task_notifications, + no_tools=no_tools, + max_no_thinking_ratio=max_no_thinking_ratio, + split_mode=split_mode, + random_strip_thinking_prob=random_strip_thinking_prob, + random_strip_thinking_seed=random_strip_thinking_seed, + n_thinking_variants=n_thinking_variants, + dump_dir=dump_dir, + dump_n_samples=dump_samples, + parse_tool_call_args=parse_tool_call_args, + ) + + # --- Distributed rank-0-only processing --- + rank = int(os.getenv("RANK", "0")) + world_size = int(os.getenv("WORLD_SIZE", "1")) + + if cache_dir is not None and world_size > 1: + done_marker = os.path.join(cache_dir, ".done") + meta_path = os.path.join(cache_dir, ".meta.json") + cache_meta = { + "version": 1, + "path": path, + "tokenizer": getattr(tokenizer, "name_or_path", None), + "process_kwargs": { + k: v + for k, v in process_kwargs.items() + if k not in ("dump_dir", "dump_n_samples") + }, + } + + def _filter_by_max_length(ds): + if max_length is None: + return ds + before = len(ds) + # Length via arrow list offsets: avoids decoding every row to + # Python lists, which for long-context datasets costs minutes of + # startup per rank while (on a validated cache) removing nothing — + # build-time _TokenizeAndMask already filtered with this max_length. + import pyarrow.compute as pc + + # ds.data is the underlying arrow table; a freshly built dataset + # carries an indices mapping (from .filter views) whose row count + # differs. Materialize the view first (no-op for load_from_disk). + if getattr(ds, "_indices", None) is not None: + ds = ds.flatten_indices() + lengths = pc.list_value_length(ds.data.column("input_ids")).to_pylist() + keep = [i for i, n in enumerate(lengths) if n <= max_length] + ds = ds.select(keep) + if len(ds) < before: + logger.info( + f"Rank {rank}: filtered {before - len(ds)} samples " + f"exceeding max_length={max_length}" + ) + if len(ds) == 0: + raise ValueError( + f"processed dataset at {cache_dir} has 0 samples after " + f"max_length={max_length} filtering" + ) + return ds + + def _load_valid_cache(): + if not os.path.exists(meta_path): + raise ValueError(f"cached dataset metadata is missing: {meta_path}") + with open(meta_path) as f: + cached_meta = json.load(f) + if cached_meta != cache_meta: + raise ValueError( + f"cached dataset metadata does not match current SWE settings: " + f"{meta_path}" + ) + dataset = load_from_disk(cache_dir) + if len(dataset) == 0: + raise ValueError(f"cached dataset is empty: {cache_dir}") + return dataset + + def _wait_for_valid_cache(): + start = time.monotonic() + last_error = None + while True: + if os.path.exists(done_marker): + try: + return _load_valid_cache() + except Exception as e: + last_error = e + elapsed = time.monotonic() - start + if elapsed > _RANK0_CACHE_TIMEOUT: + raise TimeoutError( + f"Waited {_RANK0_CACHE_TIMEOUT}s for rank 0 to rebuild " + f"a valid dataset cache at {cache_dir}. Last error: {last_error}" + ) + time.sleep(_RANK0_CACHE_POLL_INTERVAL) + + # Fast path: cache from a previous run (or rank 0 already finished). + if os.path.exists(done_marker): + if rank == 0: + try: + logger.info( + f"Rank {rank}: loading cached processed dataset from {cache_dir}" + ) + dataset = _load_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + except Exception as e: + logger.warning( + "Rank 0: invalid processed dataset cache at %s (%s); " + "rebuilding it.", + cache_dir, + e, + ) + shutil.rmtree(cache_dir, ignore_errors=True) + else: + try: + logger.info( + f"Rank {rank}: loading cached processed dataset from {cache_dir}" + ) + dataset = _load_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info(f"Final dataset: {len(dataset)} samples") + return dataset + except Exception as e: + logger.warning( + "Rank %d: cached processed dataset at %s is not usable " + "(%s); waiting for rank 0 to rebuild it.", + rank, + cache_dir, + e, + ) + dataset = _wait_for_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info( + f"Rank {rank}: loaded rebuilt dataset ({len(dataset)} samples)" + ) + return dataset + + if rank == 0: + # Rank 0: do the heavy processing and save for other ranks. + dataset = _process_swe_sft(path, tokenizer, **process_kwargs) + if len(dataset) == 0: + raise RuntimeError( + "SWE SFT preprocessing produced 0 samples; refusing to cache " + "an empty processed_dataset." + ) + shutil.rmtree(cache_dir, ignore_errors=True) + os.makedirs(cache_dir, exist_ok=True) + dataset.save_to_disk(cache_dir) + with open(meta_path, "w") as f: + json.dump(cache_meta, f, sort_keys=True) + # Write marker AFTER save completes so readers see a consistent dir. + with open(done_marker, "w") as f: + f.write(str(len(dataset))) + logger.info( + f"Rank 0: saved processed dataset " + f"({len(dataset)} samples) to {cache_dir}" + ) + dataset = _filter_by_max_length(dataset) + return dataset + else: + # Other ranks: wait for rank 0, then load with meta validation so a + # cache rebuilt for different settings (or mid-rmtree) is never + # silently loaded as this rank's dataset. + logger.info(f"Rank {rank}: waiting for rank 0 to process dataset...") + dataset = _wait_for_valid_cache() + dataset = _filter_by_max_length(dataset) + logger.info(f"Rank {rank}: loaded cached dataset ({len(dataset)} samples)") + return dataset + + # --- Non-distributed or no cache_dir: process in current process --- + return _process_swe_sft(path, tokenizer, **process_kwargs) diff --git a/areal/dataset/swe_sft/tokenization.py b/areal/dataset/swe_sft/tokenization.py new file mode 100644 index 0000000000..5258e4adbe --- /dev/null +++ b/areal/dataset/swe_sft/tokenization.py @@ -0,0 +1,539 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Chat-template handling, tokenization, loss masking, and sample dumps.""" + +import json +import os +import random +import re + +from datasets import Dataset + +from areal.utils import logging + +logger = logging.getLogger("SWESFTDataset") + +DATASET_NUM_PROC = 1 + + +# -- Chat template patch (runtime, no file modification) -------- + +# Both Bailing and Qwen3 templates have ``ns.last_query_index`` logic +# that prevents ```` rendering for assistant turns BEFORE the +# last user message, AND discards inline empty ``\n`` +# extracted from content. +# +# This breaks trajectory-mode training: +# - Multi-user trajectories: turns before the last user msg lack +# - Empty ensure_thinking via inline gets stripped +# +# The patch below handles both Bailing (`ASSISTANT` style) +# and Qwen3 (`<|im_start|>assistant` style) templates: +# 1. Adds ``had_think_tags`` detection so empty ```` survives. +# 2. Removes the ``ns.last_query_index`` gate so all assistant turns +# render ```` uniformly when think intent is detected. +# +# Applied at runtime via ``tokenizer.chat_template = patched`` — the +# original template file on disk is never modified. + +_BAILING_OLD_BLOCK = ( + "{%- if loop.index0 > ns.last_query_index %}\n" + " {%- if reasoning_content != '' %}\n" + " {{- 'ASSISTANT\\n' + '\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- 'ASSISTANT\\n' + content }}\n" + " {%- endif %}\n" + " {%- else %}\n" + " {{- 'ASSISTANT\\n' + content }}\n" + " {%- endif %}" +) +_BAILING_NEW_BLOCK = ( + "{%- if reasoning_content != '' or had_think_tags %}\n" + " {{- 'ASSISTANT\\n' + '\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- 'ASSISTANT\\n' + content }}\n" + " {%- endif %}" +) + +# Qwen3 uses `loop.last or (not loop.last and reasoning_content)` so the +# last turn always renders even with empty reasoning. We +# preserve `loop.last` and add `had_think_tags` for inline-empty support. +_QWEN3_OLD_BLOCK = ( + "{%- if loop.index0 > ns.last_query_index %}\n" + " {%- if loop.last or (not loop.last and reasoning_content) %}\n" + " {{- '<|im_start|>' + message.role + '\\n\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" + " {%- endif %}\n" + " {%- else %}\n" + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" + " {%- endif %}" +) +_QWEN3_NEW_BLOCK = ( + "{%- if loop.last or reasoning_content != '' or had_think_tags %}\n" + " {{- '<|im_start|>' + message.role + '\\n\\n'" + " + reasoning_content.strip('\\n') + '\\n\\n\\n'" + " + content.lstrip('\\n') }}\n" + " {%- else %}\n" + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n" + " {%- endif %}" +) + +_OLD_DETECT = "{%- set reasoning_content = '' %}" +_NEW_DETECT = ( + "{%- set reasoning_content = '' %}\n" + " {%- set had_think_tags = ('' in content) %}" +) + + +def _patch_chat_template_for_training(tokenizer): + """Patch Bailing/Qwen3 chat templates to render ```` uniformly. + + Detects template family by matching known render blocks: + - Bailing: ``ASSISTANT`` markers + - Qwen3: ``<|im_start|>assistant`` markers + + Other templates (e.g. plain ChatML without ``last_query_index``) + are left unchanged. If the template has ``last_query_index`` but + neither known block matches, logs a warning. + """ + template = getattr(tokenizer, "chat_template", None) + if not template or "last_query_index" not in template: + return + + if _BAILING_OLD_BLOCK in template: + family = "Bailing" + patched = template.replace(_BAILING_OLD_BLOCK, _BAILING_NEW_BLOCK) + elif _QWEN3_OLD_BLOCK in template: + family = "Qwen3" + patched = template.replace(_QWEN3_OLD_BLOCK, _QWEN3_NEW_BLOCK) + else: + # Reaching here means the template family needs the training patch + # (it gates rendering on last_query_index) but the verbatim block no + # longer matches — most likely an upstream template revision. Failing + # loudly beats silently training on data whose blocks the + # stock template strips (see _clean_message / ensure_thinking). + # + # Escape hatch for uses that do not depend on think normalization + # (e.g. precision-alignment forward dumps): set + # AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE=1 to proceed with a warning. + if os.environ.get("AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE", ""): + logger.warning( + "Chat template has last_query_index but matches neither known " + "render block; proceeding UNPATCHED because " + "AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE is set. Empty " + "blocks may be discarded by the stock template." + ) + return + raise ValueError( + "Chat template has last_query_index but matches neither the known " + "Bailing nor Qwen3 render block; the training patch cannot be " + "applied. Without it, empty blocks are discarded and " + "multi-turn thinking renders inconsistently. Update " + "_BAILING_OLD_BLOCK/_QWEN3_OLD_BLOCK for this template revision, " + "or set AREAL_SWE_ALLOW_UNPATCHED_TEMPLATE=1 if think " + "normalization is irrelevant for this run." + ) + + if _OLD_DETECT not in patched: + raise ValueError( + "Chat template render block matched but the reasoning_content " + "detect line did not; had_think_tags would be undefined and " + "empty blocks would silently vanish. Update " + "_OLD_DETECT/_NEW_DETECT for this template revision." + ) + patched = patched.replace(_OLD_DETECT, _NEW_DETECT) + + tokenizer.chat_template = patched + logger.info( + f"Patched {family} chat template for training: removed " + "last_query_index gate, added had_think_tags detection." + ) + + +_TEMPLATE_PATTERNS = [ + # ChatML (Qwen, etc.): <|im_start|>assistant\n ... <|im_end|> + (r"<\|im_start\|>assistant\n", r"<\|im_end\|>"), + # Llama 3: <|start_header_id|>assistant<|end_header_id|>\n\n ... <|eot_id|> + (r"<\|start_header_id\|>assistant<\|end_header_id\|>\n\n", r"<\|eot_id\|>"), + # GLM: <|assistant|> ... (ends at next <|user|>, <|observation|>, or end of string) + (r"<\|assistant\|>", r"(?=<\|user\|>|<\|observation\|>|\Z)"), +] + + +def _parse_tool_call_arguments(messages): + """Parse JSON-string arguments in tool_calls to dicts. + + OpenAI returns tool_call arguments as JSON strings, but some chat + templates (e.g. GLM-4.x / GLM-5.x) expect parsed dicts. Most other + templates (Qwen / ChatML, Llama 3, Bailing, ...) accept the standard + OpenAI string form, so this conversion must be opt-in. + """ + patched = [] + for m in messages: + tool_calls = m.get("tool_calls") + if not tool_calls: + patched.append(m) + continue + new_tcs = [] + for tc in tool_calls: + fn = tc.get("function", tc) + args = fn.get("arguments") + if isinstance(args, str): + try: + parsed = json.loads(args) + except (json.JSONDecodeError, TypeError): + parsed = args + fn = {**fn, "arguments": parsed} + tc = {**tc, "function": fn} if "function" in tc else fn + new_tcs.append(tc) + patched.append({**m, "tool_calls": new_tcs}) + return patched + + +def _render_tokenize_mask( + messages, + tokenizer, + assistant_pattern, + tools=None, + *, + split_mode="pair", + error_indices=None, + parse_tool_call_args=False, +): + """Render, tokenize, and build loss_mask for a message list. + + In **pair mode** (default), only the **last** assistant turn gets + ``loss_mask=1``. In **trajectory mode**, **all** assistant turns + get ``loss_mask=1`` except those at indices in *error_indices*. + + When *parse_tool_call_args* is True, JSON-string ``tool_calls`` arguments + are converted to dicts before rendering (required by GLM chat templates; + other templates such as Qwen / Llama / Bailing must keep the OpenAI + string form). + + Returns: + Tuple of ``(full_text, input_ids, loss_mask, offset_mapping)``, or + ``None`` if ``apply_chat_template`` fails. + """ + # 1) Render the full template text. + try: + kwargs = {"tokenize": False} + if tools is not None: + kwargs["tools"] = tools + if parse_tool_call_args: + messages = _parse_tool_call_arguments(messages) + full_text = tokenizer.apply_chat_template(messages, **kwargs) + except Exception as e: + logger.warning( + "apply_chat_template failed: %s. Skipping sample.", + e, + ) + return None + + # 2) Tokenize with offset mapping so we can map char→token. + encoding = tokenizer( + full_text, add_special_tokens=False, return_offsets_mapping=True + ) + input_ids = encoding["input_ids"] + offset_mapping = encoding["offset_mapping"] + + # 3) Build loss_mask. + loss_mask = [0] * len(input_ids) + + if split_mode == "trajectory": + # Trajectory mode: mask ALL assistant segments, skip error_indices. + skip = set(error_indices) if error_indices else set() + matches = list(assistant_pattern.finditer(full_text)) + + # Verify regex matches correspond 1:1 to assistant messages. + n_asst = sum(1 for m in messages if m.get("role") == "assistant") + if len(matches) != n_asst: + # Fail closed: a spurious match (e.g. a tool output quoting the + # chat-template header literal) would otherwise put loss on + # user/tool tokens and silently defeat error masking. + logger.warning( + "Segment count mismatch: %d assistant messages but %d regex " + "matches in rendered text. Dropping this sample.", + n_asst, + len(matches), + ) + return None + + for seg_idx, m in enumerate(matches): + if seg_idx in skip: + continue + rs, re_ = m.start(1), m.end(0) + for tok_idx, (cs, ce) in enumerate(offset_mapping): + if ce > rs and cs < re_: + loss_mask[tok_idx] = 1 + else: + # Pair mode: mask only the LAST assistant segment. + last_match = None + for m in assistant_pattern.finditer(full_text): + last_match = m + if last_match is not None: + rs, re_ = last_match.start(1), last_match.end(0) + for tok_idx, (cs, ce) in enumerate(offset_mapping): + if ce > rs and cs < re_: + loss_mask[tok_idx] = 1 + else: + # Loss lands nowhere; the SFT loss path tolerates all-zero masks + # (kept, not dropped, to preserve cache compatibility) but this + # always signals template/pattern drift worth investigating. + logger.warning( + "No assistant segment matched the template pattern; sample " + "keeps an all-zero loss_mask." + ) + + return full_text, input_ids, loss_mask, offset_mapping + + +class _TokenizeAndMask: + """Picklable callable for ``Dataset.map(num_proc=N)``.""" + + def __init__( + self, + tokenizer, + assistant_pattern, + max_length=None, + *, + split_mode="pair", + parse_tool_call_args=False, + ): + self.tokenizer = tokenizer + self.assistant_pattern = assistant_pattern + self.max_length = max_length + self.split_mode = split_mode + self.parse_tool_call_args = parse_tool_call_args + + def __call__(self, sample): + error_indices = ( + sample.get("error_indices", []) if self.split_mode == "trajectory" else None + ) + tools_json = sample.get("tools_json") + tools = json.loads(tools_json) if tools_json else None + result = _render_tokenize_mask( + sample["messages"], + self.tokenizer, + self.assistant_pattern, + tools, + split_mode=self.split_mode, + error_indices=error_indices, + parse_tool_call_args=self.parse_tool_call_args, + ) + if result is None: + return {"input_ids": [], "loss_mask": []} + + _full_text, input_ids, loss_mask, _offset_mapping = result + + # Early exit: overlength or empty → return empty so a single + # filter pass removes it together with template-failure empties. + if self.max_length is not None and len(input_ids) > self.max_length: + return {"input_ids": [], "loss_mask": []} + + return {"input_ids": input_ids, "loss_mask": loss_mask} + + +def _detect_template_pattern(tokenizer, tools=None): + """Detect the assistant role delimiter used by this tokenizer's template. + + When *tools* is provided the probe is rendered with ``tools=`` so that + the detected delimiters match the actual training text (some templates + alter the system block when tools are present). + + Strategy: + 1. Try known ``_TEMPLATE_PATTERNS`` (fast, battle-tested). + 2. Fall back to double-probe diff: render the template with a known + marker and with empty content, then diff the two strings to extract + the exact header and end-of-turn delimiters. + + Raises: + ValueError: If both strategies fail to detect a usable pattern. + """ + _PROBE_CONTENT = "PROBE_MARKER" + + extra_kwargs = {} + if tools is not None: + extra_kwargs["tools"] = tools + + probe_msgs = [ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": _PROBE_CONTENT}, + ] + probe_text = tokenizer.apply_chat_template( + probe_msgs, tokenize=False, **extra_kwargs + ) + + # --- Strategy 1: known patterns --- + for hdr_re, eot_re in _TEMPLATE_PATTERNS: + if re.search(hdr_re, probe_text): + pattern = re.compile(hdr_re + r"(.*?)" + eot_re, re.DOTALL) + logger.info( + f"Detected template style (known pattern): " + f"header_re={hdr_re!r}, eot_re={eot_re!r}" + ) + return pattern + + # --- Strategy 2: double-probe diff --- + try: + probe_empty = [ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": ""}, + ] + text_empty = tokenizer.apply_chat_template( + probe_empty, tokenize=False, **extra_kwargs + ) + + marker_idx = probe_text.index(_PROBE_CONTENT) + header = probe_text[:marker_idx] + tail = probe_text[marker_idx + len(_PROBE_CONTENT) :] + + if text_empty == header + tail: + # Extract the assistant-specific header by removing the shared + # user-only prefix. + user_only = tokenizer.apply_chat_template( + [{"role": "user", "content": "x"}], + tokenize=False, + **extra_kwargs, + ) + asst_header = header[len(user_only) :] + # end-of-turn delimiter: strip leading newlines, then take + # up to the first newline (or the full string if none). + eot_stripped = tail.lstrip("\n") + eot = eot_stripped.split("\n")[0] if "\n" in eot_stripped else eot_stripped + + if asst_header and eot: + hdr_re = re.escape(asst_header) + eot_re = re.escape(eot) + pattern = re.compile(hdr_re + r"(.*?)" + eot_re, re.DOTALL) + logger.info( + f"Detected template style (probe diff): " + f"header={asst_header!r}, eot={eot!r}" + ) + return pattern + except (ValueError, IndexError): + pass # PROBE_CONTENT not found in rendered text, skip + + raise ValueError( + "Could not detect chat template assistant delimiters. " + "Unable to build a reliable loss mask. " + f"Probe text: {probe_text[:200]!r}" + ) + + +def _dump_samples( + samples, + tokenizer, + assistant_pattern, + tools_list, + dump_dir, + n_samples, + *, + split_mode="pair", + error_indices_list=None, + parse_tool_call_args=False, +): + """Dump sampled message lists as ``.txt`` + ``.json`` for inspection. + + Args: + samples: List of message-list samples (pairs or full trajectories). + tokenizer: Tokenizer with ``apply_chat_template`` support. + assistant_pattern: Compiled regex from ``_detect_template_pattern``. + tools_list: Per-sample tool definitions (parallel to *samples*), + or ``None`` when no tools are available. + dump_dir: Directory to write files into (created if needed). + n_samples: Number of random samples to dump. ``-1`` dumps all. + split_mode: ``"trajectory"`` for trajectory-mode loss masking. + error_indices_list: Per-sample error segment indices (trajectory mode). + """ + import random as _random + + os.makedirs(dump_dir, exist_ok=True) + + if n_samples == -1 or n_samples >= len(samples): + indices = list(range(len(samples))) + else: + indices = sorted(_random.sample(range(len(samples)), n_samples)) + + n_written = 0 + for i in indices: + sample = samples[i] + sample_tools = tools_list[i] if tools_list else None + err_idxs = ( + error_indices_list[i] + if split_mode == "trajectory" and error_indices_list + else None + ) + + result = _render_tokenize_mask( + sample, + tokenizer, + assistant_pattern, + sample_tools, + split_mode=split_mode, + error_indices=err_idxs, + parse_tool_call_args=parse_tool_call_args, + ) + if result is None: + continue + + full_text, input_ids, loss_mask, offset_mapping = result + n_loss = sum(loss_mask) + base = os.path.join(dump_dir, f"sample_{i}") + + # --- .txt --- + with open(base + ".txt", "w", encoding="utf-8") as fout: + fout.write( + f"Sample {i}: {len(sample)} messages, " + f"{len(input_ids)} tokens, loss=1: {n_loss}\n" + ) + fout.write(f"Last msg role: {sample[-1]['role']}\n") + fout.write(f"{'=' * 72}\n\n") + + fout.write("--- Rendered Text ---\n") + fout.write(full_text) + fout.write("\n\n") + + fout.write("--- Token / Loss Mask ---\n") + fout.write(f"{'Idx':>6} | {'TokenID':>8} | Loss | Token Text\n") + fout.write(f"{'-' * 6}-+-{'-' * 8}-+------+{'-' * 40}\n") + for t in range(len(input_ids)): + cs, ce = offset_mapping[t] + tok_text = repr(full_text[cs:ce]) + fout.write( + f"{t:>6} | {input_ids[t]:>8} | {loss_mask[t]:>4} | {tok_text}\n" + ) + + # --- .json --- + tokens_list = [] + for t in range(len(input_ids)): + cs, ce = offset_mapping[t] + tokens_list.append( + { + "idx": t, + "token_id": input_ids[t], + "text": full_text[cs:ce], + "loss": loss_mask[t], + } + ) + record = { + "sample_index": i, + "n_messages": len(sample), + "n_tokens": len(input_ids), + "n_loss_tokens": n_loss, + "rendered_text": full_text, + "tokens": tokens_list, + } + with open(base + ".json", "w", encoding="utf-8") as fout: + json.dump(record, fout, ensure_ascii=False) + + n_written += 1 + + logger.info(f"Dumped {n_written} samples to {dump_dir}/") diff --git a/tests/test_swe_sft_cache.py b/tests/test_swe_sft_cache.py index 0ac582fc25..42b3711eb2 100644 --- a/tests/test_swe_sft_cache.py +++ b/tests/test_swe_sft_cache.py @@ -22,21 +22,29 @@ def _load_swe_sft_module(): areal_module = types.ModuleType("areal") dataset_module = types.ModuleType("areal.dataset") + dataset_module.__path__ = [] + swe_package = types.ModuleType("areal.dataset.swe_sft") + swe_package.__path__ = [] utils_module = types.ModuleType("areal.utils") utils_module.logging = logging areal_module.dataset = dataset_module areal_module.utils = utils_module sys.modules["areal"] = areal_module sys.modules["areal.dataset"] = dataset_module + sys.modules["areal.dataset.swe_sft"] = swe_package sys.modules["areal.utils"] = utils_module - path = Path(__file__).parents[1] / "areal" / "dataset" / "swe_sft.py" - spec = importlib.util.spec_from_file_location("areal.dataset.swe_sft", path) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules["areal.dataset.swe_sft"] = module + package_path = Path(__file__).parents[1] / "areal" / "dataset" / "swe_sft" try: - spec.loader.exec_module(module) + for name in ("messages", "tokenization", "pipeline"): + full_name = f"areal.dataset.swe_sft.{name}" + spec = importlib.util.spec_from_file_location( + full_name, package_path / f"{name}.py" + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[full_name] = module + spec.loader.exec_module(module) finally: for name in list(sys.modules): if name == "areal" or name.startswith("areal."): From 723b49dae78df518d19c0825235233e17ff92559 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Wed, 12 Aug 2026 14:07:00 +0800 Subject: [PATCH 3/8] fix(dataset): align SWE SFT with adaptive templates Port the follow-up SWE SFT fixes from swe-dev so Bailing V3 adaptive chat templates use structural assistant masks and consistent thinking modes. Signed-off-by: chucai.dzq --- areal/dataset/swe_sft/__main__.py | 1 + areal/dataset/swe_sft/messages.py | 54 +++-- areal/dataset/swe_sft/pipeline.py | 2 +- areal/dataset/swe_sft/tokenization.py | 184 +++++++++++++++-- areal/models/mcore/bailing_v3.py | 2 + areal/models/mcore/bailing_v3_mla.py | 2 + tests/test_bailing_v3_kda_cp_helpers.py | 3 +- tests/test_swe_sft_dataset.py | 262 ++++++++++++++++++++++++ 8 files changed, 461 insertions(+), 49 deletions(-) create mode 100644 tests/test_swe_sft_dataset.py diff --git a/areal/dataset/swe_sft/__main__.py b/areal/dataset/swe_sft/__main__.py index 2b114a7db0..e6adcf390d 100644 --- a/areal/dataset/swe_sft/__main__.py +++ b/areal/dataset/swe_sft/__main__.py @@ -23,6 +23,7 @@ _tokenize_samples, ) + def main(): import argparse import sys diff --git a/areal/dataset/swe_sft/messages.py b/areal/dataset/swe_sft/messages.py index 6b24b767e0..541f861a0e 100644 --- a/areal/dataset/swe_sft/messages.py +++ b/areal/dataset/swe_sft/messages.py @@ -53,8 +53,12 @@ def _iter_jsonl_records(path): nested (``conversations`` wrapper) vs flat format auto-detection via ``_extract_messages``. Records with empty messages are skipped. - Warns about multi-user trajectories which break think-tag rendering - in templates with ``ns.last_query_index`` logic (e.g. Bailing). + Reports multi-user trajectories for visibility. These are now + **handled** in this pipeline: ``_patch_chat_template_for_training`` + removes the ``ns.last_query_index`` gate at tokenization time, so + ```` renders for every assistant turn regardless of how many + user messages precede it. The count below only matters if the data + is tokenized through a pipeline that does NOT apply that patch. """ record_idx = 0 n_multi_user = 0 @@ -72,19 +76,19 @@ def _iter_jsonl_records(path): if n_user > 1: n_multi_user += 1 if n_multi_user <= 3: - logger.warning( - "Record %d has %d user messages. Templates with " - "ns.last_query_index logic (e.g. Bailing) will NOT " - "render for assistant turns before the last " - "user message. Consider filtering!", + logger.info( + "Record %d has %d user messages. Handled by " + "_patch_chat_template_for_training (last_query_index " + "gate removed) so all assistant turns render ; " + "only an issue if tokenized via an unpatched template.", record_idx, n_user, ) yield record_idx, messages, record_tools if n_multi_user > 0: - logger.warning( - "Total %d/%d records have multiple user messages. " - "These may produce no-think training signal.", + logger.info( + "Total %d/%d records have multiple user messages " + "(handled by the chat-template patch at tokenization time).", n_multi_user, record_idx, ) @@ -159,9 +163,9 @@ def _clean_message(msg, strip_thinking=True, ensure_thinking=False): ensure_thinking: If True, inject inline ``\n`` on assistant turns that lack a thinking block (either inline or in ``reasoning_content``). Requires the - patched Bailing template (via - ``_patch_chat_template_for_training``) which detects - ``had_think_tags`` and preserves empty think blocks. + patched template (Bailing / Qwen3 / other ``last_query_index`` + families, via ``_patch_chat_template_for_training``) which + detects ``had_think_tags`` and preserves empty think blocks. """ cleaned = {"role": msg["role"]} @@ -238,8 +242,9 @@ def _clean_message(msg, strip_thinking=True, ensure_thinking=False): # ``reasoning_content='\n'`` approach. # # Requires ``_patch_chat_template_for_training`` to have been called - # on the tokenizer, otherwise the stock Bailing template will extract - # and discard the empty ```` block. + # on the tokenizer, otherwise the stock template (Bailing / Qwen3 / + # other ``last_query_index`` families) will extract and discard the + # empty ```` block. if ensure_thinking and msg["role"] == "assistant" and not has_thinking: cur_content = cleaned.get("content") if cur_content is None or cur_content == "": @@ -313,13 +318,11 @@ def _is_bare_text_tool_call(msg): def _msg_has_thinking(msg): - """True if assistant *msg* has thinking content (inline or reasoning_content).""" + """True if assistant *msg* has non-empty thinking content.""" if msg.get("role") != "assistant": return False - content = msg.get("content") or "" - normalized = _THINK_OPEN_RE.sub("", content) - normalized = _THINK_CLOSE_RE.sub("", normalized) - if _THINK_RE.search(normalized): + normalized = _normalize_thinking_tags(msg.get("content") or "") + if any(match.group(1).strip() for match in _THINK_RE.finditer(normalized)): return True rc = msg.get("reasoning_content") or "" return bool(rc.strip()) @@ -381,16 +384,7 @@ def _classify_pair(pair): if target.get("role") != "assistant": return "pure_text" - content = target.get("content") or "" - rc = target.get("reasoning_content") or "" - # Require non-empty think content: pair cleaning runs BEFORE balancing - # and (with ensure_thinking) injects an empty \n into - # every no-think target, so a bare regex hit would classify everything - # as "thinking" and silently disable max_no_thinking_ratio. - _m = _THINK_RE.search(content) - has_thinking = bool(_m and _m.group(1).strip()) or bool(rc.strip()) - - if has_thinking: + if _msg_has_thinking(target): return "thinking" if target.get("tool_calls"): return "no_thinking_tool_call" diff --git a/areal/dataset/swe_sft/pipeline.py b/areal/dataset/swe_sft/pipeline.py index 144dd780f5..8cc16a0d64 100644 --- a/areal/dataset/swe_sft/pipeline.py +++ b/areal/dataset/swe_sft/pipeline.py @@ -25,10 +25,10 @@ ) from .tokenization import ( DATASET_NUM_PROC, - _TokenizeAndMask, _detect_template_pattern, _dump_samples, _patch_chat_template_for_training, + _TokenizeAndMask, ) logger = logging.getLogger("SWESFTDataset") diff --git a/areal/dataset/swe_sft/tokenization.py b/areal/dataset/swe_sft/tokenization.py index 5258e4adbe..12ea3aa204 100644 --- a/areal/dataset/swe_sft/tokenization.py +++ b/areal/dataset/swe_sft/tokenization.py @@ -4,13 +4,12 @@ import json import os -import random import re -from datasets import Dataset - from areal.utils import logging +from .messages import _msg_has_thinking + logger = logging.getLogger("SWESFTDataset") DATASET_NUM_PROC = 1 @@ -18,20 +17,22 @@ # -- Chat template patch (runtime, no file modification) -------- -# Both Bailing and Qwen3 templates have ``ns.last_query_index`` logic -# that prevents ```` rendering for assistant turns BEFORE the -# last user message, AND discards inline empty ``\n`` -# extracted from content. +# Bailing / Qwen3 (and other ``ns.last_query_index`` families) prevent +# ```` rendering for assistant turns BEFORE the last user message, AND +# discard inline empty ``\n`` extracted from content. # -# This breaks trajectory-mode training: +# Without patching this breaks trajectory-mode training: # - Multi-user trajectories: turns before the last user msg lack # - Empty ensure_thinking via inline gets stripped # -# The patch below handles both Bailing (`ASSISTANT` style) -# and Qwen3 (`<|im_start|>assistant` style) templates: +# The patch below handles Bailing (`ASSISTANT` style) and Qwen3 +# (`<|im_start|>assistant` style) via exact block replacement: # 1. Adds ``had_think_tags`` detection so empty ```` survives. -# 2. Removes the ``ns.last_query_index`` gate so all assistant turns -# render ```` uniformly when think intent is detected. +# 2. Removes the ``ns.last_query_index`` gate so all assistant turns render +# ```` uniformly when think intent is detected — which is why +# multi-user trajectories are now handled (no think loss). +# An unrecognized family (e.g. a future ring3.0 revision) is skipped with a +# warning; add its exact block here rather than blindly neutralizing. # # Applied at runtime via ``tokenizer.chat_template = patched`` — the # original template file on disk is never modified. @@ -91,22 +92,98 @@ " {%- set had_think_tags = ('' in content) %}" ) +_BAILING_V3_ASSISTANT_START = ( + '{%- elif message.role == "assistant" %}\n' + " {%- set reasoning_content = '' %}" +) +_BAILING_V3_ASSISTANT_START_WITH_GENERATION = ( + '{%- elif message.role == "assistant" %}\n' + " {{- 'ASSISTANT' }}\n" + " {%- generation %}\n" + " {%- set reasoning_content = '' %}" +) +_BAILING_V3_ASSISTANT_END = ( + " {{- '<|role_end|>' }}\n {%- elif message.role == \"tool\" %}" +) +_BAILING_V3_ASSISTANT_END_WITH_GENERATION = ( + " {{- '<|role_end|>' }}\n" + " {%- endgeneration %}\n" + ' {%- elif message.role == "tool" %}' +) + + +def _add_bailing_v3_generation_tags(template): + """Mark Bailing V3 assistant bodies without changing rendered text.""" + if "preserved_thinking = true" not in template: + return None + if _BAILING_V3_ASSISTANT_START not in template: + return None + if template.count(_BAILING_V3_ASSISTANT_END) != 1: + return None + + patched = template.replace( + _BAILING_V3_ASSISTANT_START, + _BAILING_V3_ASSISTANT_START_WITH_GENERATION, + 1, + ).replace( + _BAILING_V3_ASSISTANT_END, + _BAILING_V3_ASSISTANT_END_WITH_GENERATION, + 1, + ) + + # The header is now emitted once, outside the tracked generation region. + thinking_header = "'ASSISTANT' + " + empty_thinking_header = "'ASSISTANT\\n' + " + if patched.count(thinking_header) != 1 or patched.count(empty_thinking_header) != 2: + return None + patched = patched.replace(thinking_header, "", 1).replace( + empty_thinking_header, + "'\\n' + ", + 2, + ) + return patched + def _patch_chat_template_for_training(tokenizer): """Patch Bailing/Qwen3 chat templates to render ```` uniformly. Detects template family by matching known render blocks: - - Bailing: ``ASSISTANT`` markers + - Bailing V2.5: ``ASSISTANT`` + ``reasoning_content != ''`` - Qwen3: ``<|im_start|>assistant`` markers - Other templates (e.g. plain ChatML without ``last_query_index``) - are left unchanged. If the template has ``last_query_index`` but - neither known block matches, logs a warning. + Bailing V3 *adaptive* (config_ling_adaptive) needs no thinking patch because + ``preserved_thinking = true`` already renders thinking uniformly. Its + assistant body is instrumented with Jinja generation tags so Transformers + can construct a structural loss mask without parsing role delimiters. + + Other templates (plain ChatML without ``last_query_index``) are left + unchanged. If the template has ``last_query_index`` but no known block + matches, logs a warning and skips — add an explicit block pattern for that + family rather than relying on a blind neutralization. """ template = getattr(tokenizer, "chat_template", None) if not template or "last_query_index" not in template: return + # Bailing V3 adaptive already preserves thinking. Add Transformers' + # generation tracking around assistant bodies so loss masks do not depend + # on delimiters that may also occur literally in message payloads. + if "preserved_thinking = true" in template and "preserved_thinking or" in template: + patched = _add_bailing_v3_generation_tags(template) + if patched is not None: + tokenizer.chat_template = patched + logger.info( + "Bailing V3 adaptive template detected: added generation " + "tracking around assistant bodies." + ) + return + logger.info( + "Bailing V3 adaptive template detected (preserved_thinking=true): " + "all assistant turns already render , but its assistant " + "block was not recognized for generation tracking." + ) + return + if _BAILING_OLD_BLOCK in template: family = "Bailing" patched = template.replace(_BAILING_OLD_BLOCK, _BAILING_NEW_BLOCK) @@ -164,6 +241,13 @@ def _patch_chat_template_for_training(tokenizer): (r"<\|start_header_id\|>assistant<\|end_header_id\|>\n\n", r"<\|eot_id\|>"), # GLM: <|assistant|> ... (ends at next <|user|>, <|observation|>, or end of string) (r"<\|assistant\|>", r"(?=<\|user\|>|<\|observation\|>|\Z)"), + # Bailing (V2.5 / V3 / V3-adaptive): ASSISTANT ... <|role_end|> + # Must be a KNOWN pattern (not probe-derived): the V3-adaptive template + # (config_ling_adaptive) renders an empty ```` for + # non-reasoning turns, which the double-probe would bake into the header and + # then fail to match real reasoning turns (``{reasoning}``). + # This header/eot captures the whole turn (think + content + tool_calls). + (r"ASSISTANT", r"<\|role_end\|>"), ] @@ -227,6 +311,19 @@ def _render_tokenize_mask( kwargs = {"tokenize": False} if tools is not None: kwargs["tools"] = tools + # Adaptive-thinking templates (Bailing V3 config_ling_adaptive) read an + # ``enable_thinking`` flag that sets the ``detailed thinking on/off`` + # system label. Match it to whether THIS rendered example actually + # contains thinking, so the switch is trained consistently + # (on<->think, off<->no-think); training an all-no-think example under + # "on" (or a thinking example under "off") would corrupt the switch. + # In trajectory mode ``messages`` is the whole trajectory (=> per-traj + # rule: on iff >=1 assistant turn thinks); in pair mode it is the pair. + # Gated on the preserved_thinking + enable_thinking signature so other + # templates (e.g. Qwen3) are left untouched. + _tmpl = getattr(tokenizer, "chat_template", None) or "" + if "enable_thinking" in _tmpl and "preserved_thinking" in _tmpl: + kwargs["enable_thinking"] = any(_msg_has_thinking(m) for m in messages) if parse_tool_call_args: messages = _parse_tool_call_arguments(messages) full_text = tokenizer.apply_chat_template(messages, **kwargs) @@ -247,6 +344,61 @@ def _render_tokenize_mask( # 3) Build loss_mask. loss_mask = [0] * len(input_ids) + # Templates instrumented with Jinja's generation extension provide + # structural assistant spans. Unlike delimiter matching, these spans cannot + # be forged by literal role markers in user/tool payloads. + if re.search(r"\{%-?\s*generation\s*-?%\}", _tmpl): + try: + tracked_kwargs = { + **kwargs, + "tokenize": True, + "return_dict": True, + "return_assistant_tokens_mask": True, + } + tracked = tokenizer.apply_chat_template(messages, **tracked_kwargs) + tracked_ids = tracked["input_ids"] + tracked_mask = list(tracked["assistant_masks"]) + if tracked_ids != input_ids or len(tracked_mask) != len(input_ids): + raise ValueError("tracked tokenization differs from rendered text") + except Exception as e: + logger.warning( + "Native assistant-mask generation failed: %s. Skipping sample.", + e, + ) + return None + + segments = [] + start = None + for idx, enabled in enumerate([*tracked_mask, 0]): + if enabled and start is None: + start = idx + elif not enabled and start is not None: + segments.append((start, idx)) + start = None + + n_asst = sum(1 for m in messages if m.get("role") == "assistant") + if len(segments) != n_asst: + logger.warning( + "Native assistant-mask segment mismatch: %d assistant messages " + "but %d tracked segments. Skipping sample.", + n_asst, + len(segments), + ) + return None + + if split_mode == "trajectory": + skip = set(error_indices) if error_indices else set() + selected = ( + segment + for seg_idx, segment in enumerate(segments) + if seg_idx not in skip + ) + else: + selected = segments[-1:] if segments else [] + for start, end in selected: + loss_mask[start:end] = [1] * (end - start) + return full_text, input_ids, loss_mask, offset_mapping + if split_mode == "trajectory": # Trajectory mode: mask ALL assistant segments, skip error_indices. skip = set(error_indices) if error_indices else set() diff --git a/areal/models/mcore/bailing_v3.py b/areal/models/mcore/bailing_v3.py index c1263bca20..59191fd7ad 100644 --- a/areal/models/mcore/bailing_v3.py +++ b/areal/models/mcore/bailing_v3.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + """BailingMoeV3ForCausalLM (Ling V3) support for megatron-core. This module provides: diff --git a/areal/models/mcore/bailing_v3_mla.py b/areal/models/mcore/bailing_v3_mla.py index efadc63dc2..81f30bc830 100644 --- a/areal/models/mcore/bailing_v3_mla.py +++ b/areal/models/mcore/bailing_v3_mla.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + """Gated MLA self-attention for BailingMoeV3 (Ling V3). v3 enables *head-wise gated attention* on top of standard MLA: a per-head sigmoid gate diff --git a/tests/test_bailing_v3_kda_cp_helpers.py b/tests/test_bailing_v3_kda_cp_helpers.py index 813f325ff5..dc83d3b66b 100644 --- a/tests/test_bailing_v3_kda_cp_helpers.py +++ b/tests/test_bailing_v3_kda_cp_helpers.py @@ -183,8 +183,7 @@ def fake_make_sharded( ) ) return { - f"{prefix}{name}": _FakeShard(tensor) - for name, tensor in state_dict.items() + f"{prefix}{name}": _FakeShard(tensor) for name, tensor in state_dict.items() } def fake_default(module, prefix="", sharded_offsets=(), metadata=None): diff --git a/tests/test_swe_sft_dataset.py b/tests/test_swe_sft_dataset.py new file mode 100644 index 0000000000..e47fe3cd47 --- /dev/null +++ b/tests/test_swe_sft_dataset.py @@ -0,0 +1,262 @@ +"""Unit tests for SWE SFT thinking-mode classification.""" + +import importlib.util +import logging +import re +import sys +import types +from pathlib import Path + +import pytest + + +def _load_swe_modules(): + areal_module = types.ModuleType("areal") + dataset_module = types.ModuleType("areal.dataset") + dataset_module.__path__ = [] + swe_package = types.ModuleType("areal.dataset.swe_sft") + swe_package.__path__ = [] + utils_module = types.ModuleType("areal.utils") + utils_module.logging = logging + areal_module.utils = utils_module + sys.modules.setdefault("areal", areal_module) + sys.modules.setdefault("areal.dataset", dataset_module) + sys.modules.setdefault("areal.dataset.swe_sft", swe_package) + sys.modules.setdefault("areal.utils", utils_module) + + package_path = Path(__file__).parents[1] / "areal" / "dataset" / "swe_sft" + loaded = [] + for name in ("messages", "tokenization"): + full_name = f"areal.dataset.swe_sft.{name}" + spec = importlib.util.spec_from_file_location( + full_name, package_path / f"{name}.py" + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[full_name] = module + spec.loader.exec_module(module) + loaded.append(module) + return loaded + + +messages, tokenization = _load_swe_modules() +_classify_pair = messages._classify_pair +_clean_message = messages._clean_message +_msg_has_thinking = messages._msg_has_thinking +_add_bailing_v3_generation_tags = tokenization._add_bailing_v3_generation_tags +_render_tokenize_mask = tokenization._render_tokenize_mask + + +class _AdaptiveTokenizer: + chat_template = "enable_thinking preserved_thinking" + + def __init__(self): + self.enable_thinking = None + + def apply_chat_template(self, messages, *, tokenize, **kwargs): + assert tokenize is False + self.enable_thinking = kwargs.get("enable_thinking") + rendered = [] + for message in messages: + content = message.get("content") or "" + if message.get("role") == "assistant": + rendered.append(f"ASSISTANT{content}<|role_end|>") + else: + rendered.append(content) + return "".join(rendered) + + def __call__(self, text, *, add_special_tokens, return_offsets_mapping): + assert add_special_tokens is False + assert return_offsets_mapping is True + return { + "input_ids": list(range(len(text))), + "offset_mapping": [(idx, idx + 1) for idx in range(len(text))], + } + + +class _NativeMaskTokenizer: + chat_template = "{% generation %}" + + @staticmethod + def _render(messages): + text = "" + mask = [] + for message in messages: + content = message.get("content") or "" + if message["role"] == "assistant": + header = "ASSISTANT" + body = content + "<|role_end|>" + text += header + body + mask.extend([0] * len(header) + [1] * len(body)) + else: + body = content + "<|role_end|>" + text += body + mask.extend([0] * len(body)) + return text, mask + + def apply_chat_template(self, messages, *, tokenize, **kwargs): + text, mask = self._render(messages) + if not tokenize: + return text + assert kwargs["return_dict"] is True + assert kwargs["return_assistant_tokens_mask"] is True + return {"input_ids": list(range(len(text))), "assistant_masks": mask} + + def __call__(self, text, *, add_special_tokens, return_offsets_mapping): + return { + "input_ids": list(range(len(text))), + "offset_mapping": [(idx, idx + 1) for idx in range(len(text))], + } + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + ({"role": "user", "content": "reasoning"}, False), + ({"role": "assistant", "content": "answer"}, False), + ({"role": "assistant", "content": "answer"}, False), + ({"role": "assistant", "content": " \n "}, False), + ( + { + "role": "assistant", + "content": "reasoning", + }, + True, + ), + ({"role": "assistant", "content": "", "reasoning_content": " \n "}, False), + ( + { + "role": "assistant", + "content": "", + "reasoning_content": "reasoning", + }, + True, + ), + ], +) +def test_msg_has_thinking_requires_non_empty_reasoning(message, expected): + # Act + result = _msg_has_thinking(message) + + # Assert + assert result is expected + + +@pytest.mark.parametrize( + ("target", "expected"), + [ + ( + { + "role": "assistant", + "content": "\n", + "tool_calls": [{"type": "function"}], + }, + "no_thinking_tool_call", + ), + ({"role": "assistant", "content": "summary"}, "pure_text"), + ( + {"role": "assistant", "content": "reasoninganswer"}, + "thinking", + ), + ], +) +def test_classify_pair_uses_non_empty_thinking_semantics(target, expected): + # Arrange + pair = [{"role": "user", "content": "task"}, target] + + # Act + result = _classify_pair(pair) + + # Assert + assert result == expected + + +def test_render_tokenize_mask_sets_adaptive_mode_per_trajectory(): + # Arrange + assistant_pattern = re.compile( + r"ASSISTANT(.*?)<\|role_end\|>", re.DOTALL + ) + empty_thinking = _clean_message( + {"role": "assistant", "content": "answer"}, + strip_thinking=False, + ensure_thinking=True, + ) + real_thinking = _clean_message( + {"role": "assistant", "content": "reasoninganswer"}, + strip_thinking=False, + ensure_thinking=True, + ) + + # Act + no_thinking_tokenizer = _AdaptiveTokenizer() + _render_tokenize_mask( + [{"role": "user", "content": "task"}, empty_thinking], + no_thinking_tokenizer, + assistant_pattern, + split_mode="trajectory", + ) + mixed_tokenizer = _AdaptiveTokenizer() + _render_tokenize_mask( + [ + {"role": "user", "content": "task"}, + empty_thinking, + real_thinking, + ], + mixed_tokenizer, + assistant_pattern, + split_mode="trajectory", + ) + + # Assert + assert no_thinking_tokenizer.enable_thinking is False + assert mixed_tokenizer.enable_thinking is True + + +def test_native_mask_ignores_literal_assistant_delimiter_and_masks_errors(): + # Arrange + injected = "quoted ASSISTANT text" + messages = [ + {"role": "user", "content": injected}, + {"role": "assistant", "content": "bad"}, + {"role": "assistant", "content": "good"}, + ] + + # Act + result = _render_tokenize_mask( + messages, + _NativeMaskTokenizer(), + re.compile(r"this fallback must not be used"), + split_mode="trajectory", + error_indices=[0], + ) + + # Assert + full_text, _, loss_mask, _ = result + injected_start = full_text.index(injected) + bad_start = full_text.index("bad") + good_start = full_text.index("good") + assert not any(loss_mask[injected_start : injected_start + len(injected)]) + assert not any(loss_mask[bad_start : bad_start + len("bad")]) + assert all(loss_mask[good_start : good_start + len("good")]) + + +def test_add_bailing_v3_generation_tags_keeps_header_outside_mask(): + # Arrange + template = r"""{% set preserved_thinking = true %} +{%- if preserved_thinking or loop.index0 > ns.last_query_index %}{% endif %} +{%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {{- 'ASSISTANT' + 'thinking' }} + {{- 'ASSISTANT\n' + 'empty' }} + {{- 'ASSISTANT\n' + 'content' }} + {{- '<|role_end|>' }} + {%- elif message.role == "tool" %}""" + + # Act + patched = _add_bailing_v3_generation_tags(template) + + # Assert + assert patched is not None + assert "{{- 'ASSISTANT' }}\n {%- generation %}" in patched + assert "{{- '<|role_end|>' }}\n {%- endgeneration %}" in patched + assert patched.count("ASSISTANT") == 1 From a29fc3c45832c5086be035349f20d59c7e13ff3f Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Mon, 17 Aug 2026 12:14:52 +0800 Subject: [PATCH 4/8] test(dataset): restore modules after SWE loader tests Prevent collection-time stubs from replacing the real areal.dataset package for subsequent data-service tests. --- tests/test_swe_sft_dataset.py | 35 +++++++++++++++++++++++++---------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/tests/test_swe_sft_dataset.py b/tests/test_swe_sft_dataset.py index e47fe3cd47..6caa7115da 100644 --- a/tests/test_swe_sft_dataset.py +++ b/tests/test_swe_sft_dataset.py @@ -11,6 +11,15 @@ def _load_swe_modules(): + saved_modules = { + name: module + for name, module in sys.modules.items() + if name == "areal" or name.startswith("areal.") + } + for name in list(sys.modules): + if name == "areal" or name.startswith("areal."): + del sys.modules[name] + areal_module = types.ModuleType("areal") dataset_module = types.ModuleType("areal.dataset") dataset_module.__path__ = [] @@ -26,16 +35,22 @@ def _load_swe_modules(): package_path = Path(__file__).parents[1] / "areal" / "dataset" / "swe_sft" loaded = [] - for name in ("messages", "tokenization"): - full_name = f"areal.dataset.swe_sft.{name}" - spec = importlib.util.spec_from_file_location( - full_name, package_path / f"{name}.py" - ) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[full_name] = module - spec.loader.exec_module(module) - loaded.append(module) + try: + for name in ("messages", "tokenization"): + full_name = f"areal.dataset.swe_sft.{name}" + spec = importlib.util.spec_from_file_location( + full_name, package_path / f"{name}.py" + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[full_name] = module + spec.loader.exec_module(module) + loaded.append(module) + finally: + for name in list(sys.modules): + if name == "areal" or name.startswith("areal."): + del sys.modules[name] + sys.modules.update(saved_modules) return loaded From 6a1af814017efdc9bf08e77e530f2a361fd0eb2c Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Mon, 17 Aug 2026 20:32:14 +0800 Subject: [PATCH 5/8] refactor(engine): remove precision dump hooks Keep Bailing V3 support focused on production training behavior by removing out-of-scope routing and log-probability dump paths. --- areal/engine/megatron_engine.py | 158 +------------------------------- 1 file changed, 5 insertions(+), 153 deletions(-) diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 7cf1e5d6b6..a5efb56fdc 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -509,14 +509,6 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec, *args, **kwargs): if self.mcore_config.use_deterministic_algorithms: set_deterministic_algorithms(self.tf_config, prebuild=True) - # Precision-alignment dumps (AReaL-friend tools/precision-alignment): - # when AREAL_DUMP_ROUTING is set, enable megatron RouterReplay - # recording so MoE expert indices can be captured during forward. - if os.environ.get("AREAL_DUMP_ROUTING", "") and hasattr( - self.tf_config, "moe_enable_routing_replay" - ): - self.tf_config.moe_enable_routing_replay = True - self.is_vision_model = is_valid_vision_model(self.hf_config.model_type) # GDN/SSM models (e.g. Qwen3.5) reject packed THD input and must run # the padded BSHD forward. Derived from model type rather than a @@ -1134,11 +1126,6 @@ def optimizer_step(self): ) def lr_scheduler_step(self): - if os.environ.get("AREAL_DUMP_ROUTING", "") or os.environ.get( - "AREAL_DUMP_LOGP", "" - ): - # Precision-alignment forward-only mode: no optimizer/scheduler. - return assert self.lr_scheduler is not None, "LR Scheduler is not initialized." self.lr_scheduler.step(1) @@ -1212,23 +1199,6 @@ def forward_step(batch_iter, model): "BSHD is supported only for text-only models such as Qwen3.5" ) - # Precision-alignment routing dump: record MoE expert indices for the - # first microbatch via megatron RouterReplay (enabled in initialize). - _routing_dump_path = os.environ.get("AREAL_DUMP_ROUTING", "") - if _routing_dump_path and not getattr(self, "_routing_dumped", False): - try: - from megatron.core.transformer.moe.router_replay import ( - RouterReplay, - RouterReplayAction, - ) - - RouterReplay.set_global_router_replay_action( - RouterReplayAction.RECORD - ) - except Exception as e: - self.logger.warning(f"[ROUTING-DUMP] RECORD setup failed: {e}") - _routing_dump_path = "" - output = packed_context_parallel_forward( model, mb_input.padded_mb, @@ -1312,76 +1282,6 @@ def forward_step(batch_iter, model): ), ) - if _routing_dump_path and not getattr(self, "_routing_dumped", False): - try: - from megatron.core.transformer.moe.router_replay import ( - RouterReplay, - ) - - recorded = RouterReplay.get_recorded_data() - if recorded and any(r is not None for r in recorded): - pp_rank = mpu.get_pipeline_model_parallel_rank() - cp_rank = mpu.get_context_parallel_rank() - if ( - mpu.get_data_parallel_rank() == 0 - and mpu.get_tensor_model_parallel_rank() == 0 - ): - expert_indices = [ - r.detach().cpu() for r in recorded if r is not None - ] - # The router sees the CP-local token shard (zigzag - # split by packed_context_parallel_forward) and, - # with TP>1, the SP-local sub-shard of that. - # Apply the same CP split here; tp_rank 0 then - # holds the first contiguous SP chunk, so the - # row-count prefix slice below stays valid. - ids_src = mb_input.padded_mb["input_ids"] - cu = mb_input.padded_mb.get("cu_seqlens") - pos_src = torch.arange( - ids_src.numel(), device=ids_src.device - ) - if cp_size > 1 and cu is not None: - ids_src = split_packed_seqs_for_context_parallel( - ids_src, cu - ) - pos_src = split_packed_seqs_for_context_parallel( - pos_src, cu - ) - ids = ids_src.detach().cpu().reshape(-1) - pos = pos_src.detach().cpu().reshape(-1) - n_rows = expert_indices[0].shape[0] - save_data = { - "expert_indices": expert_indices, - "input_ids": ids[:n_rows], - # Canonical (padded packed-sequence) position - # of each recorded row so the merge tool can - # restore zigzag CP order; without it CP>1 - # shards cannot be mapped back. - "positions": pos[:n_rows], - "cp_size": cp_size, - "pp_rank": pp_rank, - "cp_rank": cp_rank, - } - if cu is not None: - save_data["padded_cu_seqlens"] = cu.detach().cpu() - orig_cu = mb_input.orig_mb.get("cu_seqlens") - if orig_cu is not None: - save_data["orig_cu_seqlens"] = orig_cu.detach().cpu() - out_file = ( - f"{_routing_dump_path}.pp{pp_rank}.cp{cp_rank}.pt" - ) - torch.save(save_data, out_file) - self.logger.info( - f"[ROUTING-DUMP] saved " - f"{len(save_data['expert_indices'])} MoE layers, " - f"pp={pp_rank} cp={cp_rank} -> {out_file}" - ) - RouterReplay.clear_global_indices() - RouterReplay.clear_global_router_replay_action() - self._routing_dumped = True - except Exception as e: - self.logger.warning(f"[ROUTING-DUMP] failed: {e}") - # Release tree attention metadata after forward pass for key in tree_attn_keys: del mb_input.padded_mb[key] @@ -1457,14 +1357,7 @@ def train_batch( if self._awex_adapter is not None: self._awex_adapter.ensure_grad_buffers() - # Precision-alignment forward-only mode: no optimizer exists (see - # _create_optimizer), so skip zero_grad/step and run forward only. - _fwd_only = bool( - os.environ.get("AREAL_DUMP_ROUTING", "") - or os.environ.get("AREAL_DUMP_LOGP", "") - ) - if not _fwd_only: - self.optimizer_zero_grad() + self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -1490,13 +1383,9 @@ def train_batch( # that extra division would shrink every gradient (and thus grad_norm and the # effective optimizer step) by `num_microbatches`. loss_multiplier = ( - float(mpu.get_data_parallel_world_size()) - if _fwd_only - else ( - mpu.get_data_parallel_world_size() - * self.optimizer.get_loss_scale().item() - * len(mb_list) - ) + mpu.get_data_parallel_world_size() + * self.optimizer.get_loss_scale().item() + * len(mb_list) ) def process_output( @@ -1514,12 +1403,9 @@ def process_output( self.forward_backward_batch( mb_list, process_output, - forward_only=_fwd_only, + forward_only=False, ) - if _fwd_only: - return {"num_micro_batches": len(mb_list.mbs)} - # Step 4: Optimizer step stats = self.optimizer_step() stats["num_micro_batches"] = len(mb_list.mbs) @@ -1936,14 +1822,6 @@ def _init_context_and_model_parallel_group(self) -> None: def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: if self.optimizer_config is None: return - if os.environ.get("AREAL_DUMP_ROUTING", "") or os.environ.get( - "AREAL_DUMP_LOGP", "" - ): - self.logger.info( - "[MegatronEngine] AREAL_DUMP_ROUTING/LOGP set, skipping optimizer " - "creation to save GPU memory (precision-alignment forward-only mode)." - ) - return assert self.model is not None and len(self.model) > 0 use_distributed_optimizer = ( @@ -3043,32 +2921,6 @@ def _compute_logprobs_and_loss( k: v for k, v in inputs.items() if not k.startswith("_cp_") } - # Precision-alignment logp dump: save final per-token logprobs for - # the first microbatch (last PP stage only; this branch already is). - _logp_dump_path = os.environ.get("AREAL_DUMP_LOGP", "") - if _logp_dump_path and not getattr(self, "_logp_dumped", False): - pp_rank = mpu.get_pipeline_model_parallel_rank() - cp_rank = mpu.get_context_parallel_rank() - if ( - mpu.get_data_parallel_rank() == 0 - and mpu.get_tensor_model_parallel_rank() == 0 - ): - save_data = { - "logprobs": logprobs.detach().cpu(), - "input_ids": inputs["input_ids"].detach().cpu(), - "pp_rank": pp_rank, - "cp_rank": cp_rank, - } - if "loss_mask" in inputs: - save_data["loss_mask"] = inputs["loss_mask"].detach().cpu() - out_file = f"{_logp_dump_path}.pp{pp_rank}.cp{cp_rank}.pt" - torch.save(save_data, out_file) - self.logger.info( - f"[LOGP-DUMP] pp={pp_rank} cp={cp_rank} " - f"logprobs={list(logprobs.shape)} -> {out_file}" - ) - self._logp_dumped = True - loss = loss_fn( logprobs, entropy, From 52678d617c97fb531c1499eef4b681fe85798643 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 18 Aug 2026 10:53:35 +0800 Subject: [PATCH 6/8] test(models): run zigzag coverage in CI Place the CP zigzag unit tests under the root test pattern used by the GCP unit-test workflow. --- tests/{models => }/test_zigzag_indices.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{models => }/test_zigzag_indices.py (100%) diff --git a/tests/models/test_zigzag_indices.py b/tests/test_zigzag_indices.py similarity index 100% rename from tests/models/test_zigzag_indices.py rename to tests/test_zigzag_indices.py From 2f0aa5cea506993a76d05288692f34e7a2971213 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Fri, 21 Aug 2026 16:52:56 +0800 Subject: [PATCH 7/8] fix: preserve Bailing V3 HF export metadata Keep runtime model metadata valid across fast, native, direct, and in-place mbridge exports while retaining the production local-source fallback. Key changes: - snapshot and validate HF config metadata before exporters overwrite it - preserve source assets and support native mbridge finalization - forward SWE preprocessing kwargs through the remote dataset path - add config round-trip, Saver, and controller regression coverage Refs: #1598 Constraint: Preserve swe-dev Bailing V3 architecture-based bridge dispatch Confidence: high Scope-risk: moderate Not-tested: Real multi-rank Bailing V3 HF save/load canary --- areal/dataset/__init__.py | 4 +- areal/engine/megatron_engine.py | 19 +- areal/models/mcore/hf_save.py | 93 ++------- areal/utils/hf_utils.py | 177 ++++++++++++++++++ areal/utils/saver.py | 14 ++ .../2026-08-21-bailing-v3-hf-export-design.md | 39 ++++ .../infra/data_service/test_trainer_compat.py | 35 ++++ tests/test_hf_config.py | 84 +++++++++ tests/test_megatron_engine.py | 46 +++++ tests/test_saver.py | 90 +++++++++ 10 files changed, 521 insertions(+), 80 deletions(-) create mode 100644 docs/plans/2026-08-21-bailing-v3-hf-export-design.md create mode 100644 tests/test_hf_config.py create mode 100644 tests/test_saver.py diff --git a/areal/dataset/__init__.py b/areal/dataset/__init__.py index 150676c400..2db3520624 100644 --- a/areal/dataset/__init__.py +++ b/areal/dataset/__init__.py @@ -196,12 +196,14 @@ def get_custom_dataset( ): from areal.infra.data_service.rdataset import RDataset + dataset_kwargs = dict(getattr(dataset_config, "dataset_kwargs", None) or {}) + dataset_kwargs.update(kwargs) return RDataset( path=dataset_config.path, type=dataset_config.type, split=split, max_length=dataset_config.max_length, - dataset_kwargs=getattr(dataset_config, "dataset_kwargs", None), + dataset_kwargs=dataset_kwargs, ) if dataset_config is not None: diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index a5efb56fdc..c46f0fcdbd 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -132,7 +132,12 @@ unpad_logits, ) from areal.utils.functional import gather_logprobs, gather_logprobs_entropy -from areal.utils.hf_utils import load_hf_processor_and_tokenizer, load_hf_tokenizer +from areal.utils.hf_utils import ( + finalize_hf_export, + load_hf_config_snapshot, + load_hf_processor_and_tokenizer, + load_hf_tokenizer, +) from areal.utils.lock import DistributedLock from areal.utils.network import find_free_ports, format_host_for_url, gethostip from areal.utils.offload import is_tms_enabled, torch_memory_saver @@ -2521,6 +2526,11 @@ def _save_model_to_hf( ) else: if self.mcore_config.use_mbridge_save: + source_config = ( + load_hf_config_snapshot(base_model_path) + if dist.get_rank() == 0 + else None + ) # when loading model using AreaL's fast hf load, the safetensor_io is never set if ( not hasattr(self.bridge, "safetensor_io") @@ -2530,6 +2540,13 @@ def _save_model_to_hf( self.config.path ) self.bridge.save_weights(models=self.model, weights_path=path) + if dist.get_rank() == 0: + finalize_hf_export( + self.bridge.hf_config, + path, + source_model_path=base_model_path, + source_config=source_config, + ) else: save_weights_to_hf_with_mbridge_fast( bridge=self.bridge, diff --git a/areal/models/mcore/hf_save.py b/areal/models/mcore/hf_save.py index f2423569e4..194159864f 100644 --- a/areal/models/mcore/hf_save.py +++ b/areal/models/mcore/hf_save.py @@ -3,7 +3,6 @@ import json import os import re -import shutil from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass @@ -26,17 +25,19 @@ from areal.infra.platforms import current_platform from areal.models.mcore.registry import unwrap_to_gpt_model from areal.utils import logging +from areal.utils.hf_utils import ( + HF_MODEL_ASSET_FILES, + copy_hf_model_assets, + finalize_hf_export, +) logger = logging.getLogger("HFSaver") +HF_MODEL_CONFIG_FILES = list(HF_MODEL_ASSET_FILES) + -HF_MODEL_CONFIG_FILES = [ - "generation_config.json", - "tokenizer_config.json", - "vocab.json", - "merges.txt", - "special_tokens_map.json", - "tokenizer.json", -] +def copy_hf_configs(src_model_dir, dst_model_dir): + """Backward-compatible wrapper for copying Hugging Face model assets.""" + copy_hf_model_assets(src_model_dir, dst_model_dir) def _maybe_convert_from_te_fp8_params( @@ -116,71 +117,6 @@ def _maybe_convert_to_torch_fp8_params( return converted_names, converted_params -def copy_hf_configs(src_model_dir, dst_model_dir): - for file in HF_MODEL_CONFIG_FILES: - try: - shutil.copy( - os.path.join(src_model_dir, file), - os.path.join(dst_model_dir, file), - ) - logger.info(f"copied {file} from {src_model_dir} to {dst_model_dir}") - except FileNotFoundError: - logger.info(f"{file} not exist in {src_model_dir} skipping.") - # Copy remote codes and chat template files - for file in os.listdir(src_model_dir): - copy = False - for prefix in ["chat_format", "configuration_", "modeling_", "tokenization_"]: - if file.startswith(prefix) and file.endswith(".py"): - copy = True - break - # Chat template files (e.g. chat_template.jinja) - if file.startswith("chat_template"): - copy = True - if copy: - shutil.copy( - os.path.join(src_model_dir, file), - os.path.join(dst_model_dir, file), - ) - logger.info(f"copied {file} from {src_model_dir} to {dst_model_dir}") - - -def _patch_saved_config(base_model_path, saved_path): - """Patch saved config.json to preserve model_type and torch_dtype. - - Some HF config classes lack a ``model_type`` class attribute, causing - ``save_pretrained()`` to lose the field (``PretrainedConfig.to_dict()`` - reads the class attribute, not the instance value). This restores - critical fields from the original model's config.json. - """ - orig_config_path = os.path.join(base_model_path, "config.json") - saved_config_path = os.path.join(saved_path, "config.json") - - if not os.path.exists(orig_config_path) or not os.path.exists(saved_config_path): - return - - with open(orig_config_path) as f: - orig_config = json.load(f) - with open(saved_config_path) as f: - saved_config = json.load(f) - - patched_fields = [] - - # Restore model_type if missing or null - if not saved_config.get("model_type") and orig_config.get("model_type"): - saved_config["model_type"] = orig_config["model_type"] - patched_fields.append(f"model_type={orig_config['model_type']}") - - # Restore torch_dtype if missing - if "torch_dtype" not in saved_config and "torch_dtype" in orig_config: - saved_config["torch_dtype"] = orig_config["torch_dtype"] - patched_fields.append(f"torch_dtype={orig_config['torch_dtype']}") - - if patched_fields: - with open(saved_config_path, "w") as f: - json.dump(saved_config, f, indent=2) - logger.info(f"Patched config.json: {', '.join(patched_fields)}") - - def _bridge_uses_stacked_experts(bridge: Bridge) -> bool: """Detect bridges whose HF format keeps experts grouped under a single stacked tensor (e.g., Qwen3-VL-MoE ``mlp.experts.gate_up_proj`` shape @@ -840,12 +776,13 @@ def _save_one_shard(x): # 7. save metadata if dist.get_rank() == 0: - bridge.hf_config.save_pretrained(weights_path) with open(os.path.join(weights_path, "model.safetensors.index.json"), "w") as f: json.dump(bin_index, f, indent=4) - if base_model_path is not None: - copy_hf_configs(base_model_path, weights_path) - _patch_saved_config(base_model_path, weights_path) + finalize_hf_export( + bridge.hf_config, + weights_path, + source_model_path=base_model_path, + ) def save_critic_value_head(models, weights_path): diff --git a/areal/utils/hf_utils.py b/areal/utils/hf_utils.py index 6d98c62f55..a1c9a045c7 100644 --- a/areal/utils/hf_utils.py +++ b/areal/utils/hf_utils.py @@ -2,7 +2,13 @@ from __future__ import annotations +import json +import os +import shutil +import tempfile +from collections.abc import Mapping from functools import lru_cache +from pathlib import Path from typing import Any, Literal, overload import transformers @@ -12,6 +18,177 @@ logger = logging.getLogger("HFUtils") +HF_MODEL_ASSET_FILES = ( + "generation_config.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "special_tokens_map.json", + "tokenizer.json", +) + + +def copy_hf_model_assets( + source_model_path: str | os.PathLike[str], + save_directory: str | os.PathLike[str], +) -> None: + """Copy tokenizer, generation, chat-template, and remote-code assets.""" + source_model_path = os.fspath(source_model_path) + save_directory = os.fspath(save_directory) + for filename in HF_MODEL_ASSET_FILES: + source = os.path.join(source_model_path, filename) + destination = os.path.join(save_directory, filename) + try: + shutil.copy(source, destination) + logger.info( + "Copied %s from %s to %s", filename, source_model_path, save_directory + ) + except FileNotFoundError: + logger.info( + "%s does not exist in %s; skipping", filename, source_model_path + ) + + for filename in os.listdir(source_model_path): + is_remote_code = filename.endswith(".py") and filename.startswith( + ("chat_format", "configuration_", "modeling_", "tokenization_") + ) + if is_remote_code or filename.startswith("chat_template"): + shutil.copy( + os.path.join(source_model_path, filename), + os.path.join(save_directory, filename), + ) + logger.info( + "Copied %s from %s to %s", + filename, + source_model_path, + save_directory, + ) + + +def save_hf_config( + config: transformers.PretrainedConfig, + save_directory: str | os.PathLike[str], + *, + source_model_path: str | os.PathLike[str] | None = None, + source_config: Mapping[str, Any] | None = None, +) -> None: + """Save a Hugging Face config without losing its runtime ``model_type``. + + ``PretrainedConfig.to_dict()`` serializes the class-level ``model_type``. + Some remote-code configs instead receive ``model_type`` as an instance field, + so a normal ``save_pretrained()`` call replaces the valid runtime value with an + empty string. Save first, then restore and validate that value in the emitted + ``config.json``. A local source config may also supply fields that Transformers + omits while re-serializing a generic config. + """ + save_directory = os.fspath(save_directory) + source_config = ( + load_hf_config_snapshot(source_model_path) + if source_config is None + else dict(source_config) + ) + + runtime_model_type = getattr(config, "model_type", None) + expected_model_type = runtime_model_type or source_config.get("model_type") + if expected_model_type is not None and not isinstance(expected_model_type, str): + raise TypeError( + "Hugging Face config model_type must be a string, got " + f"{type(expected_model_type).__name__}." + ) + + # Snapshot source fields before save_pretrained() because source and + # destination may be the same directory. + config.save_pretrained(save_directory) + + saved_config_path = Path(save_directory) / "config.json" + with saved_config_path.open() as f: + saved_config = json.load(f) + + patched_fields: list[str] = [] + if expected_model_type and saved_config.get("model_type") != expected_model_type: + saved_config["model_type"] = expected_model_type + patched_fields.append(f"model_type={expected_model_type}") + + if "torch_dtype" not in saved_config and "torch_dtype" in source_config: + saved_config["torch_dtype"] = source_config["torch_dtype"] + patched_fields.append(f"torch_dtype={source_config['torch_dtype']}") + + if patched_fields: + with tempfile.NamedTemporaryFile( + mode="w", + dir=save_directory, + prefix=".config.json.", + suffix=".tmp", + delete=False, + ) as tmp: + json.dump(saved_config, tmp, indent=2) + tmp.write("\n") + tmp_path = tmp.name + os.replace(tmp_path, saved_config_path) + logger.info("Patched config.json: %s", ", ".join(patched_fields)) + + with saved_config_path.open() as f: + persisted_config = json.load(f) + if ( + expected_model_type + and persisted_config.get("model_type") != expected_model_type + ): + raise RuntimeError( + "Saved Hugging Face config has an invalid model_type: expected " + f"{expected_model_type!r}, got {persisted_config.get('model_type')!r}." + ) + + +def finalize_hf_export( + config: transformers.PretrainedConfig, + save_directory: str | os.PathLike[str], + *, + source_model_path: str | os.PathLike[str] | None = None, + source_config: Mapping[str, Any] | None = None, +) -> None: + """Finalize an HF export with source assets and a validated config.""" + save_directory = os.fspath(save_directory) + os.makedirs(save_directory, exist_ok=True) + local_source_path: str | None = None + if source_model_path is not None: + candidate = os.fspath(source_model_path) + if not os.path.isdir(candidate): + logger.warning( + "Cannot copy source Hugging Face assets: model path is not a " + "local directory: %s", + candidate, + ) + else: + local_source_path = candidate + if os.path.samefile(candidate, save_directory): + logger.warning( + "Skipping source Hugging Face asset copy because the source " + "and checkpoint directories are the same: %s", + save_directory, + ) + else: + copy_hf_model_assets(candidate, save_directory) + + save_hf_config( + config, + save_directory, + source_model_path=local_source_path, + source_config=source_config, + ) + + +def load_hf_config_snapshot( + source_model_path: str | os.PathLike[str] | None, +) -> dict[str, Any]: + """Read a source config before an exporter can overwrite it in place.""" + if source_model_path is None: + return {} + source_config_path = Path(source_model_path) / "config.json" + if not source_config_path.is_file(): + return {} + with source_config_path.open() as f: + return json.load(f) + @overload def apply_chat_template( diff --git a/areal/utils/saver.py b/areal/utils/saver.py index d6d2b02d67..f03c2a719d 100644 --- a/areal/utils/saver.py +++ b/areal/utils/saver.py @@ -153,6 +153,20 @@ def save( if self._should_use_async(engine): self._async_save(engine, path, name, tokenizer, processor) else: + if not base_model_path: + # Megatron's HF exporter needs the source directory to retain + # remote model code and patch critical config fields. Infer it + # here so every periodic trainer save gets the same behavior. + engine_config = getattr(engine, "config", None) + configured_model_path = getattr(engine_config, "path", None) + if configured_model_path and os.path.isdir(configured_model_path): + base_model_path = os.fspath(configured_model_path) + elif configured_model_path: + logger.warning( + "Cannot copy source HuggingFace assets for checkpoint save: " + "engine model path is not a local directory: %s", + configured_model_path, + ) meta = SaveLoadMeta( path=path, weight_format="hf", diff --git a/docs/plans/2026-08-21-bailing-v3-hf-export-design.md b/docs/plans/2026-08-21-bailing-v3-hf-export-design.md new file mode 100644 index 0000000000..2900f5b43b --- /dev/null +++ b/docs/plans/2026-08-21-bailing-v3-hf-export-design.md @@ -0,0 +1,39 @@ +# Bailing V3 HF Export Hardening + +## Context + +The Bailing V3 bridge must dispatch by `architectures` because V3 and Bailing V2.5 +share `model_type="bailing_hybrid"`. The production `swe-dev` implementation therefore +loads V3 through a generic `PretrainedConfig` and constructs `BailingV3Bridge` +explicitly. That behavior is retained: changing registration or relying only on +`AutoConfig` would risk routing V2.5 checkpoints through the V3 bridge. + +The generic config keeps the source `model_type` on the instance, but Transformers +serializes the class-level value. For config classes without their own class-level +`model_type`, `save_pretrained()` can therefore write an empty or missing value. +Internal `swe-dev` later mitigated normal periodic saves by forwarding the local source +checkpoint directory, but direct saves, disk weight updates, and mbridge's native saver +still bypass that mitigation. + +## Design + +Keep model loading and weight conversion unchanged. After every mbridge HF config save, +run one shared finalization step that restores the original instance `model_type` in the +written `config.json`. When a local `base_model_path` is available, the same step also +copies tokenizer, generation, chat-template, and remote-code assets and preserves source +config fields. Port the production `Saver` fallback that derives `base_model_path` from +`engine.config.path`; this remains useful for source assets but is no longer required for +`model_type` correctness. + +The SWE entrypoint also passes preprocessing controls as call-time keyword arguments. +The single-controller `RDataset` branch will merge those arguments with +`dataset_config.dataset_kwargs`, giving explicit call-time arguments precedence, so its +behavior matches the direct SPMD loader. + +## Verification + +Add CPU tests for a remote-style config class without a class-level `model_type`, verify +the exported JSON, and reload it through `AutoConfig`. Cover the production Saver +fallback and both mbridge export finalization paths without requiring GPUs. Add a +single-controller test proving SWE preprocessing options reach `RDataset`. Existing +Bailing V3 KDA, MLA, MoE, and weight-layout code remains untouched. diff --git a/tests/infra/data_service/test_trainer_compat.py b/tests/infra/data_service/test_trainer_compat.py index 5a895f5809..f16b004307 100644 --- a/tests/infra/data_service/test_trainer_compat.py +++ b/tests/infra/data_service/test_trainer_compat.py @@ -161,6 +161,41 @@ def _fake_custom_dataset(**_kwargs): assert dataset is sentinel + def test_get_custom_dataset_forwards_kwargs_to_remote_dataset(self, monkeypatch): + from areal.api.cli_args import TrainDatasetConfig + from areal.dataset import get_custom_dataset + + monkeypatch.setenv("AREAL_SPMD_MODE", "0") + configured_kwargs = { + "num_proc": 4, + "filter_errors": False, + "configured_only": "preserved", + } + cfg = TrainDatasetConfig( + path="swe-data.jsonl", + type="sft", + dataset_kwargs=configured_kwargs, + ) + + dataset = get_custom_dataset( + split="train", + dataset_config=cfg, + split_mode="trajectory", + filter_errors=True, + random_strip_thinking_prob=0.5, + cache_dir="/tmp/swe-cache", + ) + + assert dataset._dataset_kwargs == { + "num_proc": 4, + "filter_errors": True, + "configured_only": "preserved", + "split_mode": "trajectory", + "random_strip_thinking_prob": 0.5, + "cache_dir": "/tmp/swe-cache", + } + assert cfg.dataset_kwargs == configured_kwargs + class TestGenericDatasetFallback: def test_none_split_uses_first_available_split(self, tmp_path: Path): diff --git a/tests/test_hf_config.py b/tests/test_hf_config.py new file mode 100644 index 0000000000..10c8b43943 --- /dev/null +++ b/tests/test_hf_config.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json + +from transformers import AutoConfig, PretrainedConfig + +from areal.utils.hf_utils import finalize_hf_export, save_hf_config + + +class RuntimeModelTypeConfig(PretrainedConfig): + """Config that deliberately inherits the empty class-level model_type.""" + + +def test_save_hf_config_preserves_instance_model_type_without_source(tmp_path): + """Runtime model_type survives even when no local base model is available.""" + config = RuntimeModelTypeConfig(architectures=["RuntimeModel"]) + config.model_type = "runtime_only" + + save_path = tmp_path / "saved" + save_hf_config(config, save_path) + + with (save_path / "config.json").open() as f: + saved_config = json.load(f) + assert saved_config["model_type"] == "runtime_only" + assert saved_config["architectures"] == ["RuntimeModel"] + + +def test_save_hf_config_supports_remote_config_round_trip(tmp_path): + """A config with no class model_type remains loadable through AutoConfig.""" + source_path = tmp_path / "source" + source_path.mkdir() + module_name = "configuration_runtime_only.py" + (source_path / module_name).write_text( + "from transformers import PretrainedConfig\n\n" + "class RuntimeOnlyConfig(PretrainedConfig):\n" + " pass\n" + ) + with (source_path / "config.json").open("w") as f: + json.dump( + { + "architectures": ["RuntimeOnlyModel"], + "auto_map": { + "AutoConfig": "configuration_runtime_only.RuntimeOnlyConfig" + }, + "model_type": "runtime_only", + "torch_dtype": "bfloat16", + }, + f, + ) + + config = PretrainedConfig.from_pretrained(source_path) + save_path = tmp_path / "saved" + (source_path / "generation_config.json").write_text('{"max_new_tokens": 1}') + (source_path / "chat_template.jinja").write_text("{{ messages }}") + finalize_hf_export(config, save_path, source_model_path=source_path) + + loaded = AutoConfig.from_pretrained(save_path, trust_remote_code=True) + + assert type(loaded).__name__ == "RuntimeOnlyConfig" + assert loaded.model_type == "runtime_only" + assert (save_path / module_name).is_file() + assert (save_path / "generation_config.json").is_file() + assert (save_path / "chat_template.jinja").is_file() + with (save_path / "config.json").open() as f: + saved_config = json.load(f) + assert saved_config["torch_dtype"] == "bfloat16" + + +def test_finalize_hf_export_snapshots_same_directory_config(tmp_path): + """In-place finalization preserves fields before save_pretrained overwrites them.""" + config = RuntimeModelTypeConfig(architectures=["RuntimeModel"]) + config.model_type = "runtime_only" + with (tmp_path / "config.json").open("w") as f: + json.dump( + {"model_type": "runtime_only", "torch_dtype": "bfloat16"}, + f, + ) + + finalize_hf_export(config, tmp_path, source_model_path=tmp_path) + + with (tmp_path / "config.json").open() as f: + saved_config = json.load(f) + assert saved_config["model_type"] == "runtime_only" + assert saved_config["torch_dtype"] == "bfloat16" diff --git a/tests/test_megatron_engine.py b/tests/test_megatron_engine.py index 576c5f2375..948b250904 100644 --- a/tests/test_megatron_engine.py +++ b/tests/test_megatron_engine.py @@ -1,7 +1,10 @@ +import json import os import time from importlib.metadata import version as get_version +from types import SimpleNamespace from typing import Any +from unittest.mock import Mock import pytest import torch @@ -103,6 +106,46 @@ def test_mark_duplicated_params_clears_tp_metadata_for_replicated_params(): assert sharded_linear.weight.tensor_model_parallel +def test_native_mbridge_save_finalizes_hf_config(monkeypatch, tmp_path): + """The native mbridge saver runs the shared rank-zero HF finalizer.""" + engine = MegatronEngine.__new__(MegatronEngine) + engine.model = [object()] + engine.bridge_cls = "mbridge" + engine.mcore_config = SimpleNamespace(use_mbridge_save=True) + engine.config = SimpleNamespace(is_critic=False, path="/models/source") + engine.bridge = Mock(safetensor_io=object()) + engine.cpu_group = object() + + finalize = Mock() + source_config = {"model_type": "runtime_only", "torch_dtype": "bfloat16"} + snapshot = Mock(return_value=source_config) + monkeypatch.setattr( + "areal.engine.megatron_engine.finalize_hf_export", + finalize, + ) + monkeypatch.setattr( + "areal.engine.megatron_engine.load_hf_config_snapshot", + snapshot, + ) + monkeypatch.setattr(dist, "get_rank", lambda: 0) + monkeypatch.setattr(dist, "barrier", lambda **_kwargs: None) + monkeypatch.setattr(current_platform, "synchronize", lambda: None) + + engine._save_model_to_hf(str(tmp_path), base_model_path=str(tmp_path)) + + snapshot.assert_called_once_with(str(tmp_path)) + engine.bridge.save_weights.assert_called_once_with( + models=engine.model, + weights_path=str(tmp_path), + ) + finalize.assert_called_once_with( + engine.bridge.hf_config, + str(tmp_path), + source_model_path=str(tmp_path), + source_config=source_config, + ) + + # Cannot use a "module" scope since process groups can only be initialized once. @pytest.fixture def engine(): @@ -171,6 +214,9 @@ def test_hf_save_load_weights(tmp_path_factory, engine, mock_input): start = time.perf_counter() engine.save(save_load_meta) logger.info(f"Save done, time cost: {time.perf_counter() - start:.4f} seconds.") + with open(path / "config.json") as f: + saved_config = json.load(f) + assert saved_config["model_type"] == engine.hf_config.model_type for name, param in engine.model.named_parameters(): param.zero_() diff --git a/tests/test_saver.py b/tests/test_saver.py new file mode 100644 index 0000000000..60f1249e7a --- /dev/null +++ b/tests/test_saver.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import Mock + +from areal.api.cli_args import SaverConfig +from areal.api.io_struct import FinetuneSpec +from areal.utils.saver import Saver + + +def _make_saver(tmp_path) -> Saver: + config = SaverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=str(tmp_path), + freq_steps=1, + mode="sync", + ) + ft_spec = FinetuneSpec( + total_train_epochs=1, + dataset_size=1, + train_batch_size=1, + ) + saver = Saver(config, ft_spec) + saver._should_use_async = Mock(return_value=False) + return saver + + +def test_save_without_base_model_path_uses_engine_config_path(tmp_path): + """Periodic HF saves retain source assets when callers omit the base path.""" + saver = _make_saver(tmp_path) + source_path = tmp_path / "source-checkpoint" + source_path.mkdir() + engine = Mock() + engine.config = SimpleNamespace(path=str(source_path)) + + saver.save(engine, epoch=0, step=0, global_step=0) + + meta = engine.save.call_args.args[0] + assert meta.base_model_path == str(source_path) + + +def test_save_with_base_model_path_preserves_explicit_override(tmp_path): + """An explicit source path takes precedence over the engine configuration.""" + saver = _make_saver(tmp_path) + engine = Mock() + engine.config = SimpleNamespace(path="/models/source-checkpoint") + + saver.save( + engine, + epoch=0, + step=0, + global_step=0, + base_model_path="/models/explicit-checkpoint", + ) + + meta = engine.save.call_args.args[0] + assert meta.base_model_path == "/models/explicit-checkpoint" + + +def test_save_with_hub_model_path_does_not_treat_it_as_local_directory(tmp_path): + """Hub model IDs do not enter the local Hugging Face asset-copy path.""" + saver = _make_saver(tmp_path) + engine = Mock() + engine.config = SimpleNamespace(path="organization/model") + + saver.save(engine, epoch=0, step=0, global_step=0) + + meta = engine.save.call_args.args[0] + assert meta.base_model_path is None + + +def test_save_with_source_matching_destination_preserves_source_path(tmp_path): + """The exporter receives the source path so it can snapshot config fields.""" + saver = _make_saver(tmp_path) + save_path = Saver.get_model_save_path( + "test_exp", + "test_trial", + str(tmp_path), + epoch=0, + step=0, + globalstep=0, + ) + engine = Mock() + engine.config = SimpleNamespace(path=save_path) + + saver.save(engine, epoch=0, step=0, global_step=0) + + meta = engine.save.call_args.args[0] + assert meta.base_model_path == str(save_path) From b54a863c3a721b1b76fc52f66a4fdcc77135269b Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Fri, 21 Aug 2026 18:59:01 +0800 Subject: [PATCH 8/8] test(engine): initialize native save fixture correctly Use the MegatronEngine backing process-group fields so the native mbridge finalization test can exercise the real cpu_group property. Refs: #1598 Constraint: Keep the production save path unchanged Confidence: high Scope-risk: narrow Not-tested: Full suite rerun pending on GCP --- .../2026-08-21-bailing-v3-hf-export-design.md | 14 +++++++------- tests/test_megatron_engine.py | 3 ++- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/docs/plans/2026-08-21-bailing-v3-hf-export-design.md b/docs/plans/2026-08-21-bailing-v3-hf-export-design.md index 2900f5b43b..66200461d1 100644 --- a/docs/plans/2026-08-21-bailing-v3-hf-export-design.md +++ b/docs/plans/2026-08-21-bailing-v3-hf-export-design.md @@ -2,11 +2,11 @@ ## Context -The Bailing V3 bridge must dispatch by `architectures` because V3 and Bailing V2.5 -share `model_type="bailing_hybrid"`. The production `swe-dev` implementation therefore -loads V3 through a generic `PretrainedConfig` and constructs `BailingV3Bridge` -explicitly. That behavior is retained: changing registration or relying only on -`AutoConfig` would risk routing V2.5 checkpoints through the V3 bridge. +The Bailing V3 bridge must dispatch by `architectures` because V3 and Bailing V2.5 share +`model_type="bailing_hybrid"`. The production `swe-dev` implementation therefore loads +V3 through a generic `PretrainedConfig` and constructs `BailingV3Bridge` explicitly. +That behavior is retained: changing registration or relying only on `AutoConfig` would +risk routing V2.5 checkpoints through the V3 bridge. The generic config keeps the source `model_type` on the instance, but Transformers serializes the class-level value. For config classes without their own class-level @@ -22,8 +22,8 @@ run one shared finalization step that restores the original instance `model_type written `config.json`. When a local `base_model_path` is available, the same step also copies tokenizer, generation, chat-template, and remote-code assets and preserves source config fields. Port the production `Saver` fallback that derives `base_model_path` from -`engine.config.path`; this remains useful for source assets but is no longer required for -`model_type` correctness. +`engine.config.path`; this remains useful for source assets but is no longer required +for `model_type` correctness. The SWE entrypoint also passes preprocessing controls as call-time keyword arguments. The single-controller `RDataset` branch will merge those arguments with diff --git a/tests/test_megatron_engine.py b/tests/test_megatron_engine.py index 948b250904..2edda20f68 100644 --- a/tests/test_megatron_engine.py +++ b/tests/test_megatron_engine.py @@ -114,7 +114,8 @@ def test_native_mbridge_save_finalizes_hf_config(monkeypatch, tmp_path): engine.mcore_config = SimpleNamespace(use_mbridge_save=True) engine.config = SimpleNamespace(is_critic=False, path="/models/source") engine.bridge = Mock(safetensor_io=object()) - engine.cpu_group = object() + engine.process_group_initialized = True + engine._cpu_group = object() finalize = Mock() source_config = {"model_type": "runtime_only", "torch_dtype": "bfloat16"}