Skip to content

GitLab

Unverified Commit 3a9cdc32 authored by Nicolas Patry's avatar Nicolas Patry Committed by GitHub
Browse files

Fixing auto bloom test. (#2699)

parent 78ce618c
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
server/text_generation_server/models/custom_modeling/bloom_modeling.py
View file @ 3a9cdc32
...@@ -377,7 +377,7 @@ class BloomAttention(nn.Module): ...@@ -377,7 +377,7 @@ class BloomAttention(nn.Module):
past_value.view(-1, *past_value.shape[-2:]), past_value.view(-1, *past_value.shape[-2:]),
) )
if CUSTOM_KERNELS_ENABLED: if CUSTOM_KERNELS_ENABLED and attention_mask.shape[-1] < 4096:
assert self.training is False, "Only foward pass was implemented" assert self.training is False, "Only foward pass was implemented"
assert ( assert (
attention_mask.shape[-1] < 4096 attention_mask.shape[-1] < 4096
...@@ -580,7 +580,7 @@ class BloomPreTrainedModel(PreTrainedModel): ...@@ -580,7 +580,7 @@ class BloomPreTrainedModel(PreTrainedModel):
@staticmethod @staticmethod
def _convert_to_bloom_cache( def _convert_to_bloom_cache(
past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]] past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]],
) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]: ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
""" """
Converts the cache to the format expected by Bloom, i.e. to tuple(tuple([batch_size * num_heads, ...])) Converts the cache to the format expected by Bloom, i.e. to tuple(tuple([batch_size * num_heads, ...]))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment