fix(server): fix bnb quantization for CausalLM models (#385)

huggingface · May 31, 2023 · 337afb2 · 337afb2
1 parent 87dc034
commit 337afb2
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 3 deletions.
diff --git a/server/text_generation_server/models/bloom.py b/server/text_generation_server/models/bloom.py
@@ -245,6 +245,8 @@ def linear(input, weight, bias):
                                 return linear
 
                             module.linear = replace_linear(state)
+                        else:
+                            tensor = tensor.to(device)
                     elif quantize == "gptq":
                         raise NotImplementedError("`gptq` is not implemented for now")
                     elif quantize is None:

diff --git a/server/text_generation_server/models/galactica.py b/server/text_generation_server/models/galactica.py
@@ -364,6 +364,8 @@ def linear(input, weight, bias):
                                 return linear
 
                             module.linear = replace_linear(state)
+                        else:
+                            tensor = tensor.to(device)
                     elif quantize == "gptq":
                         raise NotImplementedError("`gptq` is not implemented for now")
                     elif quantize is None:

diff --git a/server/text_generation_server/models/gpt_neox.py b/server/text_generation_server/models/gpt_neox.py
@@ -210,6 +210,8 @@ def linear(input, weight, bias):
                                 return linear
 
                             module.linear = replace_linear(state)
+                        else:
+                            tensor = tensor.to(device)
                     elif quantize == "gptq":
                         raise NotImplementedError("`gptq` is not implemented for now")
                     elif quantize is None:

diff --git a/server/text_generation_server/models/opt.py b/server/text_generation_server/models/opt.py
@@ -166,7 +166,7 @@ def load_weights(
 
                     tensor = tensor.contiguous().to(dtype)
 
-                    if quantize:
+                    if quantize == "bitsandbytes":
                         if not HAS_BITS_AND_BYTES:
                             raise ImportError(
                                 "bitsandbytes is not available on your machine either because it is not installed "
@@ -216,9 +216,14 @@ def linear(input, weight, bias):
                                 return linear
 
                             module.linear = replace_linear(state)
-
                         else:
                             tensor = tensor.to(device)
+                    elif quantize == "gptq":
+                        raise NotImplementedError("`gptq` is not implemented for now")
+                    elif quantize is None:
+                        tensor = tensor.to(device)
+                    else:
+                        raise ValueError(f"Unexpected quantize `{quantize}`")
 
                     module._parameters[param_name] = tensor
                     if name == "model.decoder.embed_tokens.weight":

diff --git a/server/text_generation_server/models/t5.py b/server/text_generation_server/models/t5.py
@@ -222,7 +222,8 @@ def linear(input, weight, bias):
                                 return linear
 
                             module.linear = replace_linear(state)
-
+                        else:
+                            tensor = tensor.to(device)
                     elif quantize == "gptq" and not module_name.endswith("wo"):
                         raise NotImplementedError("`gptq` is not implemented for now")
                     elif quantize is None or module_name.endswith("wo"):