From 2efb3638b2f4f3ed7553c5ee373fca70546d0a90 Mon Sep 17 00:00:00 2001
From: wiketool <64351271+wiketool@users.noreply.github.com>
Date: Thu, 25 Jun 2026 19:39:42 +0800
Subject: [PATCH] Include vocab sizes in EAGLE3 vocab mapping cache key

---
 scripts/train_eagle3.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/scripts/train_eagle3.py b/scripts/train_eagle3.py
index f25b6cb18..a4032095f 100644
--- a/scripts/train_eagle3.py
+++ b/scripts/train_eagle3.py
@@ -560,13 +560,19 @@ def build_dataloaders(
     )
 
     # convert to dataloader
-    cache_params_string = (
+    dataset_cache_params_string = (
         f"{args.train_data_path}-"
         f"{args.max_length}-"
         f"{args.chat_template}-"
         f"{args.target_model_path}"  # Tokenizer may also different
     )
-    cache_key = hashlib.md5(cache_params_string.encode()).hexdigest()
+    vocab_cache_params_string = (
+        f"{dataset_cache_params_string}-"
+        f"{draft_model_config.draft_vocab_size}-"
+        f"{draft_model_config.vocab_size}"
+    )
+    cache_key = hashlib.md5(dataset_cache_params_string.encode()).hexdigest()
+    vocab_cache_key = hashlib.md5(vocab_cache_params_string.encode()).hexdigest()
     train_dataset = Dataset.from_generator(
         generator=safe_conversations_generator,
         gen_kwargs={"file_path": args.train_data_path},
@@ -593,7 +599,7 @@ def build_dataloaders(
             target_vocab_size=draft_model_config.vocab_size,
             draft_vocab_size=draft_model_config.draft_vocab_size,
             cache_dir=os.path.join(args.cache_dir, "vocab_mapping"),
-            cache_key=cache_key,
+            cache_key=vocab_cache_key,
         )
 
         if not is_online: