diff --git a/examples/deepseek_v2/convert_checkpoint.py b/examples/deepseek_v2/convert_checkpoint.py index e45ea15d58e8..7cbc7935d225 100755 --- a/examples/deepseek_v2/convert_checkpoint.py +++ b/examples/deepseek_v2/convert_checkpoint.py @@ -23,7 +23,8 @@ from tensorrt_llm.layers import MoeConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models import DeepseekV2ForCausalLM -from tensorrt_llm.models.deepseek_v2.convert import load_hf_deepseek +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo def parse_arguments(): @@ -58,7 +59,7 @@ def parse_arguments(): parser.add_argument('--load_model_on_cpu', default=False, action="store_true", - help='Choose to load HF cpkt into GPU') + help='Choose to load HF cpkt into CPU') parser.add_argument( '--use_parallel_embedding', action="store_true", @@ -75,6 +76,25 @@ def parse_arguments(): 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' 'To shard it along hidden dimension, set embedding_sharding_dim=1' 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') + + parser.add_argument( + '--use_weight_only', + default=False, + action="store_true", + help='Quantize weights for the various GEMMs to INT4/INT8.' + 'See --weight_only_precision to set the precision') + parser.add_argument( + '--weight_only_precision', + const='int8', + type=str, + nargs='?', + default='int8', + choices=['int8', 'int4'], + help= + 'Define the precision for the weights when using weight-only quantization.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) + parser.add_argument('--output_dir', type=str, default='trtllm_checkpoint', @@ -125,12 +145,29 @@ def parse_arguments(): return args +def precision_to_config(precision, quant_config) -> QuantConfig: + '''update config dict for weight-only quantization + ''' + quant_config = QuantConfig() + precision_to_algo = {'int8': QuantAlgo.W8A16, 'int4': QuantAlgo.W4A16} + quant_config.quant_algo = precision_to_algo.get(precision) + return quant_config + + +def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: + '''return config dict with quantization info based on the command line args + ''' + quant_config = QuantConfig() + if args.use_weight_only: + quant_config = precision_to_config(args.weight_only_precision, + quant_config) + return quant_config + + def args_to_build_options(args): return { 'use_parallel_embedding': args.use_parallel_embedding, 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin, 'load_model_on_cpu': args.load_model_on_cpu } @@ -163,8 +200,7 @@ def convert_and_save_hf(args): override_fields = {} override_fields.update(args_to_build_options(args)) - load_model_on_cpu = args.load_model_on_cpu - hf_model = load_hf_deepseek(model_dir, load_model_on_cpu) + quant_config = args_to_quant_config(args) def convert_and_save_rank(args, rank): mapping = Mapping(world_size=world_size, @@ -175,7 +211,11 @@ def convert_and_save_rank(args, rank): moe_ep_size=args.moe_ep_size) deepseekv2 = DeepseekV2ForCausalLM.from_hugging_face( - hf_model, args.model_dir, args.dtype, mapping, **override_fields) + model_dir, + args.dtype, + mapping=mapping, + quant_config=quant_config, + **override_fields) deepseekv2.save_checkpoint(args.output_dir, save_config=(rank == 0)) del deepseekv2 @@ -187,7 +227,6 @@ def main(): print(tensorrt_llm.__version__) args = parse_arguments() - args.tp_size * args.pp_size if (args.moe_tp_size == -1 and args.moe_ep_size == -1): # moe default to tp-only args.moe_tp_size = args.tp_size diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md new file mode 100644 index 000000000000..6d6857d306e8 --- /dev/null +++ b/examples/deepseek_v3/README.md @@ -0,0 +1,170 @@ +# Deepseek-v3 + +This document shows how to build and run the `DeepSeek-v3` model in TensorRT-LLM. + + +***The current branch is a preview version only for supporting DeepSeek-V3, which will be part of the TRT-LLM official release in 0.18+ version.*** + + +- [Deepseek-V3](#deepseek-v3) + - [Support Matrix](#support-matrix) + - [Prerequisite](#prerequisite) + - [Hardware](#hardware) + - [Overview](#overview) + - [Usage](#usage) + - [Build TensorRT engine(s)](#build-tensorrt-engines) + +## Support Matrix + +| Model | FP16 | BF16 | FP8 | W8A16 | W4A16 | TP | EP | IB | +| :------------- | :---: | :---: | :---: | :-----: | :-----: | :-----: | :-----: | :-----: | +| DeepSeek-V3 | Y | Y | . | Y | Y | Y | Y | Y | + + +- W8A16: INT8 Weight-Only +- W4A16: INT4 Weight-Only +- TP: Tensor Parallel +- EP: Expert Parallel +- IB: Inflight Batching +- FP8: Support for FP8 is currently in progress and will be released soon + +***Please Note:*** +- Prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues. +- Although Int8 and Int4 weight-only quantization can reduce the amount of total GPU memory needed, they may affect the accuracy to a certain degree. + + +## Prerequisite + +First, please download DeepSeek-V3 weights from HF https://huggingface.co/deepseek-ai/DeepSeek-V3-Base. + +```bash +git lfs install +git clone https://huggingface.co/deepseek-ai/DeepSeek-V3-Base +``` + +## Hardware + +The DeepSeek-V3 model requires at least 32x80G GPU memory, model contains 660B parameters, roughly 1.3TB memory (with BF16 precision). + +***Caution: Current TRT-LLM MLA kernel only supports Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in the future release.*** + +## Overview + +The TensorRT-LLM DeepSeek-V3 implementation can be found in [tensorrt_llm/models/deepseek_v2/model.py](../../tensorrt_llm/models/deepseek_v2/model.py). The TensorRT-LLM Deepseek-V3 example code is located in [`example/deepseek_v3`](./). There is one main file: + +* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the DeepSeek-V3 model into tensorrt-llm checkpoint format. + +In addition, there are three shared files in the parent folder [`examples`](../) can be used for inference and evaluation: + +* [`../run.py`](../run.py) to run the model inference output by giving an input text. + + +## Usage + +The TensorRT-LLM DeepSeek-V3 example code is located at [examples/deepseek_v3](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. + +### Build TensorRT engine(s) + +Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. + +Firstly, convert FP8 weights to BF16: +```bash +git clone https://github.com/deepseek-ai/DeepSeek-V3.git +cd DeepSeek-V3/inferece/ +python fp8_cast_bf16.py --input-fp8-hf-path /path/to/DeepSeek-V3 --output-bf16-hf-path /path/to/deepseek-v3-bf16 +``` + +Secondly, the BF16 checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. + +```bash +# Convert Deepseek-v3 HF weights to TensorRT-LLM checkpoint in BF16. +python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --dtype bfloat16 \ + --tp_size 32 \ + --workers 8 # using multiple workers can accelerate the conversion process + + +# Use Weight-Only Int8 quantization +python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --dtype bfloat16 \ + --tp_size 32 \ + --use_weight_only \ + --weight_only_precision int8 \ + --workers 8 + +# Use Weight-Only Int4 quantization +python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --dtype bfloat16 \ + --tp_size 32 \ + --use_weight_only \ + --weight_only_precision int4 \ + --workers 8 +``` +We observed the checkpoint conversion time took hours, while using a significant amount of CPU memory, please adjust the `--workers` parameter to balance your time and memory consumption. + + +After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. + +```bash +# Build engine +trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --output_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --gpt_attention_plugin bfloat16 \ + --gemm_plugin bfloat16 \ + --max_batch_size 4 \ + --max_seq_len 4096 \ + --max_input_len 2048 \ + --use_paged_context_fmha enable \ + --workers 8 +``` + +***Caution: `--max_batch_size` and `--max_seq_len` are the main factors to determine how many GPU memory will be used during runtime, so later when try to run e.g., `summarize.py` or `mmlu.py` or `gptManagerBenchmark.cpp`may need adjust `--max_batch_size` and `--max_seq_len` accordingly to avoid OOM.(meaning rebuild TensorRT engine with smaller `--max_batch_size` and `--max_seq_len` if needed based on GPU memory size), there is beautiful technical log perf-best-practices.md (https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md) explained the mechanism.*** + +Test the engine with [run.py](../run.py) script: +``` +# run.sh +python3 ../run.py --input_text "Today is a nice day." \ + --max_output_len 30 \ + --tokenizer_dir ./DeepSeek-V3 \ + --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --top_p 0.95 \ + --temperature 0.3 + + +``` +For multi-nodes inference, let's take Slurm as an example using above command (run.sh): + +```bash +srun -N 4 -w node-[1-4] --gres=gpu:8 --ntasks-per-node 8 \ + --container-image tensorrt_llm/release:latest \ + --container-mounts ${PWD}:/workspace \ + sh /workspace/command/run.sh +``` + +and the output will be like: + +``` +... +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +Input [Text 0]: "Today is a nice day." +Output [Text 0 Beam 0]: " I am going to the park with my friends. We are going to play soccer. We are going" +``` + diff --git a/examples/deepseek_v3/convert_checkpoint.py b/examples/deepseek_v3/convert_checkpoint.py new file mode 100644 index 000000000000..7cbc7935d225 --- /dev/null +++ b/examples/deepseek_v3/convert_checkpoint.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import os +import time +import traceback +from concurrent.futures import ThreadPoolExecutor, as_completed + +import tensorrt_llm +from tensorrt_llm._utils import release_gc +from tensorrt_llm.layers import MoeConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models import DeepseekV2ForCausalLM +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--model_dir', type=str, default=None, required=True) + parser.add_argument('--tp_size', + type=int, + default=1, + help='N-way tensor parallelism size') + parser.add_argument('--pp_size', + type=int, + default=1, + help='N-way pipeline parallelism size') + parser.add_argument( + '--moe_tp_size', + type=int, + default=-1, + help= + 'N-way tensor parallelism size for MoE, default is tp_size, which will do tp-only for MoE' + ) + parser.add_argument( + '--moe_ep_size', + type=int, + default=-1, + help= + 'N-way expert parallelism size for MoE, default is 1, which will do tp-only for MoE' + ) + parser.add_argument('--dtype', + type=str, + default='float16', + choices=['float32', 'bfloat16', 'float16']) + parser.add_argument('--load_model_on_cpu', + default=False, + action="store_true", + help='Choose to load HF cpkt into CPU') + parser.add_argument( + '--use_parallel_embedding', + action="store_true", + default=False, + help= + 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' + ) + parser.add_argument( + '--embedding_sharding_dim', + type=int, + default=0, + choices=[0, 1], + help= + 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' + 'To shard it along hidden dimension, set embedding_sharding_dim=1' + 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') + + parser.add_argument( + '--use_weight_only', + default=False, + action="store_true", + help='Quantize weights for the various GEMMs to INT4/INT8.' + 'See --weight_only_precision to set the precision') + parser.add_argument( + '--weight_only_precision', + const='int8', + type=str, + nargs='?', + default='int8', + choices=['int8', 'int4'], + help= + 'Define the precision for the weights when using weight-only quantization.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) + + parser.add_argument('--output_dir', + type=str, + default='trtllm_checkpoint', + required=True, + help='The path to save the TensorRT-LLM checkpoint') + parser.add_argument( + '--workers', + type=int, + default=1, + help='The number of workers for converting checkpoint in parallel') + parser.add_argument( + '--moe_num_experts', + type=int, + default=0, + help='Specify the number of experts to use for MOE layers') + parser.add_argument( + '--moe_top_k', + type=int, + default=0, + help= + 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' + ) + parser.add_argument( + '--moe_renorm_mode', + type=int, + default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, + help= + 'Controls renormalization after gate logits. Check layers/moe.py for accepted values' + ) + parser.add_argument( + '--save_config_only', + action="store_true", + default=False, + help= + 'Only save the model config w/o read and converting weights, be careful, this is for debug only' + ) + parser.add_argument( + '--disable_weight_only_quant_plugin', + default=False, + action="store_true", + help= + 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' + 'You must also use --use_weight_only for that argument to have an impact' + ) + # Add quantization related feature later + args = parser.parse_args() + + return args + + +def precision_to_config(precision, quant_config) -> QuantConfig: + '''update config dict for weight-only quantization + ''' + quant_config = QuantConfig() + precision_to_algo = {'int8': QuantAlgo.W8A16, 'int4': QuantAlgo.W4A16} + quant_config.quant_algo = precision_to_algo.get(precision) + return quant_config + + +def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: + '''return config dict with quantization info based on the command line args + ''' + quant_config = QuantConfig() + if args.use_weight_only: + quant_config = precision_to_config(args.weight_only_precision, + quant_config) + return quant_config + + +def args_to_build_options(args): + return { + 'use_parallel_embedding': args.use_parallel_embedding, + 'embedding_sharding_dim': args.embedding_sharding_dim, + 'load_model_on_cpu': args.load_model_on_cpu + } + + +def execute(workers, func, args): + if workers == 1: + for rank, f in enumerate(func): + f(args, rank) + else: + with ThreadPoolExecutor(max_workers=workers) as p: + futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] + 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 convert_and_save_hf(args): + model_dir = args.model_dir + world_size = args.tp_size * args.pp_size + # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. + # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, + # before the refactor is done. + override_fields = {} + override_fields.update(args_to_build_options(args)) + + quant_config = args_to_quant_config(args) + + def convert_and_save_rank(args, rank): + mapping = Mapping(world_size=world_size, + rank=rank, + tp_size=args.tp_size, + pp_size=args.pp_size, + moe_tp_size=args.moe_tp_size, + moe_ep_size=args.moe_ep_size) + + deepseekv2 = DeepseekV2ForCausalLM.from_hugging_face( + model_dir, + args.dtype, + mapping=mapping, + quant_config=quant_config, + **override_fields) + deepseekv2.save_checkpoint(args.output_dir, save_config=(rank == 0)) + del deepseekv2 + + execute(args.workers, [convert_and_save_rank] * world_size, args) + release_gc() + + +def main(): + print(tensorrt_llm.__version__) + args = parse_arguments() + + if (args.moe_tp_size == -1 and args.moe_ep_size == -1): + # moe default to tp-only + args.moe_tp_size = args.tp_size + args.moe_ep_size = 1 + elif (args.moe_tp_size == -1): + args.moe_tp_size = args.tp_size // args.moe_ep_size + elif (args.moe_ep_size == -1): + args.moe_ep_size = args.tp_size // args.moe_tp_size + assert (args.moe_tp_size * args.moe_ep_size == args.tp_size + ), "moe_tp_size * moe_ep_size must equal to tp_size" + + tik = time.time() + if not os.path.exists(args.output_dir): + os.makedirs(args.output_dir) + assert args.model_dir is not None + convert_and_save_hf(args) + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + print(f'Total time of converting checkpoints: {t}') + + +if __name__ == '__main__': + main() diff --git a/tensorrt_llm/layers/attention.py b/tensorrt_llm/layers/attention.py index 871ccbb119dd..4bfc86ca8c43 100755 --- a/tensorrt_llm/layers/attention.py +++ b/tensorrt_llm/layers/attention.py @@ -1913,11 +1913,6 @@ def __init__( max_position_embeddings=1024, rotary_embedding_base=10000.0, rotary_embedding_scaling=None, - rotary_embedding_beta_fast=32, - rotary_embedding_beta_slow=1, - rotary_embedding_mscale=1, - rotary_embedding_mscale_all_dim=0, - rotary_embedding_origin_max_position=4096, rotary_scaling=None, tp_group=None, tp_size=1, @@ -1944,6 +1939,7 @@ def __init__( enable_qkv=False) self.tp_size = tp_size + self.tp_rank = tp_rank if q_lora_rank is None: self.q_lora_rank = hidden_size @@ -1965,7 +1961,7 @@ def yarn_get_mscale(scale=1, mscale=1): return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 - assert self.rotary_scaling is not None + # assert self.rotary_scaling is not None if self.rotary_scaling is not None: mscale_all_dim = self.rotary_scaling.get("mscale_all_dim", 0) scaling_factor = self.rotary_scaling["factor"] @@ -1973,12 +1969,27 @@ def yarn_get_mscale(scale=1, mscale=1): mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) self.q_scaling = 1.0 / (mscale * mscale) - embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_deepseek_attention_plugin( - self.max_position_embeddings, self.qk_rope_head_dim, - self.rotary_embedding_base, self.rotary_scaling["factor"], - rotary_embedding_origin_max_position, rotary_embedding_beta_fast, - rotary_embedding_beta_slow, rotary_embedding_mscale, - rotary_embedding_mscale_all_dim) + embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_deepseek_attention_plugin( + self.max_position_embeddings, self.qk_rope_head_dim, + self.rotary_embedding_base, self.rotary_scaling["factor"], + rotary_scaling['original_max_position_embeddings'], + self.rotary_scaling["beta_fast"], + self.rotary_scaling["beta_slow"], self.rotary_scaling["mscale"], + mscale_all_dim) + else: + embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions( + self.max_position_embeddings, self.qk_rope_head_dim, + self.rotary_embedding_base) + embed_positions_for_gpt_attention = embed_positions_for_gpt_attention.squeeze( + 0) + sin, cos = np.split(embed_positions_for_gpt_attention, 2, 1) + cos_embed = np.expand_dims(np.concatenate((cos, cos), axis=1), + axis=2) + sin_embed = np.expand_dims(np.concatenate((sin, sin), axis=1), + axis=2) + embed_positions_for_gpt_attention = np.concatenate( + (cos_embed, sin_embed), axis=-1).reshape(1, -1) + self.register_parameter( 'embed_positions_for_gpt_attention', Parameter(embed_positions_for_gpt_attention, dtype='float32')) @@ -2205,3 +2216,88 @@ def forward(self, return (context, past_key_value) else: return context + + def postprocess(self, tllm_key, weights, **kwargs): + + def split(v, tp_size, idx, dim=0): + if tp_size == 1: + return v + if len(v.shape) == 1: + return torch.chunk(v, tp_size)[idx].contiguous() + else: + return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() + + if tllm_key.endswith("kv_b_proj"): + kv_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.v_head_dim + ]) + splited_kv_b_proj = split(kv_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + k_nope_weight, v_weight = splited_kv_b_proj.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=1, + ) + kv_b_proj_weight = torch.concat([ + k_nope_weight.reshape( + self.num_attention_heads * self.qk_nope_head_dim, + self.kv_lora_rank), + v_weight.reshape(self.num_attention_heads * self.v_head_dim, + self.kv_lora_rank) + ], + dim=0) + return {tllm_key: kv_b_proj_weight} + elif tllm_key.endswith("q_b_proj"): + q_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.qk_rope_head_dim + ]) + splited_q_b_proj = split(q_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + q_b_proj_weight = splited_q_b_proj.reshape( + self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim), + self.q_lora_rank) + return {tllm_key: q_b_proj_weight} + elif tllm_key.endswith("fused_q_proj"): + assert isinstance(weights, list) and len(weights) == 2 + q_b_proj = weights[0].unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.qk_rope_head_dim + ]) + splited_q_b_proj = split(q_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + kv_b_proj = weights[1].unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.v_head_dim + ]) + splited_kv_b_proj = split(kv_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + q_nope_weight, q_pe_weight = splited_q_b_proj.split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], + dim=1, + ) + k_nope_weight, _ = splited_kv_b_proj.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=1, + ) + fused_q_nope_weight = torch.einsum( + 'hdq,hdk->hkq', + q_nope_weight, + k_nope_weight, + ) + fused_q_weight = torch.cat( + [fused_q_nope_weight, q_pe_weight], + dim=1, + ).flatten(start_dim=0, end_dim=1) + return {tllm_key: fused_q_weight} + else: + return {tllm_key: weights} diff --git a/tensorrt_llm/layers/linear.py b/tensorrt_llm/layers/linear.py index 9224c066dacf..899ad446b8c3 100644 --- a/tensorrt_llm/layers/linear.py +++ b/tensorrt_llm/layers/linear.py @@ -426,6 +426,11 @@ def postprocess(self, tllm_key, weights, **kwargs): weights = w.reshape(-1, self.in_features) # Weight else: weights = w.reshape(-1) # Bias + elif tllm_key.endswith("fused_a.weight") and isinstance(weights, list): + weights = torch.cat( + [weights[0], weights[1]], + dim=0, + ) weights = weights.to(str_dtype_to_torch(self.dtype)) return {tllm_key: weights} diff --git a/tensorrt_llm/layers/moe.py b/tensorrt_llm/layers/moe.py index 2dfabe559c98..c1e83c933207 100755 --- a/tensorrt_llm/layers/moe.py +++ b/tensorrt_llm/layers/moe.py @@ -64,6 +64,11 @@ class ExpertScaleNormalizationMode(IntEnum): DEVICE_LIMITED = 3 DEVICE_LIMITED_RENORM = 4 + class TopKMethod(IntEnum): + GREEDY = 0, + GROUP_LIMITED_GREEDY = 1, + NOAUX_TC = 2 + num_experts: int = 0 shared_expert_intermediate_size: int = 0 @@ -71,6 +76,7 @@ class ExpertScaleNormalizationMode(IntEnum): normalization_mode: ExpertScaleNormalizationMode = ExpertScaleNormalizationMode.RENORMALIZE sparse_mixer_epsilon: float = 0.01 tp_mode: int = 0 + topk_method: TopKMethod = TopKMethod.GREEDY device_limited_n_group: int = 0 device_limited_topk_group: int = 0 @@ -499,6 +505,10 @@ def __init__(self, "router": "gate" } + if moe_config.topk_method == MoeConfig.TopKMethod.NOAUX_TC: + self.e_score_correction_bias = Parameter(shape=(self.num_experts, ), + dtype=trt.float32) + self.init_experts() self.max_low_rank = None @@ -542,6 +552,38 @@ def group_limited_greedy(self, logits): self.moe_config.device_limited_routed_scaling_factor return scores + def noaux_tc(self, logits): + n_group = self.moe_config.device_limited_n_group + scores = sigmoid(logits) + scores_with_bias = scores + self.e_score_correction_bias.value + scores_shape = [ + shape(scores_with_bias, i) for i in range(scores_with_bias.ndim()) + ] + group_scores = sum(topk(scores_with_bias.view( + concat(scores_shape[:-1] + [n_group, scores_shape[-1] // n_group])), + k=2, + dim=-1)[0], + dim=-1) + _, group_idx = topk(group_scores, + k=self.moe_config.device_limited_topk_group, + dim=-1) + group_mask = scatter(group_scores * 0, -1, group_idx, + group_idx.cast(group_scores.dtype) * 0 + 1) + score_mask = expand( + unsqueeze(group_mask, -1), + concat(scores_shape[:-1] + [n_group, scores_shape[-1] // n_group]), + ).view(concat(scores_shape)) + scores_with_bias = scores_with_bias * score_mask + + _, topk_idx = topk(scores_with_bias, k=self.moe_config.top_k, dim=-1) + new_mask = scatter(scores * 0, -1, topk_idx, + topk_idx.cast(scores.dtype) * 0 + 1) + scores = scores * new_mask + score_sum = sum(scores, dim=-1, keepdim=True) + 1e-20 + scores = scores / score_sum * \ + self.moe_config.device_limited_routed_scaling_factor + return scores + def forward(self, hidden_states, finished=None, @@ -556,11 +598,10 @@ def forward(self, 0, "moe_router") routing_input = cast(hidden_states, trt.float32) routing = self.router(routing_input, moe_router_lora_params) - if self.moe_config.normalization_mode in [ - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED, - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - ]: + if self.moe_config.topk_method == MoeConfig.TopKMethod.GROUP_LIMITED_GREEDY: routing = self.group_limited_greedy(routing) + elif self.moe_config.topk_method == MoeConfig.TopKMethod.NOAUX_TC: + routing = self.noaux_tc(routing) output = self.forward_experts(hidden_states, routing, finished, lora_layer_params, side_stream_id) if side_stream_id != SideStreamIDType.disable: diff --git a/tensorrt_llm/models/__init__.py b/tensorrt_llm/models/__init__.py index 816f52354b0e..c3b5231e6827 100755 --- a/tensorrt_llm/models/__init__.py +++ b/tensorrt_llm/models/__init__.py @@ -70,6 +70,7 @@ 'DeepseekForCausalLM', 'FalconConfig', 'DeepseekV2ForCausalLM', + 'DeepSeekV2Config' 'FalconForCausalLM', 'FalconModel', 'GPTConfig', @@ -181,6 +182,7 @@ 'DeepseekForCausalLM': DeepseekForCausalLM, 'DeciLMForCausalLM': DeciLMForCausalLM, 'DeepseekV2ForCausalLM': DeepseekV2ForCausalLM, + 'DeepseekV3ForCausalLM': DeepseekV2ForCausalLM, 'EagleForCausalLM': EagleForCausalLM, 'CohereForCausalLM': CohereForCausalLM, 'MLLaMAModel': MLLaMAForCausalLM, # For modelopt diff --git a/tensorrt_llm/models/deepseek_v2/config.py b/tensorrt_llm/models/deepseek_v2/config.py new file mode 100644 index 000000000000..676418dacfc3 --- /dev/null +++ b/tensorrt_llm/models/deepseek_v2/config.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional, Union + +from transformers import AutoConfig + +from ...layers import MoeConfig +from ...mapping import Mapping +from ..modeling_utils import PretrainedConfig, QuantConfig + + +class DeepSeekV2Config(PretrainedConfig): + + def __init__(self, + *, + mlp_bias: bool = False, + attn_bias: bool = False, + rotary_base: float = 10000.0, + rotary_scaling: Optional[dict] = None, + residual_mlp: bool = False, + disable_weight_only_quant_plugin: bool = False, + moe: Optional[Union[MoeConfig, dict]] = None, + remove_duplicated_kv_heads: bool = False, + **kwargs): + self.mlp_bias = mlp_bias + self.attn_bias = attn_bias + self.rotary_base = rotary_base + self.rotary_scaling = rotary_scaling + self.residual_mlp = residual_mlp + self.disable_weight_only_quant_plugin = disable_weight_only_quant_plugin + if isinstance(moe, dict): + moe = MoeConfig.from_dict(moe) + assert isinstance(moe, MoeConfig) + self.moe = moe.validate() + self.remove_duplicated_kv_heads = remove_duplicated_kv_heads + self.fc_after_embed = False + self.use_input_layernorm_in_first_layer = True + self.use_last_layernorm = True + self.layer_idx_offset = 0 + + super().__init__(**kwargs) + + def to_dict(self): + output = super().to_dict() + # Serialize the fields added in LLaMAConfig + output['mlp_bias'] = self.mlp_bias + output['attn_bias'] = self.attn_bias + output['rotary_base'] = self.rotary_base + output['rotary_scaling'] = self.rotary_scaling + output['residual_mlp'] = self.residual_mlp + output[ + 'disable_weight_only_quant_plugin'] = self.disable_weight_only_quant_plugin + output['fc_after_embed'] = self.fc_after_embed + output[ + 'use_input_layernorm_in_first_layer'] = self.use_input_layernorm_in_first_layer + output['use_last_layernorm'] = self.use_last_layernorm + output['layer_idx_offset'] = self.layer_idx_offset + output['moe'] = self.moe.to_dict() + return output + + @classmethod + def from_hugging_face(cls, + model_dir: str, + dtype: str = 'auto', + mapping: Optional[Mapping] = None, + quant_config: Optional[QuantConfig] = None, + **kwargs): + trust_remote_code = kwargs.pop('trust_remote_code', True) + hf_config = AutoConfig.from_pretrained( + model_dir, trust_remote_code=trust_remote_code) + + moe_routed_scaling_factor = hf_config.routed_scaling_factor + moe_top_k = hf_config.num_experts_per_tok + assert moe_routed_scaling_factor > 0, 'routed_scaling_factor should be greater than 0' + if hf_config.topk_method == 'group_limited_greedy': + moe_topk_method = MoeConfig.TopKMethod.GROUP_LIMITED_GREEDY + if moe_top_k > 1 and hf_config.norm_topk_prob: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM + else: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED + elif hf_config.topk_method == 'greedy': + moe_topk_method = MoeConfig.TopKMethod.GREEDY + assert moe_routed_scaling_factor == 1.0, 'The combination of topk_method == greedy and routed_scaling_factor != 1.0 is not supported' + if moe_top_k > 1 and hf_config.norm_topk_prob: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE + else: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.NONE + elif hf_config.topk_method == 'noaux_tc': + moe_topk_method = MoeConfig.TopKMethod.NOAUX_TC + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED + else: + raise AssertionError( + f'Unsupported topk_method in hf_config: {hf_config.topk_method}' + ) + + rotary_scaling = None + if hf_config.rope_scaling: + rotary_scaling = { + 'beta_fast': + hf_config.rope_scaling['beta_fast'], + 'beta_slow': + hf_config.rope_scaling['beta_slow'], + 'factor': + hf_config.rope_scaling['factor'], + 'mscale': + hf_config.rope_scaling['mscale'], + 'mscale_all_dim': + hf_config.rope_scaling['mscale_all_dim'], + 'original_max_position_embeddings': + hf_config.rope_scaling['original_max_position_embeddings'], + 'type': + 'yarn', + } + + moe_config = MoeConfig( + num_experts=hf_config.n_routed_experts, + shared_expert_intermediate_size=hf_config.n_shared_experts * + hf_config.moe_intermediate_size, + top_k=moe_top_k, + normalization_mode=moe_renorm_mode, + device_limited_n_group=hf_config.n_group, + device_limited_topk_group=hf_config.topk_group, + device_limited_routed_scaling_factor=moe_routed_scaling_factor, + topk_method=moe_topk_method) + moe_config.validate() + + return cls( + architecture=hf_config.architectures[0], + dtype=dtype, + num_hidden_layers=hf_config.num_hidden_layers, + num_attention_heads=hf_config.num_attention_heads, + hidden_size=hf_config.hidden_size, + intermediate_size=hf_config.intermediate_size, + num_key_value_heads=hf_config.num_key_value_heads, + vocab_size=hf_config.vocab_size, + position_embedding_type='rope_gpt_neox', + max_position_embeddings=hf_config.max_position_embeddings, + hidden_act='swiglu', + norm_epsilon=hf_config.rms_norm_eps, + rotary_base=hf_config.rope_theta, + rotary_scaling=rotary_scaling, # TODO: modify this + moe_inter_size=hf_config.moe_intermediate_size, + moe=moe_config, + mapping=mapping, + quantization=quant_config, + kv_lora_rank=hf_config.kv_lora_rank, + q_lora_rank=hf_config.q_lora_rank, + qk_nope_head_dim=hf_config.qk_nope_head_dim, + qk_rope_head_dim=hf_config.qk_rope_head_dim, + v_head_dim=hf_config.v_head_dim, + topk_method=hf_config.topk_method, + first_k_dense_replace=hf_config.first_k_dense_replace, + moe_layer_freq=hf_config.moe_layer_freq, + coring_func=hf_config.scoring_func, + fp8_format=False, + **kwargs) diff --git a/tensorrt_llm/models/deepseek_v2/convert.py b/tensorrt_llm/models/deepseek_v2/convert.py index e9222e6ac23c..158b527fc70f 100755 --- a/tensorrt_llm/models/deepseek_v2/convert.py +++ b/tensorrt_llm/models/deepseek_v2/convert.py @@ -18,141 +18,16 @@ import torch from transformers import AutoConfig, AutoModelForCausalLM -from tensorrt_llm.layers import MoeConfig - from ..._utils import pad_vocab_size, release_gc -from ...mapping import Mapping +from ...layers import MoeConfig +from ...quantization import QuantAlgo from ..convert_utils import get_tllm_linear_weight +from .config import DeepSeekV2Config # `Override num_hidden_layers` used for reduce number of hidden layers in DeepseekV2ForCausalLM for debug purpose OVERRIDE_HIDDEN_LAYERS = None # 2 -## Convert config parameters to dict -def create_trt_config_from_hf(model_dir, - dtype, - mapping: Mapping, - override_fields: dict = {}): - config = {} - assert isinstance(model_dir, str) - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - # Override num_hidden_layers - if OVERRIDE_HIDDEN_LAYERS is not None: - hf_config.num_hidden_layers = OVERRIDE_HIDDEN_LAYERS - print( - f'Override hidden layers to {hf_config.num_hidden_layers} for DeepseekV2ForCausalLM' - ) - dtype = dtype - n_layer = hf_config.num_hidden_layers - n_head = hf_config.num_attention_heads - n_embd = hf_config.hidden_size - inter_size = hf_config.intermediate_size - n_kv_head = hf_config.num_key_value_heads - vocab_size = hf_config.vocab_size - n_positions = hf_config.max_position_embeddings - hidden_act = 'swiglu' # TRT-LLM request make gated activation explicit for MOE implementation - rotary_base = hf_config.rope_theta - rms_norm_eps = hf_config.rms_norm_eps - rotary_scaling_beta_fast = hf_config.rope_scaling['beta_fast'] - rotary_scaling_beta_slow = hf_config.rope_scaling['beta_slow'] - rotary_scaling_factor = hf_config.rope_scaling['factor'] - rotary_scaling_mscale = hf_config.rope_scaling['mscale'] - rotary_scaling_mscale_all_dim = hf_config.rope_scaling['mscale_all_dim'] - rotary_scaling_original_max_position_embeddings = hf_config.rope_scaling[ - 'original_max_position_embeddings'] - rotary_scaling_type = 'yarn' - kv_lora_rank = hf_config.kv_lora_rank - q_lora_rank = hf_config.q_lora_rank - qk_nope_head_dim = hf_config.qk_nope_head_dim - qk_rope_head_dim = hf_config.qk_rope_head_dim - v_head_dim = hf_config.v_head_dim - moe_num_experts = hf_config.n_routed_experts - moe_inter_size = hf_config.moe_intermediate_size - moe_num_shared_experts = hf_config.n_shared_experts - moe_top_k = hf_config.num_experts_per_tok - moe_n_group = hf_config.n_group - moe_topk_group = hf_config.topk_group - moe_routed_scaling_factor = hf_config.routed_scaling_factor - assert moe_routed_scaling_factor > 0, 'routed_scaling_factor should be greater than 0' - if hf_config.topk_method == 'group_limited_greedy': - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED - elif hf_config.topk_method == 'greedy': - assert moe_routed_scaling_factor == 1.0, 'The combination of topk_method == greedy and routed_scaling_factor != 1.0 is not supported' - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.NONE - else: - raise AssertionError( - 'Unsupported topk_method in hf_config: {hf_config.topk_method}') - - config = { - 'architecture': 'DeepseekV2ForCausalLM', - 'dtype': dtype, - 'logits_type': 'float32', - 'num_hidden_layers': n_layer, - 'num_attention_heads': n_head, - 'hidden_size': n_embd, - 'intermediate_size': inter_size, - 'num_key_value_heads': n_kv_head, - 'vocab_size': vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': n_positions, - 'hidden_act': hidden_act, - 'rotary_base': rotary_base, - 'norm_epsilon': rms_norm_eps, - 'rotary_scaling': { - 'beta_fast': rotary_scaling_beta_fast, - 'beta_slow': rotary_scaling_beta_slow, - 'factor': rotary_scaling_factor, - 'mscale': rotary_scaling_mscale, - 'mscale_all_dim': rotary_scaling_mscale_all_dim, - 'original_max_position_embeddings': - rotary_scaling_original_max_position_embeddings, - 'type': rotary_scaling_type, - }, - 'mapping': { - 'world_size': mapping.tp_size * mapping.pp_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - 'moe_tp_size': mapping.moe_tp_size, - 'moe_ep_size': mapping.moe_ep_size, - }, - 'kv_lora_rank': kv_lora_rank, - 'q_lora_rank': q_lora_rank, - 'qk_nope_head_dim': qk_nope_head_dim, - 'qk_rope_head_dim': qk_rope_head_dim, - 'v_head_dim': v_head_dim, - 'moe_num_experts': moe_num_experts, - 'moe_inter_size': moe_inter_size, - 'moe_num_shared_experts': moe_num_shared_experts, - 'moe_top_k': moe_top_k, - 'moe_renorm_mode': moe_renorm_mode, - 'moe_n_group': moe_n_group, - 'moe_topk_group': moe_topk_group, - 'moe_routed_scaling_factor': moe_routed_scaling_factor, - } - - config.update(override_fields) - - moe_config = MoeConfig( - num_experts=config['moe_num_experts'], - shared_expert_intermediate_size=config['moe_num_shared_experts'] * - config['moe_inter_size'], - top_k=config['moe_top_k'], - normalization_mode=config['moe_renorm_mode'], - device_limited_n_group=config['moe_n_group'], - device_limited_topk_group=config['moe_topk_group'], - device_limited_routed_scaling_factor=config['moe_routed_scaling_factor'] - ) - moe_config.validate() - - return config - - ## Get HF model def load_hf_deepseek(model_dir, load_model_on_cpu=False): hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) @@ -225,42 +100,42 @@ def get_param_weight(weight, prefix): return results -def convert_deepseekv2(hf_model, - config, - mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0): +def load_weights_from_hf_model(hf_model, + config: DeepSeekV2Config, + use_parallel_embedding=False, + sharding_dim=0): + quant_algo = config.quantization.quant_algo + use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] + if quant_algo == QuantAlgo.W8A16: + plugin_weight_only_quant_type = torch.int8 + elif quant_algo == QuantAlgo.W4A16: + plugin_weight_only_quant_type = torch.quint4x2 + else: + plugin_weight_only_quant_type = None + use_gemm_woq_plugin = True weights = {} tik = time.time() model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - moe_config = MoeConfig( - num_experts=config['moe_num_experts'], - shared_expert_intermediate_size=config['moe_num_shared_experts'] * - config['moe_inter_size'], - top_k=config['moe_top_k'], - normalization_mode=config['moe_renorm_mode'], - device_limited_n_group=config['moe_n_group'], - device_limited_topk_group=config['moe_topk_group'], - device_limited_routed_scaling_factor=config['moe_routed_scaling_factor'] - ) - - layers_range = mapping.pp_layers(config['num_hidden_layers']) + dtype = getattr(torch, config.dtype) + + mapping = config.mapping + moe_config = config.moe + first_k_dense_replace = config.first_k_dense_replace + layers_range = mapping.pp_layers(config.num_hidden_layers) def convert_layer(l): prefix = f'model.layers.{l}.' trtllm_prex = f'transformer.layers.{l - layers_range[0]}.' # Fuse matrices for compression # Split matrices for decompression - q_lora_rank = config['q_lora_rank'] - kv_lora_rank = config['kv_lora_rank'] - num_heads = config['num_attention_heads'] - qk_nope_head_dim = config['qk_nope_head_dim'] - qk_rope_head_dim = config['qk_rope_head_dim'] - v_head_dim = config['v_head_dim'] - hidden_size = config['hidden_size'] + q_lora_rank = config.q_lora_rank + kv_lora_rank = config.kv_lora_rank + num_heads = config.num_attention_heads + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + hidden_size = config.hidden_size if q_lora_rank is not None: q_a_proj_weight = get_weight(model_params, @@ -392,15 +267,19 @@ def convert_layer(l): get_param_weight(kv_b_proj_weight, trtllm_prex + 'attention.kv_b_proj')) weights.update( - get_tllm_linear_weight(o_proj_weight, - trtllm_prex + 'attention.dense.')) + get_tllm_linear_weight( + o_proj_weight, + trtllm_prex + 'attention.dense.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) if q_lora_rank is not None: weights.update( get_tllm_linear_weight(q_a_layernorm_weight, trtllm_prex + 'attention.q_a_layernorm.')) - if moe_config.has_moe() and l > 0: + if moe_config.has_moe() and l >= first_k_dense_replace: rank_experts = list(range(moe_config.num_experts)) if mapping.has_moe_ep(): rank_experts = mapping.ep_experts(moe_config.num_experts) @@ -439,14 +318,22 @@ def convert_layer(l): moe_experts_down_proj_weights = get_weight( model_params, prefix + 'mlp.experts.down_proj', dtype) weights.update( - get_tllm_linear_weight(moe_experts_down_proj_weights, - trtllm_prex + 'mlp.proj.')) + get_tllm_linear_weight( + moe_experts_down_proj_weights, + trtllm_prex + 'mlp.proj.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) ## mlp.experts.up_gate.weight moe_experts_up_gate_proj_weights = get_weight( model_params, prefix + 'mlp.experts.up_gate_proj', dtype) weights.update( - get_tllm_linear_weight(moe_experts_up_gate_proj_weights, - trtllm_prex + 'mlp.fc.')) + get_tllm_linear_weight( + moe_experts_up_gate_proj_weights, + trtllm_prex + 'mlp.fc.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) ## MOE hardcoded routing_input into trt.float32, please refer to moe.py line 397 moe_experts_gate_weights = get_weight(model_params, prefix + 'mlp.gate', @@ -455,6 +342,15 @@ def convert_layer(l): get_tllm_linear_weight(moe_experts_gate_weights, trtllm_prex + 'mlp.router.')) + if moe_config.topk_method == MoeConfig.TopKMethod.NOAUX_TC: + e_score_correction_bias = get_weight( + model_params, prefix + 'mlp.gate.e_score_correction_bias', + torch.float32, '') + weights.update( + get_param_weight( + e_score_correction_bias, + trtllm_prex + 'mlp.e_score_correction_bias')) + if moe_config.shared_expert_intermediate_size > 0: shared_moe_up_proj_weights = get_weight( model_params, prefix + 'mlp.shared_experts.up_proj', dtype) @@ -487,13 +383,21 @@ def convert_layer(l): weights.update( get_tllm_linear_weight( shared_moe_gate_up_proj_weights, - trtllm_prex + 'mlp.shared_expert.fc.')) + trtllm_prex + 'mlp.shared_expert.fc.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type= + plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) ## mlp.shared_experts.down_proj.weight weights.update( get_tllm_linear_weight( shared_moe_down_proj_weights, - trtllm_prex + 'mlp.shared_expert.proj.')) + trtllm_prex + 'mlp.shared_expert.proj.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type= + plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) else: ## Current MLP layer is only one, if it goes large consider to do fuse @@ -504,7 +408,12 @@ def convert_layer(l): mapping.tp_rank, dim=0) weights.update( - get_tllm_linear_weight(split_gate, trtllm_prex + 'mlp.gate.')) + get_tllm_linear_weight( + split_gate, + trtllm_prex + 'mlp.gate.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', dtype) @@ -513,7 +422,12 @@ def convert_layer(l): mapping.tp_rank, dim=0) weights.update( - get_tllm_linear_weight(split_fc, trtllm_prex + 'mlp.fc.')) + get_tllm_linear_weight( + split_fc, + trtllm_prex + 'mlp.fc.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', dtype) @@ -522,7 +436,12 @@ def convert_layer(l): mapping.tp_rank, dim=1) weights.update( - get_tllm_linear_weight(split_proj, trtllm_prex + 'mlp.proj.')) + get_tllm_linear_weight( + split_proj, + trtllm_prex + 'mlp.proj.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) # Layer norms do not use tensor parallelism input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', @@ -540,11 +459,11 @@ def convert_layer(l): if hf_model.config.tie_word_embeddings: # lm_head.weight has the same weights as embedding if mapping.is_last_pp_rank(): - if config['vocab_size'] % mapping.tp_size != 0: + if config.vocab_size % mapping.tp_size != 0: # padding - vocab_size_padded = pad_vocab_size(config['vocab_size'], + vocab_size_padded = pad_vocab_size(config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config['vocab_size'] + pad_width = vocab_size_padded - config.vocab_size v = torch.nn.functional.pad(v, (0, 0, 0, pad_width), 'constant', 0) weights['lm_head.weight'] = split(v, mapping.tp_size, @@ -559,11 +478,11 @@ def convert_layer(l): lm_head_weights = get_weight(model_params, 'lm_head', dtype) if mapping.is_last_pp_rank(): - if config['vocab_size'] % mapping.tp_size != 0: + if config.vocab_size % mapping.tp_size != 0: # padding - vocab_size_padded = pad_vocab_size(config['vocab_size'], + vocab_size_padded = pad_vocab_size(config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config['vocab_size'] + pad_width = vocab_size_padded - config.vocab_size lm_head_weights = torch.nn.functional.pad(lm_head_weights, (0, 0, 0, pad_width), 'constant', diff --git a/tensorrt_llm/models/deepseek_v2/model.py b/tensorrt_llm/models/deepseek_v2/model.py index 16b13ad6883c..4fa7798b3263 100755 --- a/tensorrt_llm/models/deepseek_v2/model.py +++ b/tensorrt_llm/models/deepseek_v2/model.py @@ -13,26 +13,31 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os from typing import Optional import torch +from tqdm import tqdm from ..._utils import pad_vocab_size, torch_dtype_to_str from ...functional import Tensor, non_gated_version, recv, send from ...layers import (MOE, AttentionMaskType, ColumnLinear, - DeepseekV2Attention, Embedding, GatedMLP, MoeConfig, + DeepseekV2Attention, Embedding, GatedMLP, PositionEmbeddingType, RmsNorm, SharedMoE) +from ...layers.moe import MOEWeightWrapper from ...mapping import Mapping from ...module import Module from ...plugin import init_all_reduce_helper +from ..model_weights_loader import ModelWeightsLoader from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) -from .convert import convert_deepseekv2, create_trt_config_from_hf + QuantConfig) +from .config import DeepSeekV2Config +from .convert import load_hf_deepseek, load_weights_from_hf_model class DeepseekV2DecoderLayer(Module): - def __init__(self, config: PretrainedConfig, layer_idx: int): + def __init__(self, config: DeepSeekV2Config, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.config = config @@ -61,13 +66,6 @@ def __init__(self, config: PretrainedConfig, layer_idx: int): position_embedding_type=PositionEmbeddingType.learned_absolute, rotary_embedding_base=config.rotary_base, rotary_embedding_scaling=None, - rotary_embedding_beta_fast=config.rotary_scaling['beta_fast'], - rotary_embedding_beta_slow=config.rotary_scaling['beta_slow'], - rotary_embedding_mscale=config.rotary_scaling['mscale'], - rotary_embedding_mscale_all_dim=config. - rotary_scaling['mscale_all_dim'], - rotary_embedding_origin_max_position=config. - rotary_scaling['original_max_position_embeddings'], rotary_scaling=config.rotary_scaling, tp_group=config.mapping.tp_group, tp_size=config.mapping.tp_size, @@ -80,21 +78,12 @@ def __init__(self, config: PretrainedConfig, layer_idx: int): ### Distinguish dense MLP and MoE MLP # dense_config = DenseConfig(intermediate_size=config.intermediate_size) - moe_config = MoeConfig( - num_experts=config.moe_num_experts, - shared_expert_intermediate_size=config.moe_num_shared_experts * - config.moe_inter_size, - top_k=config.moe_top_k, - normalization_mode=config.moe_renorm_mode, - device_limited_n_group=config.moe_n_group, - device_limited_topk_group=config.moe_topk_group, - device_limited_routed_scaling_factor=config. - moe_routed_scaling_factor) + moe_config = config.moe # layer_config = LayerMLPConfig(config=[dense_config, moe_config], moe_layer_idx_min=0, # moe_layer_idx_max=config.num_hidden_layers, # total_num_layers=config.num_hidden_layers) - if moe_config.num_experts > 0 and layer_idx > 0: + if moe_config.num_experts > 0 and layer_idx >= config.first_k_dense_replace: hidden_act = config.hidden_act mlp_hidden_size = config.moe_inter_size mlp_kwargs = {'moe_config': moe_config, 'mapping': config.mapping} @@ -118,6 +107,7 @@ def __init__(self, config: PretrainedConfig, layer_idx: int): bias=False, tp_group=config.mapping.tp_group, tp_size=config.mapping.tp_size, + quant_mode=config.quant_mode, **mlp_kwargs) ### Pose layernorm in Deepseek v2 is same as Llama @@ -159,7 +149,7 @@ def forward(self, class DeepseekV2Model(Module): - def __init__(self, config: PretrainedConfig) -> None: + def __init__(self, config: DeepSeekV2Config) -> None: super().__init__() init_all_reduce_helper() # enable use_customer_all_reduce self.dtype = config.dtype @@ -221,8 +211,9 @@ def forward(self, class DeepseekV2ForCausalLM(DecoderModelForCausalLM): + config_class = DeepSeekV2Config - def __init__(self, config: PretrainedConfig): + def __init__(self, config: DeepSeekV2Config): transformer = DeepseekV2Model(config) vocab_size_padded = pad_vocab_size(config.vocab_size, config.mapping.tp_size) @@ -241,22 +232,89 @@ def __init__(self, config: PretrainedConfig): @classmethod def from_hugging_face(cls, - hf_model, model_dir, dtype: str = 'auto', mapping: Optional[Mapping] = None, - override_fields={}, + quant_config: Optional[QuantConfig] = None, **kwargs): - assert hf_model is not None + load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) if mapping is None: mapping = Mapping() - config = create_trt_config_from_hf(model_dir, - dtype, - mapping=mapping, - override_fields=override_fields) - print(config) - pretrained_config = PretrainedConfig.from_dict(config) - pretrained_config.set_rank(mapping.rank) # TODO:remove this hack + config = DeepSeekV2Config.from_hugging_face(model_dir, + dtype=dtype, + mapping=mapping, + quant_config=quant_config, + **kwargs) + if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: + if config.q_lora_rank is not None: # Deepseek-V2&V3 + custom_dict = { + "fused_a": ["q_a_proj", "kv_a_proj_with_mqa"], + "q_a_layernrom": "q_a_layernorm", + "kv_a_layernorm": "kv_a_layernorm", + "q_b_proj": "q_b_proj.weight", + "kv_b_proj": "kv_b_proj.weight", + "fused_q_proj": ["q_b_proj.weight", "kv_b_proj.weight"], + "shared_expert": "shared_experts", + "e_score_correction_bias": "gate.e_score_correction_bias", + } + else: # Deepseek-V2-Lite + custom_dict = { + "fused_a": "kv_a_proj_with_mqa", + "kv_a_layernorm": "kv_a_layernorm", + "q_b_proj": "q_proj.weight", + "kv_b_proj": "kv_b_proj.weight", + "fused_q_proj": ["q_proj.weight", "kv_b_proj.weight"], + "shared_expert": "shared_experts", + "e_score_correction_bias": "gate.e_score_correction_bias", + } + + loader = ModelWeightsLoader(model_dir, custom_dict) + model = cls(config) + for tllm_key, _ in model.named_parameters(): + sub_module = model + for attr in tllm_key.split(".")[:-1]: + sub_module = getattr(sub_module, attr) + if "router" in tllm_key or isinstance(sub_module, + MOEWeightWrapper): + sub_module_dic = sub_module.tllm_to_externel_key_dict + sub_module_dic["mlp"] = "mlp" + if "fc" in sub_module_dic.keys(): + sub_module_dic["fc"] = [ + hf_keyword.replace("w1", "gate_proj") + for hf_keyword in sub_module_dic["fc"] + ] + sub_module_dic["fc"] = [ + hf_keyword.replace("w3", "up_proj") + for hf_keyword in sub_module_dic["fc"] + ] + if "proj" in sub_module_dic.keys(): + sub_module_dic["proj"] = [ + hf_keyword.replace("w2", "down_proj") + for hf_keyword in sub_module_dic["proj"] + ] + sub_module.tllm_to_externel_key_dict = sub_module_dic + + def concat_gate_up_proj(weights): + return torch.cat(weights, dim=-2) + + loader.update_key_mapping(model) + tllm_weights = {} + for tllm_key, _ in tqdm(model.named_parameters()): + if tllm_key.endswith("shared_expert.fc.weight"): + updated_map = loader.tllm_to_externel_key_dict + updated_map["fc"] = ["up_proj", "gate_proj"] + loader.tllm_to_externel_key_dict = updated_map + tllm_weights.update( + loader.load(tllm_key, concat_gate_up_proj)) + else: + tllm_weights.update(loader.load(tllm_key)) + loader.fill(tllm_weights) + else: + hf_model = load_hf_deepseek(model_dir, load_model_on_cpu) + weights = load_weights_from_hf_model(hf_model, config) + model = cls(config) + model.load(weights) + return model if dtype == 'auto': dtype = getattr(config, 'torch_dtype', None) diff --git a/tensorrt_llm/models/model_weights_loader.py b/tensorrt_llm/models/model_weights_loader.py index 737dfc6edae1..2c822cb2a10a 100644 --- a/tensorrt_llm/models/model_weights_loader.py +++ b/tensorrt_llm/models/model_weights_loader.py @@ -286,7 +286,6 @@ def load(self, ] else: v = self.load_tensor(external_key, tp_size, tp_dim, tp_rank) - if preprocess is not None: v = preprocess(v) @@ -309,7 +308,7 @@ def load(self, weight_dict = {tllm_key: v} for k, v in weight_dict.items(): - if not v.is_contiguous(): + if v is not None and not v.is_contiguous(): weight_dict[k] = v.contiguous() return weight_dict @@ -353,6 +352,8 @@ def fill(self, weights): for tllm_key, param in self.model.named_parameters(): if param.is_buffer: continue + if weights[tllm_key] is None: + continue w_shape = weights[tllm_key].shape if w_shape != param.shape: logger.warning( diff --git a/tensorrt_llm/quantization/quantize.py b/tensorrt_llm/quantization/quantize.py index dc54899debbc..17efc23a9b00 100644 --- a/tensorrt_llm/quantization/quantize.py +++ b/tensorrt_llm/quantization/quantize.py @@ -35,6 +35,7 @@ def quantize_layers( '*position_embedding', '*block_embedding', '*shared_expert_gate', + '*fused_a', ] for name, module, parent in model.named_modules_with_parent():