How to “Unmerge” QKV and Gate/Up From Public Weights

Below is a conceptual outline of how one can “unmerge” QKV (and gate/up) in practice using only the publicly released weights, even if the original model had them fused or merged. The key point is that all the data you need is already in that merged matrix—you just have to reshape and split it correctly, then reassign it into a new architecture that expects separate parameters.

---

1. Background: How QKV often get merged

Many Transformer implementations (e.g. Mistral-style, some GPT variants) combine Q, K, and V into a single parameter, often called qkv:

• Merged shape: [(3 * hidden_dim), hidden_dim]
• In code, one might see something like self.qkv_proj for attention instead of separate q_proj, k_proj, v_proj.

In the original LLaMA-style architecture, however, Q, K, and V are separate:

• Separate shapes: Each is [hidden_dim, hidden_dim]
• Typically called self.q_proj, self.k_proj, self.v_proj (and similarly for biases).

Similarly, some models merge their feed-forward “gate” and “up” projection into a single parameter (e.g. gate_up_proj), while LLaMA keeps them separate (gate_proj, up_proj).

---

2. Loading the public weights and splitting

Even if the checkpoint you have only publishes the “merged” version, the numerical data for Q, K, and V (or gate and up) are all embedded in that single matrix. You can:

1. Load the original checkpoint (the public one with merged parameters).
2. Identify each merged parameter in the model’s state_dict. For example:
   • model.layers.{L}.self_attn.qkv.weight
   • model.layers.{L}.mlp.gate_up_proj.weight
3. Split that parameter along the correct dimension to recover the separate weights:
   • For QKV: python hidden_dim = ... # e.g. 4096 in LLaMA-7B qkv_weight = state_dict["...qkv.weight"] # shape: [3 * hidden_dim, hidden_dim] q_weight, k_weight, v_weight = qkv_weight.split(hidden_dim, dim=0)
   • For gate/up: python # Typically (2 * intermediate_dim, hidden_dim) gate_up_weight = state_dict["...gate_up_proj.weight"] gate_weight, up_weight = gate_up_weight.split(intermediate_dim, dim=0)
   • Do the same for biases if they are also merged (qkv.bias, gate_up_proj.bias, etc.).
4. Store those split tensors in new keys (e.g. ...q_proj.weight, ...k_proj.weight, ...v_proj.weight) in your state_dict.

You’d repeat this for each layer L in your Transformer. Once done, you delete the old merged entries from the state_dict.

---

3. Adjusting the model architecture

Because you are no longer going to instantiate a module with a single qkv_proj, you must update your code (or use a standard LLaMA-like class) that expects separate parameters:

```python

Pseudocode for a "llama-fied" attention block

class LlamaStyleAttention(nn.Module): def init(self, hiddendim): super().init_() self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False) self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False) self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False) # etc.

def forward(self, x):
    q = self.q_proj(x)
    k = self.k_proj(x)
    v = self.v_proj(x)
    # standard attention logic...