Skip to content

Commit

Permalink
Activation based merging - copied over from wip-zipit branch
Browse files Browse the repository at this point in the history
  • Loading branch information
@metric-space
metric-space committed Jul 10, 2024
1 parent 3eb146f commit 377d6f6
Show file tree
Hide file tree
Showing 7 changed files with 976 additions and 101 deletions.
64 changes: 50 additions & 14 deletions mergekit/_data/architectures/gpt2.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@
"pre_weights": [
{
"name": "wte.weight",
"is_embed": true
"is_embed": true,
"output_space": "h_0"
},
{
"name": "wpe.weight"
"name": "wpe.weight",
"output_space": "h_0"
}
],
"post_weights": [
Expand All @@ -24,40 +26,74 @@
"layer_templates": {
"weights": [
{
"name": "h.${layer_index}.attn.c_attn.weight"
"name": "h.${layer_index}.ln_1.weight",
"input_space": "h_${layer_index}"
},
{
"name": "h.${layer_index}.attn.c_attn.bias"
"name": "h.${layer_index}.ln_1.bias",
"input_space": "h_${layer_index}"
},
{
"name": "h.${layer_index}.attn.c_proj.weight"
"name": "h.${layer_index}.attn.c_attn.weight",
"input_space": "h_${layer_index}",
"output_space": "qkv_${layer_index}"
},
{
"name": "h.${layer_index}.attn.c_proj.bias"
"name": "h.${layer_index}.attn.c_attn.bias",
"input_space": "qkv_${layer_index}"
},
{
"name": "h.${layer_index}.ln_1.weight"
"name": "h.${layer_index}.attn.c_proj.weight",
"input_space": "qkv_${layer_index}",
"output_space": "attn_out_${layer_index}"
},
{
"name": "h.${layer_index}.ln_1.bias"
"name": "h.${layer_index}.attn.c_proj.bias",
"input_space": "attn_out_${layer_index}"
},
{
"name": "h.${layer_index}.ln_2.weight"
"name": "h.${layer_index}.ln_2.weight",
"input_space": "post_attn_${layer_index}"
},
{
"name": "h.${layer_index}.ln_2.bias"
"name": "h.${layer_index}.ln_2.bias",
"input_space": "post_attn_${layer_index}"
},
{
"name": "h.${layer_index}.mlp.c_proj.weight"
"name": "h.${layer_index}.mlp.c_fc.weight",
"input_space": "post_attn_${layer_index}",
"output_space": "mlp_${layer_index}"
},
{
"name": "h.${layer_index}.mlp.c_proj.bias"
"name": "h.${layer_index}.mlp.c_fc.bias",
"input_space": "mlp_${layer_index}"
},
{
"name": "h.${layer_index}.mlp.c_fc.weight"
"name": "h.${layer_index}.mlp.c_proj.weight",
"input_space": "mlp_${layer_index}",
"output_space": "post_mlp_${layer_index}"
},
{
"name": "h.${layer_index}.mlp.c_fc.bias"
"name": "h.${layer_index}.mlp.c_proj.bias",
"input_space": "post_mlp_${layer_index}"
}
],
"procedural_spaces": [
{
"name": "post_attn_${layer_index}",
"type": "residual",
"inputs": [
"h_${layer_index}",
"attn_out_${layer_index}"
]
},
{
"name": "h_${layer_index+1}",
"type": "residual",
"inputs": [
"post_mlp_${layer_index}",
"post_attn_${layer_index}"
]
}
]
}
Expand Down
72 changes: 38 additions & 34 deletions mergekit/_data/architectures/llama.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,84 +8,88 @@
{
"name": "model.embed_tokens.weight",
"is_embed": true,
"output_space": "h_0"
"output_space": "running_residual"
}
],
"num_layers_config_key": "num_hidden_layers",
"layer_templates": {
"weights": [
{
"name": "model.layers.${layer_index}.input_layernorm.weight",
"input_space": "h_${layer_index}"
"input_space": "running_residual",
"is_vector": true
},
{
"name": "model.layers.${layer_index}.self_attn.q_proj.weight",
"input_space": "h_${layer_index}",
"output_space": "attn_qk_${layer_index}"
"input_space": "running_residual",
"output_space": "attn_qk_${layer_index}",
"head_split": "output",
"is_kq": true
},
{
"name": "model.layers.${layer_index}.self_attn.k_proj.weight",
"input_space": "h_${layer_index}",
"output_space": "attn_qk_${layer_index}"
"input_space": "running_residual",
"output_space": "attn_qk_${layer_index}",
"head_split": "output",
"is_kq": true
},
{
"name": "model.layers.${layer_index}.self_attn.v_proj.weight",
"input_space": "h_${layer_index}",
"output_space": "attn_v_${layer_index}"
"input_space": "running_residual",
"output_space": "attn_v_${layer_index}",
"head_split": "output"
},
{
"name": "model.layers.${layer_index}.self_attn.o_proj.weight",
"input_space": "attn_v_${layer_index}",
"output_space": "post_attn_${layer_index}"
"output_space": "running_residual",
"head_split": "input"
},
{
"name": "model.layers.${layer_index}.post_attention_layernorm.weight",
"input_space": "h_a_${layer_index}"
"input_space": "running_residual",
"is_vector": true
},
{
"name": "model.layers.${layer_index}.mlp.up_proj.weight",
"input_space": "h_a_${layer_index}",
"input_space": "running_residual",
"output_space": "up_${layer_index}"
},
{
"name": "model.layers.${layer_index}.mlp.gate_proj.weight",
"input_space": "h_a_${layer_index}",
"input_space": "running_residual",
"output_space": "up_${layer_index}"
},
{
"name": "model.layers.${layer_index}.mlp.down_proj.weight",
"input_space": "up_${layer_index}",
"output_space": "post_mlp_${layer_index}"
}
],
"procedural_spaces": [
{
"name": "h_a_${layer_index}",
"type": "residual",
"inputs": [
"h_${layer_index}",
"post_attn_${layer_index}"
]
},
{
"name": "h_${layer_index+1}",
"type": "residual",
"inputs": [
"h_a_${layer_index}",
"post_mlp_${layer_index}"
]
"output_space": "running_residual"
}
]
},
"post_weights": [
{
"name": "model.norm.weight",
"input_space": "h_${num_layers}"
"input_space": "running_residual",
"is_vector": true
},
{
"name": "lm_head.weight",
"input_space": "h_${num_layers}",
"is_embed": true
"input_space": "running_residual",
"is_lm_head": true,
"aliases": [
"model.lm_head.weight"
]
}
],
"procedural_spaces": [
{
"type": "kv_expand",
"name":"attn_qkv_${layer_index}",
"inputs": [
"attn_qk_${layer_index}",
"attn_v_${layer_index}"
]
}
]
}
Loading

0 comments on commit 377d6f6

Please sign in to comment.