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

This allows for overriding default block configs for certain layers. This should contain two sub configs: order and overrides. order specifies the order of different kinds of layers (default refers to a layer that does not apply any overrides). For each kind of layer, specify the overrides in the overrides config. For example, to specify this model (https://research.character.ai/optimizing-inference/) , the following config will be needed:

model:
    ...
    (usual model configs)
    ...
    block_overrides:
        order:
        - name: default
        - 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
          repeat: 2
        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

Also prints the following log summarizing the network:

INFO: llmfoundry.models.mpt.modeling_mpt: The following is a summary of overrides per layer.
  idx  name                        overrides
-----  --------------------------  ----------------------------------------------------------
    0  default                     []
    1  sliding_window_layer        [{'sliding_window_size': 1024}]
    2  sliding_window_layer_reuse  [{'sliding_window_size': 1024}, {'reuse_kv_layer_idx': 1}]
    3  sliding_window_layer        [{'sliding_window_size': 1024}]
    4  sliding_window_layer_reuse  [{'sliding_window_size': 1024}, {'reuse_kv_layer_idx': 3}]
    5  sliding_window_layer_reuse  [{'sliding_window_size': 1024}, {'reuse_kv_layer_idx': 3}]
    6  reuse_kv_layer              [{'reuse_kv_layer_idx': 0}]
    7  sliding_window_layer        [{'sliding_window_size': 1024}]
    8  sliding_window_layer_reuse  [{'sliding_window_size': 1024}, {'reuse_kv_layer_idx': 7}]
    9  sliding_window_layer        [{'sliding_window_size': 1024}]
   10  sliding_window_layer_reuse  [{'sliding_window_size': 1024}, {'reuse_kv_layer_idx': 9}]
   11  sliding_window_layer_reuse  [{'sliding_window_size': 1024}, {'reuse_kv_layer_idx': 9}]
   12  reuse_kv_layer              [{'reuse_kv_layer_idx': 0}]

Note that the table above prints the absolute layer index for reuse_kv_layer_idx

mosaicml / llm-foundry

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