wehub-resource-sync
584 lines
27 KiB
Python
584 lines
27 KiB
Python
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
|
|
|
|
import gc
|
|
from collections import defaultdict
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from pprint import pprint
|
|
|
|
import torch
|
|
from lightning.fabric.utilities.load import _NotYetLoadedTensor as NotYetLoadedTensor
|
|
|
|
from litgpt import Config
|
|
from litgpt.scripts.convert_hf_checkpoint import layer_template, load_param
|
|
from litgpt.utils import extend_checkpoint_dir, incremental_save, lazy_load
|
|
|
|
|
|
def copy_weights_falcon(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "transformer.word_embeddings.weight",
|
|
"transformer.h.{}.attn.qkv.weight": "transformer.h.{}.self_attention.query_key_value.weight",
|
|
"transformer.h.{}.attn.proj.weight": "transformer.h.{}.self_attention.dense.weight",
|
|
"transformer.h.{}.mlp.fc.weight": "transformer.h.{}.mlp.dense_h_to_4h.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "transformer.h.{}.mlp.dense_4h_to_h.weight",
|
|
"transformer.ln_f.bias": "transformer.ln_f.bias",
|
|
"transformer.ln_f.weight": "transformer.ln_f.weight",
|
|
"lm_head.weight": "lm_head.weight",
|
|
}
|
|
# the original model definition is different for each size
|
|
if "7b" in config.name:
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.norm_1.bias": "transformer.h.{}.input_layernorm.bias",
|
|
"transformer.h.{}.norm_1.weight": "transformer.h.{}.input_layernorm.weight",
|
|
}
|
|
)
|
|
elif "40b" in config.name or "180B" in config.name:
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.norm_1.bias": "transformer.h.{}.ln_attn.bias",
|
|
"transformer.h.{}.norm_1.weight": "transformer.h.{}.ln_attn.weight",
|
|
"transformer.h.{}.norm_2.bias": "transformer.h.{}.ln_mlp.bias",
|
|
"transformer.h.{}.norm_2.weight": "transformer.h.{}.ln_mlp.weight",
|
|
}
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
for from_name, param in lit_weights.items():
|
|
name_template, layer_idx = layer_template(from_name)
|
|
to_name = weight_map[name_template].format(layer_idx)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith((".attn.qkv.weight", ".attn.qkv.bias")):
|
|
# Reassemble [q, q, ..., k, k, ..., v, v, ...] --> [q, k, v, q, k, v, ...]
|
|
param = qkv_reassemble(param, config)
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def copy_weights_gpt_neox(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "gpt_neox.embed_in.weight",
|
|
"transformer.h.{}.norm_1.bias": "gpt_neox.layers.{}.input_layernorm.bias",
|
|
"transformer.h.{}.norm_1.weight": "gpt_neox.layers.{}.input_layernorm.weight",
|
|
"transformer.h.{}.attn.qkv.bias": "gpt_neox.layers.{}.attention.query_key_value.bias",
|
|
"transformer.h.{}.attn.qkv.weight": "gpt_neox.layers.{}.attention.query_key_value.weight",
|
|
"transformer.h.{}.attn.proj.bias": "gpt_neox.layers.{}.attention.dense.bias",
|
|
"transformer.h.{}.attn.proj.weight": "gpt_neox.layers.{}.attention.dense.weight",
|
|
"transformer.h.{}.norm_2.bias": "gpt_neox.layers.{}.post_attention_layernorm.bias",
|
|
"transformer.h.{}.norm_2.weight": "gpt_neox.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.mlp.fc.bias": "gpt_neox.layers.{}.mlp.dense_h_to_4h.bias",
|
|
"transformer.h.{}.mlp.fc.weight": "gpt_neox.layers.{}.mlp.dense_h_to_4h.weight",
|
|
"transformer.h.{}.mlp.proj.bias": "gpt_neox.layers.{}.mlp.dense_4h_to_h.bias",
|
|
"transformer.h.{}.mlp.proj.weight": "gpt_neox.layers.{}.mlp.dense_4h_to_h.weight",
|
|
"transformer.ln_f.bias": "gpt_neox.final_layer_norm.bias",
|
|
"transformer.ln_f.weight": "gpt_neox.final_layer_norm.weight",
|
|
"lm_head.weight": "embed_out.weight",
|
|
}
|
|
|
|
for from_name, param in lit_weights.items():
|
|
name_template, layer_idx = layer_template(from_name)
|
|
to_name = weight_map[name_template].format(layer_idx)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith((".attn.qkv.weight", ".attn.qkv.bias")):
|
|
# Reassemble [q, q, ..., k, k, ..., v, v, ...] --> [q, k, v, q, k, v, ...]
|
|
param = qkv_reassemble(param, config)
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def copy_weights_llama(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
untie_weights: bool = False,
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "model.embed_tokens.weight",
|
|
"transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
|
|
"transformer.h.{}.norm_1.bias": "model.layers.{}.input_layernorm.bias",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
|
|
"transformer.h.{}.norm_2.weight": "model.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.norm_2.bias": "model.layers.{}.post_attention_layernorm.bias",
|
|
"transformer.ln_f.weight": "model.norm.weight",
|
|
"transformer.ln_f.bias": "model.norm.bias",
|
|
"lm_head.weight": "lm_head.weight",
|
|
}
|
|
if config.mlp_class_name == "LLaMAMoE":
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.mlp.gate.weight": "model.layers.{}.block_sparse_moe.gate.weight",
|
|
"transformer.h.{}.mlp.experts.{}.fc_1.weight": "model.layers.{}.block_sparse_moe.experts.{}.w1.weight",
|
|
"transformer.h.{}.mlp.experts.{}.fc_2.weight": "model.layers.{}.block_sparse_moe.experts.{}.w3.weight",
|
|
"transformer.h.{}.mlp.experts.{}.proj.weight": "model.layers.{}.block_sparse_moe.experts.{}.w2.weight",
|
|
}
|
|
)
|
|
elif config.mlp_class_name in ("LLaMAMLP", "GemmaMLP"):
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight",
|
|
"transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight",
|
|
}
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
for from_name, param in lit_weights.items():
|
|
if from_name == "lm_head.weight" and untie_weights:
|
|
continue
|
|
name_template, *ids = layer_template(from_name, num_matches=2)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith(".attn.qkv.weight"):
|
|
to_names = (
|
|
"model.layers.{}.self_attn.q_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.k_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.v_proj.weight".format(*ids),
|
|
)
|
|
params = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
else:
|
|
to_names = (weight_map[name_template].format(*ids),)
|
|
params = (param,)
|
|
|
|
for to_name, param in zip(to_names, params):
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def copy_weights_gemma_2(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
untie_weights: bool = True,
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "model.embed_tokens.weight",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
|
|
"transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight",
|
|
"transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight",
|
|
"transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
|
|
"transformer.h.{}.post_attention_norm.weight": "model.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.norm_2.weight": "model.layers.{}.pre_feedforward_layernorm.weight",
|
|
"transformer.h.{}.post_mlp_norm.weight": "model.layers.{}.post_feedforward_layernorm.weight",
|
|
"transformer.ln_f.weight": "model.norm.weight",
|
|
"lm_head.weight": "lm_head.weight",
|
|
}
|
|
|
|
for from_name, param in lit_weights.items():
|
|
if from_name == "lm_head.weight" and untie_weights:
|
|
continue
|
|
name_template, *ids = layer_template(from_name, num_matches=2)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith(".attn.qkv.weight"):
|
|
to_names = (
|
|
"model.layers.{}.self_attn.q_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.k_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.v_proj.weight".format(*ids),
|
|
)
|
|
params = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
else:
|
|
to_names = (weight_map[name_template].format(*ids),)
|
|
params = (param,)
|
|
|
|
for to_name, param in zip(to_names, params):
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def copy_weights_gemma_3(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
untie_weights: bool = True,
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "model.embed_tokens.weight",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
|
|
"transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight",
|
|
"transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight",
|
|
"transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
|
|
"transformer.h.{}.post_attention_norm.weight": "model.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.norm_2.weight": "model.layers.{}.pre_feedforward_layernorm.weight",
|
|
"transformer.h.{}.post_mlp_norm.weight": "model.layers.{}.post_feedforward_layernorm.weight",
|
|
"transformer.ln_f.weight": "model.norm.weight",
|
|
"lm_head.weight": "lm_head.weight",
|
|
"transformer.h.{}.attn.norm_q.weight": "model.layers.{}.self_attn.q_norm.weight",
|
|
"transformer.h.{}.attn.norm_k.weight": "model.layers.{}.self_attn.k_norm.weight",
|
|
}
|
|
|
|
for from_name, param in lit_weights.items():
|
|
if from_name == "lm_head.weight" and untie_weights:
|
|
continue
|
|
name_template, *ids = layer_template(from_name, num_matches=2)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith(".attn.qkv.weight"):
|
|
to_names = (
|
|
"model.layers.{}.self_attn.q_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.k_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.v_proj.weight".format(*ids),
|
|
)
|
|
params = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
else:
|
|
to_names = (weight_map[name_template].format(*ids),)
|
|
params = (param,)
|
|
|
|
for to_name, param in zip(to_names, params):
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def copy_weights_phi(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "model.embed_tokens.weight",
|
|
"transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
|
|
"transformer.h.{}.norm_1.bias": "model.layers.{}.input_layernorm.bias",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.dense.weight",
|
|
"transformer.h.{}.attn.proj.bias": "model.layers.{}.self_attn.dense.bias",
|
|
"transformer.h.{}.mlp.fc.weight": "model.layers.{}.mlp.fc1.weight",
|
|
"transformer.h.{}.mlp.fc.bias": "model.layers.{}.mlp.fc1.bias",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.fc2.weight",
|
|
"transformer.h.{}.mlp.proj.bias": "model.layers.{}.mlp.fc2.bias",
|
|
"transformer.ln_f.weight": "model.final_layernorm.weight",
|
|
"transformer.ln_f.bias": "model.final_layernorm.bias",
|
|
"lm_head.weight": "lm_head.weight",
|
|
"lm_head.bias": "lm_head.bias",
|
|
}
|
|
if config.name.lower().startswith(("phi-3", "phi-4")):
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.attn.qkv.weight": "model.layers.{}.self_attn.qkv_proj.weight",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
|
|
"transformer.h.{}.norm_2.weight": "model.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight",
|
|
"transformer.ln_f.weight": "model.norm.weight",
|
|
}
|
|
)
|
|
gate_up_proj_weights = defaultdict(dict)
|
|
|
|
for from_name, param in lit_weights.items():
|
|
if from_name == "lm_head.weight" and config.name.startswith("Phi-4"):
|
|
continue
|
|
name_template, layer_idx = layer_template(from_name)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith((".attn.qkv.weight", ".attn.qkv.bias")):
|
|
if config.name.lower().startswith(("phi-3", "phi-4")):
|
|
to_names = (weight_map[name_template].format(layer_idx),)
|
|
params = (param,)
|
|
else:
|
|
weight_type = from_name.split(".")[-1] # weight or bias
|
|
to_names = (
|
|
f"model.layers.{{}}.self_attn.q_proj.{weight_type}".format(layer_idx),
|
|
f"model.layers.{{}}.self_attn.k_proj.{weight_type}".format(layer_idx),
|
|
f"model.layers.{{}}.self_attn.v_proj.{weight_type}".format(layer_idx),
|
|
)
|
|
params = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
elif from_name.endswith((".fc_1.weight", ".fc_2.weight")):
|
|
weight = load_param(param, from_name, None)
|
|
weight_name = from_name.split(".")[-2]
|
|
gate_up_proj_weights[layer_idx][weight_name] = weight
|
|
else:
|
|
to_names = (weight_map[name_template].format(layer_idx),)
|
|
params = (param,)
|
|
|
|
for to_name, param in zip(to_names, params):
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
if config.name.lower().startswith(("phi-3", "phi-4")):
|
|
for layer_idx in list(gate_up_proj_weights):
|
|
fc_1_weight = gate_up_proj_weights[layer_idx]["fc_1"]
|
|
fc_2_weight = gate_up_proj_weights[layer_idx]["fc_2"]
|
|
weight = torch.concat([fc_1_weight, fc_2_weight], dim=0)
|
|
layer_name = f"model.layers.{layer_idx}.mlp.gate_up_proj.weight"
|
|
state_dict[layer_name] = weight
|
|
del gate_up_proj_weights[layer_idx]
|
|
|
|
|
|
def copy_weights_qwen_2_5(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
untie_weights: bool = False,
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "model.embed_tokens.weight",
|
|
"transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
|
|
"transformer.h.{}.norm_2.weight": "model.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
|
|
"transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight",
|
|
"transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight",
|
|
"transformer.ln_f.weight": "model.norm.weight",
|
|
"lm_head.weight": "lm_head.weight",
|
|
}
|
|
|
|
for from_name, param in lit_weights.items():
|
|
if from_name == "lm_head.weight" and untie_weights:
|
|
continue
|
|
name_template, *ids = layer_template(from_name, num_matches=2)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith((".attn.qkv.weight", ".attn.qkv.bias")):
|
|
weight_type = from_name.split(".")[-1] # weight or bias
|
|
to_names = (
|
|
"model.layers.{}.self_attn.q_proj.{}".format(*ids, weight_type),
|
|
"model.layers.{}.self_attn.k_proj.{}".format(*ids, weight_type),
|
|
"model.layers.{}.self_attn.v_proj.{}".format(*ids, weight_type),
|
|
)
|
|
params = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
else:
|
|
to_names = (weight_map[name_template].format(*ids),)
|
|
params = (param,)
|
|
|
|
for to_name, param in zip(to_names, params):
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def copy_weights_olmo2(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
untie_weights: bool = False,
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "model.embed_tokens.weight",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
|
|
"transformer.h.{}.attn.norm_q.weight": "model.layers.{}.self_attn.q_norm.weight",
|
|
"transformer.h.{}.attn.norm_k.weight": "model.layers.{}.self_attn.k_norm.weight",
|
|
"transformer.h.{}.norm_2.weight": "model.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.norm_2.bias": "model.layers.{}.post_attention_layernorm.bias",
|
|
"transformer.h.{}.post_mlp_norm.weight": "model.layers.{}.post_feedforward_layernorm.weight",
|
|
"transformer.ln_f.weight": "model.norm.weight",
|
|
"transformer.ln_f.bias": "model.norm.bias",
|
|
"lm_head.weight": "lm_head.weight",
|
|
}
|
|
if config.mlp_class_name in ("LLaMAMLP", "GemmaMLP"):
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight",
|
|
"transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight",
|
|
}
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
for from_name, param in lit_weights.items():
|
|
if from_name == "lm_head.weight" and untie_weights:
|
|
continue
|
|
name_template, *ids = layer_template(from_name, num_matches=2)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith(".attn.qkv.weight"):
|
|
to_names = (
|
|
"model.layers.{}.self_attn.q_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.k_proj.weight".format(*ids),
|
|
"model.layers.{}.self_attn.v_proj.weight".format(*ids),
|
|
)
|
|
params = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
else:
|
|
to_names = (weight_map[name_template].format(*ids),)
|
|
params = (param,)
|
|
|
|
for to_name, param in zip(to_names, params):
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def copy_weights_qwen_3(
|
|
config: Config,
|
|
state_dict: dict[str, torch.Tensor],
|
|
lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
|
|
untie_weights: bool = False,
|
|
saver: incremental_save | None = None,
|
|
) -> None:
|
|
weight_map = {
|
|
"transformer.wte.weight": "model.embed_tokens.weight",
|
|
"transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
|
|
"transformer.h.{}.norm_2.weight": "model.layers.{}.post_attention_layernorm.weight",
|
|
"transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
|
|
"transformer.h.{}.attn.norm_q.weight": "model.layers.{}.self_attn.q_norm.weight",
|
|
"transformer.h.{}.attn.norm_k.weight": "model.layers.{}.self_attn.k_norm.weight",
|
|
"transformer.ln_f.weight": "model.norm.weight",
|
|
"lm_head.weight": "lm_head.weight",
|
|
}
|
|
if config.mlp_class_name == "LLaMAMoE":
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.mlp.gate.weight": "model.layers.{}.mlp.gate.weight",
|
|
"transformer.h.{}.mlp.experts.{}.fc_1.weight": "model.layers.{}.mlp.experts.{}.gate_proj.weight",
|
|
"transformer.h.{}.mlp.experts.{}.fc_2.weight": "model.layers.{}.mlp.experts.{}.up_proj.weight",
|
|
"transformer.h.{}.mlp.experts.{}.proj.weight": "model.layers.{}.mlp.experts.{}.down_proj.weight",
|
|
}
|
|
)
|
|
elif config.mlp_class_name == "LLaMAMLP":
|
|
weight_map.update(
|
|
{
|
|
"transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight",
|
|
"transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight",
|
|
"transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight",
|
|
}
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
for from_name, param in lit_weights.items():
|
|
if from_name == "lm_head.weight" and untie_weights:
|
|
continue
|
|
name_template, *ids = layer_template(from_name, num_matches=2)
|
|
param = load_param(param, from_name, None)
|
|
if from_name.endswith(".attn.qkv.weight"):
|
|
weight_type = from_name.split(".")[-1] # weight or bias
|
|
to_names = (
|
|
"model.layers.{}.self_attn.q_proj.{}".format(*ids, weight_type),
|
|
"model.layers.{}.self_attn.k_proj.{}".format(*ids, weight_type),
|
|
"model.layers.{}.self_attn.v_proj.{}".format(*ids, weight_type),
|
|
)
|
|
params = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
else:
|
|
to_names = (weight_map[name_template].format(*ids),)
|
|
params = (param,)
|
|
|
|
for to_name, param in zip(to_names, params):
|
|
if saver is not None:
|
|
param = saver.store_early(param)
|
|
state_dict[to_name] = param
|
|
|
|
|
|
def qkv_reassemble(param: torch.Tensor | NotYetLoadedTensor, config: Config) -> torch.Tensor:
|
|
"""Reassemble from a normal to an interleaved placement in a QKV matrix.
|
|
[Q, Q, ..., K, K, ..., V, V, ...] --> [Q, K, V, Q, K, V, ...]
|
|
"""
|
|
q, k, v = param.split(
|
|
(
|
|
config.n_head * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
config.n_query_groups * config.head_size,
|
|
)
|
|
)
|
|
qs = q.split(config.n_head // config.n_query_groups * config.head_size)
|
|
ks = k.split(config.head_size)
|
|
vs = v.split(config.head_size)
|
|
interleaved = [t for group in zip(qs, ks, vs) for t in group]
|
|
return torch.cat(interleaved)
|
|
|
|
|
|
def check_conversion_supported(lit_weights: dict[str, torch.Tensor]) -> None:
|
|
if any("lora" in wn for wn in lit_weights):
|
|
raise ValueError("Checkpoints with LoRA weights cannot be converted. Call `scripts/merge_lora.py` first.")
|
|
if any("adapter" in wn or "gating_factor" in wn for wn in lit_weights):
|
|
raise NotImplementedError("Converting adapter models is not supported.")
|
|
|
|
|
|
@torch.inference_mode()
|
|
def convert_lit_checkpoint(checkpoint_dir: Path, output_dir: Path) -> None:
|
|
"""Convert a LitGPT trained checkpoint into a Hugging Face Transformers checkpoint."""
|
|
checkpoint_dir = extend_checkpoint_dir(checkpoint_dir)
|
|
pprint(locals())
|
|
|
|
config = Config.from_file(checkpoint_dir / "model_config.yaml")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
output_path = output_dir / "model.pth"
|
|
|
|
if "falcon" in config.name:
|
|
copy_fn = partial(copy_weights_falcon, config)
|
|
elif config.name.startswith("Gemma-2"):
|
|
copy_fn = partial(copy_weights_gemma_2, config)
|
|
elif config.name.startswith("Gemma-3"):
|
|
copy_fn = partial(copy_weights_gemma_3, config)
|
|
elif config.name.lower().startswith("phi"):
|
|
copy_fn = partial(copy_weights_phi, config)
|
|
elif config.name.lower().startswith(("qwen2.5", "qwq")):
|
|
copy_fn = partial(copy_weights_qwen_2_5, config)
|
|
elif config.name.lower().startswith("olmo-2-"):
|
|
copy_fn = partial(copy_weights_olmo2, config)
|
|
elif config.name.lower().startswith("qwen3"):
|
|
copy_fn = partial(copy_weights_qwen_3, config)
|
|
elif config.mlp_class_name in ("LLaMAMLP", "GemmaMLP", "LLaMAMoE"):
|
|
untie_weights = "Gemma" in config.name
|
|
copy_fn = partial(copy_weights_llama, config, untie_weights=untie_weights)
|
|
else:
|
|
copy_fn = partial(copy_weights_gpt_neox, config)
|
|
|
|
# initialize a new empty state dict to hold our new weights
|
|
sd = {}
|
|
with incremental_save(output_path) as saver:
|
|
lit_weights = lazy_load(checkpoint_dir / "lit_model.pth")
|
|
lit_weights = lit_weights.get("model", lit_weights)
|
|
check_conversion_supported(lit_weights)
|
|
copy_fn(sd, lit_weights, saver=saver)
|
|
gc.collect()
|
|
saver.save(sd)
|