From cc07a8460465228c435dcee323d9c93ef91d3f7a Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 7 May 2024 12:11:43 +0000 Subject: [PATCH] Fixes. --- server/text_generation_server/layers/tensor_parallel.py | 2 +- .../models/custom_modeling/mamba_modeling.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/server/text_generation_server/layers/tensor_parallel.py b/server/text_generation_server/layers/tensor_parallel.py index ec4ecf38..34b9c51e 100644 --- a/server/text_generation_server/layers/tensor_parallel.py +++ b/server/text_generation_server/layers/tensor_parallel.py @@ -1,7 +1,7 @@ import torch from torch.nn import functional as F from typing import List -from text_generation_server.layers.linear import get_linear +from text_generation_server.layers.linear import get_linear, FastLinear class SuperLayer(torch.nn.Module): diff --git a/server/text_generation_server/models/custom_modeling/mamba_modeling.py b/server/text_generation_server/models/custom_modeling/mamba_modeling.py index ed638804..293051c2 100644 --- a/server/text_generation_server/models/custom_modeling/mamba_modeling.py +++ b/server/text_generation_server/models/custom_modeling/mamba_modeling.py @@ -11,9 +11,9 @@ import torch.nn.functional as F from text_generation_server.layers import ( SpeculativeHead, TensorParallelEmbedding, - FastRMSNorm, FastLinear, ) +from text_generation_server.layers.layernorm import FastRMSNorm from einops import rearrange from causal_conv1d import causal_conv1d_fn, causal_conv1d_update