From 49374d7977ba0eb0818a69ea390d682eed81cb1c Mon Sep 17 00:00:00 2001 From: thyways <2484113689@qq.com> Date: Fri, 5 Jun 2026 20:06:17 +0000 Subject: [PATCH 1/2] feat(speculative): add reasoning mode control for EAGLE/P-EAGLE/DFlash training Add --reasoning {none,save,disable} flag to regenerate.py for controlling whether target model reasoning content is preserved or suppressed during data regeneration. Add mask_reasoning_content option to EAGLE/P-EAGLE/DFlash training recipes to exclude reasoning traces from the loss mask. Co-authored-by: khazic Signed-off-by: thyways <2484113689@qq.com> Signed-off-by: khazic --- examples/speculative/dflash/qwen3_dflash.yaml | 1 + .../eagle1/qwen3_eagle1_perfectblend.yaml | 1 + .../eagle2/qwen3_eagle2_perfectblend.yaml | 1 + .../eagle3/qwen3_eagle3_perfectblend.yaml | 1 + .../p-eagle/qwen_peagle_perfectblend.yaml | 1 + .../components/datasets/llm/eagle3.py | 2 ++ .../components/speculative/regenerate.py | 32 +++++++++++++++---- nemo_automodel/recipes/llm/train_dflash.py | 2 ++ nemo_automodel/recipes/llm/train_eagle1.py | 2 ++ nemo_automodel/recipes/llm/train_eagle3.py | 2 ++ .../unit_tests/speculative/test_regenerate.py | 8 ++++- 11 files changed, 46 insertions(+), 7 deletions(-) diff --git a/examples/speculative/dflash/qwen3_dflash.yaml b/examples/speculative/dflash/qwen3_dflash.yaml index c5eb601681..943a891bed 100644 --- a/examples/speculative/dflash/qwen3_dflash.yaml +++ b/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 ... reasoning traces from loss # --- checkpoint cadence (independent; the fully-trained model is always saved) --- # ckpt_every_steps: 1000 # also save every N optimizer steps diff --git a/examples/speculative/eagle1/qwen3_eagle1_perfectblend.yaml b/examples/speculative/eagle1/qwen3_eagle1_perfectblend.yaml index 61ed7f4068..65e9e750af 100644 --- a/examples/speculative/eagle1/qwen3_eagle1_perfectblend.yaml +++ b/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 ... reasoning traces from loss optimizer: lr: 1.0e-4 diff --git a/examples/speculative/eagle2/qwen3_eagle2_perfectblend.yaml b/examples/speculative/eagle2/qwen3_eagle2_perfectblend.yaml index 33e771cca5..4873dfab23 100644 --- a/examples/speculative/eagle2/qwen3_eagle2_perfectblend.yaml +++ b/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 ... reasoning traces from loss optimizer: lr: 1.0e-4 diff --git a/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml b/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml index 5913192df3..8ea009bae1 100644 --- a/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml +++ b/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml @@ -23,6 +23,7 @@ recipe_args: shuffle_seed: 42 log_every_steps: 10 max_grad_norm: 1.0 + # mask_reasoning_content: true # exclude ... reasoning traces from loss optimizer: lr: 1.0e-4 diff --git a/examples/speculative/p-eagle/qwen_peagle_perfectblend.yaml b/examples/speculative/p-eagle/qwen_peagle_perfectblend.yaml index e69eec1a11..c5b08665c9 100644 --- a/examples/speculative/p-eagle/qwen_peagle_perfectblend.yaml +++ b/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 ... reasoning traces from loss # ---------------------------------------------------------------------- # P-EAGLE knobs. diff --git a/nemo_automodel/components/datasets/llm/eagle3.py b/nemo_automodel/components/datasets/llm/eagle3.py index 3f6d8d45f2..7ada36b4c1 100644 --- a/nemo_automodel/components/datasets/llm/eagle3.py +++ b/nemo_automodel/components/datasets/llm/eagle3.py @@ -50,6 +50,7 @@ def build_eagle3_dataloader( split: str | None = None, distributed: bool = False, shuffle_seed: int | None = 42, + mask_reasoning_content: bool = False, ) -> DataLoader: """Build a dataloader backed by the repo's chat formatting utilities.""" dataset = ChatDataset( @@ -61,6 +62,7 @@ def build_eagle3_dataloader( truncation=True, shuffle_seed=shuffle_seed, unshifted=True, + mask_reasoning_content=mask_reasoning_content, ) sampler = DistributedSampler(dataset, shuffle=shuffle) if distributed else None diff --git a/nemo_automodel/components/speculative/regenerate.py b/nemo_automodel/components/speculative/regenerate.py index 00aaef0e1c..623f03b921 100644 --- a/nemo_automodel/components/speculative/regenerate.py +++ b/nemo_automodel/components/speculative/regenerate.py @@ -86,6 +86,7 @@ class GenerationConfig: max_new_tokens: int temperature: float top_p: float + reasoning: str = "none" def _build_manifest(args: argparse.Namespace) -> dict[str, Any]: @@ -110,6 +111,7 @@ def _build_manifest(args: argparse.Namespace) -> dict[str, Any]: "max_new_tokens": args.max_new_tokens, "temperature": args.temperature, "top_p": args.top_p, + "reasoning": args.reasoning, } @@ -264,8 +266,8 @@ async def _chat_completion( *, timeout_s: float, max_retries: int, -) -> str: - """POST ``payload`` to ``url`` and return the assistant text, with bounded retries.""" +) -> dict[str, Any]: + """POST ``payload`` to ``url`` and return the assistant message dict, with bounded retries.""" aiohttp = _import_aiohttp() last_err: Exception | None = None for attempt in range(max_retries + 1): @@ -276,7 +278,7 @@ async def _chat_completion( raise RuntimeError(f"HTTP {resp.status} from {url}: {text[:200]}") resp.raise_for_status() data = await resp.json() - return data["choices"][0]["message"]["content"] + return data["choices"][0]["message"] except Exception as exc: # noqa: BLE001 -- retry any transport / 5xx error last_err = exc if attempt == max_retries: @@ -303,15 +305,22 @@ async def _regenerate_one( max_retries: int, ) -> list[dict[str, Any]]: """Call the target server once and return ``prompt + [assistant]``.""" - payload = { + payload: dict[str, Any] = { "model": gen_cfg.model, "messages": prompt, "max_tokens": gen_cfg.max_new_tokens, "temperature": gen_cfg.temperature, "top_p": gen_cfg.top_p, } - content = await _chat_completion(session, url, payload, timeout_s=timeout_s, max_retries=max_retries) - return [*prompt, {"role": "assistant", "content": content}] + if gen_cfg.reasoning == "disable": + payload["extra_body"] = {"chat_template_kwargs": {"enable_thinking": False}} + msg = await _chat_completion(session, url, payload, timeout_s=timeout_s, max_retries=max_retries) + resp_msg: dict[str, Any] = {"role": "assistant", "content": msg.get("content", "")} + if gen_cfg.reasoning == "save": + reasoning_content = msg.get("reasoning_content") + if reasoning_content: + resp_msg["reasoning_content"] = reasoning_content + return [*prompt, resp_msg] async def _process_shard( @@ -391,6 +400,7 @@ async def _run(args: argparse.Namespace) -> int: max_new_tokens=args.max_new_tokens, temperature=args.temperature, top_p=args.top_p, + reasoning=args.reasoning, ) url = args.target_server.rstrip("/") + "/chat/completions" timeout = aiohttp.ClientTimeout(total=args.timeout_s) @@ -494,6 +504,16 @@ def _build_parser() -> argparse.ArgumentParser: parser.add_argument("--top-p", type=float, default=1.0) parser.add_argument("--timeout-s", type=float, default=600.0, help="Per-request timeout in seconds.") parser.add_argument("--max-retries", type=int, default=3, help="Retries on 5xx / 429 / transport errors.") + parser.add_argument( + "--reasoning", + choices=["none", "save", "disable"], + default="none", + help=( + "Reasoning mode: 'none' for standard models (default), " + "'save' to store reasoning_content from the response, " + "or 'disable' to suppress thinking via extra_body chat_template_kwargs." + ), + ) parser.add_argument( "--resume", action="store_true", help="Skip shard indices whose parquet file already exists in --output-dir." ) diff --git a/nemo_automodel/recipes/llm/train_dflash.py b/nemo_automodel/recipes/llm/train_dflash.py index 50333b9251..f5c1471e42 100644 --- a/nemo_automodel/recipes/llm/train_dflash.py +++ b/nemo_automodel/recipes/llm/train_dflash.py @@ -141,6 +141,7 @@ def setup(self): split=recipe_cfg.get("train_split", None), distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), + mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), ) self.val_dataloader = None if recipe_cfg.get("val_data_path", None): @@ -154,6 +155,7 @@ def setup(self): split=recipe_cfg.get("val_split", None), distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), + mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), ) # DFlash draft config: a small non-causal Qwen3 stack that reuses the diff --git a/nemo_automodel/recipes/llm/train_eagle1.py b/nemo_automodel/recipes/llm/train_eagle1.py index ec45719c95..9508feb95f 100644 --- a/nemo_automodel/recipes/llm/train_eagle1.py +++ b/nemo_automodel/recipes/llm/train_eagle1.py @@ -185,6 +185,7 @@ def setup(self): split=recipe_cfg.get("train_split", None), distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), + mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), ) self.val_dataloader = None if recipe_cfg.get("val_data_path", None): @@ -198,6 +199,7 @@ def setup(self): split=recipe_cfg.get("val_split", None), distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), + mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), ) draft_config = target_config.to_dict() diff --git a/nemo_automodel/recipes/llm/train_eagle3.py b/nemo_automodel/recipes/llm/train_eagle3.py index 18e37daaf7..3fce8f1c76 100644 --- a/nemo_automodel/recipes/llm/train_eagle3.py +++ b/nemo_automodel/recipes/llm/train_eagle3.py @@ -361,6 +361,7 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): split=recipe_cfg.get("train_split", None), distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), + mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), ) self.val_dataloader = None if recipe_cfg.get("val_data_path", None): @@ -374,6 +375,7 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): split=recipe_cfg.get("val_split", None), distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), + mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), ) special_token_ids = [ diff --git a/tests/unit_tests/speculative/test_regenerate.py b/tests/unit_tests/speculative/test_regenerate.py index 840412481e..ac7b38e37c 100644 --- a/tests/unit_tests/speculative/test_regenerate.py +++ b/tests/unit_tests/speculative/test_regenerate.py @@ -152,6 +152,7 @@ def test_resume_manifest_mismatch_raises(tmp_path: Path): max_new_tokens=128, temperature=0.0, top_p=1.0, + reasoning="none", ) manifest = _build_manifest(args) (tmp_path / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") @@ -180,6 +181,7 @@ def test_resume_without_manifest_and_existing_shards_raises(tmp_path: Path): "max_new_tokens": 128, "temperature": 0.0, "top_p": 1.0, + "reasoning": "none", } with pytest.raises(ValueError, match="manifest.json is missing"): _ensure_manifest_compatible(tmp_path, manifest, resume=True, existing_shards={0}) @@ -201,6 +203,7 @@ def test_fresh_run_with_existing_shards_refuses_to_clobber(tmp_path: Path): max_new_tokens=128, temperature=0.0, top_p=1.0, + reasoning="none", ) with pytest.raises(ValueError, match="already contains"): _ensure_manifest_compatible( @@ -226,6 +229,7 @@ def test_fresh_run_into_empty_dir_writes_manifest(tmp_path: Path): max_new_tokens=128, temperature=0.0, top_p=1.0, + reasoning="none", ) manifest = _build_manifest(args) _ensure_manifest_compatible(tmp_path, manifest, resume=False, existing_shards=set()) @@ -248,6 +252,7 @@ def test_build_manifest_excludes_self_referential_and_operational_fields(): shard_size=1000, max_new_tokens=128, temperature=0.0, + reasoning="none", top_p=1.0, # Operational knobs that intentionally do NOT belong in the manifest. concurrency=99, @@ -417,6 +422,7 @@ def _run_args(tmp_path: Path, *, resume: bool, shard_size: int = 2) -> SimpleNam max_retries=0, resume=resume, log_level="INFO", + reasoning="none", ) @@ -524,7 +530,7 @@ async def fast_sleep(delay: float) -> None: result = asyncio.run(_chat_completion(session, "http://stub/completions", {}, timeout_s=1.0, max_retries=3)) - assert result == "ok" + assert result == {"role": "assistant", "content": "ok"} assert session.call_count == 3 assert len(sleep_calls) == 2 # slept once after each failed attempt assert sleep_calls[0] == 1.0 # 2**0 = 1 From 160078c2c0d71c3176580f4579b87f84363b22da Mon Sep 17 00:00:00 2001 From: thyways <2484113689@qq.com> Date: Sun, 7 Jun 2026 05:09:35 +0000 Subject: [PATCH 2/2] feat(speculative): add EAGLE-3 sequence packing for draft training Pack variable-length chat samples into fixed-width rows for EAGLE-3 training, removing the per-sample padding waste of the default max_length path. Documents within a row attend block-causally: the target uses a 4D block-causal mask (SDPA) and the draft uses varlen FlashAttention-2; cross-document TTT supervision is gated by doc_remaining so deeper steps never leak across boundaries. Opt-in via packed_sequence_size > 0, colocated target backend only. Covered by unit tests plus an FA2-vs-eager parity test. Co-authored-by: khazic Signed-off-by: thyways <2484113689@qq.com> Signed-off-by: khazic --- .../eagle3/llama_eagle3_perfectblend.yaml | 5 + .../eagle3/qwen3_eagle3_perfectblend.yaml | 20 +- .../components/datasets/llm/eagle3.py | 153 +++++- .../datasets/llm/packed_sequence.py | 30 ++ .../components/speculative/eagle/core.py | 22 +- .../speculative/eagle/draft_llama.py | 193 ++++++-- .../components/speculative/eagle/target.py | 65 ++- nemo_automodel/recipes/llm/train_eagle3.py | 20 + .../test_eagle3_packing_fa2_parity.py | 297 ++++++++++++ .../speculative/test_eagle3_packing.py | 446 ++++++++++++++++++ 10 files changed, 1180 insertions(+), 71 deletions(-) create mode 100644 tests/functional_tests/speculative/test_eagle3_packing_fa2_parity.py create mode 100644 tests/unit_tests/speculative/test_eagle3_packing.py diff --git a/examples/speculative/eagle3/llama_eagle3_perfectblend.yaml b/examples/speculative/eagle3/llama_eagle3_perfectblend.yaml index 5fa010b6fe..daefab1a3e 100644 --- a/examples/speculative/eagle3/llama_eagle3_perfectblend.yaml +++ b/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 diff --git a/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml b/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml index 8ea009bae1..6aae63e6d2 100644 --- a/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml +++ b/examples/speculative/eagle3/qwen3_eagle3_perfectblend.yaml @@ -6,18 +6,19 @@ 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 @@ -25,6 +26,11 @@ recipe_args: max_grad_norm: 1.0 # mask_reasoning_content: true # exclude ... 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 betas: [0.9, 0.95] @@ -32,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 diff --git a/nemo_automodel/components/datasets/llm/eagle3.py b/nemo_automodel/components/datasets/llm/eagle3.py index 7ada36b4c1..6ec682683e 100644 --- a/nemo_automodel/components/datasets/llm/eagle3.py +++ b/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, @@ -51,19 +151,46 @@ def build_eagle3_dataloader( 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, - mask_reasoning_content=mask_reasoning_content, - ) + """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 @@ -92,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, ) diff --git a/nemo_automodel/components/datasets/llm/packed_sequence.py b/nemo_automodel/components/datasets/llm/packed_sequence.py index 4a46ef3367..83052082af 100644 --- a/nemo_automodel/components/datasets/llm/packed_sequence.py +++ b/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) diff --git a/nemo_automodel/components/speculative/eagle/core.py b/nemo_automodel/components/speculative/eagle/core.py index e35faa3016..f62fcb3a56 100644 --- a/nemo_automodel/components/speculative/eagle/core.py +++ b/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() diff --git a/nemo_automodel/components/speculative/eagle/draft_llama.py b/nemo_automodel/components/speculative/eagle/draft_llama.py index 77ac9b13a1..d0833e230a 100644 --- a/nemo_automodel/components/speculative/eagle/draft_llama.py +++ b/nemo_automodel/components/speculative/eagle/draft_llama.py @@ -99,6 +99,7 @@ import torch.nn as nn from transformers import PretrainedConfig, PreTrainedModel +from nemo_automodel.components.datasets.llm.packed_sequence import build_block_causal_additive_mask from nemo_automodel.components.models.common import initialize_rms_norm_module from nemo_automodel.components.models.llama.rope_utils import ( LlamaRotaryEmbedding, @@ -115,25 +116,27 @@ logger = logging.getLogger(__name__) -def _load_flash_attn_func() -> tuple[bool, object | None]: +def _load_flash_attn_func() -> tuple[bool, object | None, object | None]: """Best-effort load of flash-attn without breaking eager-only users. ``safe_import_from`` already handles missing modules and missing symbols, but some broken ``flash-attn`` installs fail with lower-level loader errors (e.g. ABI / shared-library issues) that should not prevent importing this - module for the eager path. + module for the eager path. Returns the dense ``flash_attn_func`` and the + ``flash_attn_varlen_func`` (used by the packed block-causal path). """ try: has_fa, flash_attn_func = safe_import_from("flash_attn", "flash_attn_func") + _, flash_attn_varlen_func = safe_import_from("flash_attn", "flash_attn_varlen_func") except Exception as exc: # pragma: no cover - depends on local flash-attn loader failures. logger.warning("Failed to import flash_attn.flash_attn_func; FlashAttention-2 path will be disabled: %s", exc) - return False, None + return False, None, None if not has_fa: - return False, None - return True, flash_attn_func + return False, None, None + return True, flash_attn_func, flash_attn_varlen_func -_HAS_FA, _flash_attn_func = _load_flash_attn_func() +_HAS_FA, _flash_attn_func, _flash_attn_varlen_func = _load_flash_attn_func() _SUPPORTED_ATTN_IMPLEMENTATIONS = ("eager", "flash_attention_2") @@ -158,6 +161,28 @@ def _is_right_padded_attention_mask(attention_mask: torch.Tensor) -> bool: return not bool((mask_bool[:, 1:] & ~mask_bool[:, :-1]).any()) +def _seq_lens_to_cu_seqlens(seq_lens: torch.Tensor, seq_length: int) -> tuple[torch.Tensor, int]: + """Build FlashAttention varlen ``cu_seqlens`` (int32) from packed ``seq_lens``. + + Documents are flattened row-major to match the varlen attention's + ``reshape(B*T, ...)`` token order. Returns ``(cu_seqlens, max_seqlen)``. + """ + doc_lens = seq_lens[seq_lens > 0] + cu_seqlens = torch.zeros(doc_lens.numel() + 1, dtype=torch.int32, device=seq_lens.device) + cu_seqlens[1:] = torch.cumsum(doc_lens, dim=0).to(torch.int32) + expected_total = seq_lens.shape[0] * seq_length + if int(cu_seqlens[-1].item()) != expected_total: + raise ValueError( + f"Packed seq_lens sum to {int(cu_seqlens[-1].item())} but expected B*T={expected_total}; " + "each row's document lengths (with trailing padding folded into the last document) " + "must sum to the packed sequence length." + ) + # ``doc_lens`` is non-empty here: an all-zero ``seq_lens`` would sum to 0, + # which the ``expected_total`` check above already rejects. + max_seqlen = int(doc_lens.max().item()) + return cu_seqlens, max_seqlen + + class Eagle3LlamaAttention(_PeagleAttentionMixin, nn.Module): """EAGLE-3 draft attention over ``[input_emb, hidden]`` 2H features. @@ -265,6 +290,8 @@ def forward( attention_mask: torch.Tensor, position_ids: torch.Tensor, cache_hidden: list[list[torch.Tensor]], + cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, ) -> torch.Tensor: batch_size, seq_len, _ = combined_states.shape q, k, v = self._project_qkv(combined_states) @@ -286,7 +313,9 @@ def forward( cache_v.append(v) if self.attn_implementation == "flash_attention_2": - attn_output = self._flash_attention_forward(q, cache_k, cache_v, step_idx, batch_size, seq_len) + attn_output = self._flash_attention_forward( + q, cache_k, cache_v, step_idx, batch_size, seq_len, cu_seqlens, max_seqlen + ) else: attn_output = self._eager_attention_forward( q, cache_k, cache_v, attention_mask, step_idx, batch_size, seq_len @@ -342,52 +371,48 @@ def _flash_attention_forward( step_idx: int, batch_size: int, seq_len: int, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, ) -> torch.Tensor: """EAGLE-3 attention via FlashAttention-2 for the T x T causal block. - FA2 covers Block 1 (full ``T x T`` causal attention against ``K_0``) and - returns the un-normalized log-sum-exp (``softmax_lse``) alongside the - per-token output. The diagonal extension columns (Block 2) for cached - steps ``i >= 1`` are computed in eager mode, then merged into a single - softmax via the log-space identity - ``lse_full = logaddexp(lse_fa, logsumexp(diag))``; the FA output is - rescaled by ``exp(lse_fa - lse_full)`` and the diagonal contribution - is added with weights ``exp(diag - lse_full)``. - - Padding handling: FA2 is invoked with ``causal=True``. For right-padded - batches, padding keys always lie strictly above the diagonal relative - to any non-padded query position, so causal masking alone yields the - same output as the eager additive padding mask at every valid query - position. Outputs at padding query positions differ, but those are - masked out at loss time. + FA2 covers Block 1 (causal attention against ``K_0``) and returns its + log-sum-exp. The diagonal Block 2 (cached steps ``i >= 1``) is computed + eagerly and merged via the log-space identity + ``lse_full = logaddexp(lse_fa, logsumexp(diag))``: the FA output is scaled + by ``exp(lse_fa - lse_full)`` and each diagonal by ``exp(diag - lse_full)``. + + With ``cu_seqlens`` (packing), Block 1 uses ``flash_attn_varlen_func`` for + document-level causal attention; the position-wise Block 2 is unchanged. """ # FA2 expects (B, T, H, D); eager cache is (B, H, T, D). k0, v0 = cache_k[0], cache_v[0] q_fa = q.transpose(1, 2).contiguous() k0_fa = k0.transpose(1, 2).contiguous() v0_fa = v0.transpose(1, 2).contiguous() - # ``softmax_lse`` is fp32 with shape (B, H, T): the log-sum-exp of the - # SCALED Block-1 logits (the FA kernel folds in ``softmax_scale``). - out_fa, lse_fa, _ = _flash_attn_func( - q_fa, - k0_fa, - v0_fa, - softmax_scale=self.scaling, - causal=True, - return_attn_probs=True, - ) - # FA output is (B, T, H, D); bring back to (B, H, T, D) for downstream merge. - attn_output_bhtd = out_fa.transpose(1, 2) + if cu_seqlens is not None: + attn_output_bhtd, lse_fa = self._flash_block1_varlen( + q_fa, k0_fa, v0_fa, cu_seqlens, max_seqlen, batch_size, seq_len + ) + else: + # ``softmax_lse`` is fp32 with shape (B, H, T): the log-sum-exp of the + # SCALED Block-1 logits (the FA kernel folds in ``softmax_scale``). + out_fa, lse_fa, _ = _flash_attn_func( + q_fa, + k0_fa, + v0_fa, + softmax_scale=self.scaling, + causal=True, + return_attn_probs=True, + ) + # FA output is (B, T, H, D); bring back to (B, H, T, D) for the merge. + attn_output_bhtd = out_fa.transpose(1, 2) if step_idx >= 1: - # Diagonal logits share the same ``self.scaling`` factor as FA's - # internal softmax, so ``lse_fa`` and ``diag_logits`` are commensurate. + # Diagonal logits use the same scaling as FA, so the LSEs are commensurate. later_k = torch.stack(cache_k[1:], dim=0) # [step_idx, B, H, T, D] diag_logits = torch.einsum("bhtd,sbhtd->bhts", q, later_k) * self.scaling - # Combine softmax in log-space: - # lse_full = log( exp(lse_fa) + sum_i exp(diag_i) ) - # = logaddexp(lse_fa, logsumexp(diag, dim=-1)) lse_fa_f32 = lse_fa.float() # [B, H, T] diag_f32 = diag_logits.float() # [B, H, T, step_idx] diag_lse = torch.logsumexp(diag_f32, dim=-1) # [B, H, T] @@ -402,6 +427,55 @@ def _flash_attention_forward( return attn_output_bhtd.transpose(1, 2).contiguous().view(batch_size, seq_len, -1) + def _flash_block1_varlen( + self, + q_fa: torch.Tensor, + k0_fa: torch.Tensor, + v0_fa: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + batch_size: int, + seq_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Document-level causal Block 1 via ``flash_attn_varlen_func``. + + Flattens ``(B, T, H, D)`` to varlen ``(total_tokens, H, D)`` and reshapes + outputs back to ``[B, H, T, D]`` / ``[B, H, T]`` for the dense-path merge. + Note varlen ``softmax_lse`` is ``[H, total_tokens]`` (head-major), unlike + the dense ``[B, H, T]`` -- hence the explicit reshape + shape check. + """ + if _flash_attn_varlen_func is None: + raise ImportError( + "Eagle3LlamaAttention: packed FlashAttention-2 requires flash_attn.flash_attn_varlen_func." + ) + num_heads, head_dim = q_fa.shape[2], q_fa.shape[3] + total_tokens = batch_size * seq_len + q_flat = q_fa.reshape(total_tokens, num_heads, head_dim) + k0_flat = k0_fa.reshape(total_tokens, num_heads, head_dim) + v0_flat = v0_fa.reshape(total_tokens, num_heads, head_dim) + out_flat, lse_flat, _ = _flash_attn_varlen_func( + q_flat, + k0_flat, + v0_flat, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + softmax_scale=self.scaling, + causal=True, + return_attn_probs=True, + ) + # out_flat: [total_tokens, H, D] -> [B, H, T, D] + attn_output_bhtd = out_flat.reshape(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) + # lse_flat: [H, total_tokens] -> [B, H, T] + if lse_flat.shape != (num_heads, total_tokens): + raise RuntimeError( + f"Unexpected varlen softmax_lse shape {tuple(lse_flat.shape)}; " + f"expected {(num_heads, total_tokens)}. Verify the installed flash-attn version." + ) + lse_fa = lse_flat.transpose(0, 1).reshape(batch_size, seq_len, num_heads).permute(0, 2, 1) + return attn_output_bhtd, lse_fa + class Eagle3LlamaMLP(nn.Module): """Standard Llama-style SwiGLU MLP on hidden-size activations.""" @@ -455,6 +529,8 @@ def forward( attention_mask: torch.Tensor, position_ids: torch.Tensor, cache_hidden: list[list[torch.Tensor]], + cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, ) -> torch.Tensor: residual = hidden_states norm_input_embeds = self.input_layernorm(input_embeds) @@ -465,6 +541,8 @@ def forward( attention_mask, position_ids, cache_hidden=cache_hidden, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, ) residual = hidden_states @@ -662,6 +740,7 @@ def forward( attention_mask: torch.Tensor, position_ids: Optional[torch.Tensor] = None, cache_hidden: Optional[list[list[torch.Tensor]]] = None, + seq_lens: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Run one full-sequence draft update step. @@ -671,28 +750,48 @@ def forward( to it. If ``None`` is passed (e.g. from a one-shot evaluation call) a fresh ``[[], []]`` is allocated locally -- step 0 of TTT is mathematically equivalent to a plain causal forward. + + ``seq_lens`` (packing) makes Block-1 attention document-level block-causal + (eager mask / FA2 varlen); callers must pass per-document ``position_ids``. """ if position_ids is None: position_ids = torch.arange(input_ids.shape[1], device=input_ids.device, dtype=torch.long).unsqueeze(0) position_ids = position_ids.expand(input_ids.shape[0], -1) if cache_hidden is None: cache_hidden = [[], []] - if self.model.layers[ - 0 - ].self_attn.attn_implementation == "flash_attention_2" and not _is_right_padded_attention_mask(attention_mask): - raise ValueError( - "LlamaEagle3DraftModel: attn_implementation='flash_attention_2' requires a right-padded " - "attention_mask (each row must be contiguous 1s followed by 0s)." - ) + is_fa2 = self.model.layers[0].self_attn.attn_implementation == "flash_attention_2" + cu_seqlens: torch.Tensor | None = None + max_seqlen: int | None = None + if seq_lens is not None: + # Packed: structure comes from seq_lens (cu_seqlens for FA2 / block-causal + # mask for eager), so the right-padding check below does not apply. + seq_length = input_ids.shape[1] + if is_fa2: + # FA2 attends document-wise through cu_seqlens and never reads the 4D + # additive mask, so skip materializing the [B, 1, T, T] block-causal mask. + causal_mask = None + cu_seqlens, max_seqlen = _seq_lens_to_cu_seqlens(seq_lens, seq_length) + else: + causal_mask = build_block_causal_additive_mask( + seq_lens, seq_length=seq_length, dtype=projected_hidden_states.dtype, device=input_ids.device + ) + else: + if is_fa2 and not _is_right_padded_attention_mask(attention_mask): + raise ValueError( + "LlamaEagle3DraftModel: attn_implementation='flash_attention_2' requires a right-padded " + "attention_mask (each row must be contiguous 1s followed by 0s)." + ) + causal_mask = _build_causal_mask(attention_mask=attention_mask, dtype=projected_hidden_states.dtype) draft_input_embeds = self.embed_input_ids(input_ids) - causal_mask = _build_causal_mask(attention_mask=attention_mask, dtype=projected_hidden_states.dtype) hidden_states = self.model.layers[0]( input_embeds=draft_input_embeds, hidden_states=projected_hidden_states, attention_mask=causal_mask, position_ids=position_ids, cache_hidden=cache_hidden, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, ) # EAGLE-3.1 ``norm_output``: route the post-norm hidden state to both # the next TTT step (fed back via ``cur_hidden_states`` in the trainer diff --git a/nemo_automodel/components/speculative/eagle/target.py b/nemo_automodel/components/speculative/eagle/target.py index 862b92b144..9d073b8e87 100644 --- a/nemo_automodel/components/speculative/eagle/target.py +++ b/nemo_automodel/components/speculative/eagle/target.py @@ -23,6 +23,7 @@ import torch import torch.nn as nn +from nemo_automodel.components.datasets.llm.packed_sequence import build_block_causal_additive_mask from nemo_automodel.components.speculative.eagle.backend import Eagle3TargetBackend @@ -60,6 +61,12 @@ class Eagle3TargetBatch: logits: torch.Tensor | None = None target_probs: torch.Tensor | None = None position_mask: torch.Tensor | None = None + # Packing metadata (None unless packing is enabled), unshifted slot frame: + # per-document position_ids / seq_lens (block-causal mask) and doc_remaining + # (gates cross-document TTT supervision). + position_ids: torch.Tensor | None = None + seq_lens: torch.Tensor | None = None + doc_remaining: torch.Tensor | None = None def __post_init__(self) -> None: has_logits = self.logits is not None @@ -85,6 +92,10 @@ def to_trainer_inputs(self) -> dict[str, torch.Tensor]: else: inputs["target_probs"] = self.target_probs inputs["position_mask"] = self.position_mask + if self.seq_lens is not None: + inputs["position_ids"] = self.position_ids + inputs["seq_lens"] = self.seq_lens + inputs["doc_remaining"] = self.doc_remaining return inputs @@ -173,8 +184,17 @@ def generate_batch( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor, + position_ids: torch.Tensor | None = None, + seq_lens: torch.Tensor | None = None, + doc_remaining: torch.Tensor | None = None, ) -> Eagle3TargetBatch: - """Run the target model and capture aux hidden states plus logits.""" + """Run the target model and capture aux hidden states plus logits. + + With ``seq_lens`` (packing), the target runs with a ``[B, 1, T, T]`` + block-causal mask and per-document ``position_ids`` so its outputs respect + document boundaries; the packing metadata is forwarded unshifted to the + trainer. ``seq_lens=None`` keeps the original 2D-mask path. + """ layers = self._get_transformer_layers() captured: dict[int, torch.Tensor] = {} handles = [] @@ -199,10 +219,48 @@ def _hook(_module, _inputs, outputs): name: False for name in ("output_hidden_states", "output_attentions", "use_cache") if name in forward_params } + # Packing isolates documents per attention backend; the mask strategy + # differs because FlashAttention has no 4D-mask code path: + # * SDPA / eager consume the [B, 1, T, T] block-causal additive mask. + # * FlashAttention infers per-document cu_seqlens from the reset points + # in a per-document ``position_ids`` and is passed ``attention_mask=None``. + # Feeding FA the 4D additive mask instead drives its unpad gather out + # of bounds: the mask flattens to B*T*T entries and the gather indexes + # a B*T-row tensor with them. transformers only packs from position_ids + # at batch size 1 (see ``_is_packed_sequence``). + target_attention_mask = attention_mask + if seq_lens is not None: + if position_ids is None or "position_ids" not in forward_params: + raise ValueError( + "EAGLE-3 sequence packing requires per-document position_ids, but none were " + "provided or the target model's forward does not accept a `position_ids` argument." + ) + extra_kwargs["position_ids"] = position_ids + attn_impl = getattr(self.model.config, "_attn_implementation", None) or "" + if "flash" in attn_impl: + if input_ids.shape[0] != 1: + raise ValueError( + "EAGLE-3 sequence packing with a FlashAttention target only supports " + f"micro_batch_size=1 (got {input_ids.shape[0]}). FlashAttention infers " + "document boundaries from per-document position_ids, which transformers " + "packs only at batch size 1. Set micro_batch_size=1 or load the target " + "with attn_implementation='sdpa'." + ) + # attention_mask=None + per-document position_ids -> FA varlen packing. + target_attention_mask = None + else: + param_dtype = next(self.model.parameters()).dtype + target_attention_mask = build_block_causal_additive_mask( + seq_lens, + seq_length=input_ids.shape[1], + dtype=param_dtype, + device=input_ids.device, + ) + try: outputs = self.model( input_ids=input_ids, - attention_mask=attention_mask, + attention_mask=target_attention_mask, **extra_kwargs, ) finally: @@ -227,4 +285,7 @@ def _hook(_module, _inputs, outputs): input_ids=shifted_input_ids, attention_mask=attention_mask, loss_mask=shifted_loss_mask, + position_ids=position_ids, + seq_lens=seq_lens, + doc_remaining=doc_remaining, ) diff --git a/nemo_automodel/recipes/llm/train_eagle3.py b/nemo_automodel/recipes/llm/train_eagle3.py index 3fce8f1c76..b1bf998e9a 100644 --- a/nemo_automodel/recipes/llm/train_eagle3.py +++ b/nemo_automodel/recipes/llm/train_eagle3.py @@ -344,6 +344,14 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): # precomputed supervision over HTTP + NCCL. No target weights are # loaded here, which frees the training GPU's memory. backend = recipe_cfg.get("target_model_backend", "colocated") + # Sequence packing is colocated-only (the remote server does not yet honor + # per-document masking). + packed_sequence_size = recipe_cfg.get("packed_sequence_size", 0) + if packed_sequence_size > 0 and backend == "remote": + raise NotImplementedError( + "packed_sequence_size > 0 is only supported with the colocated target backend; " + "the remote backend does not yet propagate per-document masking." + ) if backend == "remote": self._setup_remote_target(recipe_cfg) elif backend == "colocated": @@ -362,6 +370,7 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), + packed_sequence_size=packed_sequence_size, ) self.val_dataloader = None if recipe_cfg.get("val_data_path", None): @@ -376,6 +385,7 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): distributed=self.dist_env.world_size > 1, shuffle_seed=recipe_cfg.get("shuffle_seed", 42), mask_reasoning_content=recipe_cfg.get("mask_reasoning_content", False), + packed_sequence_size=packed_sequence_size, ) special_token_ids = [ @@ -536,6 +546,14 @@ def _forward_batch(self, batch, target_batch=None): if target_batch is not None: return self.trainer_module(**target_batch.to_trainer_inputs()) batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()} + # Sequence-packing metadata (present only when packed_sequence_size > 0). + packing_kwargs = {} + if "seq_lens" in batch: + packing_kwargs = { + "position_ids": batch["position_ids"], + "seq_lens": batch["seq_lens"], + "doc_remaining": batch["doc_remaining"], + } if self.target_wrapper is None: # Offline cache: the supervision is already in the batch. batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()} @@ -546,11 +564,13 @@ def _forward_batch(self, batch, target_batch=None): aux_hidden_states=batch["aux_hidden_states"], target_probs=batch["target_probs"], position_mask=batch["position_mask"], + **packing_kwargs, ) target_batch = self.target_wrapper.generate_batch( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], loss_mask=batch["loss_mask"], + **packing_kwargs, ) return self.trainer_module(**target_batch.to_trainer_inputs()) diff --git a/tests/functional_tests/speculative/test_eagle3_packing_fa2_parity.py b/tests/functional_tests/speculative/test_eagle3_packing_fa2_parity.py new file mode 100644 index 0000000000..c08be7762f --- /dev/null +++ b/tests/functional_tests/speculative/test_eagle3_packing_fa2_parity.py @@ -0,0 +1,297 @@ +# 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. + +"""FlashAttention-2 EAGLE-3 packing parity + speed check on GPU. + +Random-initialises four documents of length 600 / 700 / 800 / 9000 (the last is +truncated to ``SEQ_LENGTH``) and trains them once each, ``micro_batch_size=1`` +(the recipe default; FA2 position-id packing requires batch size 1), under two +layouts: + + * no packing: four ``[1, SEQ_LENGTH]`` rows, one padded document each; + * packing: greedily packed into ``[1, SEQ_LENGTH]`` rows (here two: the three + short docs share one row, the truncated long doc fills another). + +The target (FA2, the path under test) isolates documents from per-document +position_ids; the draft runs eager (the recipe default) in fp32, fed the target's +bf16 aux/logits cast up at the boundary. + +Both layouts supervise the identical (document, position, TTT-step) triples, so +the valid-token-weighted global loss and accumulated draft gradients match within +bf16 FA2 tolerance, while packing drops the padding compute and halves the steps. +Reports the speedup and the loss / gradient deltas. + +Run directly for the numbers (``pytest -s`` or ``python``); needs one GPU with a +working flash-attn build. +""" + +from __future__ import annotations + +import importlib.util +import time + +import pytest +import torch + +from nemo_automodel.components.datasets.llm.eagle3 import _pack_collate, build_packed_eagle3_dataset +from nemo_automodel.components.speculative.eagle.core import Eagle3TrainerModule +from nemo_automodel.components.speculative.eagle.draft_llama import LlamaEagle3DraftModel +from nemo_automodel.components.speculative.eagle.target import HFEagle3TargetModel + +_HAS_FA = importlib.util.find_spec("flash_attn") is not None + +from transformers import LlamaConfig, LlamaForCausalLM + +SEQ_LENGTH = 4096 +# Four rollouts of the same ascending length profile (9000 is truncated to +# SEQ_LENGTH), enough documents to exercise multi-row packing. +DOC_LENS = [ + 600, + 700, + 800, + 900, + 1000, + 1100, + 1200, + 1300, + 1400, + 1500, + 1600, + 1700, + 1800, + 1900, + 2000, + 3000, + 4000, + 5000, + 6000, + 7000, + 8000, + 9000, +] * 4 +HIDDEN = 512 +VOCAB = 2048 +TARGET_LAYERS = 8 # deep enough for the default aux ids [1, 3, 4] +TTT_STEPS = 7 +# Target runs the FA2 path under test in bf16 (FA2 has no fp32 kernel); the draft +# stays fp32 so its fp32-RoPE q/k stay dtype-consistent and parity isn't muddied +# by extra draft-side bf16 noise. Target aux/logits are cast to fp32 at the boundary. +TARGET_DTYPE = torch.bfloat16 +DRAFT_DTYPE = torch.float32 + + +def _make_documents() -> list[dict[str, list[int]]]: + """Four random documents (input_ids + all-ones loss_mask), pre-truncated to T.""" + torch.manual_seed(7) + docs = [] + for length in DOC_LENS: + eff = min(length, SEQ_LENGTH) + ids = torch.randint(0, VOCAB, (eff,)).tolist() + docs.append({"input_ids": ids, "loss_mask": [1] * eff}) + return docs + + +def _build_target() -> HFEagle3TargetModel: + config = LlamaConfig( + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=TARGET_LAYERS, + num_attention_heads=8, + num_key_value_heads=8, + vocab_size=VOCAB, + max_position_embeddings=SEQ_LENGTH, + attn_implementation="flash_attention_2", + ) + target = LlamaForCausalLM(config).to(device="cuda", dtype=TARGET_DTYPE).eval() + target.requires_grad_(False) + return HFEagle3TargetModel(target) + + +def _build_trainer() -> Eagle3TrainerModule: + """Identical draft init for both layouts (fixed seed); eager fp32 draft.""" + torch.manual_seed(123) + draft_config = LlamaConfig( + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=1, + num_attention_heads=8, + num_key_value_heads=8, + vocab_size=VOCAB, + max_position_embeddings=SEQ_LENGTH, + ) + draft_config.draft_vocab_size = VOCAB # full vocab -> every position supervised + draft_config.target_hidden_size = HIDDEN + # Draft stays on eager (the recipe default): its FA2 path upcasts q/k via fp32 + # RoPE, which then mismatches the bf16 value cache. The target is the FA2 path + # under test here. + draft_config.attn_implementation = "eager" + draft = LlamaEagle3DraftModel(draft_config).to(device="cuda", dtype=DRAFT_DTYPE) + selected_token_ids = torch.arange(VOCAB, dtype=torch.long, device="cuda") + selected_token_mask = torch.ones(VOCAB, dtype=torch.bool, device="cuda") + return Eagle3TrainerModule( + draft, + selected_token_ids=selected_token_ids, + selected_token_mask=selected_token_mask, + ttt_steps=TTT_STEPS, + ).to("cuda") + + +def _to_cuda(batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {k: v.to("cuda", non_blocking=True) for k, v in batch.items()} + + +def _padded_batches(docs: list[dict[str, list[int]]]) -> list[dict[str, torch.Tensor]]: + """Layout A: one padded ``[1, SEQ_LENGTH]`` row per document.""" + batches = [] + for doc in docs: + ids = doc["input_ids"] + input_ids = torch.zeros(1, SEQ_LENGTH, dtype=torch.long) + loss_mask = torch.zeros(1, SEQ_LENGTH, dtype=torch.long) + attention_mask = torch.zeros(1, SEQ_LENGTH, dtype=torch.long) + input_ids[0, : len(ids)] = torch.tensor(ids, dtype=torch.long) + loss_mask[0, : len(ids)] = 1 + attention_mask[0, : len(ids)] = 1 + batches.append(_to_cuda({"input_ids": input_ids, "attention_mask": attention_mask, "loss_mask": loss_mask})) + return batches + + +def _packed_batches(docs: list[dict[str, list[int]]]) -> list[dict[str, torch.Tensor]]: + """Layout B: greedily packed rows, one ``[1, SEQ_LENGTH]`` batch per pack.""" + packs = build_packed_eagle3_dataset(docs, packed_sequence_size=SEQ_LENGTH, pad_token_id=0) + return [_to_cuda(_pack_collate([pack])) for pack in packs] + + +def _step(target_wrapper, trainer, batch, *, packed: bool): + """One micro-batch training step; seq_lens present -> packed path.""" + kwargs = {} + if packed: + kwargs = { + "position_ids": batch["position_ids"], + "seq_lens": batch["seq_lens"], + "doc_remaining": batch["doc_remaining"], + } + target_batch = target_wrapper.generate_batch( + input_ids=batch["input_ids"], + attention_mask=batch["attention_mask"], + loss_mask=batch["loss_mask"], + **kwargs, + ) + inputs = target_batch.to_trainer_inputs() + # Bridge the bf16 target -> fp32 draft: cast the float supervision tensors. + for key in ("aux_hidden_states", "target_logits", "target_probs"): + if inputs.get(key) is not None: + inputs[key] = inputs[key].float() + return trainer(**inputs) + + +def _run_layout(target_wrapper, trainer, batches, *, packed: bool): + """Accumulate a valid-token-weighted global loss + gradients over all batches. + + Each micro-batch returns a mean loss over its own valid tokens; backprop + ``loss * valid`` so ``.grad`` accumulates the gradient of the summed CE, then + divide by the global valid count. With identical supervision the two layouts + therefore land on the same global loss and gradients. + """ + trainer.zero_grad(set_to_none=True) + total_ce = 0.0 + total_valid = 0 + for batch in batches: + metrics = _step(target_wrapper, trainer, batch, packed=packed) + valid = metrics.valid_tokens + (metrics.loss * valid).backward() + total_ce += (metrics.loss.detach() * valid).item() + total_valid += int(valid.item()) + global_loss = total_ce / max(total_valid, 1) + grads = { + n: (p.grad.detach() / max(total_valid, 1)).clone() for n, p in trainer.named_parameters() if p.grad is not None + } + return global_loss, grads, total_valid + + +def _time_layout(target_wrapper, trainer, batches, *, packed: bool, warmup: int = 2, iters: int = 8) -> float: + """Median wall-clock seconds for one full pass (fwd+bwd) over all the docs.""" + timings = [] + for it in range(warmup + iters): + trainer.zero_grad(set_to_none=True) + torch.cuda.synchronize() + t0 = time.perf_counter() + for batch in batches: + metrics = _step(target_wrapper, trainer, batch, packed=packed) + metrics.loss.backward() + torch.cuda.synchronize() + if it >= warmup: + timings.append(time.perf_counter() - t0) + timings.sort() + return timings[len(timings) // 2] + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.skipif(not _HAS_FA, reason="requires flash-attn") +def test_eagle3_packing_fa2_parity_and_speed(): + docs = _make_documents() + target_wrapper = _build_target() + + padded = _padded_batches(docs) + packed = _packed_batches(docs) + + # --- Parity: fresh (identically-initialised) trainer per layout. --- + trainer_a = _build_trainer() + loss_a, grads_a, valid_a = _run_layout(target_wrapper, trainer_a, padded, packed=False) + + trainer_b = _build_trainer() + loss_b, grads_b, valid_b = _run_layout(target_wrapper, trainer_b, packed, packed=True) + + assert valid_a == valid_b, f"valid-token counts differ: {valid_a} vs {valid_b}" + assert set(grads_a) == set(grads_b) + + loss_abs = abs(loss_a - loss_b) + loss_rel = loss_abs / max(abs(loss_a), 1e-8) + + max_grad_rel = 0.0 + max_grad_name = "" + for name in grads_a: + ga, gb = grads_a[name], grads_b[name] + denom = ga.abs().max().clamp_min(1e-6) + rel = ((ga - gb).abs().max() / denom).item() + if rel > max_grad_rel: + max_grad_rel, max_grad_name = rel, name + + # --- Speed: median full-pass wall-clock for each layout. --- + t_padded = _time_layout(target_wrapper, trainer_a, padded, packed=False) + t_packed = _time_layout(target_wrapper, trainer_b, packed, packed=True) + + real_tokens = sum(min(length, SEQ_LENGTH) for length in DOC_LENS) + + print("\n=== EAGLE-3 FA2 packing parity & speed (micro_batch_size=1) ===") + print(f"docs (truncated) : {[min(length, SEQ_LENGTH) for length in DOC_LENS]} (real tokens={real_tokens})") + print(f"no-packing steps : {len(padded)} rows x {SEQ_LENGTH} = {len(padded) * SEQ_LENGTH} tok (padding waste)") + print(f"packing steps : {len(packed)} rows x {SEQ_LENGTH} = {len(packed) * SEQ_LENGTH} tok") + print(f"valid supervised toks : {valid_a} (both layouts)") + print(f"loss no-packing : {loss_a:.6f}") + print(f"loss packing : {loss_b:.6f}") + print(f"loss |abs| / rel : {loss_abs:.3e} / {loss_rel:.3e}") + print(f"max grad rel diff : {max_grad_rel:.3e} ({max_grad_name})") + print(f"full-pass no-packing : {t_padded * 1e3:.2f} ms ({len(padded)} steps)") + print(f"full-pass packing : {t_packed * 1e3:.2f} ms ({len(packed)} steps)") + print(f"speedup (padded/packed): {t_padded / t_packed:.2f}x") + + # bf16 FA2 target: parity is approximate. Tolerances are loose but would blow + # up by orders of magnitude on a real bug (e.g. cross-document leakage). + assert loss_rel < 5e-2, f"loss relative diff too large: {loss_rel:.3e}" + assert max_grad_rel < 2e-1, f"grad relative diff too large: {max_grad_rel:.3e} ({max_grad_name})" + + +if __name__ == "__main__": + test_eagle3_packing_fa2_parity_and_speed() diff --git a/tests/unit_tests/speculative/test_eagle3_packing.py b/tests/unit_tests/speculative/test_eagle3_packing.py new file mode 100644 index 0000000000..23004d9344 --- /dev/null +++ b/tests/unit_tests/speculative/test_eagle3_packing.py @@ -0,0 +1,446 @@ +# 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. + +"""CPU tests for EAGLE-3 sequence packing (dataset, masks, trainer integration).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from nemo_automodel.components.datasets.llm.eagle3 import ( + _pack_collate, + build_packed_eagle3_dataset, +) +from nemo_automodel.components.datasets.llm.packed_sequence import build_block_causal_additive_mask +from nemo_automodel.components.speculative.eagle.core import Eagle3TrainerModule +from nemo_automodel.components.speculative.eagle.draft_llama import ( + LlamaEagle3DraftModel, + _seq_lens_to_cu_seqlens, +) +from nemo_automodel.components.speculative.eagle.target import HFEagle3TargetModel + + +def _src(samples: list[tuple[list[int], list[int]]]) -> list[dict]: + """Build a tiny in-memory source dataset of (input_ids, loss_mask) pairs.""" + return [{"input_ids": ids, "loss_mask": loss} for ids, loss in samples] + + +def test_build_packed_dataset_structure(): + """Packing resets position_ids per doc, sums seq_lens to T, and builds doc_remaining.""" + # Two short docs that fit one row of width 8, plus a third that starts a new row. + src = _src( + [ + ([10, 11, 12], [0, 1, 1]), # doc A, len 3 + ([20, 21], [1, 1]), # doc B, len 2 -> A+B = 5 <= 8, same pack + ([30, 31, 32, 33, 34, 35, 36], [1, 1, 1, 1, 1, 1, 1]), # doc C len 7 -> new pack + ] + ) + packs = build_packed_eagle3_dataset(src, packed_sequence_size=8, pad_token_id=0) + assert len(packs) == 2 + + pack0 = packs[0] + # input_ids: A(3) + B(2) + pad(3) == 8 + assert pack0["input_ids"] == [10, 11, 12, 20, 21, 0, 0, 0] + # position_ids reset per doc, pad continues the last doc (clamped). + assert pack0["position_ids"][:5] == [0, 1, 2, 0, 1] + # seq_lens: trailing pad folded into the final doc -> [3, 2+3] sums to 8. + assert pack0["seq_lens"] == [3, 5] + assert sum(pack0["seq_lens"]) == 8 + # doc_remaining: tokens after each slot within its (real) doc; pad -> 0. + assert pack0["doc_remaining"] == [2, 1, 0, 1, 0, 0, 0, 0] + # attention_mask: 1 for the 5 real tokens, 0 for the 3 pad. + assert pack0["attention_mask"] == [1, 1, 1, 1, 1, 0, 0, 0] + + pack1 = packs[1] + assert pack1["input_ids"] == [30, 31, 32, 33, 34, 35, 36, 0] + assert pack1["seq_lens"] == [8] # len 7 + 1 pad folded + assert pack1["doc_remaining"] == [6, 5, 4, 3, 2, 1, 0, 0] + + +def test_pack_collate_pads_ragged_seq_lens(): + """_pack_collate stacks fixed-width fields and 0-pads ragged seq_lens.""" + src = _src([([1, 2, 3], [1, 1, 1]), ([4, 5], [1, 1]), ([6, 7, 8, 9], [1, 1, 1, 1])]) + packs = build_packed_eagle3_dataset(src, packed_sequence_size=6, pad_token_id=0) + batch = _pack_collate(packs) + bsz = len(packs) + assert batch["input_ids"].shape == (bsz, 6) + assert batch["doc_remaining"].shape == (bsz, 6) + assert batch["position_ids"].shape == (bsz, 6) + # seq_lens padded to [B, max_docs]; each row's nonzero entries sum to 6. + assert batch["seq_lens"].shape[0] == bsz + assert torch.all(batch["seq_lens"].sum(dim=1) == 6) + + +def test_block_causal_additive_mask(): + """Block-causal mask is in-doc lower-triangular and blocks cross-document attention.""" + seq_lens = torch.tensor([[3, 2]], dtype=torch.long) # docs over T=5 + mask = build_block_causal_additive_mask(seq_lens, seq_length=5, dtype=torch.float32, device=torch.device("cpu")) + assert mask.shape == (1, 1, 5, 5) + neg = torch.finfo(torch.float32).min + m = mask[0, 0] + # Doc A (0..2): causal within doc. + assert m[0, 0] == 0 and m[2, 0] == 0 and m[2, 2] == 0 + assert m[0, 1] == neg # future within doc A is masked + # Doc B (3..4): causal within doc. + assert m[3, 3] == 0 and m[4, 3] == 0 and m[4, 4] == 0 + # Cross-document is masked both directions. + assert m[3, 0] == neg and m[3, 2] == neg # doc B query cannot see doc A + assert m[2, 3] == neg # doc A query cannot see doc B + + +def test_seq_lens_to_cu_seqlens(): + """cu_seqlens is int32, monotonic, row-major, and sums to B*T; max_seqlen is the longest doc.""" + seq_lens = torch.tensor([[3, 2], [4, 1]], dtype=torch.long) # B=2, T=5 + cu, max_seqlen = _seq_lens_to_cu_seqlens(seq_lens, seq_length=5) + assert cu.dtype == torch.int32 + assert cu.tolist() == [0, 3, 5, 9, 10] # row-major doc lens [3,2,4,1] + assert int(cu[-1]) == 2 * 5 + assert max_seqlen == 4 + + +def _build_tiny_draft_model(attn_implementation: str = "eager") -> LlamaEagle3DraftModel: + config = LlamaConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + max_position_embeddings=64, + ) + config.torch_dtype = torch.float32 + config.draft_vocab_size = 16 + config.target_hidden_size = 32 + config.attn_implementation = attn_implementation + return LlamaEagle3DraftModel(config).to(torch.float32) + + +def test_packed_draft_attention_isolates_documents(): + """Block-causal draft forward must isolate documents: perturbing only doc B's + inputs leaves doc A's output bit-identical.""" + torch.manual_seed(0) + draft = _build_tiny_draft_model("eager") + draft.eval() + + seq_len = 6 + seq_lens = torch.tensor([[3, 3]], dtype=torch.long) # doc A: 0..2, doc B: 3..5 + position_ids = torch.tensor([[0, 1, 2, 0, 1, 2]], dtype=torch.long) + input_ids = torch.randint(0, 16, (1, seq_len)) + projected = torch.randn(1, seq_len, draft.config.hidden_size) + + def run(ids, proj): + return draft( + input_ids=ids, + projected_hidden_states=proj, + attention_mask=torch.ones(1, seq_len, dtype=torch.long), + position_ids=position_ids, + cache_hidden=[[], []], + seq_lens=seq_lens, + ) + + out_ref = run(input_ids, projected) + + # Perturb only document B (slots 3..5). + ids_b = input_ids.clone() + ids_b[:, 3:] = torch.randint(0, 16, (1, 3)) + proj_b = projected.clone() + proj_b[:, 3:] = torch.randn(1, 3, draft.config.hidden_size) + out_perturbed = run(ids_b, proj_b) + + # Document A (slots 0..2) output must be unchanged; document B may differ. + torch.testing.assert_close(out_ref[:, :3], out_perturbed[:, :3]) + assert not torch.allclose(out_ref[:, 3:], out_perturbed[:, 3:]) + + +def test_packed_trainer_forward_runs_and_backprops(): + """End-to-end packed trainer forward produces a finite loss and non-NaN grads.""" + torch.manual_seed(0) + draft = _build_tiny_draft_model("eager") + config = draft.config + selected_token_ids = torch.arange(config.draft_vocab_size, dtype=torch.long) + selected_token_mask = torch.zeros(config.vocab_size, dtype=torch.bool) + selected_token_mask[selected_token_ids] = True + + trainer = Eagle3TrainerModule( + draft, + selected_token_ids=selected_token_ids, + selected_token_mask=selected_token_mask, + ttt_steps=3, + ) + + src = _src( + [ + ([3, 4, 5, 6], [0, 1, 1, 1]), + ([7, 8, 9], [0, 1, 1]), + ([10, 11, 12, 13, 14], [0, 1, 1, 1, 1]), + ] + ) + packs = build_packed_eagle3_dataset(src, packed_sequence_size=8, pad_token_id=0) + batch = _pack_collate(packs) + bsz, seq_len = batch["input_ids"].shape + + target_logits = torch.randn(bsz, seq_len, config.vocab_size) + aux_hidden_states = torch.randn(bsz, seq_len, config.hidden_size * 3) + + metrics = trainer( + input_ids=batch["input_ids"], + attention_mask=batch["attention_mask"], + loss_mask=batch["loss_mask"], + aux_hidden_states=aux_hidden_states, + target_logits=target_logits, + position_ids=batch["position_ids"], + seq_lens=batch["seq_lens"], + doc_remaining=batch["doc_remaining"], + ) + assert torch.isfinite(metrics.loss) + metrics.loss.backward() + grads = [p.grad for p in trainer.parameters() if p.grad is not None] + assert grads, "expected at least one parameter to receive a gradient" + assert all(torch.isfinite(g).all() for g in grads) + + +def test_packing_matches_padding_loss_and_grads(): + """Golden parity: packing N docs into one row == N padded single-doc rows. + + Both layouts run through the same target wrapper + draft trainer: + * no packing: ``[D, T]``, each doc ``L`` real tokens padded to ``T``; + * packing: the same docs as ``[1, T]`` (``D * L == T``) with per-document + position_ids, block-causal attention, and ``doc_remaining`` gating. + + They supervise the identical (doc, position, TTT step) triples against + identical targets, so loss and every gradient match (CPU/fp32, tight tol). + Scaling L/T up and the tiny target to Qwen3-8B on GPU is a drop-in change. + """ + hidden = 32 + vocab = 128 + num_docs = 4 + doc_len = 16 + total = num_docs * doc_len # packed row width T + + target_config = LlamaConfig( + hidden_size=hidden, + intermediate_size=64, + num_hidden_layers=8, # deep enough for the default aux ids [1, 3, 4] + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=vocab, + max_position_embeddings=total, + ) + torch.manual_seed(1) + target = LlamaForCausalLM(target_config).to(torch.float32).eval() + target_wrapper = HFEagle3TargetModel(target) + + draft_config = LlamaConfig( + hidden_size=hidden, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=vocab, + max_position_embeddings=total, + ) + draft_config.draft_vocab_size = vocab # full vocab so every position is supervised + draft_config.target_hidden_size = hidden + draft_config.attn_implementation = "eager" + + def build_trainer(): + torch.manual_seed(123) # identical draft init for both layouts + draft = LlamaEagle3DraftModel(draft_config).to(torch.float32) + selected_token_ids = torch.arange(vocab, dtype=torch.long) + selected_token_mask = torch.ones(vocab, dtype=torch.bool) + return Eagle3TrainerModule( + draft, + selected_token_ids=selected_token_ids, + selected_token_mask=selected_token_mask, + ttt_steps=3, + ) + + torch.manual_seed(7) + docs = [torch.randint(0, vocab, (doc_len,)) for _ in range(num_docs)] + + # Layout A: D padded single-document rows. + ids_a = torch.zeros(num_docs, total, dtype=torch.long) + loss_a = torch.zeros(num_docs, total, dtype=torch.long) + attn_a = torch.zeros(num_docs, total, dtype=torch.long) + for d in range(num_docs): + ids_a[d, :doc_len] = docs[d] + loss_a[d, :doc_len] = 1 + attn_a[d, :doc_len] = 1 + trainer_a = build_trainer() + batch_a = target_wrapper.generate_batch(input_ids=ids_a, attention_mask=attn_a, loss_mask=loss_a) + metrics_a = trainer_a(**batch_a.to_trainer_inputs()) + metrics_a.loss.backward() + grads_a = {n: p.grad.clone() for n, p in trainer_a.named_parameters() if p.grad is not None} + + # Layout B: the same documents packed into one row (no padding). + ids_b = torch.cat(docs).unsqueeze(0) + loss_b = torch.ones(1, total, dtype=torch.long) + attn_b = torch.ones(1, total, dtype=torch.long) + position_ids = torch.cat([torch.arange(doc_len) for _ in range(num_docs)]).unsqueeze(0) + doc_remaining = torch.cat([torch.arange(doc_len - 1, -1, -1) for _ in range(num_docs)]).unsqueeze(0) + seq_lens = torch.tensor([[doc_len] * num_docs], dtype=torch.long) + trainer_b = build_trainer() + batch_b = target_wrapper.generate_batch( + input_ids=ids_b, + attention_mask=attn_b, + loss_mask=loss_b, + position_ids=position_ids, + seq_lens=seq_lens, + doc_remaining=doc_remaining, + ) + metrics_b = trainer_b(**batch_b.to_trainer_inputs()) + metrics_b.loss.backward() + grads_b = {n: p.grad.clone() for n, p in trainer_b.named_parameters() if p.grad is not None} + + # Identical supervised-token count, loss, and gradients (tight CPU/fp32 tol). + assert metrics_a.valid_tokens.item() == metrics_b.valid_tokens.item() + torch.testing.assert_close(metrics_a.loss, metrics_b.loss, rtol=1e-4, atol=1e-5) + assert set(grads_a) == set(grads_b) + for name in grads_a: + torch.testing.assert_close(grads_a[name], grads_b[name], rtol=1e-4, atol=1e-5, msg=f"grad mismatch: {name}") + + +def test_doc_remaining_gating_masks_cross_document_supervision(): + """At TTT step k, supervision is dropped where k >= doc_remaining (cross-doc target).""" + torch.manual_seed(0) + draft = _build_tiny_draft_model("eager") + config = draft.config + selected_token_ids = torch.arange(config.draft_vocab_size, dtype=torch.long) + selected_token_mask = torch.zeros(config.vocab_size, dtype=torch.bool) + selected_token_mask[selected_token_ids] = True + trainer = Eagle3TrainerModule( + draft, + selected_token_ids=selected_token_ids, + selected_token_mask=selected_token_mask, + ttt_steps=4, + ) + + # Single row, two docs of length 4 each (no padding): T=8. + seq_len = 8 + seq_lens = torch.tensor([[4, 4]], dtype=torch.long) + position_ids = torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3]], dtype=torch.long) + doc_remaining = torch.tensor([[3, 2, 1, 0, 3, 2, 1, 0]], dtype=torch.long) + input_ids = torch.randint(0, 16, (1, seq_len)) + # Supervise every position; bias the target argmax into the draft vocab so + # ``selected_token_mask`` always passes and masking is driven purely by the + # loss_mask shift + doc_remaining gating. + loss_mask = torch.ones(1, seq_len, dtype=torch.long) + target_logits = torch.randn(1, seq_len, config.vocab_size) + target_logits[..., : config.draft_vocab_size] += 50.0 + aux_hidden_states = torch.randn(1, seq_len, config.hidden_size * 3) + + def valid_count(doc_rem): + return trainer( + input_ids=input_ids, + attention_mask=torch.ones(1, seq_len, dtype=torch.long), + loss_mask=loss_mask, + aux_hidden_states=aux_hidden_states, + target_logits=target_logits, + position_ids=position_ids, + seq_lens=seq_lens, + doc_remaining=doc_rem, + ).valid_tokens.item() + + # Without gating the loss_mask only shrinks via the per-step left-shift + # (zero-filled tail): step k supervises 8-k slots -> 8+7+6+5 = 26. + assert valid_count(None) == 26 + # With gating, slot t at step k is kept only while k < doc_remaining[t]: + # step0: dr>0 at {0,1,2,4,5,6} -> 6 + # step1: dr>1 at {0,1,4,5} -> 4 + # step2: dr>2 at {0,4} -> 2 + # step3: dr>3 -> 0 + # total 12, strictly fewer than 26 -- every cross-document target is dropped. + assert valid_count(doc_remaining) == 12 + + +class _RecordingFlashTarget(torch.nn.Module): + """Minimal causal-LM stub reporting ``_attn_implementation='flash_attention_2'``. + + Records the ``attention_mask`` / ``position_ids`` it is forwarded so packing's + FlashAttention dispatch can be asserted on CPU (real FA kernels need a GPU). + """ + + def __init__(self, num_layers: int = 8, hidden: int = 16, vocab: int = 32): + super().__init__() + self.config = SimpleNamespace(_attn_implementation="flash_attention_2", num_hidden_layers=num_layers) + self.embed_tokens = torch.nn.Embedding(vocab, hidden) + # _get_transformer_layers() looks for ``self.model.layers``; Identity layers + # let the aux forward-hooks fire without a real attention implementation. + self.model = SimpleNamespace(layers=torch.nn.ModuleList(torch.nn.Identity() for _ in range(num_layers))) + self.received: dict = {} + + def get_input_embeddings(self) -> torch.nn.Embedding: + return self.embed_tokens + + def forward(self, input_ids, attention_mask=None, position_ids=None, **kwargs): + self.received = {"attention_mask": attention_mask, "position_ids": position_ids} + hidden = self.embed_tokens(input_ids) + for layer in self.model.layers: + hidden = layer(hidden) + logits = torch.randn(input_ids.shape[0], input_ids.shape[1], self.embed_tokens.num_embeddings) + return SimpleNamespace(logits=logits) + + +def test_packed_flash_target_passes_position_ids_not_4d_mask(): + """FlashAttention packing must forward ``attention_mask=None`` + per-doc + position_ids (FA infers cu_seqlens from them); a 4D mask would blow up its + unpad gather.""" + target = _RecordingFlashTarget() + wrapper = HFEagle3TargetModel(target) + + seq_len = 6 + input_ids = torch.randint(0, 32, (1, seq_len)) + position_ids = torch.tensor([[0, 1, 2, 0, 1, 2]], dtype=torch.long) + seq_lens = torch.tensor([[3, 3]], dtype=torch.long) + doc_remaining = torch.tensor([[2, 1, 0, 2, 1, 0]], dtype=torch.long) + + wrapper.generate_batch( + input_ids=input_ids, + attention_mask=torch.ones(1, seq_len, dtype=torch.long), + loss_mask=torch.ones(1, seq_len, dtype=torch.long), + position_ids=position_ids, + seq_lens=seq_lens, + doc_remaining=doc_remaining, + ) + + # The target saw no explicit mask (FA handles causality) and the per-doc positions. + assert target.received["attention_mask"] is None + torch.testing.assert_close(target.received["position_ids"], position_ids) + + +def test_packed_flash_target_rejects_batch_gt_1(): + """transformers only packs from position_ids at batch size 1, so a FA target + must reject micro_batch_size > 1 instead of silently leaking across documents.""" + target = _RecordingFlashTarget() + wrapper = HFEagle3TargetModel(target) + + seq_len = 6 + input_ids = torch.randint(0, 32, (2, seq_len)) + position_ids = torch.tensor([[0, 1, 2, 0, 1, 2], [0, 1, 2, 0, 1, 2]], dtype=torch.long) + seq_lens = torch.tensor([[3, 3], [3, 3]], dtype=torch.long) + doc_remaining = torch.tensor([[2, 1, 0, 2, 1, 0], [2, 1, 0, 2, 1, 0]], dtype=torch.long) + + with pytest.raises(ValueError, match="micro_batch_size=1"): + wrapper.generate_batch( + input_ids=input_ids, + attention_mask=torch.ones(2, seq_len, dtype=torch.long), + loss_mask=torch.ones(2, seq_len, dtype=torch.long), + position_ids=position_ids, + seq_lens=seq_lens, + doc_remaining=doc_remaining, + )