huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing mamba by using the transformers version.

Browse Source
This commit is contained in:
Nicolas Patry 2024-09-25 03:37:12 +02:00
parent 9d7a95b24b
commit cd355d08a9
No known key found for this signature in database
GPG Key ID: D2920555C90F704C
4 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
integration-tests/models/test_mamba.py
View File

@ -3,7 +3,7 @@ import pytest
@pytest.fixture(scope="module") @pytest.fixture(scope="module")
def fused_kernel_mamba_handle(launcher): def fused_kernel_mamba_handle(launcher):
with launcher("state-spaces/mamba-130m", num_shard=1) as handle: with launcher("state-spaces/mamba-130m-hf", num_shard=1) as handle:
yield handle yield handle

1
router/src/config.rs
View File

@ -145,6 +145,7 @@ pub enum Config {
LlavaNext(LlavaNext), LlavaNext(LlavaNext),
ClipVisionModel(ClipVisionModel), ClipVisionModel(ClipVisionModel),
Mistral, Mistral,
Mamba,
Idefics, Idefics,
Mllama, Mllama,
Idefics2(Idefics2), Idefics2(Idefics2),

4
server/text_generation_server/models/__init__.py
View File

@ -226,7 +226,7 @@ class ModelType(enum.Enum):
"url": "https://huggingface.co/databricks/dbrx-instruct", "url": "https://huggingface.co/databricks/dbrx-instruct",
} }
MAMBA = { MAMBA = {
"type": "ssm", "type": "mamba",
"name": "Mamba", "name": "Mamba",
"url": "https://huggingface.co/state-spaces/mamba-2.8b-slimpj", "url": "https://huggingface.co/state-spaces/mamba-2.8b-slimpj",
} }
@ -555,7 +555,7 @@ def get_model(
# TODO: fix how we determine model type for Mamba # TODO: fix how we determine model type for Mamba
if "ssm_cfg" in config_dict: if "ssm_cfg" in config_dict:
# *only happens in Mamba case # *only happens in Mamba case
model_type = "ssm" model_type = "mamba"
else: else:
raise RuntimeError( raise RuntimeError(
f"Could not determine model type for {model_id} revision {revision}" f"Could not determine model type for {model_id} revision {revision}"

10
server/text_generation_server/models/custom_modeling/mamba_modeling.py
View File

@ -196,7 +196,10 @@ class MambaModel(nn.Module):
def __init__(self, config, weights): def __init__(self, config, weights):
super().__init__() super().__init__()
prefix = "backbone" prefix = "backbone"
self.embed_tokens = TensorParallelEmbedding(f"{prefix}.embedding", weights) try:
self.embed_tokens = TensorParallelEmbedding(f"{prefix}.embeddings", weights)
except RuntimeError:
self.embed_tokens = TensorParallelEmbedding(f"{prefix}.embedding", weights)
self.blocks = nn.ModuleList( self.blocks = nn.ModuleList(
[ [
ResidualBlock(f"{prefix}.layers.{i}", config, weights, layer_id=i) ResidualBlock(f"{prefix}.layers.{i}", config, weights, layer_id=i)
@ -206,7 +209,10 @@ class MambaModel(nn.Module):
self.norm_f = FastRMSNorm.load( self.norm_f = FastRMSNorm.load(
f"{prefix}.norm_f", weights, eps=config.layer_norm_epsilon f"{prefix}.norm_f", weights, eps=config.layer_norm_epsilon
) )
self.lm_head = SpeculativeHead.load(config, f"{prefix}.embedding", weights) try:
self.lm_head = SpeculativeHead.load(config, f"{prefix}.embeddings", weights)
except RuntimeError:
self.lm_head = SpeculativeHead.load(config, f"{prefix}.embeddings", weights)
self.config = config self.config = config
def forward( def forward(