Added support for bloom-560m model

transformer_lens/HookedTransformer.py

@@ -334,6 +334,10 @@ def input_to_embed(
             # keys and queries. See HookedTransformerConfig for details
             residual = embed
             shortformer_pos_embed = None
+        #TODO: alibi embedding doesnt do anything
+        elif self.cfg.positional_embedding_type == 'alibi':
+            residual = embed 
+            shortformer_pos_embed = None
         else:
             raise ValueError(
                 f"Invalid positional_embedding_type passed in {self.cfg.positional_embedding_type}"

transformer_lens/HookedTransformerConfig.py

@@ -195,6 +195,11 @@ class HookedTransformerConfig:
     dtype: torch.dtype = torch.float32
     tokenizer_prepends_bos: Optional[bool] = None
 
+    # bloom flags
+    post_embedding_layer_norm: bool = False
+
+
+
     def __post_init__(self):
         if self.n_heads == -1:
             self.n_heads = self.d_model // self.d_head

transformer_lens/components.py

@@ -8,6 +8,7 @@
 from typing import Dict, Optional, Tuple, Union
 
 import einops
+import math
 import numpy as np
 import torch
 import torch.nn as nn
@@ -32,12 +33,17 @@ def __init__(self, cfg: Union[Dict, HookedTransformerConfig]):
         self.W_E: Float[torch.Tensor, "d_vocab d_model"] = nn.Parameter(
             torch.empty(self.cfg.d_vocab, self.cfg.d_model, dtype=cfg.dtype)
         )
+        # bloom needs post embedding layer norm
+        if cfg.post_embedding_layer_norm: 
+            self.ln = LayerNorm(cfg)
 
     def forward(
         self, tokens: Int[torch.Tensor, "batch pos"]
     ) -> Float[torch.Tensor, "batch pos d_model"]:
         # If A has shape [a, b] and B has shape [c, d], then A[:, B] has shape [a, c, d]
         # B acts as a tensor of indices into the second dimension (so >=0 and <b)
+        if self.cfg.post_embedding_layer_norm:
+            return self.ln(self.W_E[tokens, :])
         return self.W_E[tokens, :]
 
 
@@ -303,7 +309,7 @@ def forward(
     ]:
         if self.cfg.dtype not in [torch.float32, torch.float64]:
             x = x.to(torch.float32)
-
+          
         x = x - x.mean(axis=-1, keepdim=True)  # [batch, pos, length]
         scale: Float[torch.Tensor, "batch pos 1"] = self.hook_scale(
             (x.pow(2).mean(-1, keepdim=True) + self.eps).sqrt()
@@ -478,6 +484,8 @@ def __init__(
             )
             self.register_buffer("rotary_sin", sin)
             self.register_buffer("rotary_cos", cos)
+
+
 
     @property
     def OV(self) -> FactoredMatrix:
@@ -533,7 +541,6 @@ def forward(
             qkv_einops_string = "batch pos head_index d_model"
         else:
             qkv_einops_string = "batch pos d_model"
-
         q = self.hook_q(
             einsum(
                 f"{qkv_einops_string}, head_index d_model d_head \
@@ -593,6 +600,27 @@ def forward(
             )
             / self.attn_scale
         )  # [batch, head_index, query_pos, key_pos]
+
+        # alibi encoding before applying causal mask
+        if self.cfg.positional_embedding_type == 'alibi':
+            #TODO: not sure about the side effect of not using standard, double check
+            batch_size = attn_scores.size(0)
+            seq_len = attn_scores.size(-2) 
+            additive_mask = torch.ones(batch_size, seq_len)
+            dtype = self.cfg.dtype if self.cfg.dtype in [torch.float32, torch.float64] else 'torch.float32'
+            alibi = self.build_alibi_tensor(
+                attention_mask=additive_mask,
+                num_heads=self.cfg.n_heads,
+                dtype=dtype
+            ).to(attn_scores.device)
+
+            # Huggingface impl uses torch.Tensor.baddbmm, with alpha = 1/sqrt(d_head), and beta=1
+            # and alibi.baddbmm(q,k) = beta * alibi + alpha * (q@k), 
+            # here the `attn_scores` is already scaled by a factor of self.attn_scale, 
+            # we only need to add alibi matrix to the result
+            assert alibi.shape == (attn_scores.size(0), attn_scores.size(1), 1, attn_scores.size(-1)), f"alibi shape {alibi.shape}, expecting {attn_scores.shape}"
+            attn_scores += alibi # [batch, head_index, query_pos, key_pos]
+
         if self.cfg.attention_dir == "causal":
             # If causal attention, we mask it to only attend backwards. If bidirectional, we don't mask.
             attn_scores = self.apply_causal_mask(
@@ -756,7 +784,43 @@ def apply_rotary(
             x_rotated = x_rot * mask_rotary_cos + x_flip * mask_rotary_sin
 
         return torch.cat([x_rotated, x_pass], dim=-1)
+    def build_alibi_tensor(
+            self,
+            attention_mask: torch.Tensor, # batch pos
+            num_heads: int,
+            dtype: torch.dtype
+    ) -> Float[torch.Tensor, "batch head_index 1 pos"]:
+        """
+        https://github.com/huggingface/transformers/blob/21dc5859421cf0d7d82d374b10f533611745a8c5/src/transformers/models/bloom/modeling_bloom.py#L86
+        Args:
+        Returns tensor shaped (batch_size * num_heads, 1, max_seq_len)
+            attention_mask (`torch.Tensor`):
+                Token-wise attention mask, this should be of shape (batch_size, max_seq_len).
+            num_heads (`int`, *required*):
+                number of heads
+            dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`):
+                dtype of the output tensor
+        """
+        batch_size, seq_length = attention_mask.shape
+        closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
+        base = torch.tensor(
+            2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+        )
+        powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)
+        slopes = torch.pow(base, powers)
 
+        if closest_power_of_2 != num_heads:
+            extra_base = torch.tensor(
+                2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+            )
+            num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
+            extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)
+            slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
+
+        arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :]
+        alibi = slopes[..., None] * arange_tensor
+        # originally it returns tensor of shape batch * head_index, 1, pos
+        return alibi.reshape(batch_size, num_heads, 1, seq_length).to(dtype)
 
 # MLP Layers
 class MLP(nn.Module):

transformer_lens/loading_from_pretrained.py

@@ -136,6 +136,7 @@
     "stabilityai/stablelm-base-alpha-7b",
     "stabilityai/stablelm-tuned-alpha-3b",
     "stabilityai/stablelm-tuned-alpha-7b",
+    "bigscience/bloom-560m",
 ]
 """Official model names for models on HuggingFace."""
 
@@ -494,6 +495,9 @@
         "stablelm-tuned-alpha-7b",
         "stablelm-tuned-7b",
     ],
+    "bigscience/bloom-560m": [
+        "bloom-560m"
+    ],
 }
 """Model aliases for models on HuggingFace."""
 
