Allows interweaving of arbitrary kinds of 'attention' layers, like sliding window, reuse prev layer kv cache etc.

llmfoundry/models/layers/attention.py

@@ -416,6 +416,7 @@ def __init__(
         device: Optional[str] = None,
         bias: bool = True,
         sliding_window_size: int = -1,
+        reuse_kv_layer_idx: Optional[int] = None,
     ):
         super().__init__()
 
@@ -428,6 +429,7 @@ def __init__(
         self.n_heads = n_heads
         self.kv_n_heads = kv_n_heads
         self.sliding_window_size = sliding_window_size
+        self.reuse_kv_layer_idx = reuse_kv_layer_idx
 
         self.head_dim = d_model // n_heads
 
@@ -458,18 +460,29 @@ def __init__(
             self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
         self.attn_dropout_p = attn_pdrop
 
-        self.Wqkv = build_fc(
-            name=fc_type_name,
-            in_features=self.d_model,
-            out_features=self.d_model + 2 * self.kv_n_heads * self.head_dim,
-            fc_kwargs=fc_type,
-        )
-        # for param init fn; enables shape based init of fused layers
-        fuse_splits = [
-            i * self.head_dim
-            for i in range(1, self.n_heads + 2 * self.kv_n_heads)
-        ]
-        self.Wqkv._fused = (0, fuse_splits)
+        if self.reuse_kv_layer_idx is None:
+            self.Wqkv = build_fc(
+                name=fc_type_name,
+                in_features=self.d_model,
+                out_features=self.d_model + 2 * self.kv_n_heads * self.head_dim,
+                fc_kwargs=fc_type,
+            )
+            # for param init fn; enables shape based init of fused layers
+            fuse_splits = [
+                i * self.head_dim
+                for i in range(1, self.n_heads + 2 * self.kv_n_heads)
+            ]
+            self.Wqkv._fused = (0, fuse_splits)
+        else:
+            self.Wq = build_fc(
+                name=fc_type_name,
+                in_features=self.d_model,
+                out_features=self.d_model,
+                fc_kwargs=fc_type,
+            )
+            # for param init fn; enables shape based init of fused layers
+            fuse_splits = [i * self.head_dim for i in range(1, self.n_heads)]
+            self.Wq._fused = (0, fuse_splits)
 
         if self.qk_ln or self.qk_gn:
             norm_size = self.head_dim if qk_gn else d_model
@@ -478,13 +491,14 @@ def __init__(
                 normalized_shape=norm_size,
                 device=device,
             )
-            if qk_ln:
-                norm_size = self.head_dim * kv_n_heads
-            self.k_ln = build_norm(
-                name=norm_type.lower(),
-                normalized_shape=norm_size,
-                device=device,
-            )
+            if self.reuse_kv_layer_idx is None:
+                if qk_ln:
+                    norm_size = self.head_dim * kv_n_heads
+                self.k_ln = build_norm(
+                    name=norm_type.lower(),
+                    normalized_shape=norm_size,
+                    device=device,
+                )
 
         self.attn_fn = attention_implementations.get(self.attn_impl)
 
@@ -507,9 +521,11 @@ def forward(
         needs_weights: bool = False,
         alibi_slopes: Optional[torch.Tensor] = None,
         flash_attn_padding_info: Optional[dict[str, torch.Tensor]] = None,
+        prev_layer_key_value: Optional[Tuple[torch.Tensor,
+                                             torch.Tensor]] = None,
     ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[
         torch.Tensor, torch.Tensor]]]:
-        query, key, value = self.get_qkv(x)
+        query, key, value = self.get_qkv(x, prev_layer_key_value)
 
         if rotary_emb_w_meta_info is not None:
             query, key, value = self._apply_rotary_embeddings(
@@ -546,6 +562,7 @@ def forward(
     def get_qkv(
         self,
         x: torch.Tensor,
+        prev_layer_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]],
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Computes and returns the query, key, and value tensors.
 
@@ -557,6 +574,27 @@ def get_qkv(
             key (torch.Tensor): The key tensor.
             value (torch.Tensor): The value tensor.
         """
+        if self.reuse_kv_layer_idx is not None:
+            if prev_layer_key_value is None:
+                raise ValueError(
+                    'prev_layer_key_value is None, cannot reuse_prev_layer_kv.',
+                )
+            key, value = prev_layer_key_value
+
+            query = self.Wq(x)
+            if self.clip_qkv:
+                query = query.clamp(min=-self.clip_qkv, max=self.clip_qkv)
+
+            if self.qk_ln or self.qk_gn:
+                # Applying layernorm to qk
+                q_shape = query.shape
+                if self.qk_gn:
+                    b, s = query.shape[:2]
+                    query = query.view(b, s, self.n_heads, -1)
+                dtype = query.dtype
+                query = self.q_ln(query).to(dtype).view(q_shape)
+            return query, key, value
+
         qkv = self.Wqkv(x)
 
         if self.clip_qkv:
@@ -591,6 +629,10 @@ def _apply_rotary_embeddings(
         key: torch.Tensor,
         value: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.reuse_kv_layer_idx is not None:
+            orig_key, orig_value = key, value
+            key, value = torch.empty_like(key), torch.empty_like(value)
+
         rotary_emb = rotary_emb_w_meta_info['rotary_emb']
         seq_len = rotary_emb_w_meta_info['seq_len']
         offset_info = rotary_emb_w_meta_info['offset_info']
@@ -602,6 +644,7 @@ def _apply_rotary_embeddings(
             value = value.view(bsz, seqlen, -1, self.head_dim)
 
             kv = torch.stack([key, value], dim=2)
+            # Note: Rotates in place (https://github.com/Dao-AILab/flash-attention/blob/320fb59487658f033f56711efd3d61b7c7a6f8f3/flash_attn/layers/rotary.py#L429)
             query, kv = rotary_emb(
                 query,
                 kv,
@@ -652,6 +695,8 @@ def _apply_rotary_embeddings(
 
         query = query.view(bsz, seqlen, -1)
         key = key.view(bsz, seqlen, -1)
+        if self.reuse_kv_layer_idx is not None:
+            return query, orig_key, orig_value  # type: ignore
         return query, key, value
 
     def get_implementation_specific_args(
@@ -705,6 +750,7 @@ def __init__(
         device: Optional[str] = None,
         bias: bool = True,
         sliding_window_size: int = -1,
+        reuse_kv_layer_idx: Optional[int] = None,
     ):
         super().__init__(
             d_model=d_model,
@@ -721,6 +767,7 @@ def __init__(
             device=device,
             bias=bias,
             sliding_window_size=sliding_window_size,
+            reuse_kv_layer_idx=reuse_kv_layer_idx,
         )
 
 
@@ -746,6 +793,7 @@ def __init__(
         device: Optional[str] = None,
         bias: bool = True,
         sliding_window_size: int = -1,
+        reuse_kv_layer_idx: Optional[int] = None,
     ):
         super().__init__(
             d_model=d_model,
@@ -762,6 +810,7 @@ def __init__(
             device=device,
             bias=bias,
             sliding_window_size=sliding_window_size,
+            reuse_kv_layer_idx=reuse_kv_layer_idx,
         )
 
 

llmfoundry/models/layers/blocks.py

@@ -158,6 +158,8 @@ def forward(
         output_attentions: bool = False,
         alibi_slopes: Optional[torch.Tensor] = None,
         flash_attn_padding_info: Optional[dict[str, torch.Tensor]] = None,
+        prev_layer_key_value: Optional[Tuple[torch.Tensor,
+                                             torch.Tensor]] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[
         torch.Tensor, torch.Tensor]]]:
         if self.fuse_norm_attn_norm:
@@ -171,6 +173,7 @@ def forward(
                 output_attentions=output_attentions,
                 alibi_slopes=alibi_slopes,
                 flash_attn_padding_info=flash_attn_padding_info,
+                prev_layer_key_value=prev_layer_key_value,
             )
         else:
             a = self.norm_1(x)
@@ -184,6 +187,7 @@ def forward(
                 needs_weights=output_attentions,
                 alibi_slopes=alibi_slopes,
                 flash_attn_padding_info=flash_attn_padding_info,
+                prev_layer_key_value=prev_layer_key_value,
             )
             x = x + self.resid_attn_dropout(b)
             m = x
@@ -308,6 +312,8 @@ def forward(
         output_attentions: bool = False,
         alibi_slopes: Optional[torch.Tensor] = None,
         flash_attn_padding_info: Optional[dict[str, torch.Tensor]] = None,
+        prev_layer_key_value: Optional[Tuple[torch.Tensor,
+                                             torch.Tensor]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor, torch.Tensor]]]:
         a = self.norm_1(x)
@@ -321,6 +327,7 @@ def forward(
             needs_weights=output_attentions,
             alibi_slopes=alibi_slopes,
             flash_attn_padding_info=flash_attn_padding_info,
+            prev_layer_key_value=prev_layer_key_value,
         )
         x = x + self.resid_attn_dropout(b)
         m = x

llmfoundry/models/mpt/configuration_mpt.py

@@ -20,6 +20,7 @@
     ffn_config_defaults,
     init_config_defaults,
 )
+from llmfoundry.utils.warnings import ExperimentalWarning
 
 
 class MPTConfig(PretrainedConfig):
@@ -48,6 +49,7 @@ def __init__(
         fc_type: Union[str, Dict] = 'torch',
         tie_word_embeddings: bool = True,
         use_pad_tok_in_ffn: bool = True,
+        block_overrides: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
     ):
         """The MPT configuration class.
@@ -117,6 +119,30 @@ def __init__(
                 also be a dictionary that specifies the fc layer name and any kwargs for the fc layer.
             tie_word_embeddings (bool): Whether to tie the input embedding and output layers.
             use_pad_tok_in_ffn (bool): Whether to forward the pad token in the feedforward networks.
+            block_overrides: This allows for overriding default block configs for certain layers. This must contain `overrides` and `order`. `order` is a nested list which describes the order of the layers. For each kind of layer, specify the `overrides` in the overrides config (default refers to a layer that does not apply any overrides).
+                To specify this model (https://research.character.ai/optimizing-inference/) , the following config will be needed:
+                    block_overrides:
+                        order:
+                        - name: default
+                        - repeat: 2
+                          order:
+                          - name: sliding_window_layer
+                          - name: sliding_window_layer_reuse
+                          - name: sliding_window_layer
+                          - repeat: 2
+                            name: sliding_window_layer_reuse
+                          - name: reuse_kv_layer
+                        overrides:
+                            sliding_window_layer:
+                                attn_config:
+                                    sliding_window_size: 1024
+                            sliding_window_layer_reuse:
+                                attn_config:
+                                    sliding_window_size: 1024
+                                    reuse_kv_layer_idx: -1 # Relative index of the layer whose kv cache to reuse
+                            reuse_kv_layer:
+                                attn_config:
+                                    reuse_kv_layer_idx: -6 # Relative index of the layer whose kv cache to reuse
         """
         self.d_model = d_model
         self.n_heads = n_heads
@@ -145,6 +171,15 @@ def __init__(
             init_config_defaults,
         )
 
+        if 'reuse_kv_layer_idx' in self.attn_config and self.attn_config[
+            'attn_impl'] == 'torch':
+            raise NotImplementedError(
+                'reusing kv cache from a previous layer is not implemented for torch attention.',
+            )
+        if block_overrides is not None:
+            self._validate_block_overrides(block_overrides)
+        self.block_overrides = block_overrides
+
         if isinstance(fc_type, str):
             fc_type = {'name': fc_type}
         self.fc_type = fc_type
@@ -169,6 +204,23 @@ def __init__(
 
         self._validate_config()
 
+    def _validate_block_overrides(self, block_overrides: Dict[str, Any]):
+        warnings.warn(ExperimentalWarning('block_overrides'))
+        if 'order' not in block_overrides:
+            raise ValueError('`order` should be defined in block_overrides',)
+        if 'overrides' not in block_overrides:
+            raise ValueError(
+                '`overrides` should be defined in block_overrides',
+            )
+        for name, override in block_overrides['overrides'].items():
+            if name == 'default':
+                raise ValueError('block overrides cannot be named "default".',)
+            if 'attn_config' in override and 'reuse_kv_layer_idx' in override[
+                'attn_config'] and self.attn_config['attn_impl'] == 'torch':
+                raise NotImplementedError(
+                    'reusing kv cache from a previous layer is not implemented for torch attention.',
+                )
+
     def _set_config_defaults(
         self,
         config: Dict[str, Any],
@@ -335,3 +387,12 @@ def _validate_config(self) -> None:
                 )
 
         self.validate_attention_config()
+
+    @property
+    def allowed_block_overrides(self):
+        return {
+            'attn_config': {
+                'sliding_window_size': None,
+                'reuse_kv_layer_idx': None,
+            },
+        }