From deed8e81540d0f139000b610ccf5ec652fcd2165 Mon Sep 17 00:00:00 2001 From: drbh Date: Wed, 7 Feb 2024 03:24:28 +0000 Subject: [PATCH] fix: adjust typos and docker build --- Dockerfile | 10 +++++----- .../models/custom_modeling/mamba_modeling.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/Dockerfile b/Dockerfile index cc1a501a..6818005f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -154,17 +154,17 @@ COPY server/Makefile-vllm Makefile # Build specific version of vllm RUN make build-vllm-cuda -# Build megablocks -FROM kernel-builder as megablocks-builder - -RUN pip install git+https://github.com/OlivierDehaene/megablocks@181709df192de9a941fdf3a641cdc65a0462996e - # Build mamba kernels FROM kernel-builder as mamba-builder WORKDIR /usr/src COPY server/Makefile-selective-scan Makefile RUN make build-all +# Build megablocks +FROM kernel-builder as megablocks-builder + +RUN pip install git+https://github.com/OlivierDehaene/megablocks@181709df192de9a941fdf3a641cdc65a0462996e + # Text Generation Inference base image FROM nvidia/cuda:12.1.0-base-ubuntu20.04 as base 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 c515a5ca..1773f04d 100644 --- a/server/text_generation_server/models/custom_modeling/mamba_modeling.py +++ b/server/text_generation_server/models/custom_modeling/mamba_modeling.py @@ -137,7 +137,7 @@ class MambaBlock(nn.Module): z = _z[i:i+1, -1, :] x_db = self.x_proj(x) dt, B, C = torch.split(x_db, [self.dt_rank, self.d_state, self.d_state], dim=-1) - df = self.dt_proj_no_bias(x) + dt = F.linear(dt, self.dt_proj.weight) y = selective_state_update( ssm_state[i:i+1,:,:], x, dt, self.negA, B, C, self.D, z=z, dt_bias=self.dt_proj.bias, dt_softplus=True )