maga_transformer/model_loader/model_weight_info.py

import functools
import gc
import logging
import os
import torch
from maga_transformer.utils.ckpt_file_info import CkptFileInfo
from typing import List, Tuple, Union, Optional, Dict, Any
from maga_transformer.utils.database import BaseDatabase, CkptDatabase
from maga_transformer.utils.model_weight import W, CkptWeightInfo, WeightStyle, choose_available, identity, tolerate_failed
from maga_transformer.config.gpt_init_model_parameters import GptInitModelParameters
from maga_transformer.model_loader.load_config import LoadConfig
from maga_transformer.model_loader.weight_module import WeightModule, AtomicWeight, CompositeWeight
from maga_transformer.model_loader.ffn_weight import FfnConfig, FfnWeight, MoeWithSharedWeight
from maga_transformer.model_loader.attn_weight import AttnConfig

class ModelWeightInfo:
    layer_weights: Union[List[WeightModule], List[List[WeightModule]]]
    weights: List[WeightModule]

    def __init__(self, weights: List[WeightModule],
                 layer_weights: Union[List[WeightModule], List[List[WeightModule]]]) -> None:
        self.weights = weights
        self.layer_weights = layer_weights
        if len(self.layer_weights) == 0:
            return


    def set_weight_dtype(self, dtype: torch.dtype):
        if self.layer_weights:
            for weight in self.layer_weights:
                weight.data_type = dtype

    def get_layer_weight_info(self, layer_id: int, name: str) -> Optional[WeightModule]:
        from collections import deque
        queue = deque(self.layer_weights[layer_id])

        while queue:
            weight = queue.popleft()
            if weight.name == name:
                return weight
            if isinstance(weight, CompositeWeight):
                queue.extend(weight.sub_weights)

        return None


    def to_quant_weight_info(self, quant_algo: Any):
        if quant_algo is None or not quant_algo.isQuant():
            raise ValueError("quant_algo is None or not quant_algo.isQuant()")
        weights = []
        if self.weights:
            for weight in self.weights:
                weights.append(weight.create(weight, quant_algo))
        layer_weights: Union[List[WeightModule], List[List[WeightModule]]] = [] if self.layer_weights else None
        if self.layer_weights:
            for weight in self.layer_weights:
                if isinstance(weight, list):
                    layer_weight = []
                    for w in weight:
                        layer_weight.append(w.create(w, quant_algo))
                    layer_weights.append(layer_weight)
                else:
                    layer_weights.append(weight.create(weight, quant_algo))

        return ModelWeightInfo(weights, layer_weights)



