feat(speculative): add EAGLE-3 sequence packing and reasoning-mode control

examples/speculative/dflash/qwen3_dflash.yaml

@@ -41,6 +41,7 @@ recipe_args:
   shuffle_seed: 42
   log_every_steps: 10
   max_grad_norm: 1.0
+  # mask_reasoning_content: true   # exclude <think>...</think> reasoning traces from loss
 
   # --- checkpoint cadence (independent; the fully-trained model is always saved) ---
   # ckpt_every_steps: 1000            # also save every N optimizer steps

examples/speculative/eagle1/qwen3_eagle1_perfectblend.yaml

@@ -24,6 +24,7 @@ recipe_args:
   shuffle_seed: 42
   log_every_steps: 10
   max_grad_norm: 1.0
+  # mask_reasoning_content: true   # exclude <think>...</think> reasoning traces from loss
 
 optimizer:
   lr: 1.0e-4

examples/speculative/eagle2/qwen3_eagle2_perfectblend.yaml

@@ -24,6 +24,7 @@ recipe_args:
   shuffle_seed: 42
   log_every_steps: 10
   max_grad_norm: 1.0
+  # mask_reasoning_content: true   # exclude <think>...</think> reasoning traces from loss
 
 optimizer:
   lr: 1.0e-4

examples/speculative/eagle3/llama_eagle3_perfectblend.yaml

@@ -26,6 +26,11 @@ recipe_args:
 
   seq_length: 2048
 
+  # Sequence packing (EAGLE-3, colocated backend only). > 0 packs variable-length
+  # samples into rows of this width with per-document block-causal attention,
+  # removing padding waste. Set to seq_length (or larger); 0 disables.
+  packed_sequence_size: 2048
+
   # 8 GPUs * micro_batch=1 * grad_accum=4 -> effective batch 32.
   micro_batch_size: 1
   grad_accumulation_steps: 4

examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml

@@ -6,23 +6,30 @@ dist_env:
 
 recipe_args:
   target_model_name_or_path: Qwen/Qwen3-8B
-  train_data_path: /path/to/train.jsonl
+  draft_attn_implementation: flash_attention_2
+  train_data_path: ./cache/dataset/perfectblend-qwen3-8b-regen-messages
   val_data_path: null
-  train_split: null
+  train_split: "train[:50]"
   val_split: null
-  output_dir: ./outputs/eagle3_qwen3_mvp
-  seq_length: 1024
+  output_dir: ./outputs/eagle3_qwen3_8b_perfectblend
+  seq_length: 4096
   micro_batch_size: 1
   grad_accumulation_steps: 1
   num_workers: 0
   num_epochs: 1
-  ttt_steps: 4
-  draft_vocab_size: 8192
+  ttt_steps: 7
+  draft_vocab_size: 32000
   freeze_embeddings: true
   trust_remote_code: false
   shuffle_seed: 42
   log_every_steps: 10
   max_grad_norm: 1.0
+  # mask_reasoning_content: true   # exclude <think>...</think> reasoning traces from loss
+
+  # Sequence packing (EAGLE-3, colocated backend only). > 0 packs variable-length
+  # samples into rows of this width with per-document block-causal attention,
+  # removing padding waste. Set to seq_length (or larger); 0 disables.
+  packed_sequence_size: 4096
 
 optimizer:
   lr: 1.0e-4
@@ -31,6 +38,6 @@ optimizer:
 
 checkpoint:
   enabled: true
-  checkpoint_dir: ./outputs/eagle3_qwen3_mvp/checkpoints
+  checkpoint_dir: ./outputs/eagle3_qwen3_8b_perfectblend/checkpoints
   model_save_format: safetensors
   save_consolidated: true

examples/speculative/p-eagle/qwen_peagle_perfectblend.yaml

@@ -58,6 +58,7 @@ recipe_args:
   shuffle_seed: 42
   log_every_steps: 50
   max_grad_norm: 1.0
