feat(datasets): add cp=1 shared-prefix prefix-tree attention for rollouts

examples/llm_finetune/qwen/qwen25_magi_prefix_tree_rollouts.yaml

@@ -0,0 +1,106 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# To run this recipe:
+#   torchrun --nproc-per-node=1 examples/llm_finetune/finetune.py --config examples/llm_finetune/qwen/qwen25_magi_prefix_tree_rollouts.yaml
+#
+# Shared-prefix RL rollout finetune with the MagiAttention (FFA) backend and a
+# block-sparse prefix-tree mask, WITHOUT context parallelism (cp_size=1). Each
+# training group is one shared prompt + N sampled completions; the collate folds
+# them into a single deduplicated flat sequence ([prompt | c_0 | c_1 | ...]) and
+# builds the prefix-tree AttnMaskSpec. Every completion attends FULL to the shared
+# prompt and CAUSAL to itself, so the prompt is encoded once instead of N times.
+# This is the verl RFC #6401 / Automodel #2385 layout (cp=1 only). Pair with the
+# Qwen3 MoE magi examples to compare against plain varlen packing.
+#
+# Notes for the magi prefix-tree backend (cp=1):
+#   * the model MUST be a registered custom model (in MODEL_ARCH_MAPPING) so
+#     backend.attn="magi" wires the custom-model attn_func, which is the only
+#     magi path that reads the out-of-band AttnMaskSpec. Qwen2 is registered;
+#     a plain HF model would route through the HF magi backend and ignore the
+#     prefix-tree spec.
+#   * local_batch_size must be 1: each folded group is already one flat sequence
+#     with no batch dim, and the prefix-tree mask is built per group.
+#   * backend.rope_fusion must be false: magi applies RoPE positionally with the
+#     per-completion position ids the collate emits (each completion continues
+#     from the prompt length), so the fused TE rope path is bypassed.
+#   * the recipe hands the per-step AttnMaskSpec to the magi attn_func out-of-band
+#     via set_active_attn_spec; no cu_seqlens / THD packer is used here.
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+step_scheduler:
+  global_batch_size: 8
+  local_batch_size: 1        # one folded rollout group per micro-step (no batch dim)
+  ckpt_every_steps: 500
+  num_epochs: 2
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 20
+
+rng:
+  _target_: nemo_automodel.components.training.rng.StatefulRNG
+  seed: 1111
+  ranked: true
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
+  pretrained_model_name_or_path: Qwen/Qwen2.5-0.5B  # registered custom model -> magi custom attn_func
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: magi              # MagiAttention Flex-Flash-Attention backend
+    linear: torch
+    rms_norm: torch
+    rope_fusion: false      # required for magi (non-fused, positional RoPE)
+    enable_hf_state_dict_adapter: true
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: checkpoints/
+  model_save_format: torch_save
+  save_consolidated: false
+
+distributed:
+  strategy: fsdp2
+  tp_size: 1
+  cp_size: 1                 # prefix-tree wiring is cp=1 only
+  pp_size: 1
+  ep_size: 1
+  sequence_parallel: false
+  activation_checkpointing: false
+
+loss_fn:
+  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
+
+dataset:
+  _target_: nemo_automodel.components.datasets.llm.mock_prefix_tree.build_mock_rollout_dataset
+  num_groups: 64
+  completions_per_group: 4
+  prompt_len: 64
+  completion_len: 32
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn: nemo_automodel.components.datasets.llm.prefix_tree.prefix_tree_collate_fn
+  shuffle: false
+
+optimizer:
+  _target_: torch.optim.Adam
+  betas: [0.9, 0.999]
+  eps: 1e-7
+  lr: 1.0e-4
+  weight_decay: 0
+  foreach: false

nemo_automodel/components/datasets/llm/mock_prefix_tree.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Deterministic mock shared-prefix rollout data for prefix-tree smoke runs."""
+
+import random
+
+
+def build_mock_rollout_dataset(
+    *,
+    num_groups: int = 16,
+    completions_per_group: int = 4,
+    prompt_len: int = 32,
+    completion_len: int = 16,
+    vocab_size: int = 1024,
+    seed: int = 0,
+) -> list[dict]:
+    """Build a deterministic mock shared-prefix rollout dataset for smoke runs.
+
+    Each group is one shared prompt with ``completions_per_group`` completions, in
+    the ``{"prompt_ids", "completions"}`` schema consumed by
+    ``prefix_tree_collate_fn``. Token ids are random in ``[2, vocab_size)``; this
+    is a pipeline smoke, not a quality dataset.
+
+    Args:
+        num_groups: number of rollout groups.
+        completions_per_group: completions (leaves) sharing each prompt.
+        prompt_len: shared prompt length per group.
+        completion_len: length of each completion.
+        vocab_size: upper bound (exclusive) for random token ids.
+        seed: RNG seed for reproducibility.
+
+    Returns:
+        A list of ``{"prompt_ids": list[int], "completions": list[list[int]]}``.
+    """
+    rng = random.Random(seed)
+
+    def _ids(n: int) -> list[int]:
+        return [rng.randint(2, vocab_size - 1) for _ in range(n)]
+
+    return [
+        {
+            "prompt_ids": _ids(prompt_len),
+            "completions": [_ids(completion_len) for _ in range(completions_per_group)],
+        }
+        for _ in range(num_groups)
+    ]

nemo_automodel/components/datasets/llm/prefix_tree.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared-prefix rollout folding for multi-turn prefix-tree attention (cp=1).
+
+Folds a group of rollouts that share one prompt prefix (one prompt -> N sampled
+completions) into a single deduplicated flat token layout plus a block-sparse
+prefix-tree attention mask (:class:`AttnMaskSpec`). The shared prompt is stored
+once; every completion attends FULL to the prompt and CAUSAL to itself.
+
+This is the verl RFC #6401 / Automodel #2385 shared-prefix RL training layout,
+restricted to the cp=1 path (no context-parallel dispatch). The produced spec is
+consumed by the magi custom-model attention backend via
+``set_active_attn_spec``; enable it with ``model.backend.attn: magi``.
+
+Branch-point note: the shared prompt's final position is stored once, so it can
+predict only one next token. The N completions diverge there, so the
+prompt -> first-completion-token transition is left unsupervised (label
+``-100`` on the last prompt token); each completion is supervised causally from
+its own first token onward. This is inherent to deduplicating the shared prefix.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+from nemo_automodel.components.distributed.magi_attn_utils import AttnMaskSpec
+
+CROSS_ENTROPY_IGNORE_IDX = -100
+
+
+@dataclass
+class FoldedRollouts:
+    """Deduplicated flat layout + prefix-tree mask for one shared-prefix group.
+
+    Attributes:
+        input_ids: flat ``[prompt | completion_0 | completion_1 | ...]`` tokens.
+        labels: pre-shifted next-token targets (the loss does not shift); the
+            prompt and each completion's last token are ``-100`` (ignored).
+        position_ids: prompt positions ``0..P-1``; each completion continues from
+            ``P`` (so RoPE sees ``prompt ++ completion``).
+        spec: the prefix-tree :class:`AttnMaskSpec` for this group.
+        sample_token_ranges: per-completion flat ``[start, end)`` ranges of its
+            nodes, for reconstructing per-sample outputs.
+    """
+
+    input_ids: list[int]
+    labels: list[int]
+    position_ids: list[int]
+    spec: AttnMaskSpec
+    sample_token_ranges: list[list[list[int]]]
+
+
+def fold_shared_prefix_rollouts(
+    prompt_ids: list[int],
+    completions: list[list[int]],
+    *,
+    ignore_idx: int = CROSS_ENTROPY_IGNORE_IDX,
+) -> FoldedRollouts:
+    """Fold one shared-prefix rollout group into a deduplicated prefix-tree layout.
+
+    Labels follow this repo's next-token convention: they are pre-shifted (the
+    loss does not shift), so position ``p`` carries the id of the token the model
+    should predict at ``p``. Within each completion, token ``t`` predicts token
+    ``t + 1``; the completion's last token has no in-layout successor and is
+    masked. The shared prompt is masked entirely, including its last position:
+    that position is the branch point where the N completions diverge, so it
+    cannot supervise any single first-completion token (the cost of deduplicating
+    the prefix is that each completion's first token is unsupervised).
+
+    Args:
+        prompt_ids: the shared prompt tokens (node 0). May be empty.
+        completions: one token-id list per sampled completion (one leaf each).
+            Must be non-empty and every completion must be non-empty.
+        ignore_idx: label value for unsupervised positions (default ``-100``).
+
+    Returns:
+        A :class:`FoldedRollouts` with the flat tokens, labels, position ids and
+        the built :class:`AttnMaskSpec`.
+
+    Raises:
+        ValueError: if ``completions`` is empty or any completion is empty.
+    """
+    if not completions:
+        raise ValueError("completions must be non-empty (need at least one rollout).")
+    if any(len(c) == 0 for c in completions):
+        raise ValueError("every completion must be non-empty.")
+
+    prompt_len = len(prompt_ids)
+
+    # Single pass builds the deduplicated flat layout: the prompt is stored once
+    # (masked, positions 0..P-1) and each completion is appended with positions
+    # continuing from P (so RoPE sees prompt ++ completion) and pre-shifted labels
+    # (token t predicts t+1; the last token is masked).
+    input_ids = list(prompt_ids)
+    labels = [ignore_idx] * prompt_len
+    position_ids = list(range(prompt_len))
+    node_lengths = [prompt_len]
+    for completion in completions:
+        input_ids.extend(completion)
+        labels.extend(completion[1:] + [ignore_idx])
+        position_ids.extend(range(prompt_len, prompt_len + len(completion)))
+        node_lengths.append(len(completion))
+
+    # Prefix tree: node 0 is the prompt, nodes 1..N are completions; each path is
+    # prompt -> completion. A bare group (empty prompt) collapses to varlen blocks:
+    # node 0 is empty, so each completion is its own causal block.
+    if prompt_len > 0:
+        sample_paths = [[0, i + 1] for i in range(len(completions))]
+    else:
+        sample_paths = [[i + 1] for i in range(len(completions))]
+
+    spec, sample_token_ranges = AttnMaskSpec.prefix_tree(node_lengths, sample_paths)
+
+    return FoldedRollouts(
+        input_ids=input_ids,
+        labels=labels,
+        position_ids=position_ids,
+        spec=spec,
+        sample_token_ranges=sample_token_ranges,
+    )
+
+
+def prefix_tree_collate_fn(batch: list[dict]) -> dict:
+    """Collate one shared-prefix rollout group into a model-ready batch (cp=1).
+
+    Folds the group with :func:`fold_shared_prefix_rollouts` and emits the flat
+    tokens plus the prefix-tree mask. Only ``local_batch_size == 1`` is supported:
+    each group already packs many completions into one flat sequence, and the mask
+    is built per group. The returned ``magi_attn_spec`` is picked up by the recipe
+    and handed to the magi attention backend via ``set_active_attn_spec``.
+
+    Args:
+        batch: a length-1 list holding one rollout group dict with keys
+            ``prompt_ids`` and ``completions``.
+
+    Returns:
+        Dict with ``input_ids``, ``labels``, ``position_ids`` (each ``[1, T]``)
+        and ``magi_attn_spec`` (the :class:`AttnMaskSpec`).
+
+    Raises:
+        ValueError: if ``batch`` does not hold exactly one rollout group.
+    """
+    if len(batch) != 1:
+        raise ValueError(f"prefix_tree_collate_fn supports local_batch_size=1 only, got {len(batch)} groups.")
+    group = batch[0]
+    folded = fold_shared_prefix_rollouts(group["prompt_ids"], group["completions"])
+    return {
+        "input_ids": torch.tensor([folded.input_ids], dtype=torch.long),
+        "labels": torch.tensor([folded.labels], dtype=torch.long),
+        "position_ids": torch.tensor([folded.position_ids], dtype=torch.long),
+        "magi_attn_spec": folded.spec,
+    }

nemo_automodel/components/distributed/magi_attn_utils.py

@@ -749,6 +749,11 @@ def prepare_llm_batch(
         Returns ``(train_ctx, batch)``. magi does its own CP, so ``train_ctx`` is
         always ``nullcontext`` (no torch-native DTensor CP context).
         """
+        # cp=1 arbitrary mask spec (e.g. prefix tree): activate the per-step spec the
+        # collate attached, out-of-band, for the magi attn_func. Setting it every step
+        # (the batch's spec or None) is self-clearing, so a stale spec never leaks into
+        # the next batch; plain batches omit "magi_attn_spec".
+        set_active_attn_spec(batch.pop("magi_attn_spec", None))
         if self.hf_dispatch:
             # HF path: dispatch the (single causal) sequence across the CP group.
             batch, _ = magi_prepare_batch(model, batch, self.cp_group)

tests/functional_tests/attention/_prefix_tree_reference.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Independent oracle mask for the prefix-tree parity checks.
+
+Shared by ``prefix_tree_flex_parity.py`` and ``prefix_tree_magi_parity.py`` so
+the two checks validate against the same reference and cannot silently diverge.
+"""
+
+import torch
+
+
+def build_reference_mask(prompt_len: int, completion_lens: list[int], total: int) -> torch.Tensor:
+    """Dense boolean attend-mask, built straight from the rollout structure.
+
+    ``allowed[q, k]`` is True iff query ``q`` may attend to key ``k`` under the
+    "prompt ++ completion as an independent causal sequence" semantics:
+      * prompt query: causal within the prompt only;
+      * completion query: FULL to the prompt, CAUSAL within its own completion.
+
+    This is built from the rollout structure, NOT from the spec under test, so a
+    wrong spec or a wrong realization of it is caught.
+    """
+    node_of = [-1] * prompt_len
+    for c, n in enumerate(completion_lens):
+        node_of.extend([c] * n)
+    assert len(node_of) == total
+
+    allowed = torch.zeros(total, total, dtype=torch.bool)
+    for q in range(total):
+        qc = node_of[q]
+        for k in range(q + 1):  # causal upper bound: never attend to the future
+            kc = node_of[k]
+            if kc == -1 or kc == qc:
+                # prompt key (full for completions, causal for prompt) or same completion.
+                allowed[q, k] = True
+    return allowed