support convert embeddings to bf16

InternLM · Dec 8, 2023 · a6c4977 · a6c4977
1 parent 1753ead
commit a6c4977
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py
@@ -558,10 +558,14 @@ def _broadcast_np(data, dtype, shape=(batch_size, )):
                 embedding_ends = [embedding_ends]
                 embeddings = [embeddings]
             # convert to lookup table type
-            # TODO bf16
             if self.tm_model.config.weight_type == 'fp32':
                 embeddings = [[x.astype(np.float32) for x in y]
                               for y in embeddings]
+            elif self.tm_model.config.weight_type == 'bf16':
+                embeddings = [[
+                    torch.from_numpy(x).bfloat16().view(torch.half).numpy()
+                    for x in y
+                ] for y in embeddings]
             else:
                 embeddings = [[x.astype(np.float16) for x in y]
                               for y in embeddings]