from __future__ import annotations import os import pathlib import typing from pathlib import Path from typing import Optional import numpy as np import torch import torch.distributed as dist import torch.nn as nn str_type_map = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16} class BaseBelleWeights: def __init__(self, head_num, size_per_head, layer_num, vocab_size, max_seq_len, tensor_para_size, pipeline_para_size, weights_data_type: typing.Union[str, np.dtype], inference_data_type: str, has_adapters: bool = False, adapter_inter_size: int = 0, gpt_with_moe: bool = False, has_positional_encoding: bool = True, has_pre_decoder_layernorm: bool = False, has_post_decoder_layernorm: bool = True, int8_mode: int = 0, inter_size: int = 0): assert(head_num % tensor_para_size == 0) if int8_mode == 1: torch_infer_dtype = str_type_map[inference_data_type] assert torch_infer_dtype == torch.float16 or torch_infer_dtype == torch.bfloat16, "Weight only quant only supported for infer type fp16 or bf16." quant = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix self.weight_transpose_calibrate_quantize = lambda x: quant( x, torch.int8) else: assert int8_mode == 0, "Invalid int8 mode for BELLE. Must be 0 or 1" self.head_num = head_num self.size_per_head = size_per_head self.layer_num = layer_num self.vocab_size = vocab_size self.max_seq_len = max_seq_len self.tensor_para_size = tensor_para_size self.pipeline_para_size = pipeline_para_size self.layers_per_device = layer_num // pipeline_para_size self.has_adapters = has_adapters self.adapter_inter_size = adapter_inter_size self.gpt_with_moe = gpt_with_moe self.has_positional_encoding = has_positional_encoding self.has_pre_decoder_layernorm = has_pre_decoder_layernorm self.has_post_decoder_layernorm = has_post_decoder_layernorm local_head_num = head_num // tensor_para_size global_head_num = head_num local_hidden_units = local_head_num * size_per_head global_hidden_units = global_head_num * size_per_head local_inter_size = local_hidden_units * 4 if inter_size != 0: assert inter_size % tensor_para_size == 0, f"inter_size({inter_size}) \% tensor_para_size({tensor_para_size}) must be 0" local_inter_size = inter_size // tensor_para_size local_adapter_inter_size = self.adapter_inter_size // tensor_para_size self.local_head_num = local_head_num self.global_head_num = global_head_num self.local_hidden_units = local_hidden_units self.global_hidden_units = global_hidden_units self.local_inter_size = local_inter_size self.int8_mode = int8_mode self.share_embed = False if isinstance(weights_data_type, str): try: weights_data_type = { "fp16": np.float16, "fp32": np.float32, "float16": np.float16, "float32": np.float32, }[weights_data_type] except KeyError: raise ValueError( f"Don't know how to interpret weights_data_type: {weights_data_type}") assert weights_data_type in [np.float32, np.float16] self.weights_data_type = weights_data_type self.inference_data_type = inference_data_type self.w = [] self.int8_w = [] self.scale = [] # Transformer blocks self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # self_layernorm_gamma self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # self_layernorm_beta self.w.extend([torch.zeros(global_hidden_units, local_hidden_units * 3, dtype=str_type_map[self.inference_data_type])] * layer_num) # self_kernel self.w.extend([torch.zeros(local_hidden_units * 3, dtype=str_type_map[self.inference_data_type])] * layer_num) # self_bias self.w.extend([torch.zeros(local_hidden_units, global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # self_output_kernel self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # self_output_bias self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # ffn_layernorm_gamma self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # ffn_layernorm_beta self.w.extend([torch.zeros(global_hidden_units, local_inter_size, dtype=str_type_map[self.inference_data_type])] * layer_num) # ffn_kernel1 self.w.extend([torch.zeros(local_inter_size, dtype=str_type_map[ self.inference_data_type])] * layer_num) # ffn_bias1 self.w.extend([torch.zeros(local_inter_size, global_hidden_units, dtype=str_type_map[self.inference_data_type])] * layer_num) # ffn_kernel2 self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # ffn_bias2 optional_adapter_offset = 0 # After Transformer blocks if self.has_pre_decoder_layernorm: self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])) # embedding layernorm gamma self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])) # embedding layernorm beta optional_adapter_offset += 2 if self.has_post_decoder_layernorm: self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])) # final layernorm gamma self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])) # final layernorm beta optional_adapter_offset += 2 if self.has_positional_encoding: self.w.append(torch.zeros(max_seq_len, global_hidden_units, dtype=str_type_map[ self.inference_data_type])) # position_encoding_table optional_adapter_offset += 1 self.pre_embed_idx = len(self.w) self.w.append(torch.zeros(vocab_size, global_hidden_units, dtype=str_type_map[self.inference_data_type])) # embedding_table self.post_embed_idx = len(self.w) self.w.append(torch.zeros(vocab_size, global_hidden_units, dtype=str_type_map[ self.inference_data_type])) # post embedding_kernel self.adapter_offset = 2 + optional_adapter_offset self.w.extend([torch.empty( 0, dtype=str_type_map[self.inference_data_type])] * layer_num) # gating_weight self.adapter_offset += layer_num # adapters if self.has_adapters: self.w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size, dtype=str_type_map[self.inference_data_type])] * layer_num) # adaptor1_kernel1 self.w.extend([torch.zeros(local_adapter_inter_size, dtype=str_type_map[ self.inference_data_type])] * layer_num) # adaptor1_bias1 self.w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units, dtype=str_type_map[self.inference_data_type])] * layer_num) # adaptor1_kernel2 self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # adaptor1_bias2 self.w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size, dtype=str_type_map[self.inference_data_type])] * layer_num) # adaptor2_kernel1 self.w.extend([torch.zeros(local_adapter_inter_size, dtype=str_type_map[ self.inference_data_type])] * layer_num) # adaptor2_bias1 self.w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units, dtype=str_type_map[self.inference_data_type])] * layer_num) # adaptor2_kernel2 self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[ self.inference_data_type])] * layer_num) # adaptor2_bias2 # Initialization self._map(lambda w: torch.nn.init.normal_(w, mean=0., std=1.)) if (self.int8_mode != 0): self.int8_w.extend([torch.zeros(global_hidden_units, local_hidden_units * 3, dtype=torch.int8)] * layer_num) # self_int8_kernel self.scale.extend([torch.zeros( local_hidden_units * 3, dtype=torch.float)] * layer_num) # self_scale self.int8_w.extend([torch.zeros(local_hidden_units, global_hidden_units, dtype=torch.int8)] * layer_num) # self_output_int8_kernel # self_output_scale self.scale.extend( [torch.zeros(global_hidden_units, dtype=torch.float)] * layer_num) self.int8_w.extend([torch.zeros(global_hidden_units, local_inter_size, dtype=torch.int8)] * layer_num) # ffn_int8_kernel1 self.scale.extend( [torch.zeros(local_inter_size, dtype=torch.float)] * layer_num) # ffn_scale1 self.int8_w.extend([torch.zeros(local_inter_size, global_hidden_units, dtype=torch.int8)] * layer_num) # ffn_int8_kernel2 self.scale.extend( [torch.zeros(global_hidden_units, dtype=torch.float)] * layer_num) # ffn_scale2 if self.has_adapters: self.int8_w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size, dtype=torch.int8)] * layer_num) # adaptor1_int8_kernel1 self.scale.extend([torch.zeros(local_adapter_inter_size, dtype=torch.float)] * layer_num) # adaptor1_scale1 self.int8_w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units, dtype=torch.int8)] * layer_num) # adaptor1_int8_kernel2 self.scale.extend([torch.zeros( global_hidden_units, dtype=torch.float)] * layer_num) # adaptor1_scale2 self.int8_w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size, dtype=torch.int8)] * layer_num) # adaptor2_int8_kernel1 self.scale.extend([torch.zeros(local_adapter_inter_size, dtype=torch.float)] * layer_num) # adaptor2_scale1 self.int8_w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units, dtype=torch.int8)] * layer_num) # adaptor2_int8_kernel2 self.scale.extend([torch.zeros( global_hidden_units, dtype=torch.float)] * layer_num) # adaptor2_scale2 def __getitem__(self, idx): return self.w[idx] def __setitem__(self, idx, val): self.w[idx] = val def __len__(self): return len(self.w) def _map(self, func): assert(self.pre_embed_idx < self.post_embed_idx, "Pre decoder embedding index should be lower than post decoder embedding index.") for i in range(len(self.w)): if isinstance(self.w[i], list): for j in range(len(self.w[i])): self.w[i][j] = func(self.w[i][j]) else: if self.share_embed and i == self.post_embed_idx: # If sharing the pre and post embedding, any mapping to # the pre decoder weight will give the same output to the # post decoder weight, so we just copy here. self.w[self.post_embed_idx] = self.w[self.pre_embed_idx] else: self.w[i] = func(self.w[i]) def _map_int8(self, func): for i in range(len(self.int8_w)): if isinstance(self.int8_w[i], list): for j in range(len(self.int8_w[i])): self.int8_w[i][j] = func(self.int8_w[i][j]) else: self.int8_w[i] = func(self.int8_w[i]) for i in range(len(self.scale)): if isinstance(self.scale[i], list): for j in range(len(self.scale[i])): self.scale[i][j] = func(self.scale[i][j]) else: self.scale[i] = func(self.scale[i]) def _map_int8_scales(self, func): for i in range(len(self.scale)): if isinstance(self.scale[i], list): for j in range(len(self.scale[i])): self.scale[i][j] = func(self.scale[i][j]) else: self.scale[i] = func(self.scale[i]) def load(self, ckpt_path, tp_rank, pipeline_para_rank): if not os.path.exists(ckpt_path): raise FileNotFoundError(f"Failed to find {ckpt_path}") w = [] type_map = {np.float32: torch.float32, np.float16: torch.float16} # Load def is_load(i): return i >= self.layers_per_device * \ pipeline_para_rank and i < self.layers_per_device * \ (pipeline_para_rank + 1) def load_to_torch(npdata: str, is_load: bool): if is_load: return torch.from_numpy(npdata).to(str_type_map[self.inference_data_type]) #return torch.from_numpy(np.fromfile(file_path, dtype=self.weights_data_type)).to(str_type_map[self.inference_data_type]) else: return torch.empty(0).to(str_type_map[self.inference_data_type]) def get_np_data(h5f, layername, layer_num, weight_type, tp_rank=None): if tp_rank is None: return [load_to_torch(h5f[f'model.layers.{i}.{layername}.{weight_type}']["weights"][:], is_load(i)) for i in range(layer_num)] else: return [load_to_torch(h5f[f'model.layers.{i}.{layername}.{weight_type}.{tp_rank}']["weights"][:], is_load(i)) for i in range(layer_num)] def get_np_data_single(h5f, layername, weight_type, is_loaded, tp_rank=None): if weight_type is None: return load_to_torch(h5f[f'model.{layername}']["weights"][:], is_loaded) if tp_rank is None: return load_to_torch(h5f[f'model.{layername}.{weight_type}']["weights"][:], is_loaded) else: return load_to_torch(h5f[f'model.{layername}.{weight_type}.{tp_rank}']["weights"][:], is_loaded) import h5py ckpt_f = h5py.File(ckpt_path, "r") w.extend(get_np_data(ckpt_f, "input_layernorm", self.layer_num, "weight")) w.extend(get_np_data(ckpt_f, "input_layernorm", self.layer_num, "bias")) w.extend(get_np_data(ckpt_f, "attention.query_key_value", self.layer_num, "weight", tp_rank)) w.extend(get_np_data(ckpt_f, "attention.query_key_value", self.layer_num, "bias", tp_rank)) w.extend(get_np_data(ckpt_f, "attention.dense", self.layer_num, "weight", tp_rank)) w.extend(get_np_data(ckpt_f, "attention.dense", self.layer_num, "bias")) w.extend(get_np_data(ckpt_f, "post_attention_layernorm", self.layer_num, "weight")) w.extend(get_np_data(ckpt_f, "post_attention_layernorm", self.layer_num, "bias")) # if moe, load "mlp.moe.experts.dense_h_to_4h" w.extend(get_np_data(ckpt_f, "mlp.dense_h_to_4h", self.layer_num, "weight", tp_rank)) w.extend(get_np_data(ckpt_f, "mlp.dense_h_to_4h", self.layer_num, "bias", tp_rank)) # if moe, load "mlp.moe.experts.dense_4h_to_h" w.extend(get_np_data(ckpt_f, "mlp.dense_4h_to_h", self.layer_num, "weight", tp_rank)) w.extend(get_np_data(ckpt_f, "mlp.dense_4h_to_h", self.layer_num, "bias")) if self.has_pre_decoder_layernorm: w.append(get_np_data_single(ckpt_f, "pre_decoder_layernorm", "weight", True)) w.append(get_np_data_single(ckpt_f, "pre_decoder_layernorm", "bias", True)) if self.has_post_decoder_layernorm: w.append(get_np_data_single(ckpt_f, "final_layernorm", "weight", True)) w.append(get_np_data_single(ckpt_f, "final_layernorm", "bias", True)) if self.has_positional_encoding: wpe = load_to_torch(get_np_data_single(ckpt_f, "wpe", weight_type=None, is_loaded=True)).reshape(-1, self.global_hidden_units) assert self.max_seq_len <= wpe.size(0), ( f"max_seq_len ({self.max_seq_len} must not exceed " f"the value of maximum sequence length during training ({wpe.size(0)})." ) w.append(wpe) w.append(get_np_data_single(ckpt_f, "wte", weight_type=None, is_loaded=True)) if "model.lm_head.weight" in ckpt_f.keys(): self.share_embed = False w.append(get_np_data_single(ckpt_f, "lm_head", "weight", True)) else: self.share_embed = True w.append(torch.empty(0).to(str_type_map[self.inference_data_type])) gate_list = [] for i in range(self.layer_num): if f"model.layers.{i}.mlp.moe.gate.wg.weight" in ckpt_f.keys(): gate_list.append(load_to_torch( f"{ckpt_path}/model.layers.{i}.mlp.moe.gate.wg.weight.bin", True)) else: gate_list.append(load_to_torch( f"{ckpt_path}/model.layers.{i}.mlp.moe.gate.wg.weight.bin", False)) w.extend(gate_list) """ if self.has_adapters: w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_h_to_4h.weight.{tp_rank}.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_h_to_4h.weight.{tp_rank}.bin") else f"{ckpt_path}/model.layers.{i}.after_attention_adapter.moe.experts.dense_h_to_4h.weight.{tp_rank}.bin", is_load(i)) for i in range(self.layer_num)]) w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_h_to_4h.bias.{tp_rank}.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_h_to_4h.bias.{tp_rank}.bin") else f"{ckpt_path}/model.layers.{i}.after_attention_adapter.moe.experts.dense_h_to_4h.bias.{tp_rank}.bin", is_load(i)) for i in range(self.layer_num)]) w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_4h_to_h.weight.{tp_rank}.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_4h_to_h.weight.{tp_rank}.bin") else f"{ckpt_path}/model.layers.{i}.after_attention_adapter.moe.experts.dense_4h_to_h.weight.{tp_rank}.bin", is_load(i)) for i in range(self.layer_num)]) w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_4h_to_h.bias.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_attention_adapter.dense_4h_to_h.bias.bin") else f"{ckpt_path}/model.layers.{i}.after_attention_adapter.moe.experts.dense_4h_to_h.bias.bin", is_load(i)) for i in range(self.layer_num)]) w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_h_to_4h.weight.{tp_rank}.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_h_to_4h.weight.{tp_rank}.bin") else f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.moe.experts.dense_h_to_4h.weight.{tp_rank}.bin", is_load(i)) for i in range(self.layer_num)]) w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_h_to_4h.bias.{tp_rank}.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_h_to_4h.bias.{tp_rank}.bin") else f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.moe.experts.dense_h_to_4h.bias.{tp_rank}.bin", is_load(i)) for i in range(self.layer_num)]) w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_4h_to_h.weight.{tp_rank}.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_4h_to_h.weight.{tp_rank}.bin") else f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.moe.experts.dense_4h_to_h.weight.{tp_rank}.bin", is_load(i)) for i in range(self.layer_num)]) w.extend([load_to_torch( f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_4h_to_h.bias.bin" if os.path.isfile(f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.dense_4h_to_h.bias.bin") else f"{ckpt_path}/model.layers.{i}.after_ffn_adapter.moe.experts.dense_4h_to_h.bias.bin", is_load(i)) for i in range(self.layer_num)]) """ assert len(self.w) == len(w) # Reshape try: for i in range(len(w)): if w[i].nelement() == self.w[i].nelement(): self.w[i] = w[i].reshape(self.w[i].shape) else: self.w[i] = w[i] except RuntimeError: raise RuntimeError( f"head_num, size_per_head, vocab_size, and max_seq_len must be the same as the ones during training " f"(idx: {i} expected shape: {self.w[i].shape} got shape: {w[i].shape})." ) # transpose calibrate quantize the kernel layer_num = self.layer_num if self.int8_mode != 0: for i in range(layer_num): self.int8_w[i + 0 * layer_num], self.scale[i + 0 * layer_num] = self.weight_transpose_calibrate_quantize(self.w[2 * layer_num + i]) self.int8_w[i + 1 * layer_num], self.scale[i + 1 * layer_num] = self.weight_transpose_calibrate_quantize(self.w[4 * layer_num + i]) self.int8_w[i + 2 * layer_num], self.scale[i + 2 * layer_num] = self.weight_transpose_calibrate_quantize(self.w[8 * layer_num + i]) self.int8_w[i + 3 * layer_num], self.scale[i + 3 * layer_num] = self.weight_transpose_calibrate_quantize(self.w[10 * layer_num + i]) # We clear the original weights since they are no longer needed if self.int8_mode == 1: self.w[2 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type]) self.w[4 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type]) self.w[8 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type]) self.w[10 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type]) if self.has_adapters: self.int8_w[i + 4 * layer_num], self.scale[i + 4 * layer_num] = self.weight_transpose_calibrate_quantize( self.w[12 * layer_num + i + self.adapter_offset]) self.int8_w[i + 5 * layer_num], self.scale[i + 5 * layer_num] = self.weight_transpose_calibrate_quantize( self.w[14 * layer_num + i + self.adapter_offset]) self.int8_w[i + 6 * layer_num], self.scale[i + 6 * layer_num] = self.weight_transpose_calibrate_quantize( self.w[16 * layer_num + i + self.adapter_offset]) self.int8_w[i + 7 * layer_num], self.scale[i + 7 * layer_num] = self.weight_transpose_calibrate_quantize( self.w[18 * layer_num + i + self.adapter_offset]) # Similar to above: if self.int8_mode == 1: self.w[12 * layer_num + i + self.adapter_offset] = torch.empty( 0).to(str_type_map[self.inference_data_type]) self.w[14 * layer_num + i + self.adapter_offset] = torch.empty( 0).to(str_type_map[self.inference_data_type]) self.w[16 * layer_num + i + self.adapter_offset] = torch.empty( 0).to(str_type_map[self.inference_data_type]) self.w[18 * layer_num + i + self.adapter_offset] = torch.empty( 0).to(str_type_map[self.inference_data_type]) return True class BaseBelleModel(nn.Module): def __init__(self, head_num, size_per_head, vocab_size, start_id, end_id, layer_num, max_seq_len: int, tensor_para_size: int, pipeline_para_size: int, lib_path: typing.Union[str, pathlib.Path], inference_data_type: str, inter_size: int = 0, # gpt_variant_params layernorm_eps: float = 1e-6, layernorm_type: typing.Literal['pre_layernorm', 'post_layernorm'] = "pre_layernorm", activation_type: str = "Gelu", gpt_with_moe: bool = False, expert_num: int = 0, moe_k: int = 0, moe_layer_index: typing.List = [], has_positional_encoding: bool = True, has_pre_decoder_layernorm: bool = False, has_post_decoder_layernorm: bool = True, has_adapters: bool = False, adapter_inter_size: int = 0, use_attention_linear_bias: bool = False, int8_mode: int = 0, weights_data_type: typing.Union[str, np.dtype] = np.float32, shared_contexts_ratio: float = 1.0): super().__init__() self.head_num = head_num self.size_per_head = size_per_head self.vocab_size = vocab_size self.start_id = start_id self.end_id = end_id self.layer_num = layer_num self.inter_size = inter_size if inter_size != 0 else 4 * \ self.head_num * self.size_per_head # gpt_variant_params self.layernorm_eps = layernorm_eps self.layernorm_type = layernorm_type self.activation_type = activation_type self.gpt_with_moe = gpt_with_moe self.expert_num = expert_num self.moe_k = moe_k self.moe_layer_index = moe_layer_index self.has_positional_encoding = has_positional_encoding self.has_pre_decoder_layernorm = has_pre_decoder_layernorm self.has_post_decoder_layernorm = has_post_decoder_layernorm self.has_adapters = has_adapters self.adapter_inter_size = adapter_inter_size self.use_attention_linear_bias = use_attention_linear_bias # multi-gpu params self.tensor_para_size = tensor_para_size self.pipeline_para_size = pipeline_para_size self.use_sparse_gemm = False self.build_model = False self.int8_mode = int8_mode self.weights_data_type = weights_data_type self.shared_contexts_ratio = shared_contexts_ratio assert torch.cuda.is_available(), "CUDA is required for this model." assert head_num % tensor_para_size == 0, "head_num must be a multiple of tensor_para_size." assert layer_num % pipeline_para_size == 0, "layer_num must be a multiple of pipeline_para_size." # Load the C++ model into Pytorch model. torch.classes.load_library(os.path.abspath(lib_path)) # Prepare weights self.weights = BaseBelleWeights(head_num, size_per_head, layer_num, vocab_size, max_seq_len, tensor_para_size, pipeline_para_size, weights_data_type=weights_data_type, inference_data_type=inference_data_type, gpt_with_moe=self.gpt_with_moe, has_positional_encoding=self.has_positional_encoding, has_pre_decoder_layernorm=self.has_pre_decoder_layernorm, has_post_decoder_layernorm=self.has_post_decoder_layernorm, has_adapters=self.has_adapters, adapter_inter_size=self.adapter_inter_size, int8_mode=int8_mode, inter_size=inter_size) # Prepare for tensor/pipeline parallel try: dist.init_process_group(backend='mpi') except: print("[INFO] WARNING: Have initialized the process group") self.rank = dist.get_rank() self.device_count = torch.cuda.device_count() self.device = self.rank % self.device_count torch.cuda.set_device(self.device) world_size = dist.get_world_size() assert world_size == tensor_para_size * \ pipeline_para_size, "tensor_para_size * pipeline_para_size must be equal to world_size." self.tensor_para_rank = self.rank % self.tensor_para_size self.pipeline_para_rank = self.rank // self.tensor_para_size def load(self, ckpt_path): is_load = self.weights.load(ckpt_path, tp_rank=self.tensor_para_rank, pipeline_para_rank=self.pipeline_para_rank) self.cuda() torch.cuda.empty_cache() # clean cache for model weight preprocessing return is_load def sparse(self): if not self.use_sparse_gemm: self.use_sparse_gemm = True def cuda(self): self.weights._map(lambda w: w.cuda(self.device)) if self.int8_mode != 0: self.weights._map_int8(lambda w: w.cuda(self.device)) if self.build_model: del self.model self.build_model = False self.model = torch.classes.FasterTransformer.GptOp( self.head_num, self.size_per_head, self.inter_size, self.layer_num, self.expert_num, self.moe_k, self.moe_layer_index, self.vocab_size, self.start_id, self.end_id, self.use_sparse_gemm, # gpt_variant_params self.layernorm_eps, self.layernorm_type, self.activation_type, self.has_positional_encoding, self.has_pre_decoder_layernorm, self.has_post_decoder_layernorm, self.has_adapters, self.adapter_inter_size, self.use_attention_linear_bias, self.weights.w) self.build_model = True def forward(self, start_ids: torch.IntTensor, start_lengths: torch.IntTensor, output_len: int, beam_width: int = 1, top_k: typing.Optional[torch.IntTensor] = None, top_p: typing.Optional[torch.FloatTensor] = None, beam_search_diversity_rate: typing.Optional[torch.FloatTensor] = None, temperature: typing.Optional[torch.FloatTensor] = None, len_penalty: typing.Optional[torch.FloatTensor] = None, repetition_penalty: typing.Optional[torch.FloatTensor] = None, presence_penalty: typing.Optional[torch.FloatTensor] = None, min_length: typing.Optional[torch.IntTensor] = None, random_seed: typing.Optional[torch.LongTensor] = None, bad_words_list: typing.Optional[torch.IntTensor] = None, return_output_length: bool = False, return_cum_log_probs: int = 0): if not self.build_model: # for the cases we don't load model self.cuda() torch.cuda.empty_cache() # clean cache for model weight preprocessing input_len = start_ids.size(1) assert input_len > 0, "input len must be larger than zero. For an unconditional case, use start_id as the first token." # Inputs to device start_ids = start_ids.cuda(self.device) start_lengths = start_lengths.cuda(self.device) # outputs: output_ids, output_lengths, output_cum_log_probs (optional) outputs = self.model.forward(start_ids, start_lengths, output_len, beam_width, # optional, can be None top_k, # optional, can be None top_p, # optional, can be None beam_search_diversity_rate, # optional, can be None temperature, # optional, can be None len_penalty, # optional, can be None repetition_penalty, # optional, can be None presence_penalty, # optional, can be None min_length, # optional, can be None random_seed, # optional, can be None bad_words_list, # optional, can be None return_cum_log_probs) # optional, can be None if return_cum_log_probs == 0: output_ids, output_lengths = outputs else: output_ids, output_lengths, output_cum_log_probs = outputs if return_output_length: if return_cum_log_probs > 0: return output_ids, output_lengths, output_cum_log_probs else: return output_ids, output_lengths else: return output_ids def set_input_tensor(self, input_tensor): """Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes from communication, not from the input, so the model's forward_step_func won't have it. This function is thus used by internal code to bypass the input provided by the forward_step_func""" self.input_tensor = input_tensor class BaseParallelBelleModel(BaseBelleModel): def cuda(self): self.weights._map(lambda w: w.cuda(self.device)) if self.int8_mode != 0: self.weights._map_int8(lambda w: w.cuda(self.device)) if self.build_model: del self.model self.build_model = False self.model = torch.classes.FasterTransformer.ParallelGptOp( self.head_num, self.size_per_head, self.inter_size, self.layer_num, self.expert_num, self.moe_k, self.moe_layer_index, self.vocab_size, self.start_id, self.end_id, self.tensor_para_size, self.pipeline_para_size, self.int8_mode, # GPT variant parameters self.layernorm_eps, self.layernorm_type, self.activation_type, self.has_positional_encoding, self.has_pre_decoder_layernorm, self.has_post_decoder_layernorm, self.has_adapters, self.adapter_inter_size, self.use_attention_linear_bias, self.weights.w, self.weights.int8_w, self.weights.scale, self.shared_contexts_ratio) self.build_model = True class BelleWeight(BaseBelleWeights): def __init__(self, head_num, size_per_head, layer_num, vocab_size, tensor_para_size, pipeline_para_size, weights_data_type, inference_data_type, int8_mode=0): super().__init__( head_num, size_per_head, layer_num, vocab_size, 0, tensor_para_size, pipeline_para_size, weights_data_type, inference_data_type, has_adapters=False, adapter_inter_size=0, has_positional_encoding=False, has_pre_decoder_layernorm=True, has_post_decoder_layernorm=True, int8_mode=int8_mode) class BelleModel(BaseParallelBelleModel): def __init__(self, head_num, size_per_head, vocab_size, start_id, end_id, layer_num, tensor_para_size: int, pipeline_para_size: int, lib_path: str | Path, inference_data_type: str, weights_data_type: str | np.dtype = np.float32, layernorm_eps: float = 1e-5, shared_contexts_ratio: float = 1.0, int8_mode: int = 0): super().__init__( head_num, size_per_head, vocab_size, start_id, end_id, layer_num, 0, tensor_para_size, pipeline_para_size, lib_path=lib_path, inference_data_type=inference_data_type, layernorm_eps=layernorm_eps, # gpt_variant_params layernorm_type="pre_layernorm", activation_type="Gelu", has_positional_encoding=False, has_pre_decoder_layernorm=True, has_post_decoder_layernorm=True, has_adapters=False, adapter_inter_size=0, use_attention_linear_bias=True, int8_mode=int8_mode, weights_data_type=weights_data_type, shared_contexts_ratio=shared_contexts_ratio) def set_input_tensor(self, input_tensor: Optional[torch.Tensor]): """Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes from communication, not from the input, so the model's forward_step_func won't have it. This function is thus used by internal code to bypass the input provided by the forward_step_func """ self.input_tensor = input_tensor