diff --git a/kda/training/data.py b/kda/training/data.py
index e23e3dc..8aa0da1 100644
--- a/kda/training/data.py
+++ b/kda/training/data.py
@@ -57,7 +57,11 @@ class HuggingFaceTokenizer:
 
     @property
     def vocab_size(self) -> int:
-        return int(len(self._tok))
+        # len(tok) 数的是去重后的 surface form; 词表有重复 piece 时 (如 Yi-6B
+        # 63992 vs 最大 id 63999) 会小于真实 id 范围, embedding 越界触发
+        # device-side assert. 以最大 id + 1 为准.
+        max_id = max(self._tok.get_vocab().values())
+        return max(int(len(self._tok)), max_id + 1)
 
     def encode(self, text: str) -> list[int]:
         return list(self._tok.encode(text, add_special_tokens=False))