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,108 @@
+# 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 routes through the HF magi backend (attn_implementation=
+#     "magi"), whose fixed attention interface cannot receive the mask spec; the
+#     recipe raises NotImplementedError in that case rather than silently
+#     dropping the prefix-tree mask.
+#   * 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,167 @@
+# 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 the prefix-tree
+structure (``node_lengths`` / ``sample_paths``). 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 collate carries
+the structure on the batch; the magi backend builds the ``AttnMaskSpec`` from it
+and activates it (the datasets layer must not import ``components.distributed``).
+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
+
+CROSS_ENTROPY_IGNORE_IDX = -100
+
+
+@dataclass
+class FoldedRollouts:
+    """Deduplicated flat layout + prefix-tree structure for one shared-prefix group.
+
+    The attention mask itself is built in the magi backend from ``node_lengths`` /
+    ``sample_paths`` (via ``AttnMaskSpec.prefix_tree``); the datasets layer must not
+    import ``components.distributed``, so it only carries the structure.
+
+    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``).
+        node_lengths: token count of each node, flat-layout order
+            ``[len(prompt), len(c_0), len(c_1), ...]``.
+        sample_paths: root -> leaf node indices per completion, e.g.
+            ``[[0, 1], [0, 2], ...]``.
+    """
+
+    input_ids: list[int]
+    labels: list[int]
+    position_ids: list[int]
+    node_lengths: list[int]
+    sample_paths: 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))]
+
+    return FoldedRollouts(
+        input_ids=input_ids,
+        labels=labels,
+        position_ids=position_ids,
+        node_lengths=node_lengths,
+        sample_paths=sample_paths,
+    )
+
+
+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 structure. Only ``local_batch_size == 1`` is
+    supported: each group already packs many completions into one flat sequence,
+    and the mask is per group. The ``prefix_tree`` entry is popped by the magi
+    backend (``MagiState.prepare_llm_batch``), which builds and activates the
+    ``AttnMaskSpec`` from it.
+
+    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 ``prefix_tree`` (``(node_lengths, sample_paths)``).
+
+    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),
+        "prefix_tree": (folded.node_lengths, folded.sample_paths),
+    }

nemo_automodel/components/distributed/magi_attn_utils.py

@@ -749,6 +749,28 @@ 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 prefix-tree mask: the datasets layer cannot import this module (component
+        # independence), so the collate attaches the tree structure and the spec is built
+        # and activated here, out-of-band, for the magi attn_func. Setting it every step
+        # (a spec or None) is self-clearing, so a stale spec never leaks into the next
+        # batch; plain batches omit "prefix_tree".
+        prefix_tree = batch.pop("prefix_tree", None)
+        if prefix_tree is not None and self.hf_dispatch:
+            # The prefix-tree mask is handed to the attn_func out-of-band (the HF
+            # attention interface has a fixed signature and cannot receive a custom
+            # mask spec argument). Only the custom-model attn_func reads it; the HF
+            # magi backend uses the plain causal varlen key from magi_prepare_batch
+            # and would silently drop the prefix-tree mask. Fail loudly instead.
+            raise NotImplementedError(
+                "The prefix-tree attention mask is only supported on the custom-model magi "
+                "backend (model.backend.attn='magi'), not the HF magi backend "
+                "(model.attn_implementation='magi'). HF's attention interface cannot receive "
+                "the out-of-band AttnMaskSpec, so the mask would be silently dropped. Load a "
+                "model registered in MODEL_ARCH_MAPPING and set model.backend.attn='magi'."
+            )
+        node_lengths, sample_paths = prefix_tree if prefix_tree is not None else (None, None)
+        spec = AttnMaskSpec.prefix_tree(node_lengths, sample_paths)[0] if prefix_tree is not None else None
+        set_active_attn_spec(spec)
         if self.hf_dispatch:
             # HF path: dispatch the (single causal) sequence across the CP group.
             batch, _ = magi_prepare_batch(model, batch, self.cp_group)