From 7a3d9cc0ab332e212f219937bb2c23bddb159a96 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Tue, 30 Jun 2026 09:15:28 +0000 Subject: [PATCH 1/5] feat(peagle): add random anchor subsampling (max_anchors) Randomly subsamples COD chain starting points when max_anchors is set, reducing depth-1+ token count while keeping depth-0 as the full sequence. This bounds VRAM for long-context training without degrading d0 accuracy. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- scripts/train.py | 5 +- src/speculators/models/dflash/core.py | 2 +- src/speculators/models/peagle/config.py | 11 +++ src/speculators/models/peagle/core.py | 2 + src/speculators/models/peagle/data.py | 11 ++- tests/unit/models/test_peagle_data.py | 95 +++++++++++++++++++++++++ 6 files changed, 121 insertions(+), 5 deletions(-) create mode 100644 tests/unit/models/test_peagle_data.py diff --git a/scripts/train.py b/scripts/train.py index e45349d20..c208bbc08 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -935,8 +935,9 @@ def parse_args(): parser.add_argument( "--max-anchors", type=int, - default=256, - help="Maximum anchor positions for DFlash training (default: 256)", + default=None, + help="Maximum anchor positions. " + "DFlash default: 3072, P-EAGLE default: None (all).", ) parser.add_argument( "--dflash-decay-gamma", diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index d18dce302..a6b19bc20 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -176,7 +176,7 @@ def _build_base_config_kwargs( "transformer_layer_config": verifier_config, "draft_vocab_size": kwargs["draft_vocab_size"], "block_size": block_size, - "max_anchors": kwargs.get("max_anchors", 3072), + "max_anchors": kwargs.get("max_anchors") or 3072, "aux_hidden_state_layer_ids": target_layer_ids, "mask_token_id": kwargs.get("mask_token_id"), "sliding_window_non_causal": kwargs.get("sliding_window_non_causal", False), diff --git a/src/speculators/models/peagle/config.py b/src/speculators/models/peagle/config.py index 6c860720a..3544010ea 100644 --- a/src/speculators/models/peagle/config.py +++ b/src/speculators/models/peagle/config.py @@ -56,6 +56,17 @@ class PEagleSpeculatorConfig(Eagle3SpeculatorConfig): description="Token ID used for padding unused positions in parallel groups", ) + max_anchors: int | None = Field( + default=None, + description=( + "Maximum number of COD chain starting points. When set, " + "randomly subsamples valid positions for depth-1+ chains. " + "Depth 0 always retains the full sequence. " + "None means use all valid positions." + ), + ge=1, + ) + # Override Eagle3 default: P-EAGLE requires trainable embeddings # (matches p-eagle-train) embed_requires_grad: bool = Field( diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index 0842bef4e..76793fc42 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -90,6 +90,7 @@ def forward( num_depths=self.num_depths, down_sample_ratio=self.down_sample_ratio, down_sample_ratio_min=self.down_sample_ratio_min, + max_anchors=self.config.max_anchors, ) total_sampled = anchor_pos.shape[0] @@ -222,6 +223,7 @@ def from_training_args( down_sample_ratio=kwargs.get("down_sample_ratio", 0.7), down_sample_ratio_min=kwargs.get("down_sample_ratio_min", 0.2), mask_token_id=kwargs.get("mask_token_id"), + max_anchors=kwargs.get("max_anchors"), speculators_config=SpeculatorsConfig( algorithm="peagle", proposal_methods=[ diff --git a/src/speculators/models/peagle/data.py b/src/speculators/models/peagle/data.py index f06b72e0a..f7678eee4 100644 --- a/src/speculators/models/peagle/data.py +++ b/src/speculators/models/peagle/data.py @@ -10,6 +10,7 @@ def generate_cod_sample_indices( down_sample_ratio: float = 0.7, down_sample_ratio_min: float = 0.2, filter_position_zero: bool = True, + max_anchors: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Generate sampling indices for parallel sequences using COD sampling. @@ -36,12 +37,18 @@ def generate_cod_sample_indices( device = loss_mask.device all_valid_indices = torch.where(loss_mask == 1)[0] + if max_anchors is not None and all_valid_indices.shape[0] > max_anchors: + perm = torch.randperm(all_valid_indices.shape[0], device=device)[:max_anchors] + seed_indices = all_valid_indices[perm].sort()[0] + else: + seed_indices = all_valid_indices + sample_indices = [torch.arange(seq_length, device=device)] n_per_depth = [seq_length] - prev_indices = all_valid_indices + prev_indices = seed_indices for d in range(1, num_depths): - valid_length = max(0, all_valid_indices.shape[0] - d) + valid_length = max(0, seed_indices.shape[0] - d) ratio = max(down_sample_ratio**d, down_sample_ratio_min) sample_size = int(valid_length * ratio) diff --git a/tests/unit/models/test_peagle_data.py b/tests/unit/models/test_peagle_data.py new file mode 100644 index 000000000..27cd2ccea --- /dev/null +++ b/tests/unit/models/test_peagle_data.py @@ -0,0 +1,95 @@ +"""Unit tests for P-EAGLE COD sampling with max_anchors.""" + +import torch + +from speculators.models.peagle.data import generate_cod_sample_indices + + +def _loss_mask(seq_length: int) -> torch.Tensor: + return torch.ones(1, seq_length, dtype=torch.float32) + + +class TestMaxAnchors: + def test_depth0_is_full_sequence(self): + """Depth 0 should always be the full sequence regardless of max_anchors.""" + seq_len = 32 + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=_loss_mask(seq_len), + num_depths=4, + max_anchors=4, + ) + depth0_positions = anchor_pos[depth == 0] + assert depth0_positions.shape[0] == seq_len + assert torch.equal(depth0_positions, torch.arange(seq_len)) + + def test_max_anchors_caps_chains(self): + """With max_anchors, depth-1+ chains should not exceed max_anchors.""" + seq_len = 64 + max_anchors = 8 + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=_loss_mask(seq_len), + num_depths=4, + max_anchors=max_anchors, + ) + depth1_count = (depth == 1).sum().item() + assert depth1_count <= max_anchors + + def test_max_anchors_preserves_full_depth0(self): + """Depth 0 count should equal seq_length even with small max_anchors.""" + seq_len = 128 + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=_loss_mask(seq_len), + num_depths=8, + max_anchors=4, + ) + assert (depth == 0).sum().item() == seq_len + + def test_max_anchors_none_uses_all(self): + """max_anchors=None should use all valid positions (default behavior).""" + seq_len = 32 + torch.manual_seed(42) + anchor_pos_none, depth_none = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=_loss_mask(seq_len), + num_depths=4, + max_anchors=None, + ) + torch.manual_seed(42) + anchor_pos_default, depth_default = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=_loss_mask(seq_len), + num_depths=4, + ) + assert torch.equal(anchor_pos_none, anchor_pos_default) + assert torch.equal(depth_none, depth_default) + + def test_max_anchors_fewer_valid_than_cap(self): + """When valid positions < max_anchors, all valid positions are used.""" + seq_len = 16 + loss_mask = torch.zeros(1, seq_len) + loss_mask[0, :5] = 1 + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + max_anchors=100, + ) + assert (depth == 0).sum().item() == seq_len + assert (depth == 1).sum().item() > 1 + + def test_max_anchors_sorted_order(self): + """Subsampled anchors should be in sorted order for causal masking.""" + seq_len = 64 + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=_loss_mask(seq_len), + num_depths=4, + max_anchors=8, + ) + for d in range(1, 4): + d_anchors = anchor_pos[depth == d] + if d_anchors.shape[0] > 1: + assert torch.all(d_anchors[1:] >= d_anchors[:-1]) From c97346151390d616dcb7a13a084461c1e97d5fce Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Wed, 1 Jul 2026 11:45:42 +0000 Subject: [PATCH 2/5] address PR #687 review: move max_anchors to trainer kwargs, default 3072 - Remove max_anchors from PEagleSpeculatorConfig (training hyperparam, not model config); pipe through get_trainer_kwargs instead - Set --max-anchors CLI default to 3072 for both algorithms - Fix DFlash falsy check to use `is None` for max_anchors defaulting - Strengthen test assertions: check all depths, verify no-cap equality Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- scripts/train.py | 5 ++--- src/speculators/models/dflash/core.py | 6 +++++- src/speculators/models/peagle/config.py | 11 ----------- src/speculators/models/peagle/core.py | 7 ++++--- tests/unit/models/test_peagle_data.py | 19 ++++++++++++++----- 5 files changed, 25 insertions(+), 23 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index f33c67917..248aeff8b 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -935,9 +935,8 @@ def parse_args(): parser.add_argument( "--max-anchors", type=int, - default=None, - help="Maximum anchor positions. " - "DFlash default: 3072, P-EAGLE default: None (all).", + default=3072, + help="Maximum anchor positions for training (default: 3072).", ) parser.add_argument( "--dflash-decay-gamma", diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index a6b19bc20..a7e2d8abb 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -176,7 +176,11 @@ def _build_base_config_kwargs( "transformer_layer_config": verifier_config, "draft_vocab_size": kwargs["draft_vocab_size"], "block_size": block_size, - "max_anchors": kwargs.get("max_anchors") or 3072, + "max_anchors": ( + 3072 + if kwargs.get("max_anchors") is None + else kwargs["max_anchors"] + ), "aux_hidden_state_layer_ids": target_layer_ids, "mask_token_id": kwargs.get("mask_token_id"), "sliding_window_non_causal": kwargs.get("sliding_window_non_causal", False), diff --git a/src/speculators/models/peagle/config.py b/src/speculators/models/peagle/config.py index 3544010ea..6c860720a 100644 --- a/src/speculators/models/peagle/config.py +++ b/src/speculators/models/peagle/config.py @@ -56,17 +56,6 @@ class PEagleSpeculatorConfig(Eagle3SpeculatorConfig): description="Token ID used for padding unused positions in parallel groups", ) - max_anchors: int | None = Field( - default=None, - description=( - "Maximum number of COD chain starting points. When set, " - "randomly subsamples valid positions for depth-1+ chains. " - "Depth 0 always retains the full sequence. " - "None means use all valid positions." - ), - ge=1, - ) - # Override Eagle3 default: P-EAGLE requires trainable embeddings # (matches p-eagle-train) embed_requires_grad: bool = Field( diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index ca70305df..a8586ee80 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -90,7 +90,7 @@ def forward( num_depths=self.num_depths, down_sample_ratio=self.down_sample_ratio, down_sample_ratio_min=self.down_sample_ratio_min, - max_anchors=self.config.max_anchors, + max_anchors=kwargs.get("max_anchors"), ) total_sampled = anchor_pos.shape[0] @@ -223,7 +223,6 @@ def from_training_args( down_sample_ratio=kwargs.get("down_sample_ratio", 0.7), down_sample_ratio_min=kwargs.get("down_sample_ratio_min", 0.2), mask_token_id=kwargs.get("mask_token_id"), - max_anchors=kwargs.get("max_anchors"), speculators_config=SpeculatorsConfig( algorithm="peagle", proposal_methods=[ @@ -255,4 +254,6 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]: Tuple of (train_call_kwargs, val_call_kwargs) """ loss_config = resolve_loss_config(kwargs["loss_fn"]) - return {"loss_config": loss_config}, {"loss_config": loss_config} + max_anchors = kwargs.get("max_anchors") + shared = {"loss_config": loss_config, "max_anchors": max_anchors} + return dict(shared), dict(shared) diff --git a/tests/unit/models/test_peagle_data.py b/tests/unit/models/test_peagle_data.py index 27cd2ccea..44a7f22c0 100644 --- a/tests/unit/models/test_peagle_data.py +++ b/tests/unit/models/test_peagle_data.py @@ -33,8 +33,8 @@ def test_max_anchors_caps_chains(self): num_depths=4, max_anchors=max_anchors, ) - depth1_count = (depth == 1).sum().item() - assert depth1_count <= max_anchors + for d in range(1, 4): + assert (depth == d).sum().item() <= max_anchors def test_max_anchors_preserves_full_depth0(self): """Depth 0 count should equal seq_length even with small max_anchors.""" @@ -71,14 +71,23 @@ def test_max_anchors_fewer_valid_than_cap(self): seq_len = 16 loss_mask = torch.zeros(1, seq_len) loss_mask[0, :5] = 1 - anchor_pos, depth = generate_cod_sample_indices( + torch.manual_seed(42) + anchor_pos_capped, depth_capped = generate_cod_sample_indices( seq_length=seq_len, loss_mask=loss_mask, num_depths=4, max_anchors=100, ) - assert (depth == 0).sum().item() == seq_len - assert (depth == 1).sum().item() > 1 + torch.manual_seed(42) + anchor_pos_none, depth_none = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + max_anchors=None, + ) + assert (depth_capped == 0).sum().item() == seq_len + assert torch.equal(anchor_pos_capped, anchor_pos_none) + assert torch.equal(depth_capped, depth_none) def test_max_anchors_sorted_order(self): """Subsampled anchors should be in sorted order for causal masking.""" From 0d92463c41c7cfbaeea89e2df5cc13ba793d9998 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Wed, 1 Jul 2026 11:47:40 +0000 Subject: [PATCH 3/5] fix(train): update --max-anchors default to 256 and help text Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- scripts/train.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/train.py b/scripts/train.py index 248aeff8b..308eba4b1 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -936,7 +936,8 @@ def parse_args(): "--max-anchors", type=int, default=3072, - help="Maximum anchor positions for training (default: 3072).", + help="Maximum anchor positions for DFlash, DSpark, " + "and P-EAGLE training (default: 3072).", ) parser.add_argument( "--dflash-decay-gamma", From 567d5153e30fe28c2fb19dbbaf0a74372ac76474 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Wed, 1 Jul 2026 12:01:17 +0000 Subject: [PATCH 4/5] style: format dflash/core.py, revert unneeded peagle fallback Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- src/speculators/models/dflash/core.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index a7e2d8abb..5695b9d29 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -177,9 +177,7 @@ def _build_base_config_kwargs( "draft_vocab_size": kwargs["draft_vocab_size"], "block_size": block_size, "max_anchors": ( - 3072 - if kwargs.get("max_anchors") is None - else kwargs["max_anchors"] + 3072 if kwargs.get("max_anchors") is None else kwargs["max_anchors"] ), "aux_hidden_state_layer_ids": target_layer_ids, "mask_token_id": kwargs.get("mask_token_id"), From 113d297347822a1d8946db11795c35272bedd1ef Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Thu, 2 Jul 2026 14:05:02 +0000 Subject: [PATCH 5/5] refactor(peagle): list max_anchors as explicit forward param Move max_anchors from **kwargs to an explicit keyword argument in PEagleDraftModel.forward() for clarity. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- src/speculators/models/peagle/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index 4e29292ab..88c84a04a 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -63,6 +63,7 @@ def forward( loss_mask: torch.Tensor | None = None, verifier_last_hidden_states: torch.Tensor | None = None, loss_config: LossConfig | None = None, + max_anchors: int | None = None, **kwargs, ): """ @@ -97,7 +98,7 @@ def forward( num_depths=self.num_depths, down_sample_ratio=self.down_sample_ratio, down_sample_ratio_min=self.down_sample_ratio_min, - max_anchors=kwargs.get("max_anchors"), + max_anchors=max_anchors, ) total_sampled = anchor_pos.shape[0]