diff --git a/scripts/train.py b/scripts/train.py index 3f4f01518..31bd6c732 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -748,8 +748,11 @@ 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 for training. " + "Defaults to 3072 for DFlash, unlimited for P-EAGLE." + ), ) parser.add_argument( "--draft-attn-impl", @@ -778,6 +781,15 @@ def parse_args(): default=0.2, help="Minimum retention ratio for COD sampling in P-EAGLE (default: 0.2)", ) + parser.add_argument( + "--max-context-window", + type=int, + default=4096, + help=( + "Hard cap on contiguous window size when max-anchors is set " + "for P-EAGLE (default: 4096)" + ), + ) parser.add_argument( "--sliding-window", type=int, diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index 504e6a7dd..f8be124c4 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -156,7 +156,9 @@ 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=( + 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 6c860720a..2f7a17c0e 100644 --- a/src/speculators/models/peagle/config.py +++ b/src/speculators/models/peagle/config.py @@ -51,6 +51,25 @@ class PEagleSpeculatorConfig(Eagle3SpeculatorConfig): le=1.0, ) + max_anchors: int | None = Field( + default=None, + description=( + "Maximum number of COD chain starting points (loss_mask=1 tokens) " + "to use during training. When set, selects a contiguous window of " + "the sequence to cap memory usage. None means use all positions." + ), + ge=1, + ) + + max_context_window: int = Field( + default=4096, + description=( + "Hard cap on the contiguous window size when max_anchors is set. " + "Prevents sparse loss masks from reinflating the window." + ), + ge=1, + ) + mask_token_id: int | None = Field( default=None, description="Token ID used for padding unused positions in parallel groups", diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index cba6e480f..161617a98 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -42,6 +42,8 @@ def __init__( self.num_depths = config.num_depths self.down_sample_ratio = config.down_sample_ratio self.down_sample_ratio_min = config.down_sample_ratio_min + self.max_anchors = config.max_anchors + self.max_context_window = config.max_context_window self.mask_token_id = config.mask_token_id # Learnable mask_hidden parameter for padding unsampled positions @@ -91,6 +93,8 @@ 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.max_anchors, + max_context_window=self.max_context_window, ) total_sampled = anchor_pos.shape[0] @@ -214,6 +218,8 @@ def from_training_args( num_depths=kwargs.get("num_depths", 8), down_sample_ratio=kwargs.get("down_sample_ratio", 0.7), down_sample_ratio_min=kwargs.get("down_sample_ratio_min", 0.2), + max_anchors=kwargs.get("max_anchors"), + max_context_window=kwargs.get("max_context_window", 4096), mask_token_id=kwargs.get("mask_token_id"), speculators_config=SpeculatorsConfig( algorithm="peagle", diff --git a/src/speculators/models/peagle/data.py b/src/speculators/models/peagle/data.py index f06b72e0a..27ee2f5a9 100644 --- a/src/speculators/models/peagle/data.py +++ b/src/speculators/models/peagle/data.py @@ -10,6 +10,8 @@ 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, + max_context_window: int = 4096, ) -> tuple[torch.Tensor, torch.Tensor]: """ Generate sampling indices for parallel sequences using COD sampling. @@ -18,6 +20,11 @@ def generate_cod_sample_indices( decay: depth 0 retains all n positions, depth 1 retains n*r positions, depth 2 retains n*r^2 positions, etc. + When max_anchors is set, selects a contiguous window of the original sequence + containing up to max_anchors valid (loss_mask=1) positions as COD starting + points. The window preserves all intervening tokens (including prompts) + to maintain context. + Args: seq_length: Length of the sequence loss_mask: Binary mask indicating valid training positions [batch, seq_len] @@ -25,6 +32,10 @@ def generate_cod_sample_indices( down_sample_ratio: Geometric decay ratio r in (0,1) down_sample_ratio_min: Minimum retention ratio floor filter_position_zero: Whether to filter out position 0 from candidates + max_anchors: Maximum number of COD chain starting points. None means + use all positions. + max_context_window: Hard cap on contiguous window size to prevent + sparse loss masks from reinflating the window. Returns: Tuple of: @@ -36,8 +47,30 @@ def generate_cod_sample_indices( device = loss_mask.device all_valid_indices = torch.where(loss_mask == 1)[0] - sample_indices = [torch.arange(seq_length, device=device)] - n_per_depth = [seq_length] + if max_anchors is not None and all_valid_indices.shape[0] > 0: + if all_valid_indices.shape[0] > max_anchors: + n_valid = all_valid_indices.shape[0] + start_idx = int(torch.randint(0, n_valid - max_anchors + 1, (1,)).item()) + selected_valid = all_valid_indices[start_idx : start_idx + max_anchors] + else: + selected_valid = all_valid_indices + + window_start = int(selected_valid[0].item()) + window_end = int(selected_valid[-1].item()) + 1 + + if (window_end - window_start) > max_context_window: + window_end = min(window_start + max_context_window, seq_length) + selected_valid = selected_valid[ + (selected_valid >= window_start) & (selected_valid < window_end) + ] + all_valid_indices = selected_valid + + sample_indices = [torch.arange(window_start, window_end, device=device)] + n_per_depth = [window_end - window_start] + else: + sample_indices = [torch.arange(seq_length, device=device)] + n_per_depth = [seq_length] + prev_indices = all_valid_indices for d in range(1, num_depths): diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 1ada0dc67..faef08b9f 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -163,6 +163,8 @@ def make_peagle_model( draft_vocab_size: int = 64, num_depths: int = 4, down_sample_ratio: float = 0.7, + max_anchors: int | None = None, + max_context_window: int = 4096, device: str = "cuda:0", dtype: torch.dtype = torch.bfloat16, ) -> PEagleDraftModel: @@ -175,6 +177,8 @@ def make_peagle_model( num_depths=num_depths, down_sample_ratio=down_sample_ratio, down_sample_ratio_min=0.2, + max_anchors=max_anchors, + max_context_window=max_context_window, mask_token_id=0, speculators_config=SpeculatorsConfig( algorithm="peagle", diff --git a/tests/integration/models/test_model_forward.py b/tests/integration/models/test_model_forward.py index a459c2df7..1a8a13e27 100644 --- a/tests/integration/models/test_model_forward.py +++ b/tests/integration/models/test_model_forward.py @@ -338,6 +338,25 @@ def test_varying_down_sample_ratio(self, down_sample_ratio): assert loss.isfinite() loss.backward() + @pytest.mark.parametrize("max_anchors", [4, 16, None]) + def test_varying_max_anchors(self, max_anchors): + model = make_peagle_model(max_anchors=max_anchors) + samples = _make_samples([128]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = model(**batch) + + assert loss.isfinite() + loss.backward() + + def test_max_context_window(self): + model = make_peagle_model(max_anchors=16, max_context_window=32) + samples = _make_samples([128]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = model(**batch) + + assert loss.isfinite() + loss.backward() + @requires_cuda @requires_transformers_version("5.2.0") diff --git a/tests/unit/models/test_peagle_data.py b/tests/unit/models/test_peagle_data.py new file mode 100644 index 000000000..d8251dfd0 --- /dev/null +++ b/tests/unit/models/test_peagle_data.py @@ -0,0 +1,59 @@ +"""Unit tests for P-EAGLE COD sampling logic.""" + +import torch + +from speculators.models.peagle.data import generate_cod_sample_indices + + +class TestMaxContextWindow: + def test_applied_when_anchors_within_limit(self): + """max_context_window caps the window even when anchors <= max_anchors.""" + seq_length = 256 + loss_mask = torch.zeros(1, seq_length) + # Place 4 valid positions spread far apart (indices 10, 80, 160, 240) + valid_positions = [10, 80, 160, 240] + for pos in valid_positions: + loss_mask[0, pos] = 1 + + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_length, + loss_mask=loss_mask, + max_anchors=8, + max_context_window=64, + ) + + depth_0_mask = depth == 0 + depth_0_positions = anchor_pos[depth_0_mask] + window_size = depth_0_positions.shape[0] + assert window_size <= 64 + + def test_applied_when_anchors_exceed_limit(self): + """max_context_window caps the window when anchors > max_anchors.""" + seq_length = 512 + loss_mask = torch.ones(1, seq_length) + + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_length, + loss_mask=loss_mask, + max_anchors=16, + max_context_window=32, + ) + + depth_0_mask = depth == 0 + depth_0_positions = anchor_pos[depth_0_mask] + window_size = depth_0_positions.shape[0] + assert window_size <= 32 + + def test_no_windowing_without_max_anchors(self): + """Without max_anchors, full sequence is used.""" + seq_length = 128 + loss_mask = torch.ones(1, seq_length) + + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_length, + loss_mask=loss_mask, + max_anchors=None, + ) + + depth_0_mask = depth == 0 + assert anchor_pos[depth_0_mask].shape[0] == seq_length