基于昇腾硬件的AI大模型仓库Python设计源码

import abc import os import sys import re import json from types import SimpleNamespace import logging as logger from collections import OrderedDict from tqdm import tqdm import torch import torch.nn as nn from transformers import AutoModelForCausalLM from megatron.core import mpu from megatron.training.arguments import validate_args from megatron.legacy.model import module from megatron.core.enums import ModelType from megatron.training.checkpointing import load_args_from_checkpoint from megatron.training.global_vars import set_args from megatron.training.checkpointing import load_checkpoint from megatron.core import tensor_parallel from modellink.utils import parse_args from modellink.training import model_provider_func_wrapper from modellink.checkpointing import load_checkpoint_wrapper logger.basicConfig(format="") logger.getLogger().setLevel(logger.INFO) load_checkpoint = load_checkpoint_wrapper(load_checkpoint) def tensor_info(tensor): shape = tensor.shape mean_val = tensor.mean().item() min_val = tensor.min().item() max_val = tensor.max().item() return f"shape: {shape} mean_val: {mean_val} min_val: {min_val} max_val: {max_val}" class ModelBase(abc.ABC): def __init__(self, args_cmd=None): self.args_cmd = args_cmd self.args = None self.args_megatron_checkpoint = None self.module = None self.module_mapping = None self.model_cfg = self.read_model_cfg() self.__register_functions() self.kwargs_idx = OrderedDict({ "vp_rank": 0, "ep_rank": 0, "tp_rank": 0, "layer_idx": 0, "expert_idx": 0 }) def update_kwargs_idx(self, **kwargs): for key in self.kwargs_idx: if key in kwargs: self.kwargs_idx[key] = kwargs[key] else: self.kwargs_idx[key] = 0 def __register_functions(self): self.get_module_mapping() def _get_obj(self, value, **kwargs): pattern = r'(\w+)(?:\[(\w+)\])?' matches = re.findall(pattern, value) self.update_kwargs_idx(**kwargs) obj = self.get_model_item(**kwargs) for attr, attr_ident in matches: if hasattr(obj, attr): obj = getattr(obj, attr) else: return None if attr_ident: if attr_ident in self.kwargs_idx: attr_idx = self.kwargs_idx[attr_ident] obj = obj[attr_idx] else: raise AssertionError(f"check {self.__class__.__name__}.module_mapping **{attr_ident}**.") return obj def _func_generator_get_module(value): def func(self, **kwargs): return _get_obj(self, value, **kwargs) return func def _func_generator_get_weight(value): def func(self, **kwargs): return _get_obj(self, value, **kwargs).weight.data return func def _func_generator_get_bias(value): def func(self, **kwargs): return _get_obj(self, value, **kwargs).bias.data return func def _func_generator_set_weight(value): def func(self, **kwargs): return _get_obj(self, value, **kwargs).weight.data.copy_(kwargs.get('data')) return func def _func_generator_set_module(value): def func(self, **kwargs): return _get_obj(self, value, **kwargs).data.copy_(kwargs.get('data')) return func def _func_generator_set_bias(value): def func(self, **kwargs): return _get_obj(self, value, **kwargs).bias.data.copy_(kwargs.get('data')) return func def _func_generator_has_module(value): def func(self, **kwargs): obj = _get_obj(self, value, **kwargs) return True if obj else False return func def _func_generator_has_bias(value): def func(self, **kwargs): bias = getattr(_get_obj(self, value, **kwargs), 'bias', None) return bias is not None return func if self.module_mapping: for key, value in self.module_mapping.items(): setattr(self, "get_" + key + "_module", _func_generator_get_module(value).__get__(self, ModelBase)) setattr(self, "set_" + key + "_module", _func_generator_set_module(value).__get__(self, ModelBase)) setattr(self, "get_" + key + "_weight", _func_generator_get_weight(value).__get__(self, ModelBase)) setattr(self, "get_" + key + "_bias", _func_generator_get_bias(value).__get__(self, ModelBase)) setattr(self, "set_" + key + "_weight", _func_generator_set_weight(value).__get__(self, ModelBase)) setattr(self, "set_" + key + "_bias", _func_generator_set_bias(value).__get__(self, ModelBase)) setattr(self, "has_" + key + "_module", _func_generator_has_module(value).__get__(self, ModelBase)) setattr(self, "has_" + key + "_bias", _func_generator_has_bias(value).__get__(self, ModelBase)) def update_module(self, src_model): self.set_preprocess_state(src_model) self.set_postprocess_state(src_model) for layer_idx in tqdm(range(self.args.num_layers), "set layer states"): self.set_layer_state(src_model, layer_idx) def set_preprocess_state(self, src_model): '''Set embedding params.''' embeddings_weight = src_model.get_embedding_word_embeddings_weight() self.set_embedding_word_embeddings_weight(data=embeddings_weight) if src_model.has_embedding_word_embeddings_norm_module(): embd_norm_weight = src_model.get_embedding_word_embeddings_norm_weight() embd_norm_bias = src_model.get_embedding_word_embeddings_norm_bias() self.set_embedding_word_embeddings_norm_weight(data=embd_norm_weight) self.set_embedding_word_embeddings_norm_bias(data=embd_norm_bias) def set_postprocess_state(self, src_model): final_layernorm_weight = src_model.get_final_layernorm_weight() self.set_final_layernorm_weight(data=final_layernorm_weight) if self.args.untie_embeddings_and_output_weights: output_layer_weight = src_model.get_output_layer_weight() self.set_output_layer_weight(data=output_layer_weight) if self.has_final_layernorm_bias(): final_layernorm_bias = src_model.get_final_layernorm_bias() self.set_final_layernorm_bias(data=final_layernorm_bias) def set_layer_state(self, src_model, layer_idx): self.set_attn_state(layer_idx, src_model) self.set_mlp_state(layer_idx, src_model) input_layernorm_weight = src_model.get_layers_input_layernorm_weight(layer_idx=layer_idx) self.set_layers_input_layernorm_weight(layer_idx=layer_idx, data=input_layernorm_weight) if self.args.post_norm: post_attn_layernorm_weight = src_model.get_layers_self_attention_post_attention_layernorm_weight( layer_idx=layer_idx) self.set_layers_self_attention_post_attention_layernorm_weight(layer_idx=layer_idx, data=post_attn_layernorm_weight) else: pre_mlp_layernorm_weight = src_model.get_layers_self_attention_pre_mlp_layernorm_weight(layer_idx=layer_idx) self.set_layers_self_attention_pre_mlp_layernorm_weight(layer_idx=layer_idx, data=pre_mlp_layernorm_weight) if self.has_layers_input_layernorm_bias(layer_idx=layer_idx):