mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
9dbc5b38ba
* Update TensorRT-LLM --------- Co-authored-by: Marks101 <markus.schnoes@gmx.de> Co-authored-by: lkm2835 <lkm2835@gmail.com>
284 lines
12 KiB
Python
284 lines
12 KiB
Python
import json
|
|
import os
|
|
import traceback
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import safetensors
|
|
from transformers import AutoModelForCausalLM
|
|
|
|
from ..._utils import pad_vocab_size
|
|
from ...functional import PositionEmbeddingType, Tensor
|
|
from ...layers import (MLP, Attention, AttentionMaskType, BlockSparseAttnParams,
|
|
Embedding, LayerNorm, ParallelLMHead, RmsNorm)
|
|
from ...lora_manager import LoraConfig, use_lora
|
|
from ...module import Module
|
|
from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM,
|
|
PretrainedConfig)
|
|
from .convert import convert_hf_config, convert_hf_weights
|
|
|
|
|
|
class Phi3DecoderLayer(Module):
|
|
|
|
def __init__(self, config: PretrainedConfig, layer_idx: int):
|
|
super().__init__()
|
|
self.config = config
|
|
self.layer_idx = layer_idx
|
|
tp_group = config.mapping.tp_group
|
|
tp_size = config.mapping.tp_size
|
|
|
|
attention_mask_type = AttentionMaskType.causal
|
|
block_sparse_attn_params = BlockSparseAttnParams()
|
|
q_scaling = 1.0
|
|
self.gegelu_limit = None
|
|
|
|
self.small_variant = config.architecture == "Phi3SmallForCausalLM"
|
|
if self.small_variant:
|
|
self.gegelu_limit = config.gegelu_limit
|
|
|
|
# MuP uses norm_factor=attention_head_size (rather than sqrt(attention_head_size))
|
|
# We achieve this using q_scaling = sqrt(attention_head_size)
|
|
hidden_size = config.hidden_size
|
|
num_attention_heads = config.num_attention_heads
|
|
attention_head_size = hidden_size / num_attention_heads
|
|
q_scaling = attention_head_size**.5
|
|
|
|
block_sparse = (
|
|
(layer_idx + 1) % config.dense_attention_every_n_layers) != 0
|
|
attention_mask_type = AttentionMaskType.blocksparse if block_sparse else AttentionMaskType.causal
|
|
|
|
block_sparse_attn_params = BlockSparseAttnParams(
|
|
config.blocksparse_block_size,
|
|
config.blocksparse_homo_head_pattern,
|
|
config.blocksparse_num_local_blocks,
|
|
config.blocksparse_vertical_stride)
|
|
|
|
self.input_layernorm = LayerNorm(
|
|
normalized_shape=config.hidden_size, dtype=config.dtype)
|
|
self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size,
|
|
dtype=config.dtype)
|
|
else:
|
|
self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size,
|
|
eps=config.norm_epsilon,
|
|
dtype=config.dtype)
|
|
self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size,
|
|
eps=config.norm_epsilon,
|
|
dtype=config.dtype)
|
|
|
|
layers_range = config.mapping.pp_layers(config.num_hidden_layers)
|
|
local_layer_idx = layer_idx - layers_range[0]
|
|
position_embedding_type = PositionEmbeddingType.rope_gpt_neox
|
|
|
|
rope_scaling_short_factors, rope_scaling_long_factors = None, None
|
|
rope_scaling_short_mscale, rope_scaling_long_mscale = None, None
|
|
original_max_position_embeddings = config.max_position_embeddings
|
|
|
|
if hasattr(config, "longrope_scaling_short_factors"):
|
|
rope_scaling_short_factors = np.asarray(
|
|
config.longrope_scaling_short_factors).astype(np.float32)
|
|
rope_scaling_long_factors = np.asarray(
|
|
config.longrope_scaling_long_factors).astype(np.float32)
|
|
|
|
original_max_position_embeddings = config.original_max_position_embeddings
|
|
position_embedding_type = PositionEmbeddingType.long_rope
|
|
|
|
if self.small_variant:
|
|
rope_scaling_short_mscale = config.longrope_short_mscale
|
|
rope_scaling_long_mscale = config.longrope_long_mscale
|
|
|
|
self.attention = Attention(
|
|
local_layer_idx=local_layer_idx,
|
|
hidden_size=config.hidden_size,
|
|
num_attention_heads=config.num_attention_heads,
|
|
num_kv_heads=config.num_key_value_heads,
|
|
position_embedding_type=position_embedding_type,
|
|
rotary_embedding_base=config.rotary_base,
|
|
max_position_embeddings=config.max_position_embeddings,
|
|
dtype=config.dtype,
|
|
attention_mask_type=attention_mask_type,
|
|
bias=self.small_variant,
|
|
q_scaling=q_scaling,
|
|
tp_group=tp_group,
|
|
tp_size=tp_size,
|
|
quant_mode=config.quant_mode,
|
|
rope_scaling_short_factors=rope_scaling_short_factors,
|
|
rope_scaling_long_factors=rope_scaling_long_factors,
|
|
rope_scaling_short_mscale=rope_scaling_short_mscale,
|
|
rope_scaling_long_mscale=rope_scaling_long_mscale,
|
|
original_max_position_embeddings=original_max_position_embeddings,
|
|
block_sparse_params=block_sparse_attn_params)
|
|
|
|
self.mlp = MLP(hidden_size=config.hidden_size,
|
|
ffn_hidden_size=config.intermediate_size,
|
|
hidden_act=config.hidden_act,
|
|
dtype=config.dtype,
|
|
tp_group=tp_group,
|
|
tp_size=tp_size,
|
|
quant_mode=config.quant_mode,
|
|
bias=self.small_variant)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: Tensor,
|
|
attention_mask=None,
|
|
use_cache=False,
|
|
kv_cache_params=None,
|
|
attention_params=None,
|
|
lora_layer_params=None,
|
|
):
|
|
|
|
input_layernorm_output = self.input_layernorm(hidden_states)
|
|
attention_output = self.attention(
|
|
input_layernorm_output,
|
|
attention_mask=attention_mask,
|
|
use_cache=use_cache,
|
|
kv_cache_params=kv_cache_params,
|
|
attention_params=attention_params,
|
|
norm_before_bmm1=not self.small_variant,
|
|
lora_layer_params=lora_layer_params,
|
|
)
|
|
|
|
if use_cache:
|
|
attention_output, presents = attention_output
|
|
|
|
post_attention_input = hidden_states + attention_output
|
|
post_attention_output = self.post_layernorm(post_attention_input)
|
|
feed_forward_hidden_states = self.mlp(
|
|
post_attention_output,
|
|
gegelu_limit=self.gegelu_limit,
|
|
lora_layer_params=lora_layer_params)
|
|
hidden_states = post_attention_input + feed_forward_hidden_states
|
|
if use_cache:
|
|
return (hidden_states, presents)
|
|
return hidden_states
|
|
|
|
|
|
class Phi3Model(Module):
|
|
|
|
def __init__(self, config: PretrainedConfig):
|
|
super().__init__()
|
|
self.vocab_embedding = Embedding(num_embeddings=config.vocab_size,
|
|
embedding_dim=config.hidden_size,
|
|
dtype=config.dtype)
|
|
|
|
self.layers = DecoderLayerList(Phi3DecoderLayer, config)
|
|
self.small_variant = config.architecture == "Phi3SmallForCausalLM"
|
|
if self.small_variant:
|
|
self.ln_f = LayerNorm(normalized_shape=config.hidden_size,
|
|
dtype=config.dtype)
|
|
self.mup_embedding_multiplier = config.mup_embedding_multiplier
|
|
else:
|
|
self.ln_f = RmsNorm(normalized_shape=config.hidden_size,
|
|
eps=config.norm_epsilon,
|
|
dtype=config.dtype)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: Tensor,
|
|
position_ids=None,
|
|
use_cache=False,
|
|
attention_mask=None,
|
|
kv_cache_params=None,
|
|
attention_params=None,
|
|
prompt_embedding_table=None,
|
|
prompt_tasks=None,
|
|
prompt_vocab_size=None,
|
|
lora_params=None,
|
|
):
|
|
args = [prompt_embedding_table, prompt_tasks, prompt_vocab_size
|
|
] if prompt_embedding_table is not None else []
|
|
hidden_states = self.vocab_embedding(input_ids, *args)
|
|
|
|
if self.small_variant and self.mup_embedding_multiplier > 0.0:
|
|
hidden_states = hidden_states * self.mup_embedding_multiplier
|
|
|
|
hidden_states = self.layers(
|
|
hidden_states,
|
|
use_cache=use_cache,
|
|
attention_mask=attention_mask,
|
|
kv_cache_params=kv_cache_params,
|
|
attention_params=attention_params,
|
|
lora_params=lora_params,
|
|
)
|
|
if use_cache:
|
|
hidden_states, presents = hidden_states
|
|
|
|
hidden_states = self.ln_f(hidden_states)
|
|
|
|
if use_cache:
|
|
return (hidden_states, tuple(presents))
|
|
return hidden_states
|
|
|
|
|
|
class Phi3ForCausalLM(DecoderModelForCausalLM):
|
|
|
|
def __init__(self, config: PretrainedConfig):
|
|
transformer = Phi3Model(config)
|
|
vocab_size_padded = pad_vocab_size(config.vocab_size,
|
|
config.mapping.tp_size)
|
|
|
|
lm_head = ParallelLMHead(config.hidden_size,
|
|
vocab_size_padded,
|
|
bias=False,
|
|
dtype=config.dtype,
|
|
tp_group=config.mapping.tp_group,
|
|
tp_size=config.mapping.tp_size,
|
|
gather_output=True)
|
|
self.trtllm_modules_to_hf_modules = {
|
|
"attn_qkv": ["qkv_proj", "query_key_value"],
|
|
"attn_dense": ["o_proj", "dense"],
|
|
"mlp_h_to_4h": ["gate_up_proj", "up_proj"],
|
|
"mlp_4h_to_h": "down_proj",
|
|
}
|
|
super().__init__(config, transformer, lm_head)
|
|
|
|
@classmethod
|
|
def convert_hf_checkpoint(cls,
|
|
hf_model_dir: str,
|
|
dtype: Optional[str] = "float16",
|
|
output_dir: Optional[str] = None,
|
|
args=None):
|
|
'''
|
|
Convert Huggingface checkpoint to TRT-LLM checkpoint
|
|
'''
|
|
|
|
hf_model = AutoModelForCausalLM.from_pretrained(hf_model_dir,
|
|
torch_dtype="auto",
|
|
trust_remote_code=True)
|
|
config = convert_hf_config(hf_model.config, dtype, args)
|
|
with open(os.path.join(output_dir, 'config.json'), 'w') as f:
|
|
json.dump(config, f, indent=4)
|
|
|
|
small_variant = config['architecture'] == "Phi3SmallForCausalLM"
|
|
|
|
def covert_and_save(rank):
|
|
weights = convert_hf_weights(hf_model, dtype, config, small_variant,
|
|
args, rank)
|
|
safetensors.torch.save_file(
|
|
weights, os.path.join(output_dir, f'rank{rank}.safetensors'))
|
|
|
|
world_size = args.tp_size * args.pp_size
|
|
if args.workers == 1:
|
|
for rank in range(world_size):
|
|
covert_and_save(rank)
|
|
else:
|
|
with ThreadPoolExecutor(max_workers=args.workers) as p:
|
|
futures = [
|
|
p.submit(covert_and_save, rank)
|
|
for rank in range(world_size)
|
|
]
|
|
exceptions = []
|
|
for future in as_completed(futures):
|
|
try:
|
|
future.result()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
exceptions.append(e)
|
|
assert len(
|
|
exceptions
|
|
) == 0, "Checkpoint conversion failed, please check error log."
|
|
|
|
def use_lora(self, lora_config: LoraConfig):
|
|
use_lora(self, lora_config, self.trtllm_modules_to_hf_modules)
|