Function convert_back

gpu/convert_safetensors.py:49–93

(
    safetensors_path: str,
    output_file: str,
    model_name: Optional[str] = None,
)

    return rearrange(w, '(h l d) i -> (h d l) i', h=config.n_local_heads, l=2)

def convert_back(
    safetensors_path: str,
    output_file: str,
    model_name: Optional[str] = None,
):
    st_dict = load_file(safetensors_path)

    cfg = ModelArgs.from_name(model_name)
    print(f"Using model configurations: {cfg}")

    recovered: dict = {}

    for layer in range(cfg.n_layer):
        base = f"model.layers.{layer}."

        wq = st_dict[f"{base}self_attn.q_proj.weight"]
        wk = st_dict[f"{base}self_attn.k_proj.weight"]
        wv = st_dict[f"{base}self_attn.v_proj.weight"]

        wq = invert_convert_q(wq, cfg)
        wk = invert_convert_k(wk, cfg)

        wqkv = torch.cat([wq, wk, wv], dim=0)
        recovered[f"layers.{layer}.attention.wqkv.weight"] = wqkv

        recovered[f"layers.{layer}.attention.wo.weight"] = st_dict[f"{base}self_attn.o_proj.weight"]

        recovered[f"layers.{layer}.attention_norm.weight"] = st_dict[f"{base}input_layernorm.weight"]
        recovered[f"layers.{layer}.ffn_norm.weight"] = st_dict[f"{base}post_attention_layernorm.weight"]
        recovered[f"layers.{layer}.attention.attn_sub_norm.weight"] = st_dict[f"{base}self_attn.attn_sub_norm.weight"]
        recovered[f"layers.{layer}.feed_forward.ffn_sub_norm.weight"] = st_dict[f"{base}mlp.ffn_sub_norm.weight"]

        gate = st_dict[f"{base}mlp.gate_proj.weight"]
        up   = st_dict[f"{base}mlp.up_proj.weight"]
        w13  = torch.cat([gate, up], dim=0)
        recovered[f"layers.{layer}.feed_forward.w13.weight"] = w13

        recovered[f"layers.{layer}.feed_forward.w2.weight"] = st_dict[f"{base}mlp.down_proj.weight"]

    recovered["tok_embeddings.weight"] = st_dict["model.embed_tokens.weight"]
    recovered["output.weight"]         = st_dict["model.embed_tokens.weight"]
    recovered["norm.weight"]           = st_dict["model.norm.weight"]

    print(f"Saving to {output_file}")
    torch.save(recovered, output_file)

if __name__ == "__main__":
    import argparse