Automatic vocabulary size padding (Lightning-AI#223)

Author: carmocca

lit_llama/model.py

@@ -5,20 +5,29 @@
 # mypy: ignore-errors
 import math
 from dataclasses import dataclass
+from typing import Optional
 
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
 from typing_extensions import Self
 
+from lit_llama.utils import find_multiple
+
+
 @dataclass
 class LLaMAConfig:
     block_size: int = 2048
     vocab_size: int = 32000
+    padded_vocab_size: Optional[int] = None
     n_layer: int = 32
     n_head: int = 32
     n_embd: int = 4096
 
+    def __post_init__(self):
+        if self.padded_vocab_size is None:
+            self.padded_vocab_size = find_multiple(self.vocab_size, 64)
+
     @classmethod
     def from_name(cls, name: str) -> Self:
         return cls(**llama_configs[name])
@@ -35,14 +44,13 @@ def from_name(cls, name: str) -> Self:
 class LLaMA(nn.Module):
     def __init__(self, config: LLaMAConfig) -> None:
         super().__init__()
-        assert config.vocab_size is not None
-        assert config.block_size is not None
+        assert config.padded_vocab_size is not None
         self.config = config
 
-        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
         self.transformer = nn.ModuleDict(
             dict(
-                wte=nn.Embedding(config.vocab_size, config.n_embd),
+                wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
                 h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
                 ln_f=RMSNorm(config.n_embd),
             )
@@ -103,7 +111,7 @@ def __init__(self, config: LLaMAConfig) -> None:
         self.n_head = config.n_head
         self.n_embd = config.n_embd
         self.block_size = config.block_size
-        self.rope_cache = None
+        self.rope_cache: Optional[torch.Tensor] = None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
@@ -150,9 +158,7 @@ def __init__(self, config: LLaMAConfig) -> None:
         super().__init__()
         hidden_dim = 4 * config.n_embd
         n_hidden = int(2 * hidden_dim / 3)
-        N = 256
-        # ensure n_hidden is multiple of N
-        n_hidden = ((n_hidden - 1) // N) * N + N
+        n_hidden = find_multiple(n_hidden, 256)
 
         self.c_fc1 = nn.Linear(config.n_embd, n_hidden, bias=False)
         self.c_fc2 = nn.Linear(config.n_embd, n_hidden, bias=False)
@@ -210,6 +216,7 @@ def build_rope_cache(seq_len: int, n_elem: int, dtype: torch.dtype, device: torc
         cache = cache.half()
     return cache
 
+
 def apply_rope(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
     x = x.transpose(1, 2)
 

lit_llama/utils.py

@@ -31,6 +31,12 @@ def llama_model_lookup(checkpoint: dict) -> str:
     return llama_model_sizes[embedding_size]
 
 
+def find_multiple(n: int, k: int) -> int:
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
 def save_model_checkpoint(fabric, model, file_path):
     """Handles boilerplate logic for retrieving and saving the state_dict.
     

tests/test_utils.py

@@ -45,3 +45,12 @@ def test_lazy_load_subclass(lit_llama):
                 actual = sd_lazy[k]
                 expected = sd[k]
                 torch.testing.assert_close(actual._load_tensor(), expected)
+
+
+def test_find_multiple(lit_llama):
+    from lit_llama.utils import find_multiple
+
+    assert find_multiple(17, 5) == 20
+    assert find_multiple(30, 7) == 35
+    assert find_multiple(10, 2) == 10
+    assert find_multiple(5, 10) == 10