Add bnb.nn.StableEmbedding for quantized training

Lightning-AI · Oct 4, 2024 · 461c33d · 461c33d
1 parent c03f3f0
commit 461c33d
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/litgpt/finetune/adapter.py b/litgpt/finetune/adapter.py
@@ -175,6 +175,12 @@ def main(
 
     if isinstance(fabric.strategy.precision, BitsandbytesPrecision):
         optimizer = instantiate_bnb_optimizer(optimizer, model.parameters())
+
+        from bitsandbytes.nn import StableEmbedding
+        old_embedding = model.transformer.wte
+        model.transformer.wte = StableEmbedding(old_embedding.num_embeddings, old_embedding.embedding_dim)
+        with torch.no_grad():
+            model.wte.weight.copy_(old_embedding.weight)
     else:
         optimizer = instantiate_torch_optimizer(optimizer, model.parameters())
 

diff --git a/litgpt/finetune/adapter_v2.py b/litgpt/finetune/adapter_v2.py
@@ -175,6 +175,12 @@ def main(
 
     if isinstance(fabric.strategy.precision, BitsandbytesPrecision):
         optimizer = instantiate_bnb_optimizer(optimizer, model.parameters())
+
+        from bitsandbytes.nn import StableEmbedding
+        old_embedding = model.transformer.wte
+        model.transformer.wte = StableEmbedding(old_embedding.num_embeddings, old_embedding.embedding_dim)
+        with torch.no_grad():
+            model.wte.weight.copy_(old_embedding.weight)
     else:
         optimizer = instantiate_torch_optimizer(optimizer, model.parameters())
 

diff --git a/litgpt/finetune/lora.py b/litgpt/finetune/lora.py
@@ -205,6 +205,13 @@ def main(
 
     if isinstance(fabric.strategy.precision, BitsandbytesPrecision):
         optimizer = instantiate_bnb_optimizer(optimizer, model.parameters())
+
+        from bitsandbytes.nn import StableEmbedding
+        old_embedding = model.transformer.wte
+        model.transformer.wte = StableEmbedding(old_embedding.num_embeddings, old_embedding.embedding_dim)
+        with torch.no_grad():
+            model.wte.weight.copy_(old_embedding.weight)
+
     else:
         optimizer = instantiate_torch_optimizer(optimizer, model.parameters())