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

support input.shape 3

Browse Source
This commit is contained in:
OlivierDehaene 2023-06-08 16:29:01 +02:00
parent b67405bd8e
commit bdd811dd83
1 changed files with 35 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -163,6 +163,7 @@ class TensorParallelHead(SuperLayer):
if world_size == 1: if world_size == 1:
return super().forward(input) return super().forward(input)
if input.shape == 2:
out_dim = self.linear.weight.shape[0] out_dim = self.linear.weight.shape[0]
if input.shape[0] == 1: if input.shape[0] == 1:
@ -196,6 +197,14 @@ class TensorParallelHead(SuperLayer):
) )
return world_out.T return world_out.T
output = super().forward(input)
world_output = [
torch.empty_like(output) for _ in range(self.process_group.size())
]
torch.distributed.all_gather(world_output, output, group=self.process_group)
world_output = torch.cat(world_output, dim=-1)
return world_output
class TensorParallelColumnLinear(SuperLayer): class TensorParallelColumnLinear(SuperLayer):
@classmethod @classmethod