diff --git a/scripts/train.py b/scripts/train.py index a934f400f..dc0893612 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1043,8 +1043,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 6ce3cfba3..63e10d250 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -159,7 +159,7 @@ def from_training_args( transformer_layer_config=verifier_config, draft_vocab_size=kwargs["draft_vocab_size"], block_size=kwargs.get("block_size", 8), - 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 7ac846d48..20274ee51 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, + ) + sink_size: int | None = Field( default=None, description=( diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index 47f9037e6..384a76e70 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] @@ -224,6 +225,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"), sink_size=kwargs.get("sink_size"), max_context_window=kwargs.get("max_context_window"), speculators_config=SpeculatorsConfig( diff --git a/src/speculators/models/peagle/data.py b/src/speculators/models/peagle/data.py index f06b72e0a..2c76e4a03 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,6 +37,10 @@ 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] + all_valid_indices = all_valid_indices[perm].sort()[0] + sample_indices = [torch.arange(seq_length, device=device)] n_per_depth = [seq_length] prev_indices = all_valid_indices diff --git a/tests/unit/models/test_peagle_data.py b/tests/unit/models/test_peagle_data.py new file mode 100644 index 000000000..6e6a0fcfb --- /dev/null +++ b/tests/unit/models/test_peagle_data.py @@ -0,0 +1,94 @@ +"""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 anchor_pos.shape[0] > 0 + + 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])