Multi-node Tensor Parallelism (#550)

Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
2024-06-17 20:41:24 -07:00
parent 53a7ebd89a
commit 09593e9bc9
10 changed files with 167 additions and 46 deletions
--- a/python/sglang/srt/model_config.py
+++ b/python/sglang/srt/model_config.py
@@ -71,7 +71,11 @@ class ModelConfig:
            return 1

        # For DBRX and MPT
-        if self.hf_config.model_type in ["dbrx", "mpt"]:
+        if self.hf_config.model_type in ["mpt"]:
+            if "kv_n_heads" in self.hf_config.attn_config:
+                return self.hf_config.attn_config["kv_n_heads"]
+            return self.hf_config.num_attention_heads
+        if self.hf_config.model_type in ["dbrx"]:
            return getattr(
                self.hf_config.attn_config,
                "kv_n_heads",