+  # mask_reasoning_content: true   # exclude <think>...</think> reasoning traces from loss
 
   # ----------------------------------------------------------------------
   # P-EAGLE knobs.

nemo_automodel/components/datasets/llm/eagle3.py

@@ -39,6 +39,106 @@ def _stack_batch(features: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
     return batch
 
 
+def build_packed_eagle3_dataset(
+    source_dataset,
+    *,
+    packed_sequence_size: int,
+    pad_token_id: int,
+) -> list[dict[str, list[int]]]:
+    """Greedily pack variable-length chat samples into rows of ``packed_sequence_size``.
+
+    Each source sample is one *document*; documents are concatenated into a
+    fixed-width row with ``position_ids`` reset per document and trailing pad
+    folded into the final document (so ``seq_lens`` sums to the row width).
+
+    Cross-document leakage at TTT boundaries is handled by ``doc_remaining[t]``
+    (real tokens after slot ``t`` within its document): the trainer supervises
+    slot ``t`` at step ``k`` to predict ``k+1`` ahead, valid iff
+    ``k < doc_remaining[t]``. This masks every cross-document / into-padding
+    supervision -- packing creates many such boundaries per row.
+
+    Returns a list of packed-row dicts with keys ``input_ids``, ``loss_mask``,
+    ``attention_mask``, ``position_ids``, ``doc_remaining`` (length
+    ``packed_sequence_size``) and ``seq_lens`` (per-document padded lengths).
+    """
+    packs: list[dict[str, list[int]]] = []
+    cur_ids: list[int] = []
+    cur_loss: list[int] = []
+    cur_pos: list[int] = []
+    cur_remaining: list[int] = []
+    cur_seq_lens: list[int] = []
+
+    def _flush() -> None:
+        nonlocal cur_ids, cur_loss, cur_pos, cur_remaining, cur_seq_lens
+        if not cur_ids:
+            return
+        valid_len = len(cur_ids)
+        num_pad = packed_sequence_size - valid_len
+        ids = cur_ids + [pad_token_id] * num_pad
+        loss = cur_loss + [0] * num_pad
+        attn = [1] * valid_len + [0] * num_pad
+        remaining = cur_remaining + [0] * num_pad
+        # Fold trailing pad into the final document (continue its position ids).
+        last_pos = cur_pos[-1] if cur_pos else -1
+        pos = cur_pos + [min(last_pos + 1 + j, packed_sequence_size - 1) for j in range(num_pad)]
+        seq_lens = list(cur_seq_lens)
+        if num_pad > 0:
+            seq_lens[-1] += num_pad
+        packs.append(
+            {
+                "input_ids": ids,
+                "loss_mask": loss,
+                "attention_mask": attn,
+                "position_ids": pos,
+                "doc_remaining": remaining,
+                "seq_lens": seq_lens,
+            }
+        )
+        cur_ids, cur_loss, cur_pos, cur_remaining, cur_seq_lens = [], [], [], [], []
+
+    for sample in source_dataset:
+        ids = list(sample["input_ids"])
+        loss = [int(bool(m)) for m in sample["loss_mask"]]
+        length = len(ids)
+        if length == 0:
+            continue
+        if length > packed_sequence_size:
+            # Guard: the source dataset's truncation should already cap at the row width.
+            ids = ids[:packed_sequence_size]
+            loss = loss[:packed_sequence_size]
+            length = packed_sequence_size
+        if cur_ids and len(cur_ids) + length > packed_sequence_size:
+            _flush()
+        cur_ids += ids
+        cur_loss += loss
+        cur_pos += list(range(length))
+        # Real tokens remaining after each slot within this document.
+        cur_remaining += list(range(length - 1, -1, -1))
+        cur_seq_lens.append(length)
+    _flush()
+
+    if not packs:
+        raise ValueError(
+            f"No packs were produced from the source dataset for packed_sequence_size="
+            f"{packed_sequence_size}. The dataset may be empty."
+        )
+    return packs
+
+
+def _pack_collate(features: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+    """Collate packed rows; ragged ``seq_lens`` is 0-padded to ``[B, max_docs]``."""
+    batch = {}
+    for key in ("input_ids", "loss_mask", "attention_mask", "position_ids", "doc_remaining"):
+        batch[key] = torch.tensor([feature[key] for feature in features], dtype=torch.long)
+    max_docs = max(len(feature["seq_lens"]) for feature in features)
+    seq_lens = torch.zeros(len(features), max_docs, dtype=torch.long)
+    for i, feature in enumerate(features):
+        row = feature["seq_lens"]
+        seq_lens[i, : len(row)] = torch.tensor(row, dtype=torch.long)
+    batch["seq_lens"] = seq_lens
+    return batch
+
+
 def build_eagle3_dataloader(
     *,
     data_path: str,
@@ -50,18 +150,47 @@ def build_eagle3_dataloader(
     split: str | None = None,
     distributed: bool = False,
     shuffle_seed: int | None = 42,
+    mask_reasoning_content: bool = False,
+    packed_sequence_size: int = 0,
 ) -> DataLoader:
-    """Build a dataloader backed by the repo's chat formatting utilities."""
-    dataset = ChatDataset(
-        data_path,
-        tokenizer=tokenizer,
-        split=split,
-        seq_length=seq_length,
-        padding="max_length",
-        truncation=True,
-        shuffle_seed=shuffle_seed,
-        unshifted=True,
-    )
+    """Build a dataloader backed by the repo's chat formatting utilities.
+
+    ``packed_sequence_size > 0`` (EAGLE-3 only) enables sequence packing (see
+    :func:`build_packed_eagle3_dataset`), removing the padding waste of the
+    default ``padding="max_length"`` path; ``== 0`` keeps the original behavior.
+    """
+    collate_fn = _stack_batch
+    if packed_sequence_size > 0:
+        # Source samples are unpadded (one document each); packing pads the row.
+        source = ChatDataset(
+            data_path,
+            tokenizer=tokenizer,
+            split=split,
+            seq_length=packed_sequence_size,
+            padding="do_not_pad",
+            truncation=True,
+            shuffle_seed=shuffle_seed,
+            unshifted=True,
+            mask_reasoning_content=mask_reasoning_content,
+        )
+        dataset = build_packed_eagle3_dataset(
+            source,
+            packed_sequence_size=packed_sequence_size,
+            pad_token_id=source.pad_token_id,
+        )
+        collate_fn = _pack_collate
+    else:
+        dataset = ChatDataset(
+            data_path,
+            tokenizer=tokenizer,
+            split=split,
+            seq_length=seq_length,
+            padding="max_length",
+            truncation=True,
+            shuffle_seed=shuffle_seed,
+            unshifted=True,
+            mask_reasoning_content=mask_reasoning_content,
+        )
     sampler = DistributedSampler(dataset, shuffle=shuffle) if distributed else None
 
     # The EAGLE recipes load the target model onto CUDA before iterating, so the
@@ -90,7 +219,7 @@ def build_eagle3_dataloader(
         shuffle=shuffle and sampler is None,
         num_workers=num_workers,
         pin_memory=torch.cuda.is_available(),
-        collate_fn=_stack_batch,
+        collate_fn=collate_fn,
         drop_last=False,
         **worker_kwargs,
     )

nemo_automodel/components/datasets/llm/packed_sequence.py

@@ -384,3 +384,33 @@ def packed_block_causal_mask(seq_lens: list[torch.Tensor]):
         _MaskType: BlockMask or Tensor if torch version < 2.5.0.
     """
     return create_block_causal_mask(seq_lens=seq_lens)
+
+
+def build_block_causal_additive_mask(
+    seq_lens: torch.Tensor,
+    *,
+    seq_length: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
+    """Build a ``[B, 1, T, T]`` additive block-causal mask directly on ``device``.
+
+    In-document causal attention is allowed (``0``); cross-document and padding
+    positions are ``finfo(dtype).min``. ``seq_lens`` is the ``[B, max_docs]``
+    0-padded per-document length tensor; each row's non-zero entries sum to
+    ``seq_length`` (trailing pad folded into the final document).
+    """
+    min_value = torch.finfo(dtype).min
+    seq_lens = seq_lens.to(device)
+    positions = torch.arange(seq_length, device=device)
+    # Per-position document id: the count of document boundaries at or before the
+    # position. 0-length padding entries leave the cumulative boundary unchanged,
+    # so they never split a real document. ``[B, T]``.
+    boundaries = seq_lens.cumsum(dim=1)  # [B, max_docs]
+    doc_id = (boundaries.unsqueeze(1) <= positions.view(1, -1, 1)).sum(dim=2)  # [B, T]
+    same_doc = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)  # [B, T, T]
+    causal = torch.tril(torch.ones(seq_length, seq_length, dtype=torch.bool, device=device))
+    # In-document lower-triangular attention is allowed; everything else is masked.
+    allowed = same_doc & causal
+    mask = torch.where(allowed, torch.zeros((), dtype=dtype, device=device), min_value)
+    return mask.unsqueeze(1)

nemo_automodel/components/speculative/eagle/core.py

@@ -92,6 +92,9 @@ def forward(
         *,
         target_probs: torch.Tensor | None = None,
         position_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        seq_lens: torch.Tensor | None = None,
+        doc_remaining: torch.Tensor | None = None,
     ) -> Eagle3StepMetrics:
         """Run the EAGLE-3 unrolled draft loss for one batch.
 
@@ -107,6 +110,11 @@ def forward(
         ``input_ids`` / ``loss_mask`` / ``position_mask`` /
         ``target_probs`` roll forward by one position per step.
 
+        Packing: ``position_ids`` / ``seq_lens`` make the draft's Block-1 attention
+        document-level block-causal, and ``doc_remaining`` gates supervision per
+        step (slot ``t`` valid at step ``k`` only while ``k < doc_remaining[t]``),
+        masking every cross-document TTT prediction.
+
         Two supervision sources are accepted: the live path passes the
         target's full-vocab ``target_logits`` and the draft distribution is
         derived here; the offline-cache path (``precompute_eagle3``) passes the
@@ -163,17 +171,27 @@ def forward(
                 input_ids=cur_input_ids,
                 projected_hidden_states=cur_hidden_states,
                 attention_mask=attention_mask,
+                position_ids=position_ids,
                 cache_hidden=cache_hidden,
+                seq_lens=seq_lens,
             )
             logits = self.draft_model.compute_logits(cur_hidden_states)
+
+            # Packing: drop supervision whose step_idx-ahead target crosses this
+            # slot's document boundary. Gate recomputed per step (depends on step_idx).
+            step_position_mask = cur_position_mask
+            if doc_remaining is not None:
+                in_doc = (step_idx < doc_remaining).unsqueeze(-1)
+                step_position_mask = cur_position_mask & in_doc
+
             step_loss = masked_soft_cross_entropy(
                 logits=logits,
                 target_probs=cur_target_probs,
-                position_mask=cur_position_mask,
+                position_mask=step_position_mask,
             )
             running_loss = running_loss + step_loss * (0.8**step_idx)
 
-            valid_mask = cur_position_mask.squeeze(-1).bool()
+            valid_mask = step_position_mask.squeeze(-1).bool()
             correct = (logits.argmax(dim=-1) == cur_target_probs.argmax(dim=-1)) & valid_mask
             running_correct = running_correct + correct.sum()
             running_valid = running_valid + valid_mask.sum()