Merge branch 'main' into mm-dataset

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Author: DarkLight1337

.github/dependabot.yml

@@ -13,7 +13,7 @@ updates:
     directory: "/"
     target-branch: "main"
     schedule:
-      interval: "weekly"
+      interval: "quarterly"
     open-pull-requests-limit: 10
     versioning-strategy: "increase"
     allow:

pyproject.toml

@@ -69,7 +69,7 @@ dev = [
     "mkdocs~=1.6.1",
     "mkdocs-material[imaging]~=9.7.6",
     "markdown~=3.10.2",
-    "pymdown-extensions~=10.21.2",
+    "pymdown-extensions~=10.21.3",
     "mkdocs-section-index~=0.3.12",
     "mkdocs-minify-plugin~=0.8.0",
     "mkdocs-api-autonav~=0.4.0",
@@ -92,7 +92,7 @@ dev = [
 
     # code quality
     "mypy~=2.1.0",
-    "ruff~=0.14.1",
+    "ruff~=0.15.13",
 
     # docs quality
     "mdformat~=1.0.0",

src/speculators/train/utils.py

@@ -156,6 +156,6 @@ def apply_fully_sharded(model: torch.nn.Module):
     for layer in model.layers:  # type: ignore[union-attr]
         fully_shard(layer, mp_policy=mp_policy)
 
-    fully_shard(model, mp_policy=mp_policy)
+    fully_shard(model)
 
     return model

tests/e2e/regression/test_training_only_acceptance.py

@@ -20,11 +20,20 @@ def _resolve_repo(repo_id: str, repo_type: str = "dataset") -> Path:
     try:
         return Path(
             snapshot_download(
-                repo_id=repo_id, repo_type=repo_type, local_files_only=True
+                repo_id=repo_id,
+                repo_type=repo_type,
+                local_files_only=True,
+                allow_patterns=["*.arrow", "*.json", "*.pt", "hidden_states/*"],
             )
         )
     except LocalEntryNotFoundError:
-        return Path(snapshot_download(repo_id=repo_id, repo_type=repo_type))
+        return Path(
+            snapshot_download(
+                repo_id=repo_id,
+                repo_type=repo_type,
+                allow_patterns=["*.arrow", "*.json", "*.pt", "hidden_states/*"],
+            )
+        )
 
 
 @requires_cadence("weekly")
@@ -41,7 +50,7 @@ def test_eagle3_qwen3_8b_sharegpt(tmp_path: Path, prompts: list[list[dict[str, s
         seq_length=8192,
         epochs=epochs,
         lr=3e-4,
-        draft_vocab_size=8192,
+        draft_vocab_size=32000,
         online=False,
         log_freq=50,
         timeout=30 * 60,  # 30 mins
@@ -71,7 +80,7 @@ def test_dflash_qwen3_8b_sharegpt(tmp_path: Path, prompts: list[list[dict[str, s
         seq_length=8192,
         epochs=epochs,
         lr=3e-4,
-        draft_vocab_size=8192,
+        draft_vocab_size=32000,
         num_layers=3,
         online=False,
         log_freq=50,
@@ -86,3 +95,39 @@ def test_dflash_qwen3_8b_sharegpt(tmp_path: Path, prompts: list[list[dict[str, s
         prompts=prompts,
         acceptance_thresholds=[0.30, 0.05, 0.001, 0, 0, 0, 0],
     )
+
+
+@requires_cadence("weekly")
+@pytest.mark.regression
+def test_peagle_qwen3_8b_sharegpt(tmp_path: Path, prompts: list[list[dict[str, str]]]):
+    save_path = tmp_path / "checkpoints"
+    hidden_states_dir = _resolve_repo("inference-optimization/Qwen3-8b-sharegpt-5k")
+    epochs = 5
+    run_training(
+        model="Qwen/Qwen3-8B",
+        speculator_type="peagle",
+        data_path=hidden_states_dir,
+        save_path=save_path,
+        seq_length=8192,
+        epochs=epochs,
+        lr=6e-4,
+        draft_vocab_size=32000,
+        num_layers=4,
+        online=False,
+        log_freq=50,
+        timeout=30 * 60,  # 30 mins
+        extra_train_args=[
+            "--no-norm-before-residual",
+            "--scheduler-type",
+            "cosine",
+        ],
+    )
+    final_checkpoint = str(save_path / str(epochs - 1))
+    run_vllm_engine(
+        model_path=final_checkpoint,
+        tmp_path=tmp_path,
+        max_tokens=512,
+        ignore_eos=True,
+        prompts=prompts,
+        acceptance_thresholds=[0.45, 0.1, 0.01],
+    )

tests/e2e/smoke/test_offline_training.py

@@ -33,14 +33,29 @@
     ("model", "dataset", "speculator_type", "extra_train_args", "target_layer_ids"),
     [
         (TEXT_MODEL, "sharegpt", "eagle3", [], None),  # Use default EAGLE layers
+        (MM_MODEL, "sharegpt4v_coco", "eagle3", [], None),  # Multimodal
         (
             TEXT_MODEL,
             "sharegpt",
             "dflash",
             ["--block-size", "8", "--max-anchors", "256", "--num-layers", "3"],
             [1, 13, 25],
         ),  # DFlash with 3 layers + verifier last layer
-        (MM_MODEL, "sharegpt4v_coco", "eagle3", [], None),  # Multimodal
+        (
+            "peagle",
+            [
+                "--num-layers",
+                "4",
+                "--num-depths",
+                "4",
+                "--down-sample-ratio",
+                "0.7",
+                "--down-sample-ratio-min",
+                "0.2",
+                "--no-norm-before-residual",
+            ],
+            None,
+        ),  # P-EAGLE with parallel multi-token prediction
     ],
 )
 def test_offline_smoke(

tests/unit/convert/test_eagle3_converter.py

@@ -260,14 +260,15 @@ def test_nm_testing_2layer_eagle3_model_config(self, tmp_path):
         checkpoint_path = "nm-testing/testing-llama3.1.8b-2layer-eagle3"
         base_model = "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic"
 
+        local_checkpoint_path = tmp_path / checkpoint_path.rsplit("/", maxsplit=1)[-1]
         converter.convert(
             checkpoint_path,
-            tmp_path / checkpoint_path.split("/")[-1],
+            local_checkpoint_path,
             base_model,
             norm_before_residual=False,
         )
 
-        config = load_checkpoint_config(tmp_path / checkpoint_path.split("/")[-1])
+        config = load_checkpoint_config(local_checkpoint_path)
 
         # Verify that num_hidden_layers is correctly set to 2
         assert config["transformer_layer_config"]["num_hidden_layers"] == 2