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

Fixing layer imports (for isinstance compat).

Browse Source
This commit is contained in:
Nicolas Patry 2023-05-15 16:46:32 +02:00
parent edc9ce9beb
commit 42d8efcb04
2 changed files with 46 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
server/text_generation_server/models/t5.py
View File

@ -18,11 +18,20 @@ from text_generation_server.utils import (
) )
from text_generation_server.utils.layers import ( from text_generation_server.utils.layers import (
FastLinear, FastLinear,
)
from transformers.models.t5.parallel_layers import (
TensorParallelRowLinear, TensorParallelRowLinear,
TensorParallelColumnLinear, TensorParallelColumnLinear,
TensorParallelEmbedding, TensorParallelEmbedding,
) )
HAS_BITS_AND_BYTES = True
try:
import bitsandbytes as bnb
from bitsandbytes.nn import Int8Params
except ImportError as e:
HAS_BITS_AND_BYTES = False
class T5Sharded(Seq2SeqLM): class T5Sharded(Seq2SeqLM):
def __init__( def __init__(

7
server/text_generation_server/utils/layers.py
View File

@ -1,7 +1,6 @@
import torch import torch
from torch import nn from torch import nn
import dropout_layer_norm
HAS_BITS_AND_BYTES = True HAS_BITS_AND_BYTES = True
try: try:
@ -182,6 +181,9 @@ class TensorParallelEmbedding(nn.Embedding):
return out return out
try:
import dropout_layer_norm
class FastLayerNorm(nn.LayerNorm): class FastLayerNorm(nn.LayerNorm):
def forward(self, hidden_states, residual=None): def forward(self, hidden_states, residual=None):
if hidden_states.shape[-1] > 8192: if hidden_states.shape[-1] > 8192:
@ -217,6 +219,9 @@ class FastLayerNorm(nn.LayerNorm):
return normed_hidden_states, residual return normed_hidden_states, residual
except ImportError:
pass
try: try:
from flash_attn.layers.rotary import RotaryEmbedding from flash_attn.layers.rotary import RotaryEmbedding