From 695a2a6d3767ab33e905268f35135bc1b5194a77 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 26 May 2026 16:44:05 +0000 Subject: [PATCH 1/3] Add forward/backward integration tests Signed-off-by: Fynn Schmitt-Ulms --- tests/integration/conftest.py | 276 ++++++++++++++++++ .../integration/models/test_dflash_forward.py | 165 +++++++++++ .../integration/models/test_eagle3_forward.py | 172 +++++++++++ .../integration/models/test_peagle_forward.py | 164 +++++++++++ 4 files changed, 777 insertions(+) create mode 100644 tests/integration/conftest.py create mode 100644 tests/integration/models/test_dflash_forward.py create mode 100644 tests/integration/models/test_eagle3_forward.py create mode 100644 tests/integration/models/test_peagle_forward.py diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 000000000..88a8584dd --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,276 @@ +"""Shared fixtures and factories for integration tests.""" + +import copy + +import pytest +import torch +from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.models.qwen3.configuration_qwen3 import Qwen3Config + +from speculators import SpeculatorsConfig, VerifierConfig +from speculators.models.dflash import DFlashSpeculatorConfig +from speculators.models.dflash.core import DFlashDraftModel +from speculators.models.eagle3 import Eagle3SpeculatorConfig +from speculators.models.eagle3.core import Eagle3DraftModel +from speculators.models.peagle.config import PEagleSpeculatorConfig +from speculators.models.peagle.core import PEagleDraftModel +from speculators.proposals.greedy import GreedyTokenProposalConfig +from speculators.train.data import create_collate_fn + +# --------------------------------------------------------------------------- +# Tiny verifier configs +# --------------------------------------------------------------------------- + +TINY_LLAMA_CONFIG = LlamaConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + max_position_embeddings=256, + rms_norm_eps=1e-6, + tie_word_embeddings=False, + _attn_implementation="eager", # type: ignore[call-arg] +) + +TINY_QWEN3_CONFIG = Qwen3Config( + vocab_size=128, + hidden_size=64, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + max_position_embeddings=256, + rms_norm_eps=1e-6, + tie_word_embeddings=False, +) + + +# --------------------------------------------------------------------------- +# Model factories +# --------------------------------------------------------------------------- + + +def _fill_nan_weights(model): + """Replace NaN-initialized weights with deterministic values.""" + with torch.no_grad(): + for param in model.parameters(): + if param.isnan().any(): + torch.nn.init.normal_(param, mean=0.0, std=0.02) + for buf in model.buffers(): + if buf.is_floating_point() and buf.isnan().any(): + buf.zero_() + + +def make_eagle3_model( + *, + draft_vocab_size: int = 64, + norm_before_residual: bool = False, + device: str = "cuda", +) -> Eagle3DraftModel: + """Create a tiny Eagle3 model with real initialized weights.""" + config = Eagle3SpeculatorConfig( + transformer_layer_config=copy.deepcopy(TINY_LLAMA_CONFIG), + draft_vocab_size=draft_vocab_size, + norm_before_residual=norm_before_residual, + embed_requires_grad=False, + speculators_config=SpeculatorsConfig( + algorithm="eagle3", + proposal_methods=[GreedyTokenProposalConfig(speculative_tokens=1)], + default_proposal_method="greedy", + verifier=VerifierConfig( + name_or_path=None, + architectures=["LlamaForCausalLM"], + ), + ), + ) + model = Eagle3DraftModel(config) + _fill_nan_weights(model) + return model.to(device) # type: ignore[arg-type] + + +def make_dflash_model( + *, + draft_vocab_size: int = 64, + block_size: int = 4, + max_anchors: int = 8, + device: str = "cuda", +) -> DFlashDraftModel: + """Create a tiny DFlash model with real initialized weights.""" + config = DFlashSpeculatorConfig( + transformer_layer_config=copy.deepcopy(TINY_QWEN3_CONFIG), + draft_vocab_size=draft_vocab_size, + block_size=block_size, + max_anchors=max_anchors, + aux_hidden_state_layer_ids=[0, 1], + mask_token_id=0, + speculators_config=SpeculatorsConfig( + algorithm="dflash", + proposal_methods=[ + GreedyTokenProposalConfig(speculative_tokens=block_size - 1) + ], + default_proposal_method="greedy", + verifier=VerifierConfig( + name_or_path=None, + architectures=["Qwen3ForCausalLM"], + ), + ), + ) + model = DFlashDraftModel(config) + _fill_nan_weights(model) + return model.to(device) # type: ignore[arg-type] + + +def make_peagle_model( + *, + draft_vocab_size: int = 64, + num_depths: int = 4, + down_sample_ratio: float = 0.7, + device: str = "cuda", +) -> PEagleDraftModel: + """Create a tiny PEagle model with real initialized weights.""" + config = PEagleSpeculatorConfig( + transformer_layer_config=copy.deepcopy(TINY_LLAMA_CONFIG), + draft_vocab_size=draft_vocab_size, + norm_before_residual=False, + embed_requires_grad=True, + num_depths=num_depths, + down_sample_ratio=down_sample_ratio, + down_sample_ratio_min=0.2, + mask_token_id=0, + speculators_config=SpeculatorsConfig( + algorithm="peagle", + proposal_methods=[GreedyTokenProposalConfig(speculative_tokens=num_depths)], + default_proposal_method="greedy", + verifier=VerifierConfig( + name_or_path=None, + architectures=["LlamaForCausalLM"], + ), + ), + ) + model = PEagleDraftModel(config) + _fill_nan_weights(model) + return model.to(device) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Synthetic data factory +# --------------------------------------------------------------------------- + + +def make_sample( + *, + seq_len: int, + hidden_size: int, + hidden_multiplier: int = 3, + vocab_size: int, + loss_mask_pattern: str = "all", + include_verifier_states: bool = True, + boundary_token_ids: list[int] | None = None, + dtype: torch.dtype = torch.float32, +) -> dict[str, torch.Tensor]: + """Generate a single synthetic sample (no batch dim, seq_len on dim 0). + + Matches the per-sample format produced by BaseDataset.__getitem__: + hidden_states: [seq_len, hidden_multiplier * hidden_size] + input_ids: [seq_len] + verifier_last_hidden_states: [seq_len, hidden_size] + loss_mask: [seq_len] + lengths: [1] + position_ids: [seq_len] + """ + hidden_states = torch.randn(seq_len, hidden_multiplier * hidden_size, dtype=dtype) + + input_ids = torch.randint(0, vocab_size, (seq_len,)) + if boundary_token_ids: + for i, tok_id in enumerate(boundary_token_ids): + if i < seq_len: + input_ids[i] = tok_id + + if loss_mask_pattern == "all": + loss_mask = torch.ones(seq_len, dtype=dtype) + elif loss_mask_pattern == "none": + loss_mask = torch.zeros(seq_len, dtype=dtype) + elif loss_mask_pattern == "random": + loss_mask = torch.randint(0, 2, (seq_len,)).to(dtype) + elif loss_mask_pattern == "alternating": + loss_mask = torch.zeros(seq_len, dtype=dtype) + loss_mask[::2] = 1.0 + else: + raise ValueError(f"Unknown loss_mask_pattern: {loss_mask_pattern}") + + result = { + "hidden_states": hidden_states, + "input_ids": input_ids, + "loss_mask": loss_mask, + "lengths": torch.tensor([seq_len], dtype=torch.long), + "position_ids": torch.arange(seq_len, dtype=torch.long), + } + + if include_verifier_states: + result["verifier_last_hidden_states"] = torch.randn( + seq_len, hidden_size, dtype=dtype + ) + + return result + + +def make_batch( + *, + max_len: int, + samples: list[dict[str, torch.Tensor]], + hidden_size: int, + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Collate a list of samples into a single batch using the real collate_fn. + + Uses ``create_collate_fn`` from ``speculators.train.data`` to pack and pad + samples exactly the way the training pipeline does. + + Args: + max_len: Target sequence length to pad/truncate to. + samples: List of per-sample dicts (from ``make_sample``). + hidden_size: Model hidden size (needed by collate for empty batches). + device: Target device for the output tensors. + + Returns: + Dict with keys matching model forward() signatures, all on ``device``. + """ + collate_fn = create_collate_fn(max_len=max_len, hidden_size=hidden_size) + batch = collate_fn(samples) + return {k: v.to(device) for k, v in batch.items()} + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +HIDDEN_SIZE = TINY_LLAMA_CONFIG.hidden_size +VOCAB_SIZE = TINY_LLAMA_CONFIG.vocab_size + + +@pytest.fixture +def eagle3_model(): + model = make_eagle3_model() + yield model + del model + torch.cuda.empty_cache() + + +@pytest.fixture +def dflash_model(): + model = make_dflash_model() + yield model + del model + torch.cuda.empty_cache() + + +@pytest.fixture +def peagle_model(): + model = make_peagle_model() + yield model + del model + torch.cuda.empty_cache() diff --git a/tests/integration/models/test_dflash_forward.py b/tests/integration/models/test_dflash_forward.py new file mode 100644 index 000000000..e000597be --- /dev/null +++ b/tests/integration/models/test_dflash_forward.py @@ -0,0 +1,165 @@ +"""Integration tests for DFlashDraftModel forward passes with real weights.""" + +import pytest +import torch + +from tests.conftest import requires_cuda +from tests.integration.conftest import ( + HIDDEN_SIZE, + VOCAB_SIZE, + make_batch, + make_dflash_model, + make_sample, +) + +MAX_LEN = 128 +NUM_TARGET_LAYERS = 2 + +LOSS_MASK_CASES = ["all", "none", "random", "alternating"] + + +def _make_samples( + seq_lengths: list[int], + loss_mask_pattern: str = "all", + vocab_size: int = VOCAB_SIZE, + boundary_token_ids: list[int] | None = None, +) -> list[dict[str, torch.Tensor]]: + return [ + make_sample( + seq_len=sl, + hidden_size=HIDDEN_SIZE, + hidden_multiplier=NUM_TARGET_LAYERS, + vocab_size=vocab_size, + loss_mask_pattern=loss_mask_pattern, + include_verifier_states=True, + boundary_token_ids=boundary_token_ids, + ) + for sl in seq_lengths + ] + + +SAMPLE_CONFIGS = [ + pytest.param([128], id="single_sample"), + pytest.param([64, 64], id="two_equal"), + pytest.param([32, 96], id="two_unequal"), + pytest.param([16] * 8, id="eight_tiny"), + pytest.param([8] * 20, id="twenty_tiny"), +] + + +@requires_cuda +class TestDFlashTraining: + """Forward pass in training mode (DFlash always requires verifier states).""" + + @pytest.mark.parametrize("seq_lengths", SAMPLE_CONFIGS) + @pytest.mark.parametrize("loss_mask_pattern", LOSS_MASK_CASES) + def test_forward_backward_produces_valid_outputs( + self, dflash_model, seq_lengths, loss_mask_pattern + ): + samples = _make_samples(seq_lengths, loss_mask_pattern=loss_mask_pattern) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = dflash_model(**batch) + + assert loss.isfinite(), f"Loss is not finite: {loss.item()}" + assert "loss_sum" in metrics + assert "loss_total" in metrics + + loss.backward() + + @pytest.mark.parametrize("block_size", [2, 4, 8]) + def test_varying_block_size(self, block_size): + model = make_dflash_model(block_size=block_size, max_anchors=4) + 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() + + @pytest.mark.parametrize("max_anchors", [2, 8, 16]) + def test_varying_max_anchors(self, max_anchors): + model = make_dflash_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() + + +@requires_cuda +class TestDFlashMultiBatch: + """Run multiple batches back-to-back through the same model to test + statefulness, cache clearing, and varying batch compositions.""" + + def test_multi_then_single(self, dflash_model): + """Multi-sample batches, then single-sample.""" + batch_configs: list[list[int]] = [ + [32, 32, 32], + [64, 64], + [128], + [16], + # Empty batch ([]) excluded: create_empty_sample dtype bug (#527) + ] + for seq_lengths in batch_configs: + samples = _make_samples(seq_lengths) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = dflash_model(**batch) + assert loss.isfinite(), f"Loss not finite for seq_lengths={seq_lengths}" + loss.backward() + + def test_alternating_batch_sizes(self, dflash_model): + """Alternate between large multi-sample and single-sample batches.""" + batch_configs: list[list[int]] = [ + [16] * 8, + [128], + [32, 32, 32, 32], + [64], + [8] * 16, + [128], + ] + for seq_lengths in batch_configs: + samples = _make_samples(seq_lengths) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = dflash_model(**batch) + assert loss.isfinite() + loss.backward() + + def test_varying_loss_masks_across_batches(self, dflash_model): + """Each batch uses a different loss mask pattern.""" + for pattern in LOSS_MASK_CASES: + samples = _make_samples([64, 64], loss_mask_pattern=pattern) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = dflash_model(**batch) + assert loss.isfinite(), f"Loss not finite for loss_mask_pattern={pattern}" + loss.backward() + + +@requires_cuda +class TestDFlashVocabBoundary: + """Tests with draft vocab mapping.""" + + @pytest.fixture + def dflash_draft_vocab_model(self): + model = make_dflash_model(draft_vocab_size=32) + t2d = torch.zeros(VOCAB_SIZE, dtype=torch.bool) + t2d[:32] = True + d2t = torch.arange(32, dtype=torch.long) + model.load_vocab_mappings(t2d.to("cuda"), d2t.to("cuda")) + yield model + del model + torch.cuda.empty_cache() + + def test_boundary_tokens(self, dflash_draft_vocab_model): + samples = _make_samples([128], vocab_size=32, boundary_token_ids=[0, 31]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = dflash_draft_vocab_model(**batch) + + assert loss.isfinite() + loss.backward() diff --git a/tests/integration/models/test_eagle3_forward.py b/tests/integration/models/test_eagle3_forward.py new file mode 100644 index 000000000..202f8b5bf --- /dev/null +++ b/tests/integration/models/test_eagle3_forward.py @@ -0,0 +1,172 @@ +"""Integration tests for Eagle3DraftModel forward passes with real weights.""" + +import pytest +import torch + +from tests.conftest import requires_cuda +from tests.integration.conftest import ( + HIDDEN_SIZE, + VOCAB_SIZE, + make_batch, + make_eagle3_model, + make_sample, +) + +MAX_LEN = 128 + +LOSS_MASK_CASES = ["all", "none", "random", "alternating"] + + +def _make_samples( + seq_lengths: list[int], + loss_mask_pattern: str = "all", + include_verifier_states: bool = True, + vocab_size: int = VOCAB_SIZE, + boundary_token_ids: list[int] | None = None, +) -> list[dict[str, torch.Tensor]]: + return [ + make_sample( + seq_len=sl, + hidden_size=HIDDEN_SIZE, + hidden_multiplier=3, + vocab_size=vocab_size, + loss_mask_pattern=loss_mask_pattern, + include_verifier_states=include_verifier_states, + boundary_token_ids=boundary_token_ids, + ) + for sl in seq_lengths + ] + + +# --------------------------------------------------------------------------- +# Sample compositions for parametrize +# --------------------------------------------------------------------------- + +SAMPLE_CONFIGS = [ + pytest.param([128], id="single_sample"), + pytest.param([64, 64], id="two_equal"), + pytest.param([32, 96], id="two_unequal"), + pytest.param([16] * 8, id="eight_tiny"), + pytest.param([8] * 20, id="twenty_tiny"), +] + + +@requires_cuda +class TestEagle3Training: + """Forward pass in training mode (with verifier_last_hidden_states).""" + + @pytest.mark.parametrize("seq_lengths", SAMPLE_CONFIGS) + @pytest.mark.parametrize("loss_mask_pattern", LOSS_MASK_CASES) + def test_forward_backward_produces_valid_outputs( + self, eagle3_model, seq_lengths, loss_mask_pattern + ): + samples = _make_samples(seq_lengths, loss_mask_pattern=loss_mask_pattern) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + ttt_steps = 2 + draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=ttt_steps) + + assert len(draft_tokens) == ttt_steps + for dt in draft_tokens: + assert dt.shape == (1, MAX_LEN) + assert dt.dtype == torch.long + + assert loss.isfinite(), f"Loss is not finite: {loss.item()}" + assert "loss_sum" in metrics + assert "loss_total" in metrics + + loss.backward() + + @pytest.mark.parametrize("ttt_steps", [1, 3, 5]) + def test_varying_ttt_steps(self, eagle3_model, ttt_steps): + samples = _make_samples([128]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=ttt_steps) + + assert len(draft_tokens) == ttt_steps + assert loss.isfinite() + + loss.backward() + + +@requires_cuda +class TestEagle3MultiBatch: + """Run multiple batches back-to-back through the same model to test + statefulness, cache clearing, and varying batch compositions.""" + + def test_multi_then_single_then_empty(self, eagle3_model): + """Multi-sample batches, then single-sample, then empty batch.""" + batch_configs: list[list[int]] = [ + [16, 16, 8, 10, 15, 12], + [32, 32, 32], + [64, 64], + [32, 3, 17], + [128], + [16], + # Empty batch ([]) excluded: create_empty_sample dtype bug (#527) + ] + torch.compiler.reset() + for seq_lengths in batch_configs: + samples = _make_samples(seq_lengths) if seq_lengths else [] + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=2) + assert loss.isfinite(), f"Loss not finite for seq_lengths={seq_lengths}" + loss.backward() + + def test_alternating_batch_sizes(self, eagle3_model): + """Alternate between large multi-sample and single-sample batches.""" + batch_configs: list[list[int]] = [ + [16] * 8, + [128], + [32, 32, 32, 32], + [64], + [8] * 16, + [128], + ] + for seq_lengths in batch_configs: + samples = _make_samples(seq_lengths) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=2) + assert loss.isfinite() + loss.backward() + + def test_varying_loss_masks_across_batches(self, eagle3_model): + """Each batch uses a different loss mask pattern.""" + for pattern in LOSS_MASK_CASES: + samples = _make_samples([64, 64], loss_mask_pattern=pattern) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=2) + assert loss.isfinite(), f"Loss not finite for loss_mask_pattern={pattern}" + loss.backward() + + +@requires_cuda +class TestEagle3VocabBoundary: + """Tests with draft vocab mapping and boundary token IDs.""" + + @pytest.fixture + def eagle3_draft_vocab_model(self): + model = make_eagle3_model(draft_vocab_size=32) + t2d = torch.zeros(VOCAB_SIZE, dtype=torch.bool) + t2d[:32] = True + d2t = torch.arange(32, dtype=torch.long) + model.load_vocab_mappings(t2d.to("cuda"), d2t.to("cuda")) + yield model + del model + torch.cuda.empty_cache() + + def test_boundary_tokens_training(self, eagle3_draft_vocab_model): + samples = _make_samples([128], vocab_size=32, boundary_token_ids=[0, 31]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = eagle3_draft_vocab_model(**batch, ttt_steps=2) + + assert len(draft_tokens) == 2 + assert loss.isfinite() + for dt in draft_tokens: + assert (dt >= 0).all() + assert (dt < 32).all() diff --git a/tests/integration/models/test_peagle_forward.py b/tests/integration/models/test_peagle_forward.py new file mode 100644 index 000000000..7094a7b57 --- /dev/null +++ b/tests/integration/models/test_peagle_forward.py @@ -0,0 +1,164 @@ +"""Integration tests for PEagleDraftModel forward passes with real weights.""" + +import pytest +import torch + +from tests.conftest import requires_cuda +from tests.integration.conftest import ( + HIDDEN_SIZE, + VOCAB_SIZE, + make_batch, + make_peagle_model, + make_sample, +) + +MAX_LEN = 128 + +LOSS_MASK_CASES = ["all", "none", "random", "alternating"] + + +def _make_samples( + seq_lengths: list[int], + loss_mask_pattern: str = "all", + vocab_size: int = VOCAB_SIZE, + boundary_token_ids: list[int] | None = None, +) -> list[dict[str, torch.Tensor]]: + return [ + make_sample( + seq_len=sl, + hidden_size=HIDDEN_SIZE, + hidden_multiplier=3, + vocab_size=vocab_size, + loss_mask_pattern=loss_mask_pattern, + include_verifier_states=True, + boundary_token_ids=boundary_token_ids, + ) + for sl in seq_lengths + ] + + +SAMPLE_CONFIGS = [ + pytest.param([128], id="single_sample"), + pytest.param([64, 64], id="two_equal"), + pytest.param([32, 96], id="two_unequal"), + pytest.param([16] * 8, id="eight_tiny"), + pytest.param([8] * 20, id="twenty_tiny"), +] + + +@requires_cuda +class TestPEagleTraining: + """Forward pass in training mode (PEagle always requires verifier states).""" + + @pytest.mark.parametrize("seq_lengths", SAMPLE_CONFIGS) + @pytest.mark.parametrize("loss_mask_pattern", LOSS_MASK_CASES) + def test_forward_backward_produces_valid_outputs( + self, peagle_model, seq_lengths, loss_mask_pattern + ): + samples = _make_samples(seq_lengths, loss_mask_pattern=loss_mask_pattern) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = peagle_model(**batch) + + assert loss.isfinite(), f"Loss is not finite: {loss.item()}" + assert "loss_sum" in metrics + assert "loss_total" in metrics + + loss.backward() + + @pytest.mark.parametrize("num_depths", [2, 4, 8]) + def test_varying_num_depths(self, num_depths): + model = make_peagle_model(num_depths=num_depths) + 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() + + @pytest.mark.parametrize("down_sample_ratio", [0.3, 0.7, 1.0]) + def test_varying_down_sample_ratio(self, down_sample_ratio): + model = make_peagle_model(down_sample_ratio=down_sample_ratio) + 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 +class TestPEagleMultiBatch: + """Run multiple batches back-to-back through the same model to test + statefulness, cache clearing, and varying batch compositions.""" + + def test_multi_then_single(self, peagle_model): + """Multi-sample batches, then single-sample.""" + batch_configs: list[list[int]] = [ + [32, 32, 32], + [64, 64], + [128], + [16], + # Empty batch ([]) excluded: create_empty_sample dtype bug (#527) + ] + for seq_lengths in batch_configs: + samples = _make_samples(seq_lengths) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = peagle_model(**batch) + assert loss.isfinite(), f"Loss not finite for seq_lengths={seq_lengths}" + loss.backward() + + def test_alternating_batch_sizes(self, peagle_model): + """Alternate between large multi-sample and single-sample batches.""" + batch_configs: list[list[int]] = [ + [16] * 8, + [128], + [32, 32, 32, 32], + [64], + [8] * 16, + [128], + ] + for seq_lengths in batch_configs: + samples = _make_samples(seq_lengths) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = peagle_model(**batch) + assert loss.isfinite() + loss.backward() + + def test_varying_loss_masks_across_batches(self, peagle_model): + """Each batch uses a different loss mask pattern.""" + for pattern in LOSS_MASK_CASES: + samples = _make_samples([64, 64], loss_mask_pattern=pattern) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = peagle_model(**batch) + assert loss.isfinite(), f"Loss not finite for loss_mask_pattern={pattern}" + loss.backward() + + +@requires_cuda +class TestPEagleVocabBoundary: + """Tests with draft vocab mapping and boundary token IDs.""" + + @pytest.fixture + def peagle_draft_vocab_model(self): + model = make_peagle_model(draft_vocab_size=32) + t2d = torch.zeros(VOCAB_SIZE, dtype=torch.bool) + t2d[:32] = True + d2t = torch.arange(32, dtype=torch.long) + model.load_vocab_mappings(t2d.to("cuda"), d2t.to("cuda")) + yield model + del model + torch.cuda.empty_cache() + + def test_boundary_tokens(self, peagle_draft_vocab_model): + samples = _make_samples([128], vocab_size=32, boundary_token_ids=[0, 31]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = peagle_draft_vocab_model(**batch) + + assert loss.isfinite() + loss.backward() From ae553cf37cd26128536b345d689e13e7cc766171 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 26 May 2026 14:47:09 -0400 Subject: [PATCH 2/3] Update tests/integration/models/test_dflash_forward.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Fynn Schmitt-Ulms --- tests/integration/models/test_dflash_forward.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/integration/models/test_dflash_forward.py b/tests/integration/models/test_dflash_forward.py index e000597be..a315b1a48 100644 --- a/tests/integration/models/test_dflash_forward.py +++ b/tests/integration/models/test_dflash_forward.py @@ -162,4 +162,7 @@ def test_boundary_tokens(self, dflash_draft_vocab_model): draft_tokens, loss, metrics = dflash_draft_vocab_model(**batch) assert loss.isfinite() + assert draft_tokens.dtype == torch.long + assert (draft_tokens >= 0).all() + assert (draft_tokens < 32).all() loss.backward() From 63c153a801bbad5b04230db4a2c337015703f775 Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Tue, 26 May 2026 19:15:19 +0000 Subject: [PATCH 3/3] Switch to parameterized tests Signed-off-by: Fynn Schmitt-Ulms --- tests/integration/conftest.py | 10 +- .../integration/models/test_dflash_forward.py | 168 ------------ .../integration/models/test_eagle3_forward.py | 172 ------------ .../integration/models/test_model_forward.py | 255 ++++++++++++++++++ .../integration/models/test_peagle_forward.py | 164 ----------- 5 files changed, 260 insertions(+), 509 deletions(-) delete mode 100644 tests/integration/models/test_dflash_forward.py delete mode 100644 tests/integration/models/test_eagle3_forward.py create mode 100644 tests/integration/models/test_model_forward.py delete mode 100644 tests/integration/models/test_peagle_forward.py diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 88a8584dd..c08d93553 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -69,7 +69,7 @@ def make_eagle3_model( *, draft_vocab_size: int = 64, norm_before_residual: bool = False, - device: str = "cuda", + device: str = "cuda:0", ) -> Eagle3DraftModel: """Create a tiny Eagle3 model with real initialized weights.""" config = Eagle3SpeculatorConfig( @@ -97,7 +97,7 @@ def make_dflash_model( draft_vocab_size: int = 64, block_size: int = 4, max_anchors: int = 8, - device: str = "cuda", + device: str = "cuda:0", ) -> DFlashDraftModel: """Create a tiny DFlash model with real initialized weights.""" config = DFlashSpeculatorConfig( @@ -105,7 +105,7 @@ def make_dflash_model( draft_vocab_size=draft_vocab_size, block_size=block_size, max_anchors=max_anchors, - aux_hidden_state_layer_ids=[0, 1], + aux_hidden_state_layer_ids=[0, 1, 2], mask_token_id=0, speculators_config=SpeculatorsConfig( algorithm="dflash", @@ -129,7 +129,7 @@ def make_peagle_model( draft_vocab_size: int = 64, num_depths: int = 4, down_sample_ratio: float = 0.7, - device: str = "cuda", + device: str = "cuda:0", ) -> PEagleDraftModel: """Create a tiny PEagle model with real initialized weights.""" config = PEagleSpeculatorConfig( @@ -223,7 +223,7 @@ def make_batch( max_len: int, samples: list[dict[str, torch.Tensor]], hidden_size: int, - device: str = "cuda", + device: str = "cuda:0", ) -> dict[str, torch.Tensor]: """Collate a list of samples into a single batch using the real collate_fn. diff --git a/tests/integration/models/test_dflash_forward.py b/tests/integration/models/test_dflash_forward.py deleted file mode 100644 index a315b1a48..000000000 --- a/tests/integration/models/test_dflash_forward.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Integration tests for DFlashDraftModel forward passes with real weights.""" - -import pytest -import torch - -from tests.conftest import requires_cuda -from tests.integration.conftest import ( - HIDDEN_SIZE, - VOCAB_SIZE, - make_batch, - make_dflash_model, - make_sample, -) - -MAX_LEN = 128 -NUM_TARGET_LAYERS = 2 - -LOSS_MASK_CASES = ["all", "none", "random", "alternating"] - - -def _make_samples( - seq_lengths: list[int], - loss_mask_pattern: str = "all", - vocab_size: int = VOCAB_SIZE, - boundary_token_ids: list[int] | None = None, -) -> list[dict[str, torch.Tensor]]: - return [ - make_sample( - seq_len=sl, - hidden_size=HIDDEN_SIZE, - hidden_multiplier=NUM_TARGET_LAYERS, - vocab_size=vocab_size, - loss_mask_pattern=loss_mask_pattern, - include_verifier_states=True, - boundary_token_ids=boundary_token_ids, - ) - for sl in seq_lengths - ] - - -SAMPLE_CONFIGS = [ - pytest.param([128], id="single_sample"), - pytest.param([64, 64], id="two_equal"), - pytest.param([32, 96], id="two_unequal"), - pytest.param([16] * 8, id="eight_tiny"), - pytest.param([8] * 20, id="twenty_tiny"), -] - - -@requires_cuda -class TestDFlashTraining: - """Forward pass in training mode (DFlash always requires verifier states).""" - - @pytest.mark.parametrize("seq_lengths", SAMPLE_CONFIGS) - @pytest.mark.parametrize("loss_mask_pattern", LOSS_MASK_CASES) - def test_forward_backward_produces_valid_outputs( - self, dflash_model, seq_lengths, loss_mask_pattern - ): - samples = _make_samples(seq_lengths, loss_mask_pattern=loss_mask_pattern) - batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) - draft_tokens, loss, metrics = dflash_model(**batch) - - assert loss.isfinite(), f"Loss is not finite: {loss.item()}" - assert "loss_sum" in metrics - assert "loss_total" in metrics - - loss.backward() - - @pytest.mark.parametrize("block_size", [2, 4, 8]) - def test_varying_block_size(self, block_size): - model = make_dflash_model(block_size=block_size, max_anchors=4) - 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() - - @pytest.mark.parametrize("max_anchors", [2, 8, 16]) - def test_varying_max_anchors(self, max_anchors): - model = make_dflash_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() - - -@requires_cuda -class TestDFlashMultiBatch: - """Run multiple batches back-to-back through the same model to test - statefulness, cache clearing, and varying batch compositions.""" - - def test_multi_then_single(self, dflash_model): - """Multi-sample batches, then single-sample.""" - batch_configs: list[list[int]] = [ - [32, 32, 32], - [64, 64], - [128], - [16], - # Empty batch ([]) excluded: create_empty_sample dtype bug (#527) - ] - for seq_lengths in batch_configs: - samples = _make_samples(seq_lengths) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = dflash_model(**batch) - assert loss.isfinite(), f"Loss not finite for seq_lengths={seq_lengths}" - loss.backward() - - def test_alternating_batch_sizes(self, dflash_model): - """Alternate between large multi-sample and single-sample batches.""" - batch_configs: list[list[int]] = [ - [16] * 8, - [128], - [32, 32, 32, 32], - [64], - [8] * 16, - [128], - ] - for seq_lengths in batch_configs: - samples = _make_samples(seq_lengths) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = dflash_model(**batch) - assert loss.isfinite() - loss.backward() - - def test_varying_loss_masks_across_batches(self, dflash_model): - """Each batch uses a different loss mask pattern.""" - for pattern in LOSS_MASK_CASES: - samples = _make_samples([64, 64], loss_mask_pattern=pattern) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = dflash_model(**batch) - assert loss.isfinite(), f"Loss not finite for loss_mask_pattern={pattern}" - loss.backward() - - -@requires_cuda -class TestDFlashVocabBoundary: - """Tests with draft vocab mapping.""" - - @pytest.fixture - def dflash_draft_vocab_model(self): - model = make_dflash_model(draft_vocab_size=32) - t2d = torch.zeros(VOCAB_SIZE, dtype=torch.bool) - t2d[:32] = True - d2t = torch.arange(32, dtype=torch.long) - model.load_vocab_mappings(t2d.to("cuda"), d2t.to("cuda")) - yield model - del model - torch.cuda.empty_cache() - - def test_boundary_tokens(self, dflash_draft_vocab_model): - samples = _make_samples([128], vocab_size=32, boundary_token_ids=[0, 31]) - batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) - draft_tokens, loss, metrics = dflash_draft_vocab_model(**batch) - - assert loss.isfinite() - assert draft_tokens.dtype == torch.long - assert (draft_tokens >= 0).all() - assert (draft_tokens < 32).all() - loss.backward() diff --git a/tests/integration/models/test_eagle3_forward.py b/tests/integration/models/test_eagle3_forward.py deleted file mode 100644 index 202f8b5bf..000000000 --- a/tests/integration/models/test_eagle3_forward.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Integration tests for Eagle3DraftModel forward passes with real weights.""" - -import pytest -import torch - -from tests.conftest import requires_cuda -from tests.integration.conftest import ( - HIDDEN_SIZE, - VOCAB_SIZE, - make_batch, - make_eagle3_model, - make_sample, -) - -MAX_LEN = 128 - -LOSS_MASK_CASES = ["all", "none", "random", "alternating"] - - -def _make_samples( - seq_lengths: list[int], - loss_mask_pattern: str = "all", - include_verifier_states: bool = True, - vocab_size: int = VOCAB_SIZE, - boundary_token_ids: list[int] | None = None, -) -> list[dict[str, torch.Tensor]]: - return [ - make_sample( - seq_len=sl, - hidden_size=HIDDEN_SIZE, - hidden_multiplier=3, - vocab_size=vocab_size, - loss_mask_pattern=loss_mask_pattern, - include_verifier_states=include_verifier_states, - boundary_token_ids=boundary_token_ids, - ) - for sl in seq_lengths - ] - - -# --------------------------------------------------------------------------- -# Sample compositions for parametrize -# --------------------------------------------------------------------------- - -SAMPLE_CONFIGS = [ - pytest.param([128], id="single_sample"), - pytest.param([64, 64], id="two_equal"), - pytest.param([32, 96], id="two_unequal"), - pytest.param([16] * 8, id="eight_tiny"), - pytest.param([8] * 20, id="twenty_tiny"), -] - - -@requires_cuda -class TestEagle3Training: - """Forward pass in training mode (with verifier_last_hidden_states).""" - - @pytest.mark.parametrize("seq_lengths", SAMPLE_CONFIGS) - @pytest.mark.parametrize("loss_mask_pattern", LOSS_MASK_CASES) - def test_forward_backward_produces_valid_outputs( - self, eagle3_model, seq_lengths, loss_mask_pattern - ): - samples = _make_samples(seq_lengths, loss_mask_pattern=loss_mask_pattern) - batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) - ttt_steps = 2 - draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=ttt_steps) - - assert len(draft_tokens) == ttt_steps - for dt in draft_tokens: - assert dt.shape == (1, MAX_LEN) - assert dt.dtype == torch.long - - assert loss.isfinite(), f"Loss is not finite: {loss.item()}" - assert "loss_sum" in metrics - assert "loss_total" in metrics - - loss.backward() - - @pytest.mark.parametrize("ttt_steps", [1, 3, 5]) - def test_varying_ttt_steps(self, eagle3_model, ttt_steps): - samples = _make_samples([128]) - batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) - draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=ttt_steps) - - assert len(draft_tokens) == ttt_steps - assert loss.isfinite() - - loss.backward() - - -@requires_cuda -class TestEagle3MultiBatch: - """Run multiple batches back-to-back through the same model to test - statefulness, cache clearing, and varying batch compositions.""" - - def test_multi_then_single_then_empty(self, eagle3_model): - """Multi-sample batches, then single-sample, then empty batch.""" - batch_configs: list[list[int]] = [ - [16, 16, 8, 10, 15, 12], - [32, 32, 32], - [64, 64], - [32, 3, 17], - [128], - [16], - # Empty batch ([]) excluded: create_empty_sample dtype bug (#527) - ] - torch.compiler.reset() - for seq_lengths in batch_configs: - samples = _make_samples(seq_lengths) if seq_lengths else [] - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=2) - assert loss.isfinite(), f"Loss not finite for seq_lengths={seq_lengths}" - loss.backward() - - def test_alternating_batch_sizes(self, eagle3_model): - """Alternate between large multi-sample and single-sample batches.""" - batch_configs: list[list[int]] = [ - [16] * 8, - [128], - [32, 32, 32, 32], - [64], - [8] * 16, - [128], - ] - for seq_lengths in batch_configs: - samples = _make_samples(seq_lengths) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=2) - assert loss.isfinite() - loss.backward() - - def test_varying_loss_masks_across_batches(self, eagle3_model): - """Each batch uses a different loss mask pattern.""" - for pattern in LOSS_MASK_CASES: - samples = _make_samples([64, 64], loss_mask_pattern=pattern) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = eagle3_model(**batch, ttt_steps=2) - assert loss.isfinite(), f"Loss not finite for loss_mask_pattern={pattern}" - loss.backward() - - -@requires_cuda -class TestEagle3VocabBoundary: - """Tests with draft vocab mapping and boundary token IDs.""" - - @pytest.fixture - def eagle3_draft_vocab_model(self): - model = make_eagle3_model(draft_vocab_size=32) - t2d = torch.zeros(VOCAB_SIZE, dtype=torch.bool) - t2d[:32] = True - d2t = torch.arange(32, dtype=torch.long) - model.load_vocab_mappings(t2d.to("cuda"), d2t.to("cuda")) - yield model - del model - torch.cuda.empty_cache() - - def test_boundary_tokens_training(self, eagle3_draft_vocab_model): - samples = _make_samples([128], vocab_size=32, boundary_token_ids=[0, 31]) - batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) - draft_tokens, loss, metrics = eagle3_draft_vocab_model(**batch, ttt_steps=2) - - assert len(draft_tokens) == 2 - assert loss.isfinite() - for dt in draft_tokens: - assert (dt >= 0).all() - assert (dt < 32).all() diff --git a/tests/integration/models/test_model_forward.py b/tests/integration/models/test_model_forward.py new file mode 100644 index 000000000..e21f994b9 --- /dev/null +++ b/tests/integration/models/test_model_forward.py @@ -0,0 +1,255 @@ +"""Integration tests for draft model forward passes with real weights. + +Covers DFlash, Eagle3, and PEagle models with shared parametrized tests +for training, multi-batch, and vocab boundary scenarios, plus model-specific +parameter variation tests. +""" + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import pytest +import torch + +from tests.conftest import requires_cuda +from tests.integration.conftest import ( + HIDDEN_SIZE, + VOCAB_SIZE, + make_batch, + make_dflash_model, + make_eagle3_model, + make_peagle_model, + make_sample, +) + +MAX_LEN = 128 +HIDDEN_MULTIPLIER = 3 +LOSS_MASK_CASES = ["all", "none", "random", "alternating"] + +SAMPLE_CONFIGS = [ + pytest.param([128], id="single_sample"), + pytest.param([64, 64], id="two_equal"), + pytest.param([32, 96], id="two_unequal"), + pytest.param([8] * 20, id="twenty_tiny"), +] + +MULTI_BATCH_CONFIGS: list[list[int]] = [ + [16, 16, 8, 10, 15, 12], + [32, 32, 32], + [64, 64], + [32, 3, 17], + [128], + [16], +] + +# --------------------------------------------------------------------------- +# Model specs +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ModelSpec: + name: str + factory: Callable[..., Any] + forward_kwargs: dict[str, Any] = field(default_factory=dict) + + +DFLASH_SPEC = ModelSpec(name="dflash", factory=make_dflash_model) +EAGLE3_SPEC = ModelSpec( + name="eagle3", factory=make_eagle3_model, forward_kwargs={"ttt_steps": 2} +) +PEAGLE_SPEC = ModelSpec(name="peagle", factory=make_peagle_model) + +ALL_SPECS = [ + pytest.param(DFLASH_SPEC, id="dflash"), + pytest.param(EAGLE3_SPEC, id="eagle3"), + pytest.param(PEAGLE_SPEC, id="peagle"), +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_samples( + seq_lengths: list[int], + loss_mask_pattern: str = "all", + vocab_size: int = VOCAB_SIZE, + boundary_token_ids: list[int] | None = None, +) -> list[dict[str, torch.Tensor]]: + return [ + make_sample( + seq_len=sl, + hidden_size=HIDDEN_SIZE, + hidden_multiplier=HIDDEN_MULTIPLIER, + vocab_size=vocab_size, + loss_mask_pattern=loss_mask_pattern, + include_verifier_states=True, + boundary_token_ids=boundary_token_ids, + ) + for sl in seq_lengths + ] + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(params=ALL_SPECS) +def model_and_spec(request): + spec: ModelSpec = request.param + model = spec.factory() + yield model, spec + del model + torch.cuda.empty_cache() + + +@pytest.fixture +def draft_vocab_model(request): + spec: ModelSpec = request.param + model = spec.factory(draft_vocab_size=32) + t2d = torch.zeros(VOCAB_SIZE, dtype=torch.bool) + t2d[:32] = True + d2t = torch.arange(32, dtype=torch.long) + model.load_vocab_mappings(t2d.to("cuda"), d2t.to("cuda")) + yield model, spec + del model + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# Shared tests +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestTraining: + """Forward + backward pass across all models.""" + + @pytest.mark.parametrize("seq_lengths", SAMPLE_CONFIGS) + @pytest.mark.parametrize("loss_mask_pattern", LOSS_MASK_CASES) + def test_forward_backward(self, model_and_spec, seq_lengths, loss_mask_pattern): + model, spec = model_and_spec + samples = _make_samples(seq_lengths, loss_mask_pattern=loss_mask_pattern) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = model(**batch, **spec.forward_kwargs) + + assert loss.isfinite(), f"Loss is not finite: {loss.item()}" + assert "loss_sum" in metrics + assert "loss_total" in metrics + loss.backward() + + +@requires_cuda +class TestMultiBatch: + """Run multiple batches back-to-back to test statefulness and cache clearing.""" + + def test_multi_then_single(self, model_and_spec): + model, spec = model_and_spec + torch.compiler.reset() + for seq_lengths in MULTI_BATCH_CONFIGS: + samples = _make_samples(seq_lengths) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = model(**batch, **spec.forward_kwargs) + assert loss.isfinite(), f"Loss not finite for seq_lengths={seq_lengths}" + loss.backward() + + def test_varying_loss_masks_across_batches(self, model_and_spec): + model, spec = model_and_spec + torch.compiler.reset() + for pattern in LOSS_MASK_CASES: + samples = _make_samples([64, 64], loss_mask_pattern=pattern) + batch = make_batch( + max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE + ) + draft_tokens, loss, metrics = model(**batch, **spec.forward_kwargs) + assert loss.isfinite(), f"Loss not finite for loss_mask_pattern={pattern}" + loss.backward() + + +@requires_cuda +class TestVocabBoundary: + """Tests with draft vocab mapping.""" + + @pytest.mark.parametrize("draft_vocab_model", ALL_SPECS, indirect=True) + def test_boundary_tokens(self, draft_vocab_model): + model, spec = draft_vocab_model + samples = _make_samples([128], vocab_size=32, boundary_token_ids=[0, 31]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = model(**batch, **spec.forward_kwargs) + + assert loss.isfinite() + loss.backward() + + +# --------------------------------------------------------------------------- +# Model-specific parameter tests +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestDFlashParams: + @pytest.mark.parametrize("block_size", [2, 4, 8]) + def test_varying_block_size(self, block_size): + model = make_dflash_model(block_size=block_size, max_anchors=4) + 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() + + @pytest.mark.parametrize("max_anchors", [2, 8, 16]) + def test_varying_max_anchors(self, max_anchors): + model = make_dflash_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() + + +@requires_cuda +class TestEagle3Params: + @pytest.mark.parametrize("ttt_steps", [1, 3, 5]) + def test_varying_ttt_steps(self, ttt_steps): + model = make_eagle3_model() + samples = _make_samples([128]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, metrics = model(**batch, ttt_steps=ttt_steps) + + assert len(draft_tokens) == ttt_steps + for dt in draft_tokens: + assert dt.shape == (1, MAX_LEN) + assert dt.dtype == torch.long + assert loss.isfinite() + loss.backward() + + +@requires_cuda +class TestPEagleParams: + @pytest.mark.parametrize("num_depths", [2, 4, 8]) + def test_varying_num_depths(self, num_depths): + model = make_peagle_model(num_depths=num_depths) + 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() + + @pytest.mark.parametrize("down_sample_ratio", [0.3, 0.7, 1.0]) + def test_varying_down_sample_ratio(self, down_sample_ratio): + model = make_peagle_model(down_sample_ratio=down_sample_ratio) + 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() diff --git a/tests/integration/models/test_peagle_forward.py b/tests/integration/models/test_peagle_forward.py deleted file mode 100644 index 7094a7b57..000000000 --- a/tests/integration/models/test_peagle_forward.py +++ /dev/null @@ -1,164 +0,0 @@ -"""Integration tests for PEagleDraftModel forward passes with real weights.""" - -import pytest -import torch - -from tests.conftest import requires_cuda -from tests.integration.conftest import ( - HIDDEN_SIZE, - VOCAB_SIZE, - make_batch, - make_peagle_model, - make_sample, -) - -MAX_LEN = 128 - -LOSS_MASK_CASES = ["all", "none", "random", "alternating"] - - -def _make_samples( - seq_lengths: list[int], - loss_mask_pattern: str = "all", - vocab_size: int = VOCAB_SIZE, - boundary_token_ids: list[int] | None = None, -) -> list[dict[str, torch.Tensor]]: - return [ - make_sample( - seq_len=sl, - hidden_size=HIDDEN_SIZE, - hidden_multiplier=3, - vocab_size=vocab_size, - loss_mask_pattern=loss_mask_pattern, - include_verifier_states=True, - boundary_token_ids=boundary_token_ids, - ) - for sl in seq_lengths - ] - - -SAMPLE_CONFIGS = [ - pytest.param([128], id="single_sample"), - pytest.param([64, 64], id="two_equal"), - pytest.param([32, 96], id="two_unequal"), - pytest.param([16] * 8, id="eight_tiny"), - pytest.param([8] * 20, id="twenty_tiny"), -] - - -@requires_cuda -class TestPEagleTraining: - """Forward pass in training mode (PEagle always requires verifier states).""" - - @pytest.mark.parametrize("seq_lengths", SAMPLE_CONFIGS) - @pytest.mark.parametrize("loss_mask_pattern", LOSS_MASK_CASES) - def test_forward_backward_produces_valid_outputs( - self, peagle_model, seq_lengths, loss_mask_pattern - ): - samples = _make_samples(seq_lengths, loss_mask_pattern=loss_mask_pattern) - batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) - draft_tokens, loss, metrics = peagle_model(**batch) - - assert loss.isfinite(), f"Loss is not finite: {loss.item()}" - assert "loss_sum" in metrics - assert "loss_total" in metrics - - loss.backward() - - @pytest.mark.parametrize("num_depths", [2, 4, 8]) - def test_varying_num_depths(self, num_depths): - model = make_peagle_model(num_depths=num_depths) - 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() - - @pytest.mark.parametrize("down_sample_ratio", [0.3, 0.7, 1.0]) - def test_varying_down_sample_ratio(self, down_sample_ratio): - model = make_peagle_model(down_sample_ratio=down_sample_ratio) - 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 -class TestPEagleMultiBatch: - """Run multiple batches back-to-back through the same model to test - statefulness, cache clearing, and varying batch compositions.""" - - def test_multi_then_single(self, peagle_model): - """Multi-sample batches, then single-sample.""" - batch_configs: list[list[int]] = [ - [32, 32, 32], - [64, 64], - [128], - [16], - # Empty batch ([]) excluded: create_empty_sample dtype bug (#527) - ] - for seq_lengths in batch_configs: - samples = _make_samples(seq_lengths) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = peagle_model(**batch) - assert loss.isfinite(), f"Loss not finite for seq_lengths={seq_lengths}" - loss.backward() - - def test_alternating_batch_sizes(self, peagle_model): - """Alternate between large multi-sample and single-sample batches.""" - batch_configs: list[list[int]] = [ - [16] * 8, - [128], - [32, 32, 32, 32], - [64], - [8] * 16, - [128], - ] - for seq_lengths in batch_configs: - samples = _make_samples(seq_lengths) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = peagle_model(**batch) - assert loss.isfinite() - loss.backward() - - def test_varying_loss_masks_across_batches(self, peagle_model): - """Each batch uses a different loss mask pattern.""" - for pattern in LOSS_MASK_CASES: - samples = _make_samples([64, 64], loss_mask_pattern=pattern) - batch = make_batch( - max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE - ) - draft_tokens, loss, metrics = peagle_model(**batch) - assert loss.isfinite(), f"Loss not finite for loss_mask_pattern={pattern}" - loss.backward() - - -@requires_cuda -class TestPEagleVocabBoundary: - """Tests with draft vocab mapping and boundary token IDs.""" - - @pytest.fixture - def peagle_draft_vocab_model(self): - model = make_peagle_model(draft_vocab_size=32) - t2d = torch.zeros(VOCAB_SIZE, dtype=torch.bool) - t2d[:32] = True - d2t = torch.arange(32, dtype=torch.long) - model.load_vocab_mappings(t2d.to("cuda"), d2t.to("cuda")) - yield model - del model - torch.cuda.empty_cache() - - def test_boundary_tokens(self, peagle_draft_vocab_model): - samples = _make_samples([128], vocab_size=32, boundary_token_ids=[0, 31]) - batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) - draft_tokens, loss, metrics = peagle_draft_vocab_model(**batch) - - assert loss.isfinite() - loss.backward()