From d91cc809f356d582721a4abe0f9b4e4c626350db Mon Sep 17 00:00:00 2001 From: hootandy321 <133196559+hootandy321@users.noreply.github.com> Date: Thu, 12 Nov 2026 02:07:08 -0500 Subject: [PATCH] feat: add Qwen3 P-EAGLE model and KL loss core --- aspec/models/peagle.py | 267 ++++++++++++++++++++++++ docs/framework/REFERENCES.md | 1 + docs/support.md | 1 + docs/user-guide/algorithms.md | 6 + docs/user-guide/peagle.md | 126 +++++++++++ tests/cpu/test_peagle.py | 381 ++++++++++++++++++++++++++++++++++ 6 files changed, 782 insertions(+) create mode 100644 aspec/models/peagle.py create mode 100644 docs/user-guide/peagle.md create mode 100644 tests/cpu/test_peagle.py diff --git a/aspec/models/peagle.py b/aspec/models/peagle.py new file mode 100644 index 0000000..c15a8c6 --- /dev/null +++ b/aspec/models/peagle.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm-project/speculators at +# 3702320434bbdaa853f0846a66e6b6b8f1d97aec: +# models/peagle/{core,attention,metrics}.py and +# models/eagle3/model_definitions.py. See licenses/Speculators-LICENSE. +# Changes: Qwen3-only model/loss prerequisite, deterministic dense depth layout, +# contiguous-document isolation, explicit aligned inputs and detached Target data. +"""Parallel P-EAGLE core; no trainer, sampler, checkpoint/export or serving adapter.""" + +import copy +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import Qwen3Config +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3DecoderLayer, + Qwen3RMSNorm, + Qwen3RotaryEmbedding, +) + + +def _segments(document_ids: torch.Tensor) -> torch.Tensor: + changes = torch.ones_like(document_ids) + changes[1:] = (document_ids[1:] != document_ids[:-1]).long() + return changes.cumsum(0) + + +def peagle_dense_layout( + document_ids: torch.Tensor, num_depths: int +) -> tuple[torch.Tensor, torch.Tensor]: + """All base positions, then every in-document (anchor, depth) pair. + + This deterministic unsampled layout is for the model/loss prerequisite; + it does not implement the reference's stochastic COD training policy. + """ + if document_ids.ndim != 1 or document_ids.dtype != torch.long or not document_ids.numel(): + raise ValueError("P-EAGLE document_ids must be a nonempty int64 vector") + if (document_ids < -1).any(): + raise ValueError("P-EAGLE uses -1 for padding and nonnegative document IDs") + if isinstance(num_depths, bool) or not isinstance(num_depths, int) or num_depths < 1: + raise ValueError("P-EAGLE num_depths must be a positive integer") + segments = _segments(document_ids) + base = torch.arange(document_ids.numel(), device=document_ids.device) + anchors, depths = [base], [torch.zeros_like(base)] + for d in range(1, min(num_depths, document_ids.numel())): + candidates = base[:-d] + valid = (document_ids[:-d] >= 0) & (segments[:-d] == segments[d:]) + candidates = candidates[valid] + anchors.append(candidates) + depths.append(torch.full_like(candidates, d)) + return torch.cat(anchors), torch.cat(depths) + + +def peagle_attention_mask( + anchor_pos: torch.Tensor, depth: torch.Tensor, document_ids: torch.Tensor +) -> torch.Tensor: + """Boolean [N,N] mask: causal base context and earlier depths of the same chain. + + Padding/cross-document slots attend only to themselves. Repeated document + labels in disjoint segments never share attention. + """ + if anchor_pos.ndim != 1 or depth.shape != anchor_pos.shape: + raise ValueError("P-EAGLE anchor/depth vectors must have the same shape") + if anchor_pos.dtype != torch.long or depth.dtype != torch.long: + raise ValueError("P-EAGLE anchor/depth vectors must use int64") + if document_ids.ndim != 1 or document_ids.dtype != torch.long: + raise ValueError("P-EAGLE document_ids must be an int64 vector") + if not (anchor_pos.device == depth.device == document_ids.device): + raise ValueError("P-EAGLE layout tensors must share a device") + positions = anchor_pos + depth + if ((anchor_pos < 0) | (depth < 0) | (positions >= document_ids.numel())).any(): + raise ValueError("P-EAGLE layout is outside the sequence") + segments = _segments(document_ids) + valid = (document_ids[anchor_pos] >= 0) & (segments[anchor_pos] == segments[positions]) + same_document = segments[anchor_pos, None] == segments[anchor_pos][None, :] + base_context = (depth[None, :] == 0) & (anchor_pos[:, None] >= anchor_pos[None, :]) + same_chain = (anchor_pos[:, None] == anchor_pos[None, :]) & ( + depth[:, None] >= depth[None, :] + ) + allowed = valid[:, None] & valid[None, :] & same_document & (base_context | same_chain) + return allowed | torch.eye(anchor_pos.numel(), dtype=torch.bool, device=anchor_pos.device) + + +class _PEagleFirstLayer(Qwen3DecoderLayer): + """Speculators' EAGLE first-layer fusion, retaining Qwen3 Q/K normalization.""" + + def __init__(self, config): + super().__init__(config, layer_idx=0) + self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + for name in ("q_proj", "k_proj", "v_proj"): + previous = getattr(self.self_attn, name) + setattr( + self.self_attn, name, + nn.Linear(2 * config.hidden_size, previous.out_features, + bias=config.attention_bias), + ) + + def forward(self, hidden_states, attention_mask=None, position_embeddings=None): + embeds, hidden = hidden_states.chunk(2, dim=-1) + fused = torch.cat((self.input_layernorm(embeds), self.hidden_norm(hidden)), dim=-1) + attended, _ = self.self_attn( + hidden_states=fused, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + hidden = hidden + attended # norm_before_residual=False in the selected recipe + return hidden + self.mlp(self.post_attention_layernorm(hidden)) + + +@dataclass +class PEagleOutput: + logits: torch.Tensor + anchor_pos: torch.Tensor + depth: torch.Tensor + valid: torch.Tensor + num_depths: int + + @property + def positions(self) -> torch.Tensor: + return self.anchor_pos + self.depth + + +class PEagleDraftModel(nn.Module): + """One packed row of aligned Qwen3 Target features, all depths in one pass. + + At aligned index i: input_ids=x[i+1], hidden_states=Target auxiliary h[i], + teacher_logits=Target logits[i+1], loss_mask=reference mask[i+1]. + Each document must be shifted separately before packing. Future depth d + at anchor a uses learned token/hidden masks and targets aligned index a+d. + See docs/user-guide/peagle.md for the scope and alignment contract. + """ + + def __init__( + self, config: Qwen3Config, *, mask_token_id: int, + num_aux_hidden_states: int = 3, num_depths: int = 4, + ): + super().__init__() + if config.model_type != "qwen3": + raise ValueError("P-EAGLE core currently supports only Qwen3") + if config._attn_implementation not in (None, "eager", "sdpa"): + raise ValueError("P-EAGLE core supports only eager or sdpa attention") + if config.attention_dropout != 0 or any( + kind != "full_attention" for kind in config.layer_types + ): + raise ValueError("P-EAGLE core requires full attention and zero dropout") + if (getattr(config, "draft_vocab_size", config.vocab_size) != config.vocab_size + or config.tie_word_embeddings + or any(getattr(config, field, False) for field in ( + "norm_before_residual", "norm_before_fc", "fc_norm", + ))): + raise ValueError("P-EAGLE core requires the selected full-vocabulary untied-head recipe") + if (config.num_attention_heads < 1 or config.num_key_value_heads < 1 + or config.num_attention_heads % config.num_key_value_heads + or config.head_dim < 1 or config.head_dim % 2): + raise ValueError("P-EAGLE core requires valid GQA geometry and an even head dimension") + if config.num_hidden_layers < 1 or config.hidden_act != "silu": + raise ValueError("P-EAGLE core requires at least one layer and SiLU") + for name, value in (("num_aux_hidden_states", num_aux_hidden_states), + ("num_depths", num_depths)): + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"P-EAGLE {name} must be a positive integer") + if isinstance(mask_token_id, bool) or not isinstance(mask_token_id, int) or not ( + 0 <= mask_token_id < config.vocab_size + ): + raise ValueError("P-EAGLE mask_token_id must be in the Target vocabulary") + self.config = copy.deepcopy(config) + self.config._attn_implementation = config._attn_implementation or "sdpa" + self.mask_token_id = mask_token_id + self.num_aux_hidden_states = num_aux_hidden_states + self.num_depths = num_depths + h = config.hidden_size + # No padding_idx: the learned mask-token row must receive gradients. + self.embed_tokens = nn.Embedding(config.vocab_size, h) + self.mask_hidden = nn.Parameter(torch.randn(1, 1, num_aux_hidden_states * h)) + self.fc = nn.Linear(num_aux_hidden_states * h, h, bias=False) + self.layers = nn.ModuleList([_PEagleFirstLayer(self.config)] + [ + Qwen3DecoderLayer(self.config, layer_idx=i) + for i in range(1, config.num_hidden_layers) + ]) + self.rotary_emb = Qwen3RotaryEmbedding(self.config) + self.norm = Qwen3RMSNorm(h, eps=config.rms_norm_eps) + self.lm_head = nn.Linear(h, config.vocab_size, bias=False) + self.apply(self._init_weights) + + def _init_weights(self, module): + if isinstance(module, (nn.Linear, nn.Embedding)): + nn.init.normal_(module.weight, std=self.config.initializer_range) + if getattr(module, "bias", None) is not None: + nn.init.zeros_(module.bias) + + def forward(self, input_ids, hidden_states, document_ids, position_ids=None) -> PEagleOutput: + if input_ids.ndim != 2 or input_ids.shape[0] != 1 or not input_ids.shape[1]: + raise ValueError("P-EAGLE core requires one nonempty packed row [1,S]") + if input_ids.dtype != torch.long or document_ids.shape != input_ids.shape: + raise ValueError("P-EAGLE requires int64 IDs and document_ids shaped [1,S]") + expected = (*input_ids.shape, self.fc.in_features) + if hidden_states.shape != expected or not hidden_states.is_floating_point(): + raise ValueError(f"P-EAGLE hidden_states must be floating point with shape {expected}") + if not (input_ids.device == hidden_states.device == document_ids.device + == self.embed_tokens.weight.device): + raise ValueError("P-EAGLE inputs and model must share a device") + if ((input_ids < 0) | (input_ids >= self.config.vocab_size)).any(): + raise ValueError("P-EAGLE input token is outside the Target vocabulary") + docs = document_ids[0] + anchors, depth = peagle_dense_layout(docs, self.num_depths) + positions = anchors + depth + if position_ids is None: + base = torch.arange(docs.numel(), device=docs.device) + starts = torch.cat((torch.ones_like(docs[:1]), (docs[1:] != docs[:-1]).long())) + position_ids = (base - torch.where(starts.bool(), base, 0).cummax(0).values)[None] + 1 + if (position_ids.shape != input_ids.shape or position_ids.dtype != torch.long + or position_ids.device != input_ids.device or (position_ids < 0).any()): + raise ValueError("P-EAGLE position_ids must be nonnegative int64 [1,S] on the input device") + is_base = (depth == 0)[None] + token_ids = torch.where(is_base, input_ids[:, positions], self.mask_token_id) + sampled_hidden = torch.where( + is_base[..., None], hidden_states.detach()[:, positions], + self.mask_hidden.to(hidden_states.dtype), + ) + projected = self.fc(sampled_hidden.to(self.fc.weight.dtype)) + states = torch.cat((self.embed_tokens(token_ids), projected), dim=-1) + allowed = peagle_attention_mask(anchors, depth, docs) + mask = states.new_zeros(allowed.shape).masked_fill(~allowed, -torch.inf)[None, None] + rope = self.rotary_emb(states, position_ids[:, positions]) + for layer in self.layers: + states = layer(states, attention_mask=mask, position_embeddings=rope) + return PEagleOutput( + self.lm_head(self.norm(states)), anchors, depth, + docs[positions] >= 0, self.num_depths, + ) + + +def peagle_kl_loss( + output: PEagleOutput, teacher_logits: torch.Tensor, loss_mask: torch.Tensor, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Forward KL(Target || Draft), mean over all supervised sampled positions. + + Targets are detached. Per-depth sums/counts are observations of independent + parallel predictions, not decoder acceptance or a joint accepted prefix. + The 1e-5 denominator matches the pinned Speculators loss reduction. + """ + if (teacher_logits.ndim != 3 or teacher_logits.shape[0] != 1 + or teacher_logits.shape[2] != output.logits.shape[2] + or loss_mask.shape != teacher_logits.shape[:2] + or teacher_logits.shape[1] != int((output.depth == 0).sum())): + raise ValueError("P-EAGLE teacher/mask must match the aligned row and full vocabulary") + if not (teacher_logits.device == loss_mask.device == output.logits.device): + raise ValueError("P-EAGLE loss tensors must share a device") + if not ((loss_mask == 0) | (loss_mask == 1)).all(): + raise ValueError("P-EAGLE loss_mask must be binary") + selected = output.valid & loss_mask[0, output.positions].bool() + logits = output.logits[0, selected].float() + targets = teacher_logits.detach()[0, output.positions[selected]].float() + logq = F.log_softmax(logits, dim=-1) + p = F.softmax(targets, dim=-1) + terms = F.kl_div(logq, p, reduction="none").sum(-1) + loss = terms.sum() / (selected.sum() + 1e-5) + with torch.no_grad(): + d = output.depth[selected] + counts = torch.zeros(output.num_depths, device=logits.device) + sums, correct = torch.zeros_like(counts), torch.zeros_like(counts) + counts.scatter_add_(0, d, torch.ones_like(terms)) + sums.scatter_add_(0, d, terms.detach()) + correct.scatter_add_(0, d, (logits.argmax(-1) == targets.argmax(-1)).float()) + return loss, {"kl_sum": sums, "count": counts, "correct": correct} diff --git a/docs/framework/REFERENCES.md b/docs/framework/REFERENCES.md index 0f33c80..af0633c 100644 --- a/docs/framework/REFERENCES.md +++ b/docs/framework/REFERENCES.md @@ -20,6 +20,7 @@ the retained source identities; applicable license and notice texts are distribu | [Transformers](https://github.com/huggingface/transformers) v4.57.1 `8cb5963cc22174954e7dca2c0a3320b7dc2f4edc` | Adapted `Qwen3MoeSparseMoeBlock` and load-balancing objective in `aspec/models/draft/moe.py` | | [DMTP](https://sylvan820.github.io/D-MTP/) | Inspiration for dual-anchor MTP hidden conditioning and zero-initialized embedding residual; independent ASpec integration reuses the retained native predictor, without copying DMTP code | | [Speculators](https://github.com/vllm-project/speculators) `7ec7aba9450b533334e8f7992baf4e93be3c1a57` | Native MTP recurrent objective and checkpoint bindings; ASpec supplies its own training/layout integration | +| [Speculators](https://github.com/vllm-project/speculators) `3702320434bbdaa853f0846a66e6b6b8f1d97aec` | Adapted P-EAGLE model/loss and Qwen3 first-layer fusion; deterministic dense layout and ASpec input/observation boundary | | [Ling-3.0-flash](https://huggingface.co/inclusionAI/Ling-3.0-flash) `fcc60a52b34c69c6271d0ac8a8cb0815020b6a43` | Adapted MoE routing and independently retained reference class bodies; Apache-2.0 source headers preserved | | [EAGLE](https://github.com/SafeAILab/EAGLE), [GPT-NeoX](https://github.com/EleutherAI/gpt-neox) and Transformers through TorchSpec | Model/data helpers in `aspec/models/draft/base.py` and `aspec/data/preprocessing.py`; original source headers retain the transitive attribution | diff --git a/docs/support.md b/docs/support.md index 58bd775..b11d542 100644 --- a/docs/support.md +++ b/docs/support.md @@ -11,6 +11,7 @@ and tests do not imply CPU training support. | --- | --- | --- | | DSpark / DFlash | Training, exact recovery, evaluation and export; Dense and selected MoE/attention structures | Objectives, feature bindings and assets remain algorithm-specific | | EAGLE3 | Recurrent TTT, full/pruned vocabulary, GQA/MHA or materialized MLA; FA4 packing/CP | No EAGLE3 MoE, anchored mode or whole-model compile | +| P-EAGLE core | Qwen3 full-vocabulary parallel model/loss, deterministic dense slots, eager/SDPA | Model/loss prerequisite only; no public trainer, COD sampler, resume, export or consumer validation; see [scope](user-guide/peagle.md) | | EAGLE4 | Dense and selected upstream MoE; Draft/candidate stages; packing/CP and native assets | Candidate Top-K forward-KL is explicitly implemented from public semantics, not an unpublished official recipe; no serving integration | | DFlash2 | Experimental full/causal/SWA training with Flex, cuDNN THD or FA4 packing/CP | Experimental objective; no claim to reproduce an official training recipe | | MTP-conditioned DSpark / DFlash | Frozen native Qwen3.5 predictor, token/hidden dual-anchor training, exact recovery and strict native assets | Dense full GQA, CP1, MBS1, REPLICATE, one node, non-packing, decay objective; no generic serving consumer | diff --git a/docs/user-guide/algorithms.md b/docs/user-guide/algorithms.md index 917f01b..1d28759 100644 --- a/docs/user-guide/algorithms.md +++ b/docs/user-guide/algorithms.md @@ -55,6 +55,12 @@ aspec --config configs/recipes/qwen3_8b_eagle3.yaml --dry-run \ TTT shifts, teacher halos and supervision remain document-local. Unsupported factorized/ambiguous heads fail export rather than producing an unfaithful serving asset. +## P-EAGLE core prerequisite + +The [P-EAGLE model/loss API](peagle.md) provides a dense unsampled Qwen3 parallel +Draft and forward-KL objective with independent numerical checks. Public training, +COD sampling, recovery, export and consumer validation remain pending. + ## EAGLE4: two stages The Draft stage trains the full recurrent model. The candidate stage freezes the entire Draft, diff --git a/docs/user-guide/peagle.md b/docs/user-guide/peagle.md new file mode 100644 index 0000000..f5b4193 --- /dev/null +++ b/docs/user-guide/peagle.md @@ -0,0 +1,126 @@ +# P-EAGLE model and loss core + +This is the first prerequisite for [ASpec #16](https://github.com/inclusionAI/ASpec/issues/16). +It provides `aspec.models.peagle.PEagleDraftModel` and `peagle_kl_loss` for one +packed row of Qwen3 features, full vocabulary, full attention, zero attention dropout, +and eager or SDPA execution. It is a model/loss API; the ASpec training CLI does not +select this algorithm yet. + +## Reference and selected recipe + +The implementation is **adapted** from +[Speculators `3702320434bbdaa853f0846a66e6b6b8f1d97aec`](https://github.com/vllm-project/speculators/tree/3702320434bbdaa853f0846a66e6b6b8f1d97aec/src/speculators/models/peagle) +and its EAGLE first-layer fusion. The selected reference is the +[Qwen3-8B example](https://github.com/vllm-project/speculators/blob/3702320434bbdaa853f0846a66e6b6b8f1d97aec/examples/train/peagle_qwen3_8b_ultrachat_online_5k.sh): +four Draft layers, four prediction depths, no normalization before the first residual, +no pre-FC normalization and no per-feature FC normalization. The example's online +training, optimizer schedule and COD sampling policy are deferred. Speculators +deprecates P-EAGLE at this revision; this API does not depend on its current main branch. +The applicable license is distributed in `licenses/Speculators-LICENSE`. + +The pinned TorchSpec EAGLE3 implementation was reviewed. ASpec retains its recurrent +Llama-style EAGLE3 path. Reusing that trainer would freeze token embeddings and feed +recurrent hidden states; its attention also omits Qwen3 Q/K normalization. This core +instead adapts the applicable Speculators first-layer fusion and reuses Transformers' +Qwen3 attention, MLP, norms, rotary embeddings and subsequent decoder layers. + +## Inputs and prediction semantics + +Before packing, shift **each document separately**, as in the reference's +`models/eagle3/data.py::shift_batch`: + +| Aligned index `i` | Value | +| --- | --- | +| `input_ids[:, i]` | Original token `x[i+1]` | +| `hidden_states[:, i]` | Concatenated selected auxiliary Target features `h[i]` | +| `teacher_logits[:, i]` | Frozen Target logits at original index `i+1` | +| `loss_mask[:, i]` | Reference loss mask at original index `i+1` | +| `position_ids[:, i]` | Document-local position of original token `x[i+1]` | + +Target logits at original index `j` predict token `x[j+1]`. Thus a slot +`(anchor=a, depth=d)` predicts `x[a+d+2]` and distills +`teacher_logits[:, a+d]`. Loss masks follow this explicit reference alignment; +callers must adapt other mask conventions before calling this API. Globally +shifting an already packed row would create invalid cross-document pairs. + +Shapes are `[1,S]` for token/document/position IDs, `[1,S,A*H]` for auxiliary +features, and `[1,S,V]` for Target logits. The caller supplies the ordered `A` +Target feature layers and matching Qwen3 hidden width `H`; `V` is the full +Target vocabulary. This core does not load, normalize or bind a real Target checkpoint. +Features must be finite and in the normalization convention of the selected reference. +Target auxiliary features and loss targets are detached. + +`document_ids=-1` marks padding. Every contiguous document segment is isolated, even +if a document label is reused later. Omitted position IDs start at one and reset +at each segment. Supply explicit IDs when retaining a document-local position offset. + +Depth zero includes every position and uses the real token and Target auxiliary +features. Greater depths use a **trainable token embedding** at `mask_token_id` +and a shared **trainable `mask_hidden`** vector in auxiliary-feature space. +One FC projection and the Draft weights are shared across depths. All embedding +rows, FC/attention/MLP/norm/head parameters and `mask_hidden` are trainable; no +frozen Target module is owned by the Draft. The mask embedding remains trainable +even when its ID equals `config.pad_token_id`. + +All slots pass through each Draft layer once. A slot attends to preceding base +context in its document and to earlier/equal depths in its own anchor chain. +It cannot read future real tokens or another chain's non-base slots. Padding +attends only to itself, keeping attention rows finite. + +The deterministic layout includes every in-document anchor/depth pair. This is an +**unsampled numerical core**, not the reference's stochastic COD sampler. Dense +attention requires quadratic memory in the total number of slots; use short inputs +for this API. This choice makes the model/loss semantics independently reviewable. + +## Loss and observations + +`peagle_kl_loss(output, teacher_logits, loss_mask)` computes forward +`KL(Target || Draft)` in FP32, with the pinned reference's reduction +`sum(selected_KL) / (selected_count + 1e-5)`. It gathers targets at +`output.anchor_pos + output.depth`. Padding and unsupervised positions are +excluded before softmax; empty supervision returns a differentiable zero. + +The returned `kl_sum`, `count` and `correct` vectors have one entry per requested +depth. Aggregate sums/counts before calculating means. These describe independent +parallel predictions; they do not measure a jointly accepted decoding prefix. + +## Minimal numerical use + +Inside the [isolated Linux environment](installation.md), this example creates +synthetic features and randomly initialized small Qwen3 weights: + +```python +import torch +from transformers import Qwen3Config +from aspec.models.peagle import PEagleDraftModel, peagle_kl_loss + +config = Qwen3Config( + vocab_size=32, hidden_size=16, intermediate_size=32, + num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, + head_dim=8, attention_dropout=0.0, +) +model = PEagleDraftModel(config, mask_token_id=0, num_depths=4) +output = model( + input_ids=torch.randint(0, 32, (1, 6)), + hidden_states=torch.randn(1, 6, 48), + document_ids=torch.zeros(1, 6, dtype=torch.long), +) +loss, observations = peagle_kl_loss( + output, torch.randn(1, 6, 32), torch.ones(1, 6), +) +loss.backward() +``` + +Run `HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 python -m pytest -q tests/cpu/test_peagle.py` +for independent scalar attention rules, explicit linear-algebra forward/backward +comparisons, target-index alignment, trainable masks, document/future isolation, +empty supervision, BF16 CPU arithmetic and strict state-dict completeness. + +## Remaining work + +COD sampling and its RNG/resume semantics, the public trainer and configuration, +Target feature/head bindings, exact training resume, asset export/strict asset loading, +proposal consumers, GPU numerical validation and per-depth held-out/acceptance +experiments remain pending. A strict PyTorch state-dict check is not an ASpec export +or exact-resume guarantee. Synthetic checks do not establish Qwen3-8B training +quality, GPU throughput, consumer compatibility or decoding acceptance. diff --git a/tests/cpu/test_peagle.py b/tests/cpu/test_peagle.py new file mode 100644 index 0000000..91e14ad --- /dev/null +++ b/tests/cpu/test_peagle.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Independent scalar/linear-algebra checks of the bounded P-EAGLE core.""" + +import copy +import math + +import pytest +import torch +import torch.nn.functional as F +from transformers import LlamaConfig, Qwen3Config + +from aspec.models.peagle import ( + PEagleDraftModel, + peagle_attention_mask, + peagle_dense_layout, + peagle_kl_loss, +) + + +def model(*, layers=2, depths=4, implementation="sdpa", **changes): + config = Qwen3Config( + vocab_size=23, hidden_size=16, intermediate_size=32, + num_hidden_layers=layers, num_attention_heads=2, num_key_value_heads=1, + head_dim=8, max_position_embeddings=64, rms_norm_eps=1e-6, + rope_theta=10000.0, attention_dropout=0.0, pad_token_id=0, + layer_types=["full_attention"] * layers, + ) + config._attn_implementation = implementation + for name, value in changes.items(): + setattr(config, name, value) + torch.manual_seed(7) + return PEagleDraftModel(config, mask_token_id=0, num_aux_hidden_states=3, num_depths=depths) + + +def example(): + generator = torch.Generator().manual_seed(9) + return dict( + input_ids=torch.tensor([[2, 4, 6, 8, 7, 5, 3, 1]]), + hidden_states=torch.randn(1, 8, 48, generator=generator, requires_grad=True), + document_ids=torch.tensor([[0, 0, 0, 1, 1, 1, -1, -1]]), + position_ids=torch.tensor([[1, 2, 3, 1, 2, 3, 0, 0]]), + ) + + +def scalar_layout(docs, depths): + slots = [(i, 0) for i in range(len(docs))] + for d in range(1, depths): + for a in range(len(docs) - d): + if docs[a] >= 0 and all(v == docs[a] for v in docs[a:a + d + 1]): + slots.append((a, d)) + return slots + + +def scalar_mask(slots, docs): + result = torch.zeros(len(slots), len(slots), dtype=torch.bool) + for qi, (a, d) in enumerate(slots): + for ki, (b, e) in enumerate(slots): + if qi == ki: + result[qi, ki] = True + elif docs[a] >= 0 and docs[b] >= 0: + connected = all(v == docs[a] for v in docs[min(a, b):max(a, b) + 1]) + result[qi, ki] = connected and ((e == 0 and b <= a) or (a == b and e <= d)) + return result + + +def rms(x, norm): + dtype = x.dtype + x = x.float() + return norm.weight * (x / torch.sqrt(x.square().mean(-1, keepdim=True) + norm.variance_epsilon)).to(dtype) + + +def linear(x, layer): + return F.linear(x, layer.weight, layer.bias) + + +def reference_forward(m, batch): + """No production layout, attention, rotary, decoder, norm or loss calls.""" + docs = batch["document_ids"][0].tolist() + slots = scalar_layout(docs, m.num_depths) + embeds = torch.stack([ + m.embed_tokens.weight[batch["input_ids"][0, a] if d == 0 else m.mask_token_id] + for a, d in slots + ])[None] + features = torch.stack([ + batch["hidden_states"].detach()[0, a] if d == 0 else m.mask_hidden[0, 0] + for a, d in slots + ])[None] + states = linear(features, m.fc) + mask = scalar_mask(slots, docs) + positions = torch.tensor([batch["position_ids"][0, a + d] for a, d in slots]) + dim = m.config.head_dim + theta = m.config.rope_parameters["rope_theta"] + inv_freq = theta ** (-torch.arange(0, dim, 2).float() / dim) + angles = positions[:, None] * inv_freq[None] + angles = torch.cat((angles, angles), -1)[None, None] + + def rotate(x): + first, second = x.chunk(2, -1) + return x * angles.cos() + torch.cat((-second, first), -1) * angles.sin() + + for i, layer in enumerate(m.layers): + residual = states + if i == 0: + attention_input = torch.cat(( + rms(embeds, layer.input_layernorm), rms(states, layer.hidden_norm), + ), -1) + else: + attention_input = rms(states, layer.input_layernorm) + attn = layer.self_attn + n = len(slots) + q = linear(attention_input, attn.q_proj).reshape(1, n, m.config.num_attention_heads, dim) + k = linear(attention_input, attn.k_proj).reshape(1, n, m.config.num_key_value_heads, dim) + v = linear(attention_input, attn.v_proj).reshape(1, n, m.config.num_key_value_heads, dim) + q = rotate(rms(q, attn.q_norm).transpose(1, 2)) + k = rotate(rms(k, attn.k_norm).transpose(1, 2)) + groups = m.config.num_attention_heads // m.config.num_key_value_heads + k = k.repeat_interleave(groups, dim=1) + v = v.transpose(1, 2).repeat_interleave(groups, dim=1) + scores = q @ k.transpose(-1, -2) / math.sqrt(dim) + probabilities = scores.masked_fill(~mask, -torch.inf).softmax(-1) + attended = (probabilities @ v).transpose(1, 2).reshape(1, n, -1) + states = residual + linear(attended, attn.o_proj) + normed = rms(states, layer.post_attention_layernorm) + mlp = layer.mlp + states = states + linear(F.silu(linear(normed, mlp.gate_proj)) * linear(normed, mlp.up_proj), mlp.down_proj) + return linear(rms(states, m.norm), m.lm_head), slots + + +def reference_loss(logits, slots, docs, teacher, loss_mask): + terms = [] + for i, (a, d) in enumerate(slots): + if docs[a] >= 0 and loss_mask[0, a + d]: + t = teacher.detach()[0, a + d].float() + logp = t - t.logsumexp(-1) + q = logits[0, i].float() + logq = q - q.logsumexp(-1) + terms.append((logp.exp() * (logp - logq)).sum()) + return torch.stack(terms).sum() / (len(terms) + 1e-5) + + +@pytest.mark.parametrize("docs", [ + [0], [-1], [0, 0, 0, 0], [0, 0, 1, 1], + [0, 0, -1, -1], [-1, 0, 0, -1], [0, 0, 1, 0, 0], [-1, -1, -1], +]) +@pytest.mark.parametrize("depths", [1, 2, 5]) +def test_layout_and_attention_match_independent_scalar_rules(docs, depths): + anchors, depth = peagle_dense_layout(torch.tensor(docs), depths) + slots = scalar_layout(docs, depths) + assert list(zip(anchors.tolist(), depth.tolist())) == slots + actual = peagle_attention_mask(anchors, depth, torch.tensor(docs)) + assert torch.equal(actual, scalar_mask(slots, docs)) + assert actual.any(-1).all() + for i, (a, _) in enumerate(slots): + if docs[a] == -1: + assert actual[i].sum() == 1 + + +@pytest.mark.parametrize("implementation", ["eager", "sdpa"]) +@pytest.mark.parametrize("layers", [1, 2]) +@pytest.mark.parametrize("depths", [1, 4]) +def test_forward_loss_and_all_parameter_gradients_match_reference(implementation, layers, depths): + actual = model(implementation=implementation, layers=layers, depths=depths) + expected = copy.deepcopy(actual) + batch = example() + logits, slots = reference_forward(expected, batch) + output = actual(**batch) + torch.testing.assert_close(output.logits, logits, rtol=2e-5, atol=2e-6) + torch.manual_seed(13) + teacher = torch.randn(1, 8, 23, requires_grad=True) + loss_mask = torch.tensor([[0, 1, 1, 0, 1, 1, 1, 1]]) + loss, metrics = peagle_kl_loss(output, teacher, loss_mask) + expected_loss = reference_loss(logits, slots, batch["document_ids"][0].tolist(), teacher, loss_mask) + torch.testing.assert_close(loss, expected_loss, rtol=2e-5, atol=2e-6) + loss.backward() + expected_loss.backward() + for (name, param), (other_name, other) in zip(actual.named_parameters(), expected.named_parameters()): + assert name == other_name + assert param.requires_grad + assert param.grad is not None, name + if other.grad is None: + assert not param.grad.any(), name + else: + torch.testing.assert_close(param.grad, other.grad, rtol=3e-4, atol=2e-6, msg=name) + assert teacher.grad is None + assert batch["hidden_states"].grad is None + counts = torch.tensor([ + sum(int(loss_mask[0, a + d]) for a, d in slots + if d == wanted and batch["document_ids"][0, a] >= 0) + for wanted in range(depths) + ]).float() + assert torch.equal(metrics["count"], counts) + expected_sums, expected_correct = torch.zeros(depths), torch.zeros(depths) + for i, (a, d) in enumerate(slots): + if batch["document_ids"][0, a] >= 0 and loss_mask[0, a + d]: + t = teacher.detach()[0, a + d].float() + logp = t - t.logsumexp(-1) + q = logits.detach()[0, i].float() + logq = q - q.logsumexp(-1) + expected_sums[d] += (logp.exp() * (logp - logq)).sum() + expected_correct[d] += (t.argmax() == q.argmax()).float() + torch.testing.assert_close(metrics["kl_sum"], expected_sums, rtol=2e-5, atol=2e-6) + assert torch.equal(metrics["correct"], expected_correct) + assert not metrics["kl_sum"].requires_grad + + +def test_future_tokens_other_documents_and_padding_do_not_leak(): + m = model() + batch = example() + before = m(**batch) + changed = {k: v.detach().clone() for k, v in batch.items()} + changed["input_ids"][0, 1:] = torch.tensor([1, 2, 3, 4, 6, 8, 9]) + changed["hidden_states"][0, 1:] = 100 * torch.randn(7, 48) + after = m(**changed) + chain_zero = before.anchor_pos == 0 + torch.testing.assert_close(before.logits[:, chain_zero], after.logits[:, chain_zero], rtol=0, atol=0) + + changed = {k: v.detach().clone() for k, v in batch.items()} + changed["input_ids"][0, :3] = 1 + changed["hidden_states"][0, :3] = 100 + after = m(**changed) + second_document = (before.anchor_pos >= 3) & (before.anchor_pos <= 5) + torch.testing.assert_close(before.logits[:, second_document], after.logits[:, second_document], rtol=0, atol=0) + + changed = {k: v.detach().clone() for k, v in batch.items()} + changed["input_ids"][0, 6:] = 22 + changed["hidden_states"][0, 6:] = -100 + after = m(**changed) + torch.testing.assert_close(before.logits[:, before.valid], after.logits[:, after.valid], rtol=0, atol=0) + + +def test_mask_token_and_hidden_parameters_update_even_when_mask_id_is_padding_id(): + m = model() + batch = example() + output = m(**batch) + torch.manual_seed(4) + loss, _ = peagle_kl_loss(output, torch.randn(1, 8, 23), torch.ones(1, 8)) + mask_before = m.mask_hidden.detach().clone() + embed_before = m.embed_tokens.weight[0].detach().clone() + optimizer = torch.optim.SGD(m.parameters(), lr=0.1) + loss.backward() + assert m.mask_hidden.grad.abs().sum() > 0 + assert m.embed_tokens.weight.grad[0].abs().sum() > 0 + optimizer.step() + assert not torch.equal(mask_before, m.mask_hidden) + assert not torch.equal(embed_before, m.embed_tokens.weight[0]) + + +@pytest.mark.parametrize("padding", [False, True]) +def test_empty_supervision_has_finite_differentiable_zero(padding): + m = model() + batch = example() + if padding: + batch["document_ids"].fill_(-1) + output = m(**batch) + teacher = torch.full((1, 8, 23), torch.nan, requires_grad=True) + loss, metrics = peagle_kl_loss(output, teacher, torch.zeros(1, 8)) + assert loss.item() == 0 and loss.requires_grad + assert torch.isfinite(output.logits).all() + assert not metrics["count"].any() + loss.backward() + assert teacher.grad is None + for param in m.parameters(): + assert param.grad is not None + assert torch.isfinite(param.grad).all() and not param.grad.any() + + +def test_unsupervised_nonfinite_teacher_values_are_excluded_before_softmax(): + m = model() + output = m(**example()) + teacher = torch.randn(1, 8, 23) + mask = torch.tensor([[0, 1, 1, 0, 1, 1, 0, 0]]) + before, _ = peagle_kl_loss(output, teacher, mask) + teacher[:, ~mask[0].bool()] = torch.nan + after, _ = peagle_kl_loss(output, teacher, mask) + torch.testing.assert_close(before, after, rtol=0, atol=0) + + +def test_reference_next_token_shift_and_depth_target_alignment(): + m = model(layers=1) + torch.manual_seed(21) + raw_ids = torch.tensor([[2, 4, 6, 8, 10]]) + raw_features = torch.randn(1, 5, 48) + raw_teacher = torch.randn(1, 5, 23) + output = m(raw_ids[:, 1:], raw_features[:, :-1], torch.zeros(1, 4, dtype=torch.long)) + loss, _ = peagle_kl_loss(output, raw_teacher[:, 1:], torch.ones(1, 4)) + terms = [] + for i, (a, d) in enumerate(zip(output.anchor_pos.tolist(), output.depth.tolist())): + source_index = a + d + 1 + target = raw_teacher[0, source_index] + p = target.softmax(-1) + terms.append((p * (target.log_softmax(-1) - output.logits[0, i].log_softmax(-1))).sum()) + torch.testing.assert_close(loss, torch.stack(terms).sum() / (len(terms) + 1e-5)) + + +def test_positions_reset_per_document_and_support_explicit_offsets(): + m = model() + batch = example() + explicit = m(**batch) + batch.pop("position_ids") + automatic = m(**batch) + torch.testing.assert_close(explicit.logits[:, explicit.valid], automatic.logits[:, automatic.valid], rtol=0, atol=0) + batch["position_ids"] = torch.tensor([[25, 26, 27, 50, 51, 52, 0, 0]]) + expected, _ = reference_forward(m, batch) + torch.testing.assert_close(m(**batch).logits, expected, rtol=2e-5, atol=2e-6) + + +@pytest.mark.parametrize("missing_key", ["mask_hidden", "embed_tokens.weight"]) +def test_strict_state_dict_requires_both_learned_masks(missing_key): + m = model() + batch = example() + state = m.state_dict() + restored = model() + restored.load_state_dict(state, strict=True) + torch.testing.assert_close(m(**batch).logits, restored(**batch).logits, rtol=0, atol=0) + missing = dict(state) + missing.pop(missing_key) + with pytest.raises(RuntimeError, match=missing_key): + restored.load_state_dict(missing, strict=True) + + +def test_bfloat16_cpu_forward_backward_is_finite(): + m = model().to(torch.bfloat16) + batch = example() + batch["hidden_states"] = batch["hidden_states"].detach().to(torch.bfloat16) + output = m(**batch) + loss, _ = peagle_kl_loss(output, torch.randn(1, 8, 23), torch.ones(1, 8)) + assert output.logits.dtype == torch.bfloat16 and torch.isfinite(loss) + loss.backward() + for param in m.parameters(): + assert param.grad is not None and torch.isfinite(param.grad).all() + + +@pytest.mark.parametrize("changes", [ + {"attention_dropout": 0.1}, {"layer_types": ["sliding_attention", "full_attention"]}, + {"hidden_act": "relu"}, {"num_hidden_layers": 0}, + {"draft_vocab_size": 12}, {"tie_word_embeddings": True}, + {"norm_before_residual": True}, {"norm_before_fc": True}, {"fc_norm": True}, + {"num_key_value_heads": 3}, {"head_dim": 7}, +]) +def test_unsupported_model_recipe_fails_closed(changes): + with pytest.raises(ValueError, match="P-EAGLE"): + model(**changes) + + +def test_other_target_and_attention_backend_fail_closed(): + with pytest.raises(ValueError, match="Qwen3"): + PEagleDraftModel(LlamaConfig(), mask_token_id=0) + with pytest.raises(ValueError, match="eager or sdpa"): + model(implementation="flash_attention_2") + + +@pytest.mark.parametrize("depths", [0, -1, 1.2, True]) +def test_invalid_depths_fail_closed(depths): + with pytest.raises(ValueError, match="positive integer"): + peagle_dense_layout(torch.tensor([0, 0]), depths) + + +@pytest.mark.parametrize("kind", ["batch", "feature", "token", "position", "document"]) +def test_malformed_inputs_fail_closed(kind): + m = model() + batch = example() + if kind == "batch": + batch["input_ids"] = batch["input_ids"].repeat(2, 1) + elif kind == "feature": + batch["hidden_states"] = batch["hidden_states"][..., :-1] + elif kind == "token": + batch["input_ids"][0, 0] = 23 + elif kind == "position": + batch["position_ids"][0, 0] = -1 + else: + batch["document_ids"][0, 0] = -2 + with pytest.raises(ValueError, match="P-EAGLE"): + m(**batch) + + +def test_wrong_loss_mask_and_vocabulary_fail_closed(): + output = model()(**example()) + with pytest.raises(ValueError, match="binary"): + peagle_kl_loss(output, torch.randn(1, 8, 23), torch.full((1, 8), 0.5)) + with pytest.raises(ValueError, match="full vocabulary"): + peagle_kl_loss(output, torch.randn(1, 8, 22), torch.ones(1, 8))