class ModelDeployWeightInfo:

    TRT_ENGINE_LAYER_WEIGHT_MAP = {
        W.pre_ln_beta : 'transformer.layers.{i}.input_layernorm.bias',
        W.pre_ln_gamma : 'transformer.layers.{i}.input_layernorm.weight',
        W.attn_qkv_w : 'transformer.layers.{i}.attention.qkv.weight',
        W.attn_qkv_b : 'transformer.layers.{i}.attention.qkv.bias',
        W.attn_qkv_s : 'transformer.layers.{i}.attention.qkv.weights_scaling_factor',

        W.attn_o_w : 'transformer.layers.{i}.attention.dense.weight',
        W.attn_o_b : 'transformer.layers.{i}.attention.dense.bias',
        W.attn_o_s : 'transformer.layers.{i}.attention.dense.weights_scaling_factor',

        W.ffn_w3 : 'transformer.layers.{i}.mlp.fc.weight',
        W.ffn_b3 : 'transformer.layers.{i}.mlp.fc.bias',
        W.ffn_s3 : 'transformer.layers.{i}.mlp.fc.weights_scaling_factor',

        W.ffn_w2 : 'transformer.layers.{i}.mlp.proj.weight',
        W.ffn_b2 : 'transformer.layers.{i}.mlp.proj.bias',
        W.ffn_s2 : 'transformer.layers.{i}.mlp.proj.weights_scaling_factor',

        W.post_ln_gamma : 'transformer.layers.{i}.post_layernorm.weight',
        W.post_ln_beta : 'transformer.layers.{i}.post_layernorm.bias',

    }

    TRT_ENGINE_LAYER_WEIGHT_MAP2 = {
        W.pre_ln_beta : 'transformer.layers.{i}.input_layernorm.bias',
        W.pre_ln_gamma : 'transformer.layers.{i}.input_layernorm.weight',
        W.attn_qkv_w : 'transformer.layers.{i}.attention.qkv.weight',
        W.attn_qkv_b : 'transformer.layers.{i}.attention.qkv.bias',
        W.attn_qkv_s : 'transformer.layers.{i}.attention.qkv.weights_scaling_factor',

        W.attn_o_w : 'transformer.layers.{i}.attention.dense.weight',
        W.attn_o_b : 'transformer.layers.{i}.attention.dense.bias',
        W.attn_o_s : 'transformer.layers.{i}.attention.dense.weights_scaling_factor',

        W.ffn_w1 : 'transformer.layers.{i}.mlp.fc.weight',
        W.ffn_b1 : 'transformer.layers.{i}.mlp.fc.bias',
        W.ffn_s1 : 'transformer.layers.{i}.mlp.fc.weights_scaling_factor',

        W.ffn_w2 : 'transformer.layers.{i}.mlp.proj.weight',
        W.ffn_b2 : 'transformer.layers.{i}.mlp.proj.bias',
        W.ffn_s2 : 'transformer.layers.{i}.mlp.proj.weights_scaling_factor',

        W.ffn_w3 : 'transformer.layers.{i}.mlp.gate.weight',
        W.ffn_b3 : 'transformer.layers.{i}.mlp.gate.bias',
        W.ffn_s3 : 'transformer.layers.{i}.mlp.gate.weights_scaling_factor',

        W.post_ln_gamma : 'transformer.layers.{i}.post_layernorm.weight',
        W.post_ln_beta : 'transformer.layers.{i}.post_layernorm.bias',

    }
    def __init__(self, config: GptInitModelParameters, tp_size: int, tp_rank: int):
        self.config = config
        self._use_qk_norm = config.use_qk_norm
        self._hidden_size = config.hidden_size
        self._inter_size = config.inter_size
        self._inter_padding_size = config.inter_padding_size
        self._moe_inter_padding_size = config.moe_inter_padding_size
        self._head_num = config.head_num
        self._head_num_kv = config.head_num_kv
        self.tp_size = tp_size
        self.tp_rank = tp_rank
        self.ep_size = config.ep_size
        self.ep_rank = config.ep_rank
        self.dp_size = config.dp_size
        self.dp_rank = config.dp_rank
        self.num_nodes: int = config.num_nodes
        self.ffn_tp_rank = config.ffn_tp_rank
        self.ffn_tp_size = config.ffn_tp_size
        self._size_per_head = config.size_per_head
        self.phy2log_ = config.phy2log
        if self._head_num_kv == -1:
            self._head_num_kv = self._head_num
        self._quant_algo = config.quant_algo
        self._num_layers = config.num_layers
        self._layer_head_num = config.layer_head_num
        self._layer_inter_padding_size = config.layer_inter_padding_size
        self._has_prefix_encoder = False
        self._is_sparse_head = config.is_sparse_head
        self._layer_head_num = config.layer_head_num
        self._src_quantization_bit = config.src_quantization_bit
        self.tp_split_emb_and_lm_head = config.tp_split_emb_and_lm_head


        self._is_gated_activation = config.gpt_init_params.isGatedActivation()
        self.expert_num_ = config.gpt_init_params.expert_num
        self.moe_n_group_ = config.moe_n_group
        self.enable_eplb_ = config.enable_eplb
        self.phy_exp_num_ = config.phy_exp_num
        self.enable_merge_w13_ = config.enable_merge_w13
        self.moe_k_      = config.gpt_init_params.moe_k
        self.moe_layer_index_ = config.gpt_init_params.moe_layer_index
        self.moe_style_ = config.gpt_init_params.moe_style
        self._moe_inter_padding_size = config.moe_inter_padding_size

        self.tie_word_embeddings = config.tie_word_embeddings
        self.need_ffn_act_scale = config.need_ffn_act_scale
        self.use_expert_attention = config.use_expert_attention
        self.weight_style = WeightStyle.RTP_LLM_STYLE if config.is_ft_style_weight else WeightStyle.NONE


        # for mla
        self.kv_lora_rank = config.kv_lora_rank
        self.nope_head_dim = config.nope_head_dim
        self.rope_head_dim = config.rope_head_dim
        self.v_head_dim = config.v_head_dim
        self.routed_scaling_factor = config.routed_scaling_factor
        self.is_ft_style_weight = config.is_ft_style_weight

        # for vit sep
        self.vit_separation = config.vit_separation

        # for eplb
        self.phy2log = config.phy2log
        # for moe
        self._use_stack_weight = False

    @property
    def support_lora(self):
        return False

    @property
    def attn_config(self):
        attn_config = AttnConfig(
            hidden_size=self._hidden_size,
            size_per_head=self._size_per_head,
            head_num=self._head_num,
            head_num_kv=self._head_num_kv
        )
        return attn_config

    @property
    def ffn_config(self):
        ffn_config = FfnConfig(
            is_gated_activation=self._is_gated_activation,
            inter_padding_size=self._inter_padding_size,
            is_moe=False
        )
        return ffn_config
            
    def get_weight_info(self) -> ModelWeightInfo:
        weight_info = self._get_weight_info()
        use_fp32 = os.environ.get("USE_FLOAT32", None) is not None
        if use_fp32:
            weight_info = weight_info.set_weight_dtype(torch.float32)

        if weight_info.layer_weights and not isinstance(weight_info.layer_weights[0], List):
            layer_weights = []
            for _ in range(self._num_layers):
                layer_weights.append(weight_info.layer_weights)
            weight_info.layer_weights = layer_weights

        if self.weight_style != WeightStyle.NONE:
            logging.info("fix weight style")
            weight_info = self._fix_weight_style_layer_weight(weight_info)
            
        if self.enable_merge_w13_:
            logging.info("fix merge_w13")
            weight_info = self._fix_merge_w1_w3(weight_info)
            
        if self._quant_algo is not None and self._quant_algo.isQuant():
            weight_info = weight_info.to_quant_weight_info(self._quant_algo)

        if self.tie_word_embeddings:
            logging.info("fix tie_word_embeddings")
            weight_info = self._fix_tie_lm_head(weight_info)
        if self._is_sparse_head:
            logging.info("Skiping load empty weight for head_num == 0")
            weight_info = self._process_sparse_weight(weight_info)

        return weight_info

    def _fix_weight_style_layer_weight(self, origin_weight_info: ModelWeightInfo):
        global_weights = []
        m1 = {
            W.embedding : 'transformer.vocab_embedding.weight',
            W.lm_head : 'lm_head.weight',
            W.final_ln_gamma : 'transformer.ln_f.weight'
        } if self.weight_style == WeightStyle.TRT_ENGINE else {}

        def __update_weight_style(weight: WeightModule, name_map: Dict[str, str]):
            if isinstance(weight , AtomicWeight):
                weight.weight_style = self.weight_style
                if weight.name in name_map:
                    if len(weight.weights) == 1:
                        weight.weights[0].name = name_map[weight.name]
                    elif weight.name in [W.attn_qkv_b, W.attn_qkv_w]:
                        weight.weights = [CkptWeightInfo(name_map[weight.name])]
                        weight.process_fun = identity
                        logging.error(f"{weight.name} have many weight, maybe cause bug {weight.weights}")
                    elif len(weight.weights) >= 2:
                        raise ValueError(f"{weight.name} should have only one or zero weight, {weight.weights}")
                    logging.info(f"update weight style for {weight.name}: {weight.weights[0].name}")
            elif isinstance(weight, CompositeWeight):
                weight.weight_style = self.weight_style
                for _, sub_weight in weight.sub_weights.items():
                    __update_weight_style(sub_weight, name_map)


        for _, weight in enumerate(origin_weight_info.weights):
            __update_weight_style(weight, m1)
            global_weights.append(weight)
        origin_weight_info.weights = global_weights

        layer_weights = []
        for weights in origin_weight_info.layer_weights:
            ffn_weight = [weight for weight in weights if weight.name == W.ffn]
            assert len(ffn_weight) == 1
            if ffn_weight[0].w1 is not None and self.weight_style == WeightStyle.TRT_ENGINE:
                m2 = self.TRT_ENGINE_LAYER_WEIGHT_MAP2
            elif self.weight_style == WeightStyle.TRT_ENGINE:
                m2 = self.TRT_ENGINE_LAYER_WEIGHT_MAP
            else:
                m2 = {}

            fix_weight = []
            for weight in weights:
                __update_weight_style(weight, m2)
                fix_weight.append(weight)
            layer_weights.append(fix_weight)

        origin_weight_info.layer_weights = layer_weights
        logging.info(f"fix weight style {origin_weight_info.layer_weights[0]}")
        return origin_weight_info

    def _fix_merge_w1_w3(self, origin_weight_info: ModelWeightInfo):
        def __update_weight_config(weight: WeightModule):
            if isinstance(weight , FfnWeight) or isinstance(weight, MoeWithSharedWeight):
                weight.config.enable_merge_w13 = True
                params = weight.extract_params(weight.__class__, weight, None)
                return weight.__class__(**params)
            else:
                return weight

        layer_weights = []
        for weights in origin_weight_info.layer_weights:
            fix_weight = []
            for weight in weights:
                fix_weight.append(__update_weight_config(weight))
            layer_weights.append(fix_weight)

        origin_weight_info.layer_weights = layer_weights
        logging.info(f"fix weight config when need_merge_w13 {origin_weight_info.layer_weights[0]}")
        return origin_weight_info

    def _fix_tie_lm_head(self, origin_weight_info: ModelWeightInfo) -> ModelWeightInfo:
        word_emb_idx = -1
        word_emb = None
        lm_head_idx = -1
        lm_head = None
        for idx, weight in enumerate(origin_weight_info.weights):
            if weight.name == W.embedding:
                word_emb_idx = idx
                word_emb = weight
            elif weight.name == W.lm_head:
                lm_head = weight
                lm_head_idx = idx
        if not lm_head or not word_emb:
            return origin_weight_info

        assert len(lm_head.weights) == 1 and len(word_emb.weights) == 1
        lm_head_ckpt_weigth_infos = [CkptWeightInfo(w.name, functools.partial(tolerate_failed, origin_func=w.merge_fun)) for w in lm_head.weights]
        lm_head_ckpt_weigth_infos.extend([CkptWeightInfo(w.name, functools.partial(tolerate_failed, origin_func=w.merge_fun)) for w in word_emb.weights])
        lm_head_merge_funcs = [lm_head.process_fun, word_emb.process_fun]
        lm_head = AtomicWeight(W.lm_head, lm_head_ckpt_weigth_infos, functools.partial(choose_available, origin_func_list = lm_head_merge_funcs))
        origin_weight_info.weights[lm_head_idx] = lm_head
        return origin_weight_info


    def _process_sparse_weight(self, origin_weight_info: ModelWeightInfo) -> ModelWeightInfo:
        if not isinstance(origin_weight_info.layer_weights[0], list):
            raise Exception("model weight use sparse config should be list(list())")
        new_layer_weights = []
        for i, layer_weight in enumerate(origin_weight_info.layer_weights):
            if self._layer_head_num[i] == 0:
                new_weights = [weight for weight in layer_weight if weight.name not in W.skip_weights_list]
            else:
                new_weights = layer_weight
            new_layer_weights.append(new_weights)
        return ModelWeightInfo(origin_weight_info.weights, new_layer_weights)

    def _get_weight_info(self) -> ModelWeightInfo:
        raise NotImplementedError()

    def create_model_weight_info(self, database: BaseDatabase) -> ModelWeightInfo:
        if isinstance(database, CkptDatabase):
            self.process_meta_from_ckpt(database.PretrainFileList)
            self.process_meta_from_ckpt(database.FinetuneFileList)
            if not self.is_ft_style_weight:
                return self.get_weight_info()
        else:
            raise Exception("Unknown database class")

    def process_meta_from_ckpt(self, ckpt_metas: List[CkptFileInfo]):
        if len(ckpt_metas) == 0:
            return
        if not self.is_ft_style_weight:
            # call subclass process_meta
            meta_dicts = [ckpt_file.get_metadata() for ckpt_file in ckpt_metas]
            weight_keys = set(functools.reduce(lambda x,y:x+y, [list(meta.keys()) for meta in meta_dicts], []))
            self._process_meta(meta_dicts, weight_keys)

    def _process_meta(self, meta_dict, weight_keys):
        pass

    def _get_layer_start_end_id(self) -> Tuple[int, int]:
        raise NotImplementedError()

    @staticmethod
    def _contains(keys: List[str], val: str):
        for key in keys:
            if val in key:
                return True
        return False

    def create_load_config(self,
                           compute_dtype: torch.dtype,
                           database: BaseDatabase,
                           exported_device: Optional[Any] = None):
        merge_lora = False
        if not self.is_ft_style_weight:
            merge_lora = database.has_lora() and bool(os.environ.get("MERGE_LORA", 1))

        if database.has_lora() and not self.support_lora:
            raise Exception(f"current weights_info: {self.__class__} not support lora, but database has lora")

        load_config = LoadConfig(
            database = database,
            num_layers = self._num_layers,
            hidden_size = self._hidden_size,
            head_num = self._head_num,
            head_num_kv = self._head_num_kv,
            size_per_head = self._size_per_head,
            use_stack_weight = self._use_stack_weight,
            need_ffn_act_scale = self.need_ffn_act_scale,
            inter_size = self._inter_size,
            moe_layer_index = self.moe_layer_index_,
            moe_n_group = self.moe_n_group_,
            inter_padding_size = self._inter_padding_size,
            moe_inter_padding_size = self._moe_inter_padding_size,
            expert_num = self.expert_num_,
            enable_eplb = self.enable_eplb_,
            phy_exp_num = self.phy_exp_num_,
            enable_merge_w13 = self.enable_merge_w13_,
            tp_size = self.tp_size,
            tp_rank = self.tp_rank,
            ep_size = self.ep_size,
            ep_rank = self.ep_rank,
            dp_size = self.dp_size,
            dp_rank = self.dp_rank,
            num_nodes = self.num_nodes,
            ffn_tp_rank = self.ffn_tp_rank,
            ffn_tp_size = self.ffn_tp_size,
            tp_split_emb_and_lm_head = self.tp_split_emb_and_lm_head,
            merge_lora = merge_lora,
            vit_separation = self.vit_separation,
            compute_dtype = compute_dtype,
            quant_algo = self._quant_algo,
            bit = self._quant_algo.getWeightBits(),
            is_ft_style_weight = self.is_ft_style_weight,
            phy2log=self.phy2log_,
            exported_device = exported_device
        )
        return load_config