@@ -721,6 +725,23 @@ def convert_hf_model_config(model_name: str, **kwargs):
             "act_fn": "gelu",
             "attention_dir": "bidirectional",
         }
+    elif architecture == 'BloomForCausalLM':
+        cfg_dict = {
+            "d_model" : hf_config.hidden_size,
+            "d_head" : hf_config.hidden_size // hf_config.n_head,
+            "n_heads" : hf_config.n_head,
+            "d_mlp": hf_config.hidden_size * 4,
+            "n_layers": hf_config.n_layer,
+            "n_ctx": 2048, # is there a variable for this?
+            "d_vocab": hf_config.vocab_size,
+            "act_fn" : "gelu_fast",
+            "eps": hf_config.layer_norm_epsilon,
+            "normalization_type": "LN", # double check this
+            "post_embedding_layer_norm": True,
+            "positional_embedding_type": 'alibi'
+        }
+
+        # print("bloom config", cfg_dict)
     else:
         raise NotImplementedError(f"{architecture} is not currently supported.")
     # All of these models use LayerNorm
@@ -1043,6 +1064,8 @@ def get_pretrained_state_dict(
             state_dict = convert_llama_weights(hf_model, cfg)
         elif cfg.original_architecture == "BertForMaskedLM":
             state_dict = convert_bert_weights(hf_model, cfg)
+        elif cfg.original_architecture == "BloomForCausalLM":
+            state_dict = convert_bloom_weights(hf_model, cfg)
         else:
             raise ValueError(
                 f"Loading weights from the architecture is not currently supported: {cfg.original_architecture}, generated from model name {cfg.model_name}. Feel free to open an issue on GitHub to request this feature."
@@ -1623,6 +1646,65 @@ def convert_bert_weights(bert, cfg: HookedTransformerConfig):
     return state_dict
 
 
+#TODO: bloom weight conversion
+def convert_bloom_weights(bloom, cfg: HookedTransformerConfig):
+    state_dict = {}
+
+    state_dict["embed.W_E"] = bloom.transformer.word_embeddings.weight
+
+    # Bloom uses post embedding layer norm
+    state_dict["embed.ln.w"] = bloom.transformer.word_embeddings_layernorm.weight
+    print(state_dict['embed.ln.w'][:5])
+    state_dict["embed.ln.b"] = bloom.transformer.word_embeddings_layernorm.bias
+
+
+
+    for l in range(cfg.n_layers):
+        state_dict[f"blocks.{l}.ln1.w"] = bloom.transformer.h[l].input_layernorm.weight
+        state_dict[f"blocks.{l}.ln1.b"] = bloom.transformer.h[l].input_layernorm.bias
+
+        # Bloom attn weight is stored as a fused matrx. BloomAttn: Linear(in=1024, out=3072)
+        # The .weight returned matrix will be in shape (3072, 1024)
+        W = bloom.transformer.h[l].self_attention.query_key_value.weight
+        # First transpose -> (1024, 3072), then split into (d_model, n_heads, 3, d_head)
+        W_split = W.T.reshape(cfg.d_model, cfg.n_heads, 3, cfg.d_head)
+
+        W_Q, W_K, W_V = W_split[...,0,:], W_split[...,1,:], W_split[...,2,:]
+        W_Q = einops.rearrange(W_Q, "m n h ->n m h", n=cfg.n_heads)
+        W_K = einops.rearrange(W_K, "m n h ->n m h", n=cfg.n_heads)
+        W_V = einops.rearrange(W_V, "m n h ->n m h", n=cfg.n_heads)
+        state_dict[f"blocks.{l}.attn.W_Q"] = W_Q
+        state_dict[f"blocks.{l}.attn.W_K"] = W_K
+        state_dict[f"blocks.{l}.attn.W_V"] = W_V
+
+        qkv_bias = bloom.transformer.h[l].self_attention.query_key_value.bias
+        qkv_bias = qkv_bias.reshape(cfg.n_heads, 3, cfg.d_head)
+
+        state_dict[f"blocks.{l}.attn.b_Q"] = qkv_bias[:, 0, :]
+        state_dict[f"blocks.{l}.attn.b_K"] = qkv_bias[:, 1, :]
+        state_dict[f"blocks.{l}.attn.b_V"] = qkv_bias[:, 2, :]
+
+        W_O = bloom.transformer.h[l].self_attention.dense.weight.T #[1024, 1024]
+        W_O = einops.rearrange(W_O, "(n h) m->n h m", n=cfg.n_heads) # [n_heads, d_head, d_model]
+        state_dict[f"blocks.{l}.attn.W_O"] = W_O
+        state_dict[f"blocks.{l}.attn.b_O"] = bloom.transformer.h[l].self_attention.dense.bias
+
+        state_dict[f"blocks.{l}.ln2.w"] = bloom.transformer.h[l].post_attention_layernorm.weight
+        state_dict[f"blocks.{l}.ln2.b"] = bloom.transformer.h[l].post_attention_layernorm.bias
+
+        W_in = bloom.transformer.h[l].mlp.dense_h_to_4h.weight.T
+        state_dict[f"blocks.{l}.mlp.W_in"] = W_in
+        state_dict[f"blocks.{l}.mlp.b_in"] = bloom.transformer.h[l].mlp.dense_h_to_4h.bias
+
+        W_out = bloom.transformer.h[l].mlp.dense_4h_to_h.weight.T
+        state_dict[f"blocks.{l}.mlp.W_out"] = W_out
+        state_dict[f"blocks.{l}.mlp.b_out"] = bloom.transformer.h[l].mlp.dense_4h_to_h.bias
+    state_dict["unembed.W_U"] = bloom.lm_head.weight.T # why transpose? cuz right mult?
+
+    state_dict["ln_final.w"] = bloom.transformer.ln_f.weight
+    state_dict["ln_final.b"] = bloom.transformer.ln_f.bias
+    return state_dict
+
 @dataclasses.dataclass
 class Config:
     d_model: int = 768
@@ -1660,3 +1742,5 @@ def get_basic_config(model_name: str, **kwargs) -> Config:
             ]
         }
     )
+
+# %%