huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-06-19 15:52:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

turn tp embed back on

Browse Source
This commit is contained in:
OlivierDehaene 2023-03-24 18:20:47 +01:00
parent fc778e46fb
commit 15a6b79c7e
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
server/text_generation_server/models/flash_neox_modeling.py
View File

@ -499,11 +499,11 @@ class FlashGPTNeoXModel(FlashGPTNeoXPreTrainedModel):
if config.vocab_size % self.tp_world_size == 0:
self.tp_embeddings = True
# if self.tp_embeddings:
# self.embed_in = TensorParallelEmbedding(
# config.vocab_size, config.hidden_size, process_group=process_group
# )
# else:
if self.tp_embeddings:
self.embed_in = TensorParallelEmbedding(
config.vocab_size, config.hidden_size, process_group=process_group
)
else:
self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
self.layers = nn.ModuleList(