class ModelWeights:
    def __init__(self, num_layers: int, device: str, dtype: torch.dtype):
        self.device = device
        self.weights: List[Dict[str, torch.Tensor]] = []
        self.global_weights: Dict[str, torch.Tensor] = {}
        self._dtype = dtype
        self.is_ft_style_weight: bool = False

        for _ in range(num_layers):
            self.weights.append({})

    def set_layer_weight(self, layer_id: int, name: str, tensor: torch.Tensor):
        self.weights[layer_id][name] = tensor
        gc.collect()

    def set_global_weight(self, name: str, tensor: torch.Tensor):
        self.global_weights[name] = tensor

    def get_global_weight(self, name: str):
        return self.global_weights.get(name, None)

    def steal_global_weight(self, name: str):
        if name not in self.global_weights:
            return None
        tensor = self.global_weights[name]
        del self.global_weights[name]
        return tensor

    @property
    def dtype(self):
        return self._dtype


    @staticmethod
    def layer_weight_prefix(tp_rank:int, dp_rank: int, ep_rank: int):
        return f"rank_{tp_rank:02d}_{dp_rank:02d}_{ep_rank:02d}.layers."

    @staticmethod
    def global_weight_prefix(tp_rank:int, dp_rank: int, ep_rank: int):
        return f"rank_{tp_rank:02d}_{dp_rank:02d}_{ep_rank:02d}.global."