A stack of BERT layers providing the backbone of Mosaic BERT. This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertAlibiEncoder`, but with substantial modifications to implement unpadding and ALiBi. Compared to the analogous Hugging Fa
| 71 | |
| 72 | |
| 73 | class BertAlibiEncoder(nn.Module): |
| 74 | """A stack of BERT layers providing the backbone of Mosaic BERT. |
| 75 | |
| 76 | This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertAlibiEncoder`, |
| 77 | but with substantial modifications to implement unpadding and ALiBi. |
| 78 | |
| 79 | Compared to the analogous Hugging Face BERT module, this module handles unpadding to reduce unnecessary computation |
| 80 | at padded tokens, and pre-computes attention biases to implement ALiBi. |
| 81 | """ |
| 82 | |
| 83 | def __init__(self, config): |
| 84 | super().__init__() |
| 85 | layer = BertAlibiLayer(config) |
| 86 | self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)]) |
| 87 | |
| 88 | self.num_attention_heads = config.num_attention_heads |
| 89 | |
| 90 | # The alibi mask will be dynamically expanded if it is too small for |
| 91 | # the input the model receives. But it generally helps to initialize it |
| 92 | # to a reasonably large size to help pre-allocate CUDA memory. |
| 93 | # The default `alibi_starting_size` is 512. |
| 94 | self._current_alibi_size = int(config.alibi_starting_size) |
| 95 | self.alibi = torch.zeros((1, self.num_attention_heads, self._current_alibi_size, self._current_alibi_size)) |
| 96 | self.rebuild_alibi_tensor(size=config.alibi_starting_size) |
| 97 | |
| 98 | def rebuild_alibi_tensor(self, size: int, device: Optional[Union[torch.device, str]] = None): |
| 99 | # Alibi |
| 100 | # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1) |
| 101 | # In the causal case, you can exploit the fact that softmax is invariant to a uniform translation |
| 102 | # of the logits, which makes the math work out *after* applying causal masking. If no causal masking |
| 103 | # will be applied, it is necessary to construct the diagonal mask. |
| 104 | n_heads = self.num_attention_heads |
| 105 | |
| 106 | def _get_alibi_head_slopes(n_heads: int) -> List[float]: |
| 107 | def get_slopes_power_of_2(n_heads: int) -> List[float]: |
| 108 | start = 2 ** (-(2 ** -(math.log2(n_heads) - 3))) |
| 109 | ratio = start |
| 110 | return [start * ratio**i for i in range(n_heads)] |
| 111 | |
| 112 | # In the paper, they only train models that have 2^a heads for some a. This function |
| 113 | # has some good properties that only occur when the input is a power of 2. To |
| 114 | # maintain that even when the number of heads is not a power of 2, we use a |
| 115 | # workaround. |
| 116 | if math.log2(n_heads).is_integer(): |
| 117 | return get_slopes_power_of_2(n_heads) |
| 118 | |
| 119 | closest_power_of_2 = 2 ** math.floor(math.log2(n_heads)) |
| 120 | slopes_a = get_slopes_power_of_2(closest_power_of_2) |
| 121 | slopes_b = _get_alibi_head_slopes(2 * closest_power_of_2) |
| 122 | slopes_b = slopes_b[0::2][: n_heads - closest_power_of_2] |
| 123 | return slopes_a + slopes_b |
| 124 | |
| 125 | context_position = torch.arange(size, device=device)[:, None] |
| 126 | memory_position = torch.arange(size, device=device)[None, :] |
| 127 | relative_position = torch.abs(memory_position - context_position) |
| 128 | # [n_heads, max_token_length, max_token_length] |
| 129 | relative_position = relative_position.unsqueeze(0).expand(n_heads, -1, -1) |
| 130 | slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device) |