From a7e54c552dab3e22ef2a11b6b07f26f22704b1e8 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 3 Jul 2026 08:27:48 +0000 Subject: [PATCH] [TRTLLM-14022][feat] BREAKING: Remove legacy TensorRT Python backend Remove the TensorRT execution backend from the Python package so that `import tensorrt_llm` and the PyTorch backend run without the `tensorrt` package installed. The `tensorrt` pip dependency itself is dropped in the follow-up C++/packaging step; `requirements.txt` is unchanged here. - Delete the TensorRT graph/build/runtime machinery: builder, network, functional graph ops, layers/, plugin/, the legacy models/ implementations, the runtime/ session and ModelRunner* classes, _tensorrt_engine, TRT quant layers and tools, and the trtllm-build/refit/prune CLIs. - Keep import paths stable: partly-TensorRT modules (functional.py, models/modeling_utils.py, _common.py, runtime/, bench/build) are gutted in place so the TensorRT-free symbols used by the PyTorch backend stay put; no `_torch/` import churn. - Drop the TensorRT public API surface (Builder/BuildConfig/TrtLlmArgs/...); `LLM(backend="tensorrt")` now raises and `trtllm-serve --backend` accepts pytorch/_autodeploy only (api_stability serve-CLI reference updated). - Remove or de-TRT the tests that exercised the deleted code (disaggregated trt-backend tests, TRT bench/openai e2e params, mixed llmapi unittest suites) and delist them from test-db/qa/waives; shared test harnesses keep their signatures for the PyTorch twins. - Port test_gather_generation_logits_cuda_graph to the PyTorch accuracy suite; repoint examples/apps to the PyTorch LLM; regenerate ruff-legacy-baseline.json. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- examples/apps/chat.py | 12 +- examples/apps/fastapi_server.py | 9 +- ruff-legacy-baseline.json | 1330 +-- setup.py | 3 - tensorrt_llm/__init__.py | 50 +- tensorrt_llm/_common.py | 227 +- tensorrt_llm/_tensorrt_engine/__init__.py | 3 - tensorrt_llm/_torch/auto_deploy/llm_args.py | 14 - .../_torch/custom_ops/torch_custom_ops.py | 3 +- .../native/bounce/gather_scatter.py | 3 +- .../disaggregation/native/bounce/impl.py | 2 +- .../_torch/disaggregation/native/transfer.py | 3 +- .../_torch/distributed/allreduce_helper.py | 155 + tensorrt_llm/_torch/distributed/ops.py | 3 +- .../checkpoints/hf/qwen3_moe_weight_mapper.py | 7 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 8 +- tensorrt_llm/_utils.py | 181 +- tensorrt_llm/bench/benchmark/__init__.py | 3 +- .../bench/benchmark/utils/asynchronous.py | 3 +- tensorrt_llm/bench/benchmark/utils/general.py | 2 +- tensorrt_llm/bench/build/build.py | 298 +- tensorrt_llm/builder.py | 1292 --- tensorrt_llm/commands/bench.py | 2 - tensorrt_llm/commands/build.py | 553 -- tensorrt_llm/commands/eval.py | 25 +- tensorrt_llm/commands/prune.py | 122 - tensorrt_llm/commands/refit.py | 183 - tensorrt_llm/commands/serve.py | 83 +- tensorrt_llm/disaggregated_params.py | 5 - tensorrt_llm/evaluate/cnn_dailymail.py | 5 +- tensorrt_llm/evaluate/covost2.py | 5 +- tensorrt_llm/evaluate/json_mode_eval.py | 5 +- tensorrt_llm/evaluate/lm_eval.py | 17 +- tensorrt_llm/evaluate/longbench_v2.py | 5 +- tensorrt_llm/evaluate/mmlu.py | 5 +- tensorrt_llm/executor/base_worker.py | 32 +- tensorrt_llm/executor/executor.py | 3 +- tensorrt_llm/executor/ray_gpu_worker.py | 5 +- tensorrt_llm/executor/rpc_worker.py | 7 +- tensorrt_llm/executor/worker.py | 7 +- tensorrt_llm/functional.py | 7922 +---------------- tensorrt_llm/graph_rewriting.py | 645 -- tensorrt_llm/layers/__init__.py | 79 - tensorrt_llm/layers/activation.py | 22 - tensorrt_llm/layers/attention.py | 2808 ------ tensorrt_llm/layers/cast.py | 29 - tensorrt_llm/layers/conv.py | 260 - tensorrt_llm/layers/embedding.py | 673 -- tensorrt_llm/layers/language_adapter.py | 94 - tensorrt_llm/layers/linear.py | 551 -- tensorrt_llm/layers/lora.py | 181 - tensorrt_llm/layers/mlp.py | 565 -- tensorrt_llm/layers/moe.py | 1553 ---- tensorrt_llm/layers/normalization.py | 341 - tensorrt_llm/layers/pooling.py | 38 - tensorrt_llm/layers/recurrent.py | 316 - tensorrt_llm/layers/ssm.py | 358 - tensorrt_llm/llmapi/__init__.py | 9 +- tensorrt_llm/llmapi/build_cache.py | 312 - tensorrt_llm/llmapi/llm.py | 259 +- tensorrt_llm/llmapi/llm_args.py | 425 +- tensorrt_llm/llmapi/llm_utils.py | 628 +- tensorrt_llm/logger.py | 33 +- tensorrt_llm/lora_helper.py | 21 - tensorrt_llm/lora_manager.py | 111 +- tensorrt_llm/models/__init__.py | 217 +- tensorrt_llm/models/baichuan/__init__.py | 14 - tensorrt_llm/models/baichuan/config.py | 92 - tensorrt_llm/models/baichuan/convert.py | 931 -- tensorrt_llm/models/baichuan/model.py | 253 - tensorrt_llm/models/bert/__init__.py | 14 - tensorrt_llm/models/bert/config.py | 117 - tensorrt_llm/models/bert/convert.py | 309 - tensorrt_llm/models/bert/model.py | 548 -- tensorrt_llm/models/bloom/__init__.py | 14 - tensorrt_llm/models/bloom/model.py | 165 - tensorrt_llm/models/chatglm/__init__.py | 14 - tensorrt_llm/models/chatglm/config.py | 182 - tensorrt_llm/models/chatglm/convert.py | 727 -- tensorrt_llm/models/chatglm/model.py | 372 - tensorrt_llm/models/clip/__init__.py | 14 - tensorrt_llm/models/clip/model.py | 213 - tensorrt_llm/models/cogvlm/__init__.py | 14 - tensorrt_llm/models/cogvlm/config.py | 44 - tensorrt_llm/models/cogvlm/convert.py | 240 - tensorrt_llm/models/cogvlm/model.py | 291 - tensorrt_llm/models/commandr/__init__.py | 14 - tensorrt_llm/models/commandr/config.py | 88 - tensorrt_llm/models/commandr/model.py | 195 - tensorrt_llm/models/dbrx/__init__.py | 14 - tensorrt_llm/models/dbrx/config.py | 59 - tensorrt_llm/models/dbrx/model.py | 188 - tensorrt_llm/models/deepseek_v1/__init__.py | 14 - tensorrt_llm/models/deepseek_v1/config.py | 105 - tensorrt_llm/models/deepseek_v1/convert.py | 290 - tensorrt_llm/models/deepseek_v1/model.py | 279 - tensorrt_llm/models/deepseek_v2/__init__.py | 14 - tensorrt_llm/models/deepseek_v2/config.py | 148 - tensorrt_llm/models/deepseek_v2/convert.py | 929 -- tensorrt_llm/models/deepseek_v2/model.py | 361 - tensorrt_llm/models/dit/__init__.py | 14 - tensorrt_llm/models/dit/model.py | 382 - tensorrt_llm/models/eagle/__init__.py | 14 - tensorrt_llm/models/eagle/config.py | 222 - tensorrt_llm/models/eagle/model.py | 1327 --- tensorrt_llm/models/enc_dec/__init__.py | 14 - tensorrt_llm/models/enc_dec/model.py | 2195 ----- tensorrt_llm/models/falcon/__init__.py | 14 - tensorrt_llm/models/falcon/config.py | 118 - tensorrt_llm/models/falcon/convert.py | 499 -- tensorrt_llm/models/falcon/model.py | 273 - tensorrt_llm/models/gemma/__init__.py | 14 - tensorrt_llm/models/gemma/config.py | 224 - tensorrt_llm/models/gemma/convert.py | 926 -- tensorrt_llm/models/gemma/model.py | 396 - tensorrt_llm/models/gemma/smoothquant.py | 848 -- tensorrt_llm/models/gemma/utils/__init__.py | 14 - tensorrt_llm/models/gemma/utils/layers.py | 39 - tensorrt_llm/models/gemma/utils/modules.py | 206 - tensorrt_llm/models/gemma/utils/params.py | 72 - .../gemma/utils/positional_embeddings.py | 92 - tensorrt_llm/models/gemma/utils/sampler.py | 191 - .../models/gemma/utils/transformer.py | 114 - tensorrt_llm/models/gemma/weight.py | 181 - tensorrt_llm/models/generation_mixin.py | 889 -- tensorrt_llm/models/gpt/__init__.py | 14 - tensorrt_llm/models/gpt/config.py | 324 - tensorrt_llm/models/gpt/convert.py | 1455 --- tensorrt_llm/models/gpt/model.py | 417 - tensorrt_llm/models/gptj/__init__.py | 14 - tensorrt_llm/models/gptj/config.py | 55 - tensorrt_llm/models/gptj/convert.py | 170 - tensorrt_llm/models/gptj/model.py | 202 - tensorrt_llm/models/gptneox/__init__.py | 14 - tensorrt_llm/models/gptneox/model.py | 147 - tensorrt_llm/models/grok/__init__.py | 14 - tensorrt_llm/models/grok/convert.py | 518 -- tensorrt_llm/models/grok/model.py | 273 - tensorrt_llm/models/grok/weight.py | 35 - tensorrt_llm/models/llama/__init__.py | 14 - tensorrt_llm/models/llama/config.py | 280 - tensorrt_llm/models/llama/convert.py | 2449 ----- tensorrt_llm/models/llama/model.py | 614 -- tensorrt_llm/models/mamba/__init__.py | 14 - tensorrt_llm/models/mamba/config.py | 313 - tensorrt_llm/models/mamba/convert.py | 245 - tensorrt_llm/models/mamba/model.py | 471 - tensorrt_llm/models/medusa/__init__.py | 14 - tensorrt_llm/models/medusa/config.py | 114 - tensorrt_llm/models/medusa/model.py | 267 - tensorrt_llm/models/medusa/weight.py | 648 -- tensorrt_llm/models/mllama/__init__.py | 14 - tensorrt_llm/models/mllama/config.py | 256 - tensorrt_llm/models/mllama/model.py | 1569 ---- tensorrt_llm/models/mmdit_sd3/__init__.py | 14 - tensorrt_llm/models/mmdit_sd3/config.py | 132 - tensorrt_llm/models/mmdit_sd3/model.py | 620 -- tensorrt_llm/models/model_weights_loader.py | 402 - tensorrt_llm/models/modeling_utils.py | 1474 +-- tensorrt_llm/models/mpt/__init__.py | 14 - tensorrt_llm/models/mpt/model.py | 176 - .../models/multimodal_encoders/__init__.py | 14 - .../models/multimodal_encoders/config.py | 116 - .../models/multimodal_encoders/model.py | 175 - tensorrt_llm/models/nemotron_nas/__init__.py | 14 - tensorrt_llm/models/nemotron_nas/config.py | 208 - tensorrt_llm/models/nemotron_nas/convert.py | 381 - .../models/nemotron_nas/layer_config.py | 86 - tensorrt_llm/models/nemotron_nas/model.py | 816 -- tensorrt_llm/models/opt/__init__.py | 14 - tensorrt_llm/models/opt/model.py | 181 - tensorrt_llm/models/phi/__init__.py | 14 - tensorrt_llm/models/phi/config.py | 87 - tensorrt_llm/models/phi/convert.py | 145 - tensorrt_llm/models/phi/model.py | 217 - tensorrt_llm/models/phi3/__init__.py | 14 - tensorrt_llm/models/phi3/config.py | 143 - tensorrt_llm/models/phi3/convert.py | 134 - tensorrt_llm/models/phi3/model.py | 316 - tensorrt_llm/models/phi3/split_weights.py | 218 - tensorrt_llm/models/qwen/__init__.py | 14 - tensorrt_llm/models/qwen/config.py | 211 - tensorrt_llm/models/qwen/convert.py | 1284 --- tensorrt_llm/models/qwen/model.py | 545 -- tensorrt_llm/models/qwen/utils.py | 116 - .../models/recurrentgemma/__init__.py | 14 - tensorrt_llm/models/recurrentgemma/model.py | 624 -- tensorrt_llm/models/redrafter/__init__.py | 14 - tensorrt_llm/models/redrafter/drafter.py | 117 - tensorrt_llm/models/redrafter/model.py | 317 - .../models/redrafter/redrafter_helper.py | 759 -- tensorrt_llm/models/stdit/__init__.py | 14 - tensorrt_llm/models/stdit/config.py | 115 - tensorrt_llm/models/stdit/model.py | 1624 ---- tensorrt_llm/models/unet/__init__.py | 14 - tensorrt_llm/models/unet/attention.py | 332 - tensorrt_llm/models/unet/embeddings.py | 117 - tensorrt_llm/models/unet/pp/attention.py | 107 - tensorrt_llm/models/unet/pp/conv2d.py | 109 - tensorrt_llm/models/unet/pp/groupnorm.py | 57 - tensorrt_llm/models/unet/pp/unet_pp.py | 131 - tensorrt_llm/models/unet/resnet.py | 254 - tensorrt_llm/models/unet/unet_2d_blocks.py | 636 -- tensorrt_llm/models/unet/unet_2d_condition.py | 249 - tensorrt_llm/models/unet/weights.py | 197 - tensorrt_llm/module.py | 288 - tensorrt_llm/network.py | 959 -- tensorrt_llm/parameter.py | 281 - tensorrt_llm/plugin/__init__.py | 23 - tensorrt_llm/plugin/plugin.py | 780 -- tensorrt_llm/profiler.py | 79 - tensorrt_llm/python_plugin.py | 578 -- tensorrt_llm/quantization/functional.py | 1444 +-- tensorrt_llm/quantization/layers.py | 3383 ------- tensorrt_llm/quantization/quantize.py | 611 -- .../quantization/quantize_by_modelopt.py | 91 - tensorrt_llm/runtime/__init__.py | 34 +- tensorrt_llm/runtime/enc_dec_model_runner.py | 537 -- tensorrt_llm/runtime/generation.py | 4785 ---------- tensorrt_llm/runtime/kv_cache_manager.py | 473 - tensorrt_llm/runtime/medusa_utils.py | 182 - .../memory_pools/memory_pools_allocator.py | 80 - tensorrt_llm/runtime/memory_pools/pool.py | 7 - .../memory_pools/pools_kv_cache_manager.py | 59 - tensorrt_llm/runtime/model_config.py | 105 + tensorrt_llm/runtime/model_runner.py | 1001 --- tensorrt_llm/runtime/model_runner_cpp.py | 1223 --- .../runtime/multimodal_model_runner.py | 2744 ------ .../runtime/processor_wrapper/__init__.py | 20 - .../mllama_processor_wrapper.py | 76 - .../processor_wrapper/processor_wrapper.py | 37 - tensorrt_llm/runtime/redrafter_utils.py | 341 - tensorrt_llm/runtime/session.py | 324 - tensorrt_llm/serialization.py | 7 +- tensorrt_llm/serve/openai_server.py | 6 +- tensorrt_llm/tools/multimodal_builder.py | 1844 ---- tensorrt_llm/tools/onnx_utils.py | 80 - tensorrt_llm/tools/plugin_gen/core.py | 713 -- tensorrt_llm/tools/plugin_gen/plugin_gen.py | 374 - tensorrt_llm/tools/plugin_gen/shape_infer.py | 355 - .../plugin_gen/templates/CMakeLists.txt.tpl | 86 - .../plugin_gen/templates/functional.py.tpl | 70 - .../tools/plugin_gen/templates/plugin.cpp.tpl | 310 - .../tools/plugin_gen/templates/plugin.h.tpl | 120 - .../plugin_gen/templates/plugin_common.cpp | 96 - .../plugin_gen/templates/plugin_common.h | 208 - .../templates/tritonPlugins.cpp.tpl | 147 - tensorrt_llm/top_model_mixin.py | 72 - .../usage/llm_args_golden_manifest.json | 2286 ----- tensorrt_llm/usage/llmapi_config.py | 3 +- tensorrt_llm/usage/usage_lib.py | 5 +- tests/integration/defs/.test_durations | 10 - tests/integration/defs/accuracy/README.md | 1 - .../defs/accuracy/accuracy_core.py | 8 +- .../integration/defs/accuracy/test_llm_api.py | 433 - .../defs/accuracy/test_llm_api_pytorch.py | 10 + ...sagg_config_ctxtp2_gentp1_trt_backend.yaml | 16 - .../disagg_config_gen_only_trt_backend.yaml | 15 - .../disagg_config_trt_backend.yaml | 18 - .../defs/disaggregated/test_disaggregated.py | 56 - .../defs/llmapi/_run_llmapi_llm.py | 26 +- .../defs/llmapi/test_llm_api_qa.py | 26 - tests/integration/defs/llmapi/test_llm_e2e.py | 177 - tests/integration/defs/test_e2e.py | 300 +- .../test_lists/qa/llm_function_core.txt | 5 +- .../integration/test_lists/test-db/l0_a10.yml | 11 - .../test_lists/test-db/l0_a100.yml | 1 - .../integration/test_lists/test-db/l0_a30.yml | 2 - .../test_lists/test-db/l0_dgx_h100.yml | 1 - .../test_lists/test-db/l0_dgx_h200.yml | 7 - .../test_lists/test-db/l0_gh200.yml | 2 - .../test_lists/test-db/l0_h100.yml | 11 - .../test_lists/test-db/l0_l40s.yml | 1 - tests/integration/test_lists/waives.txt | 5 - tests/microbenchmarks/README.md | 26 - tests/microbenchmarks/all_reduce.py | 3 +- tests/microbenchmarks/build_time_benchmark.py | 314 - tests/microbenchmarks/build_time_dashboard.py | 199 - tests/microbenchmarks/minimax_all_reduce.py | 2 +- .../references/trtllm_serve_cli.yaml | 2 +- tests/unittest/conftest.py | 14 +- tests/unittest/dynamo/test_imports.py | 1 - tests/unittest/llmapi/apps/_test_llm_chat.py | 8 +- .../unittest/llmapi/apps/_test_llm_server.py | 7 +- .../unittest/llmapi/apps/_test_openai_chat.py | 10 +- .../llmapi/apps/_test_openai_completions.py | 38 +- .../apps/_test_openai_consistent_chat.py | 235 - .../unittest/llmapi/apps/_test_openai_misc.py | 10 +- .../llmapi/apps/_test_openai_multi_chat.py | 146 - .../llmapi/apps/_test_openai_multi_gpu.py | 2 +- .../llmapi/apps/_test_openai_reasoning.py | 5 - .../apps/_test_trtllm_serve_top_logprobs.py | 24 +- tests/unittest/llmapi/run_llm.py | 27 +- .../unittest/llmapi/run_llm_with_postproc.py | 12 +- tests/unittest/llmapi/test_grpc.py | 6 +- tests/unittest/llmapi/test_llm.py | 1573 +--- tests/unittest/llmapi/test_llm_args.py | 278 +- tests/unittest/llmapi/test_llm_download.py | 4 +- tests/unittest/llmapi/test_llm_quant.py | 72 +- tests/unittest/llmapi/test_llm_telemetry.py | 49 - tests/unittest/llmapi/test_llm_utils.py | 48 +- .../unittest/llmapi/test_memory_profiling.py | 18 +- tests/unittest/tools/plugin_gen/__init__.py | 0 .../tools/plugin_gen/kernel_config.py | 49 - .../usage/test_llmapi_config_capture.py | 50 - tests/unittest/utils/test_logger.py | 4 - tests/unittest/utils/util.py | 141 +- 307 files changed, 1204 insertions(+), 95346 deletions(-) delete mode 100644 tensorrt_llm/_tensorrt_engine/__init__.py create mode 100644 tensorrt_llm/_torch/distributed/allreduce_helper.py delete mode 100644 tensorrt_llm/builder.py delete mode 100644 tensorrt_llm/commands/build.py delete mode 100644 tensorrt_llm/commands/prune.py delete mode 100644 tensorrt_llm/commands/refit.py mode change 100755 => 100644 tensorrt_llm/functional.py delete mode 100644 tensorrt_llm/graph_rewriting.py delete mode 100755 tensorrt_llm/layers/__init__.py delete mode 100644 tensorrt_llm/layers/activation.py delete mode 100755 tensorrt_llm/layers/attention.py delete mode 100644 tensorrt_llm/layers/cast.py delete mode 100644 tensorrt_llm/layers/conv.py delete mode 100644 tensorrt_llm/layers/embedding.py delete mode 100755 tensorrt_llm/layers/language_adapter.py delete mode 100644 tensorrt_llm/layers/linear.py delete mode 100644 tensorrt_llm/layers/lora.py delete mode 100644 tensorrt_llm/layers/mlp.py delete mode 100755 tensorrt_llm/layers/moe.py delete mode 100644 tensorrt_llm/layers/normalization.py delete mode 100644 tensorrt_llm/layers/pooling.py delete mode 100644 tensorrt_llm/layers/recurrent.py delete mode 100644 tensorrt_llm/layers/ssm.py delete mode 100644 tensorrt_llm/llmapi/build_cache.py delete mode 100644 tensorrt_llm/models/baichuan/__init__.py delete mode 100644 tensorrt_llm/models/baichuan/config.py delete mode 100644 tensorrt_llm/models/baichuan/convert.py delete mode 100644 tensorrt_llm/models/baichuan/model.py delete mode 100644 tensorrt_llm/models/bert/__init__.py delete mode 100644 tensorrt_llm/models/bert/config.py delete mode 100644 tensorrt_llm/models/bert/convert.py delete mode 100644 tensorrt_llm/models/bert/model.py delete mode 100644 tensorrt_llm/models/bloom/__init__.py delete mode 100644 tensorrt_llm/models/bloom/model.py delete mode 100644 tensorrt_llm/models/chatglm/__init__.py delete mode 100644 tensorrt_llm/models/chatglm/config.py delete mode 100644 tensorrt_llm/models/chatglm/convert.py delete mode 100644 tensorrt_llm/models/chatglm/model.py delete mode 100644 tensorrt_llm/models/clip/__init__.py delete mode 100644 tensorrt_llm/models/clip/model.py delete mode 100644 tensorrt_llm/models/cogvlm/__init__.py delete mode 100644 tensorrt_llm/models/cogvlm/config.py delete mode 100644 tensorrt_llm/models/cogvlm/convert.py delete mode 100644 tensorrt_llm/models/cogvlm/model.py delete mode 100644 tensorrt_llm/models/commandr/__init__.py delete mode 100644 tensorrt_llm/models/commandr/config.py delete mode 100644 tensorrt_llm/models/commandr/model.py delete mode 100644 tensorrt_llm/models/dbrx/__init__.py delete mode 100644 tensorrt_llm/models/dbrx/config.py delete mode 100644 tensorrt_llm/models/dbrx/model.py delete mode 100644 tensorrt_llm/models/deepseek_v1/__init__.py delete mode 100755 tensorrt_llm/models/deepseek_v1/config.py delete mode 100755 tensorrt_llm/models/deepseek_v1/convert.py delete mode 100755 tensorrt_llm/models/deepseek_v1/model.py delete mode 100644 tensorrt_llm/models/deepseek_v2/__init__.py delete mode 100644 tensorrt_llm/models/deepseek_v2/config.py delete mode 100755 tensorrt_llm/models/deepseek_v2/convert.py delete mode 100755 tensorrt_llm/models/deepseek_v2/model.py delete mode 100644 tensorrt_llm/models/dit/__init__.py delete mode 100644 tensorrt_llm/models/dit/model.py delete mode 100644 tensorrt_llm/models/eagle/__init__.py delete mode 100644 tensorrt_llm/models/eagle/config.py delete mode 100644 tensorrt_llm/models/eagle/model.py delete mode 100644 tensorrt_llm/models/enc_dec/__init__.py delete mode 100644 tensorrt_llm/models/enc_dec/model.py delete mode 100644 tensorrt_llm/models/falcon/__init__.py delete mode 100644 tensorrt_llm/models/falcon/config.py delete mode 100644 tensorrt_llm/models/falcon/convert.py delete mode 100644 tensorrt_llm/models/falcon/model.py delete mode 100644 tensorrt_llm/models/gemma/__init__.py delete mode 100644 tensorrt_llm/models/gemma/config.py delete mode 100644 tensorrt_llm/models/gemma/convert.py delete mode 100644 tensorrt_llm/models/gemma/model.py delete mode 100644 tensorrt_llm/models/gemma/smoothquant.py delete mode 100644 tensorrt_llm/models/gemma/utils/__init__.py delete mode 100644 tensorrt_llm/models/gemma/utils/layers.py delete mode 100644 tensorrt_llm/models/gemma/utils/modules.py delete mode 100644 tensorrt_llm/models/gemma/utils/params.py delete mode 100644 tensorrt_llm/models/gemma/utils/positional_embeddings.py delete mode 100644 tensorrt_llm/models/gemma/utils/sampler.py delete mode 100644 tensorrt_llm/models/gemma/utils/transformer.py delete mode 100644 tensorrt_llm/models/gemma/weight.py delete mode 100644 tensorrt_llm/models/generation_mixin.py delete mode 100644 tensorrt_llm/models/gpt/__init__.py delete mode 100644 tensorrt_llm/models/gpt/config.py delete mode 100644 tensorrt_llm/models/gpt/convert.py delete mode 100644 tensorrt_llm/models/gpt/model.py delete mode 100644 tensorrt_llm/models/gptj/__init__.py delete mode 100644 tensorrt_llm/models/gptj/config.py delete mode 100644 tensorrt_llm/models/gptj/convert.py delete mode 100644 tensorrt_llm/models/gptj/model.py delete mode 100644 tensorrt_llm/models/gptneox/__init__.py delete mode 100644 tensorrt_llm/models/gptneox/model.py delete mode 100644 tensorrt_llm/models/grok/__init__.py delete mode 100644 tensorrt_llm/models/grok/convert.py delete mode 100644 tensorrt_llm/models/grok/model.py delete mode 100644 tensorrt_llm/models/grok/weight.py delete mode 100644 tensorrt_llm/models/llama/__init__.py delete mode 100644 tensorrt_llm/models/llama/config.py delete mode 100644 tensorrt_llm/models/llama/convert.py delete mode 100644 tensorrt_llm/models/llama/model.py delete mode 100644 tensorrt_llm/models/mamba/__init__.py delete mode 100644 tensorrt_llm/models/mamba/config.py delete mode 100644 tensorrt_llm/models/mamba/convert.py delete mode 100644 tensorrt_llm/models/mamba/model.py delete mode 100644 tensorrt_llm/models/medusa/__init__.py delete mode 100644 tensorrt_llm/models/medusa/config.py delete mode 100644 tensorrt_llm/models/medusa/model.py delete mode 100644 tensorrt_llm/models/medusa/weight.py delete mode 100644 tensorrt_llm/models/mllama/__init__.py delete mode 100644 tensorrt_llm/models/mllama/config.py delete mode 100644 tensorrt_llm/models/mllama/model.py delete mode 100644 tensorrt_llm/models/mmdit_sd3/__init__.py delete mode 100644 tensorrt_llm/models/mmdit_sd3/config.py delete mode 100644 tensorrt_llm/models/mmdit_sd3/model.py delete mode 100644 tensorrt_llm/models/model_weights_loader.py delete mode 100644 tensorrt_llm/models/mpt/__init__.py delete mode 100644 tensorrt_llm/models/mpt/model.py delete mode 100644 tensorrt_llm/models/multimodal_encoders/__init__.py delete mode 100644 tensorrt_llm/models/multimodal_encoders/config.py delete mode 100644 tensorrt_llm/models/multimodal_encoders/model.py delete mode 100644 tensorrt_llm/models/nemotron_nas/__init__.py delete mode 100644 tensorrt_llm/models/nemotron_nas/config.py delete mode 100644 tensorrt_llm/models/nemotron_nas/convert.py delete mode 100644 tensorrt_llm/models/nemotron_nas/layer_config.py delete mode 100644 tensorrt_llm/models/nemotron_nas/model.py delete mode 100644 tensorrt_llm/models/opt/__init__.py delete mode 100644 tensorrt_llm/models/opt/model.py delete mode 100644 tensorrt_llm/models/phi/__init__.py delete mode 100644 tensorrt_llm/models/phi/config.py delete mode 100644 tensorrt_llm/models/phi/convert.py delete mode 100644 tensorrt_llm/models/phi/model.py delete mode 100644 tensorrt_llm/models/phi3/__init__.py delete mode 100644 tensorrt_llm/models/phi3/config.py delete mode 100644 tensorrt_llm/models/phi3/convert.py delete mode 100644 tensorrt_llm/models/phi3/model.py delete mode 100644 tensorrt_llm/models/phi3/split_weights.py delete mode 100644 tensorrt_llm/models/qwen/__init__.py delete mode 100644 tensorrt_llm/models/qwen/config.py delete mode 100644 tensorrt_llm/models/qwen/convert.py delete mode 100644 tensorrt_llm/models/qwen/model.py delete mode 100644 tensorrt_llm/models/qwen/utils.py delete mode 100644 tensorrt_llm/models/recurrentgemma/__init__.py delete mode 100644 tensorrt_llm/models/recurrentgemma/model.py delete mode 100644 tensorrt_llm/models/redrafter/__init__.py delete mode 100644 tensorrt_llm/models/redrafter/drafter.py delete mode 100644 tensorrt_llm/models/redrafter/model.py delete mode 100644 tensorrt_llm/models/redrafter/redrafter_helper.py delete mode 100644 tensorrt_llm/models/stdit/__init__.py delete mode 100644 tensorrt_llm/models/stdit/config.py delete mode 100644 tensorrt_llm/models/stdit/model.py delete mode 100644 tensorrt_llm/models/unet/__init__.py delete mode 100644 tensorrt_llm/models/unet/attention.py delete mode 100644 tensorrt_llm/models/unet/embeddings.py delete mode 100755 tensorrt_llm/models/unet/pp/attention.py delete mode 100755 tensorrt_llm/models/unet/pp/conv2d.py delete mode 100755 tensorrt_llm/models/unet/pp/groupnorm.py delete mode 100755 tensorrt_llm/models/unet/pp/unet_pp.py delete mode 100644 tensorrt_llm/models/unet/resnet.py delete mode 100644 tensorrt_llm/models/unet/unet_2d_blocks.py delete mode 100644 tensorrt_llm/models/unet/unet_2d_condition.py delete mode 100644 tensorrt_llm/models/unet/weights.py delete mode 100644 tensorrt_llm/module.py delete mode 100644 tensorrt_llm/network.py delete mode 100644 tensorrt_llm/parameter.py delete mode 100644 tensorrt_llm/plugin/__init__.py delete mode 100644 tensorrt_llm/plugin/plugin.py delete mode 100644 tensorrt_llm/python_plugin.py delete mode 100644 tensorrt_llm/quantization/layers.py delete mode 100644 tensorrt_llm/quantization/quantize.py delete mode 100644 tensorrt_llm/runtime/enc_dec_model_runner.py delete mode 100755 tensorrt_llm/runtime/generation.py delete mode 100644 tensorrt_llm/runtime/kv_cache_manager.py delete mode 100644 tensorrt_llm/runtime/medusa_utils.py delete mode 100644 tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py delete mode 100644 tensorrt_llm/runtime/memory_pools/pool.py delete mode 100644 tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py create mode 100644 tensorrt_llm/runtime/model_config.py delete mode 100644 tensorrt_llm/runtime/model_runner.py delete mode 100644 tensorrt_llm/runtime/model_runner_cpp.py delete mode 100644 tensorrt_llm/runtime/multimodal_model_runner.py delete mode 100644 tensorrt_llm/runtime/processor_wrapper/__init__.py delete mode 100644 tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py delete mode 100644 tensorrt_llm/runtime/processor_wrapper/processor_wrapper.py delete mode 100644 tensorrt_llm/runtime/redrafter_utils.py delete mode 100644 tensorrt_llm/runtime/session.py delete mode 100644 tensorrt_llm/tools/multimodal_builder.py delete mode 100644 tensorrt_llm/tools/onnx_utils.py delete mode 100644 tensorrt_llm/tools/plugin_gen/core.py delete mode 100644 tensorrt_llm/tools/plugin_gen/plugin_gen.py delete mode 100644 tensorrt_llm/tools/plugin_gen/shape_infer.py delete mode 100644 tensorrt_llm/tools/plugin_gen/templates/CMakeLists.txt.tpl delete mode 100644 tensorrt_llm/tools/plugin_gen/templates/functional.py.tpl delete mode 100644 tensorrt_llm/tools/plugin_gen/templates/plugin.cpp.tpl delete mode 100644 tensorrt_llm/tools/plugin_gen/templates/plugin.h.tpl delete mode 100644 tensorrt_llm/tools/plugin_gen/templates/plugin_common.cpp delete mode 100644 tensorrt_llm/tools/plugin_gen/templates/plugin_common.h delete mode 100644 tensorrt_llm/tools/plugin_gen/templates/tritonPlugins.cpp.tpl delete mode 100644 tensorrt_llm/top_model_mixin.py delete mode 100644 tests/integration/defs/accuracy/test_llm_api.py delete mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml delete mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml delete mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml delete mode 100644 tests/integration/defs/llmapi/test_llm_e2e.py delete mode 100644 tests/microbenchmarks/README.md delete mode 100644 tests/microbenchmarks/build_time_benchmark.py delete mode 100644 tests/microbenchmarks/build_time_dashboard.py delete mode 100644 tests/unittest/llmapi/apps/_test_openai_consistent_chat.py delete mode 100644 tests/unittest/llmapi/apps/_test_openai_multi_chat.py delete mode 100644 tests/unittest/tools/plugin_gen/__init__.py delete mode 100644 tests/unittest/tools/plugin_gen/kernel_config.py diff --git a/examples/apps/chat.py b/examples/apps/chat.py index 0e3f2d928158..661949a3a607 100755 --- a/examples/apps/chat.py +++ b/examples/apps/chat.py @@ -5,8 +5,8 @@ import colorama from transformers import AutoTokenizer, PreTrainedTokenizer -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import BuildConfig, KvCacheConfig, SamplingParams +from tensorrt_llm import LLM +from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams class LlmConsole(code.InteractiveConsole): @@ -72,10 +72,6 @@ def main(model: str, tokenizer: str, tp_size: int): free_gpu_memory_fraction=0.8) kv_cache_config.enable_block_reuse = True - build_config = BuildConfig(max_batch_size=1, - max_input_len=6000, - max_num_tokens=10240) - sampling_params = SamplingParams(max_tokens=100, temperature=0.5, top_p=0.95, @@ -83,7 +79,9 @@ def main(model: str, tokenizer: str, tp_size: int): llm = LLM(model, tokenizer, - build_config=build_config, + max_batch_size=1, + max_input_len=6000, + max_num_tokens=10240, kv_cache_config=kv_cache_config, tensor_parallel_size=tp_size) diff --git a/examples/apps/fastapi_server.py b/examples/apps/fastapi_server.py index 510b281a701b..e1d90c37934e 100755 --- a/examples/apps/fastapi_server.py +++ b/examples/apps/fastapi_server.py @@ -18,9 +18,9 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, Response, StreamingResponse -from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm import LLM from tensorrt_llm.executor import CppExecutorError, RequestError -from tensorrt_llm.llmapi import BuildConfig, KvCacheConfig, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams TIMEOUT_KEEP_ALIVE = 5 # seconds. @@ -120,8 +120,6 @@ def entrypoint(model_dir: str, port = port or 8000 logging.info(f"Starting server at {host}:{port}") - build_config = BuildConfig(max_batch_size=10, max_beam_width=max_beam_width) - kv_cache_config = KvCacheConfig( free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction) @@ -129,7 +127,8 @@ def entrypoint(model_dir: str, tokenizer, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, - build_config=build_config, + max_batch_size=10, + max_beam_width=max_beam_width, kv_cache_config=kv_cache_config) server = LlmServer(llm=llm) diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index ab0274cacc59..accac89d425e 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -1,8 +1,8 @@ { "_meta": { "generated_by": "scripts/legacy_utils.py lint-update-violations", - "total_violations": 5338, - "total_files": 492 + "total_violations": 2414, + "total_files": 276 }, ".github/scripts/label_community_user.py": { "D212": 1 @@ -65,7 +65,7 @@ }, "cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py": { "F403": 1, - "F405": 184 + "F405": 22 }, "docs/source/helper.py": { "D205": 2, @@ -144,22 +144,11 @@ "D212": 1, "F821": 1 }, - "examples/medusa/convert_checkpoint.py": { - "F821": 1 - }, "examples/mmlu.py": { "D205": 1, "D415": 1, "E741": 1 }, - "examples/models/contrib/bloom/convert_checkpoint.py": { - "D200": 2, - "D202": 2, - "D210": 3, - "D212": 2, - "D415": 2, - "E741": 1 - }, "examples/models/contrib/chatglm-6b/tokenization_chatglm.py": { "D205": 4, "D210": 4, @@ -179,55 +168,6 @@ "D212": 3, "D415": 3 }, - "examples/models/contrib/cogvlm/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/contrib/dbrx/convert_checkpoint.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 1 - }, - "examples/models/contrib/dit/sample.py": { - "D200": 1, - "D205": 1, - "D212": 1, - "D415": 1 - }, - "examples/models/contrib/dit/utils_modelopt.py": { - "D200": 1, - "D212": 1 - }, - "examples/models/contrib/gptj/convert_checkpoint.py": { - "F821": 1 - }, - "examples/models/contrib/gptneox/convert_checkpoint.py": { - "D202": 1, - "D210": 2, - "D415": 1, - "E741": 2 - }, - "examples/models/contrib/grok/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/contrib/mmdit/sample.py": { - "D200": 1 - }, - "examples/models/contrib/mpt/convert_checkpoint.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 2 - }, - "examples/models/contrib/opt/convert_checkpoint.py": { - "E741": 1 - }, "examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py": { "D202": 1, "D205": 6, @@ -235,122 +175,9 @@ "D414": 1, "D415": 2 }, - "examples/models/contrib/stdit/pipeline_tllm.py": { - "D200": 1, - "D205": 1, - "D212": 1, - "D415": 1 - }, - "examples/models/contrib/stdit/text_encoder.py": { - "D200": 1, - "D212": 1 - }, - "examples/models/contrib/stdit/utils.py": { - "D200": 2, - "D205": 2, - "D212": 4, - "D415": 2 - }, - "examples/models/contrib/stdit/vae.py": { - "D415": 1 - }, - "examples/models/contrib/stdit/video_transforms.py": { - "D200": 2, - "D205": 17, - "D212": 19, - "D410": 8, - "D411": 9, - "D415": 18 - }, - "examples/models/core/bert/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/commandr/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/enc_dec/helper.py": { - "D205": 1, - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/enc_dec/run.py": { - "D202": 1, - "D208": 18, - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/glm-4-9b/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/glm-4-9b/tokenization_chatglm.py": { - "D200": 1, - "D205": 1, - "D210": 2, - "D212": 4, - "D415": 3, - "F821": 1 - }, - "examples/models/core/gpt/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/internlm2/convert_checkpoint.py": { - "D210": 1, - "D415": 1, - "E741": 1 - }, - "examples/models/core/llama/convert_checkpoint.py": { - "D200": 2, - "D300": 2, - "D403": 2, - "D415": 2, - "E722": 1 - }, - "examples/models/core/mamba/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "F821": 6 - }, - "examples/models/core/mllama/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "E722": 1 - }, "examples/models/core/multimodal/run.py": { "E731": 1 }, - "examples/models/core/phi/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/qwen/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "E722": 1 - }, "examples/models/core/qwen2audio/run.py": { "D200": 1, "D212": 1, @@ -366,22 +193,6 @@ "examples/models/core/qwenvl/run_chat.py": { "E722": 1 }, - "examples/models/core/whisper/convert_checkpoint.py": { - "D415": 1 - }, - "examples/models/core/whisper/run.py": { - "E712": 1 - }, - "examples/models/core/whisper/whisper_utils.py": { - "D200": 1, - "D202": 1, - "D205": 3, - "D212": 4, - "D403": 1, - "D411": 2, - "D415": 4, - "D416": 2 - }, "examples/ngram/run_dtm_ngram.py": { "D200": 1, "D205": 1, @@ -423,12 +234,6 @@ "E741": 1, "F601": 2 }, - "examples/redrafter/convert_checkpoint.py": { - "D205": 1, - "D212": 1, - "D415": 1, - "E722": 1 - }, "examples/run.py": { "E711": 4 }, @@ -452,9 +257,6 @@ "D212": 1, "E741": 1 }, - "scripts/check_test_list.py": { - "D212": 1 - }, "scripts/dco_check.py": { "D212": 1 }, @@ -485,42 +287,23 @@ "tensorrt_llm/__init__.py": { "D202": 1, "D212": 1, - "E402": 27 - }, - "tensorrt_llm/_torch/attention_backend/sparse/dsa.py": { - "F821": 2 + "E402": 20 }, "tensorrt_llm/_torch/attention_backend/sparse/kernel.py": { "E731": 3 }, "tensorrt_llm/_torch/attention_backend/sparse/rocket.py": { - "E712": 2, - "F821": 4 - }, - "tensorrt_llm/_torch/attention_backend/sparse/utils.py": { - "F821": 4 + "E712": 2 }, "tensorrt_llm/_torch/attention_backend/trtllm.py": { "E712": 1 }, - "tensorrt_llm/_torch/attention_backend/utils.py": { - "F821": 2 - }, - "tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py": { - "E402": 1 - }, "tensorrt_llm/_torch/compilation/patterns/residual_add_norm.py": { "E402": 1 }, "tensorrt_llm/_torch/compilation/piecewise_optimizer.py": { "E711": 2 }, - "tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py": { - "E741": 11 - }, - "tensorrt_llm/_torch/custom_ops/torch_custom_ops.py": { - "F821": 2 - }, "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py": { "E741": 3 }, @@ -528,10 +311,6 @@ "E712": 1, "E731": 1 }, - "tensorrt_llm/_torch/model_config.py": { - "F811": 1, - "F821": 4 - }, "tensorrt_llm/_torch/models/checkpoints/auto_mapper.py": { "F821": 1 }, @@ -562,12 +341,6 @@ "tensorrt_llm/_torch/models/modeling_nemotron.py": { "E731": 1 }, - "tensorrt_llm/_torch/models/modeling_qwen2vl.py": { - "F811": 1 - }, - "tensorrt_llm/_torch/models/modeling_qwen3_next.py": { - "F821": 1 - }, "tensorrt_llm/_torch/models/modeling_siglip.py": { "F811": 1 }, @@ -575,9 +348,6 @@ "E731": 1, "F821": 5 }, - "tensorrt_llm/_torch/modules/attention.py": { - "E731": 2 - }, "tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py": { "F821": 2 }, @@ -588,7 +358,7 @@ "F821": 1 }, "tensorrt_llm/_torch/modules/fused_moe/interface.py": { - "E402": 4 + "E402": 3 }, "tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py": { "E712": 2 @@ -611,29 +381,12 @@ "tensorrt_llm/_torch/modules/mamba/ssd_state_passing.py": { "E731": 1 }, - "tensorrt_llm/_torch/pyexecutor/_util.py": { - "F811": 1 - }, - "tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py": { - "F821": 1 - }, - "tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py": { - "F821": 1 - }, - "tensorrt_llm/_torch/pyexecutor/model_engine.py": { - "F821": 1 - }, "tensorrt_llm/_torch/pyexecutor/model_loader.py": { - "F811": 1, - "F821": 13 + "F811": 1 }, "tensorrt_llm/_torch/pyexecutor/py_executor.py": { "E712": 1, - "E721": 1, - "F821": 2 - }, - "tensorrt_llm/_torch/pyexecutor/resource_manager.py": { - "F821": 2 + "E721": 1 }, "tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py": { "F821": 1 @@ -641,9 +394,6 @@ "tensorrt_llm/_torch/speculative/auto_heuristic.py": { "F821": 1 }, - "tensorrt_llm/_torch/speculative/interface.py": { - "F821": 1 - }, "tensorrt_llm/_torch/speculative/ngram.py": { "E741": 1 }, @@ -663,16 +413,12 @@ "D415": 1 }, "tensorrt_llm/bench/build/build.py": { - "D202": 1, - "D205": 2, - "D208": 1, - "D210": 3, + "D205": 1, + "D210": 2, "D411": 1 }, "tensorrt_llm/bench/build/dataclasses.py": { - "D205": 1, - "D208": 1, - "D210": 5 + "D210": 3 }, "tensorrt_llm/bench/build/tuning.py": { "D205": 2, @@ -686,26 +432,6 @@ "D200": 6, "D212": 6 }, - "tensorrt_llm/builder.py": { - "D205": 8, - "D208": 23, - "D210": 1, - "D212": 4, - "D300": 7, - "D403": 1, - "D415": 4, - "E712": 2 - }, - "tensorrt_llm/commands/prune.py": { - "D200": 1, - "D212": 1, - "D300": 1 - }, - "tensorrt_llm/commands/refit.py": { - "D200": 1, - "D212": 1, - "D300": 1 - }, "tensorrt_llm/commands/serve.py": { "D202": 2, "D212": 2, @@ -725,20 +451,17 @@ }, "tensorrt_llm/executor/base_worker.py": { "D200": 1, - "D202": 3, - "D205": 1, - "D208": 2, - "D210": 8, + "D202": 2, + "D210": 7, "D212": 1, - "D300": 5, + "D300": 4, "D415": 2 }, "tensorrt_llm/executor/executor.py": { - "D205": 7, + "D205": 3, "D210": 1, - "D212": 4, "D410": 3, - "D411": 7 + "D411": 3 }, "tensorrt_llm/executor/ipc.py": { "D202": 2, @@ -778,14 +501,8 @@ }, "tensorrt_llm/executor/result.py": { "D200": 3, - "D202": 1, "D205": 1, - "D210": 3, - "D212": 5, - "D300": 3, - "F402": 1, - "F811": 1, - "F821": 1 + "F402": 1 }, "tensorrt_llm/executor/rpc/rpc_client.py": { "D200": 1, @@ -820,36 +537,16 @@ "D300": 6, "F811": 1 }, - "tensorrt_llm/functional.py": { - "D200": 42, - "D202": 6, - "D205": 20, - "D207": 1, - "D208": 6, - "D210": 2, - "D212": 142, - "D214": 3, - "D300": 144, - "D301": 2, - "D411": 8, - "D415": 9 - }, - "tensorrt_llm/inputs/evs.py": { - "D205": 1, - "D212": 2 - }, "tensorrt_llm/inputs/multimodal.py": { - "D202": 1, "D415": 1 }, "tensorrt_llm/inputs/registry.py": { - "D200": 3, - "D205": 6, - "D212": 17, + "D200": 2, + "D205": 5, + "D212": 13, "D301": 1, "D405": 1, - "D411": 1, - "D415": 3, + "D415": 2, "E722": 2 }, "tensorrt_llm/inputs/utils.py": { @@ -859,54 +556,6 @@ "D214": 1, "E731": 2 }, - "tensorrt_llm/layers/attention.py": { - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/layers/embedding.py": { - "D200": 1, - "D205": 3, - "D212": 6, - "D415": 1, - "D416": 1 - }, - "tensorrt_llm/layers/language_adapter.py": { - "D205": 1, - "D212": 1 - }, - "tensorrt_llm/layers/linear.py": { - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "E711": 2 - }, - "tensorrt_llm/layers/moe.py": { - "D200": 1, - "D212": 1, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/layers/recurrent.py": { - "D205": 1, - "D212": 1, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/layers/ssm.py": { - "D205": 3, - "D212": 3, - "D300": 3, - "D415": 3 - }, - "tensorrt_llm/llmapi/build_cache.py": { - "D200": 9, - "D210": 2, - "D212": 13, - "D300": 12, - "D415": 8, - "E722": 1 - }, "tensorrt_llm/llmapi/disagg_utils.py": { "D205": 2, "D210": 1, @@ -917,31 +566,14 @@ }, "tensorrt_llm/llmapi/llm.py": { "D200": 2, - "D202": 1, - "D205": 4, - "D207": 2, - "D212": 2, - "D300": 4, - "D410": 2, - "D411": 2 + "D205": 4 }, "tensorrt_llm/llmapi/llm_args.py": { - "D200": 23, - "D202": 1, "D205": 11, - "D208": 2, - "D210": 7, - "D212": 37, - "D300": 8, - "D415": 2 + "D212": 1 }, "tensorrt_llm/llmapi/llm_utils.py": { - "D200": 2, - "D205": 4, - "D209": 1, - "D210": 9, - "D212": 3, - "D300": 12 + "D200": 1 }, "tensorrt_llm/llmapi/mgmn_leader_node.py": { "D205": 2, @@ -949,9 +581,7 @@ "D300": 2 }, "tensorrt_llm/llmapi/mm_encoder.py": { - "D200": 1, - "D205": 3, - "D207": 1, + "D205": 2, "D410": 1, "D411": 1 }, @@ -979,9 +609,9 @@ "D202": 1, "D205": 5, "D209": 1, - "D210": 14, + "D210": 13, "D212": 5, - "D300": 8, + "D300": 7, "E722": 1 }, "tensorrt_llm/mapping.py": { @@ -990,69 +620,6 @@ "D300": 1, "D415": 4 }, - "tensorrt_llm/models/baichuan/config.py": { - "F821": 2 - }, - "tensorrt_llm/models/baichuan/convert.py": { - "D202": 1, - "D205": 1, - "D208": 17, - "D212": 1, - "D415": 1, - "E731": 1, - "E741": 3 - }, - "tensorrt_llm/models/baichuan/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/bert/convert.py": { - "D200": 4, - "D208": 1, - "D212": 5, - "D300": 1, - "D403": 4, - "D415": 5 - }, - "tensorrt_llm/models/bert/model.py": { - "D200": 2, - "D202": 1, - "D205": 1, - "D212": 3, - "D300": 1, - "D415": 2 - }, - "tensorrt_llm/models/chatglm/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/chatglm/convert.py": { - "D200": 2, - "D208": 1, - "D212": 1, - "D300": 1, - "D415": 1, - "E741": 1 - }, - "tensorrt_llm/models/chatglm/model.py": { - "D200": 2, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/cogvlm/convert.py": { - "E741": 1 - }, - "tensorrt_llm/models/commandr/model.py": { - "D200": 1, - "D202": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, "tensorrt_llm/models/convert_utils.py": { "D200": 2, "D202": 1, @@ -1066,343 +633,6 @@ "E731": 1, "F821": 1 }, - "tensorrt_llm/models/deepseek_v1/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/deepseek_v1/convert.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 2 - }, - "tensorrt_llm/models/deepseek_v1/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/deepseek_v2/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/deepseek_v2/convert.py": { - "E741": 4 - }, - "tensorrt_llm/models/deepseek_v2/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/dit/model.py": { - "D200": 2, - "D205": 2, - "D208": 2, - "D212": 3, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/models/eagle/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/eagle/model.py": { - "D202": 3, - "D205": 4, - "D212": 5, - "D300": 4, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/enc_dec/model.py": { - "D202": 2, - "D205": 2, - "D208": 4, - "D300": 2, - "E712": 1 - }, - "tensorrt_llm/models/falcon/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/falcon/convert.py": { - "D202": 1, - "D210": 2, - "D415": 1, - "E741": 2 - }, - "tensorrt_llm/models/falcon/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/gemma/config.py": { - "D202": 1, - "D415": 1 - }, - "tensorrt_llm/models/gemma/convert.py": { - "D205": 1, - "D212": 1, - "D415": 1 - }, - "tensorrt_llm/models/gemma/smoothquant.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 1 - }, - "tensorrt_llm/models/gemma/utils/modules.py": { - "D200": 1 - }, - "tensorrt_llm/models/gemma/utils/params.py": { - "D207": 6 - }, - "tensorrt_llm/models/gemma/utils/positional_embeddings.py": { - "D200": 1 - }, - "tensorrt_llm/models/gemma/weight.py": { - "D200": 1, - "D212": 1, - "E741": 3 - }, - "tensorrt_llm/models/generation_mixin.py": { - "E712": 2 - }, - "tensorrt_llm/models/gpt/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/gpt/convert.py": { - "D202": 1, - "D205": 1, - "D212": 1, - "E741": 2 - }, - "tensorrt_llm/models/gpt/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/gptj/config.py": { - "D200": 1, - "D212": 1, - "F811": 1, - "F821": 1 - }, - "tensorrt_llm/models/gptj/convert.py": { - "E741": 1 - }, - "tensorrt_llm/models/gptj/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/grok/convert.py": { - "D200": 2, - "D208": 1, - "D210": 1, - "D212": 1, - "D300": 2, - "D415": 2, - "E741": 2 - }, - "tensorrt_llm/models/llama/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/llama/convert.py": { - "D200": 1, - "D205": 1, - "D208": 4, - "D210": 1, - "D212": 2, - "D300": 3, - "D415": 1, - "E741": 5 - }, - "tensorrt_llm/models/llama/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/models/mamba/convert.py": { - "D210": 1, - "E741": 1 - }, - "tensorrt_llm/models/mamba/model.py": { - "D205": 1, - "D208": 2, - "D300": 1, - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/models/medusa/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/medusa/model.py": { - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/models/medusa/weight.py": { - "E741": 2 - }, - "tensorrt_llm/models/mllama/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/mllama/model.py": { - "D200": 1, - "D202": 1, - "D205": 1, - "D208": 2, - "D210": 1, - "D300": 2, - "D415": 1, - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/models/mmdit_sd3/model.py": { - "D200": 1 - }, - "tensorrt_llm/models/model_weights_loader.py": { - "D415": 2, - "E722": 1 - }, - "tensorrt_llm/models/modeling_utils.py": { - "D200": 3, - "D202": 1, - "D205": 4, - "D208": 4, - "D210": 3, - "D212": 2, - "D300": 5, - "D415": 4, - "E721": 2, - "E722": 1 - }, - "tensorrt_llm/models/multimodal_encoders/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/multimodal_encoders/model.py": { - "D200": 1, - "D202": 1, - "D205": 1, - "D208": 2, - "D210": 1, - "D300": 2, - "D415": 1 - }, - "tensorrt_llm/models/nemotron_nas/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/nemotron_nas/convert.py": { - "D200": 1, - "D415": 1, - "E741": 1, - "F821": 7 - }, - "tensorrt_llm/models/nemotron_nas/model.py": { - "D200": 2, - "D202": 1, - "D205": 6, - "D212": 8, - "D415": 6, - "F821": 1 - }, - "tensorrt_llm/models/phi/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/phi/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/phi3/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/phi3/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/qwen/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/qwen/convert.py": { - "D200": 1, - "D208": 1, - "D212": 1, - "D300": 1, - "D415": 1, - "E741": 2 - }, - "tensorrt_llm/models/qwen/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/recurrentgemma/model.py": { - "D205": 1, - "D208": 2, - "D300": 1, - "E712": 1 - }, - "tensorrt_llm/models/redrafter/model.py": { - "D200": 1, - "D202": 2, - "D205": 2, - "D212": 3, - "D300": 1, - "D415": 3 - }, - "tensorrt_llm/models/redrafter/redrafter_helper.py": { - "D200": 5, - "D202": 2, - "D205": 12, - "D208": 7, - "D212": 17, - "D300": 9, - "D403": 1, - "D410": 2, - "D411": 5, - "D415": 14 - }, - "tensorrt_llm/models/stdit/model.py": { - "D200": 1, - "D202": 1, - "D205": 1, - "D208": 2, - "D300": 1 - }, - "tensorrt_llm/models/unet/embeddings.py": { - "D205": 1, - "D212": 1 - }, - "tensorrt_llm/network.py": { - "D200": 5, - "D205": 1, - "D208": 4, - "D210": 1, - "D212": 13, - "D300": 14, - "D415": 2, - "E731": 1, - "F811": 1, - "F821": 10 - }, - "tensorrt_llm/plugin/plugin.py": { - "D205": 2, - "D208": 10, - "D212": 1, - "D415": 1, - "E712": 1, - "E721": 1, - "F821": 5 - }, - "tensorrt_llm/quantization/functional.py": { - "D205": 4, - "D212": 4, - "D300": 4, - "D411": 2, - "D415": 2 - }, - "tensorrt_llm/quantization/layers.py": { - "D200": 2, - "D205": 1, - "D208": 8, - "D212": 3, - "D415": 1, - "E712": 2 - }, "tensorrt_llm/quantization/quantize_by_modelopt.py": { "D200": 1, "D205": 1, @@ -1427,81 +657,10 @@ "D205": 2, "D212": 3, "D403": 2, - "D415": 2 - }, - "tensorrt_llm/runtime/__init__.py": { - "E402": 7 - }, - "tensorrt_llm/runtime/enc_dec_model_runner.py": { - "E731": 1 - }, - "tensorrt_llm/runtime/generation.py": { - "D200": 4, - "D202": 1, - "D205": 3, - "D208": 2, - "D212": 6, - "D300": 3, - "D403": 2, - "D415": 3, - "E711": 5, - "F403": 1, - "F405": 9 - }, - "tensorrt_llm/runtime/kv_cache_manager.py": { - "D200": 8, - "D202": 1, - "D205": 6, - "D212": 14, - "D415": 10, - "E712": 1, - "E741": 1 - }, - "tensorrt_llm/runtime/medusa_utils.py": { - "D200": 1, - "D212": 1 - }, - "tensorrt_llm/runtime/model_runner.py": { - "D200": 1, - "D205": 3, - "D212": 7, - "D410": 2, - "D411": 2, - "E741": 1 - }, - "tensorrt_llm/runtime/model_runner_cpp.py": { - "D200": 1, - "D205": 1, - "D212": 3, - "D410": 2, - "D411": 2, - "E711": 1 - }, - "tensorrt_llm/runtime/multimodal_model_runner.py": { - "D202": 1, - "D208": 27, - "D212": 6, - "D214": 9, - "D410": 2, - "D411": 3, - "E741": 2 - }, - "tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py": { - "D200": 1, - "D212": 1, - "D300": 1, "D415": 1 }, - "tensorrt_llm/runtime/session.py": { - "D200": 4, - "D202": 1, - "D205": 6, - "D210": 1, - "D212": 6, - "D300": 11, - "D403": 2, - "D415": 8, - "E741": 3 + "tensorrt_llm/runtime/__init__.py": { + "E402": 1 }, "tensorrt_llm/scaffolding/contrib/DeepConf/deep_conf_controller.py": { "E712": 1 @@ -1547,10 +706,6 @@ "PLE0302": 1 }, "tensorrt_llm/scaffolding/task.py": { - "F811": 1, - "F821": 1 - }, - "tensorrt_llm/scaffolding/task_collection.py": { "F821": 1 }, "tensorrt_llm/scaffolding/worker.py": { @@ -1571,11 +726,6 @@ "D210": 1, "F811": 1 }, - "tensorrt_llm/serve/openai_server.py": { - "D205": 1, - "D212": 2, - "F821": 2 - }, "tensorrt_llm/serve/responses_utils.py": { "D200": 2, "D205": 4, @@ -1633,35 +783,10 @@ }, "tensorrt_llm/tokenizer/tokenizer.py": { "D202": 1, - "D205": 1, "D209": 1, "D210": 3, "D300": 5 }, - "tensorrt_llm/tools/plugin_gen/core.py": { - "D200": 4, - "D210": 6, - "D212": 8, - "D300": 14, - "D415": 2, - "F821": 1 - }, - "tensorrt_llm/tools/plugin_gen/plugin_gen.py": { - "D200": 10, - "D202": 1, - "D212": 10, - "D300": 10, - "D403": 2, - "D415": 3 - }, - "tensorrt_llm/tools/plugin_gen/shape_infer.py": { - "D200": 1, - "D212": 2, - "D300": 2, - "D403": 1, - "D415": 1, - "E731": 1 - }, "tensorrt_llm/tools/ppl.py": { "D200": 1, "D212": 1 @@ -1674,29 +799,14 @@ "D403": 9, "D415": 11 }, - "tests/integration/defs/accuracy/accuracy_core.py": { - "D205": 1, - "D212": 1 - }, "tests/integration/defs/accuracy/test_disaggregated_serving.py": { "D212": 1, "F601": 1 }, - "tests/integration/defs/accuracy/test_llm_api.py": { - "D300": 1, - "D415": 2 - }, - "tests/integration/defs/accuracy/test_llm_api_autodeploy.py": { - "D205": 1, - "D209": 1, - "D415": 1, - "F811": 2 - }, "tests/integration/defs/accuracy/test_llm_api_pytorch.py": { "D212": 1, "D300": 3, - "D415": 3, - "E402": 7 + "D415": 3 }, "tests/integration/defs/common.py": { "D200": 1, @@ -1710,17 +820,16 @@ }, "tests/integration/defs/conftest.py": { "D200": 3, - "D202": 10, - "D205": 4, + "D202": 9, + "D205": 3, "D209": 1, "D212": 5, "D214": 1, - "D300": 96, - "D403": 52, - "D415": 96, + "D300": 92, + "D403": 51, + "D415": 92, "E722": 2, - "E741": 1, - "F811": 1 + "E741": 1 }, "tests/integration/defs/cpp/conftest.py": { "E402": 2 @@ -1730,7 +839,6 @@ "E741": 2 }, "tests/integration/defs/disaggregated/test_disaggregated.py": { - "D202": 2, "D205": 3, "D209": 1, "D212": 5, @@ -1766,12 +874,6 @@ "tests/integration/defs/examples/test_commandr.py": { "D300": 1 }, - "tests/integration/defs/examples/test_gemma.py": { - "D202": 1, - "D300": 6, - "D403": 5, - "D415": 5 - }, "tests/integration/defs/examples/test_gpt.py": { "D202": 3, "D300": 5, @@ -1787,11 +889,10 @@ "D415": 1 }, "tests/integration/defs/examples/test_llama.py": { - "D202": 2, - "D300": 6, - "D403": 4, - "D415": 4, - "E712": 1 + "D202": 2, + "D300": 4, + "D403": 2, + "D415": 2 }, "tests/integration/defs/examples/test_mamba.py": { "D300": 2, @@ -1816,9 +917,8 @@ }, "tests/integration/defs/examples/test_phi.py": { "D202": 1, - "D300": 3, - "D403": 1, - "D415": 3 + "D300": 2, + "D415": 2 }, "tests/integration/defs/examples/test_qwen.py": { "D202": 1, @@ -1833,11 +933,6 @@ "D300": 1, "D415": 1 }, - "tests/integration/defs/examples/test_recurrentgemma.py": { - "D202": 1, - "D300": 2, - "D415": 2 - }, "tests/integration/defs/llmapi/test_llm_api_qa.py": { "D200": 1, "D212": 1, @@ -1874,7 +969,7 @@ "D212": 3 }, "tests/integration/defs/perf/open_search_db_utils.py": { - "D200": 2, + "D200": 1, "D212": 4, "E402": 1 }, @@ -1941,14 +1036,8 @@ }, "tests/integration/defs/test_e2e.py": { "D200": 5, - "D202": 3, - "D205": 3, - "D210": 1, - "D212": 8, - "D300": 9, - "D403": 4, - "D415": 10, - "F811": 2 + "D415": 6, + "F811": 1 }, "tests/integration/defs/test_list_parser.py": { "D202": 5, @@ -1970,9 +1059,6 @@ "D212": 2, "D415": 2 }, - "tests/integration/defs/triton_server/common.py": { - "E402": 1 - }, "tests/integration/defs/triton_server/conftest.py": { "D200": 3, "D202": 1, @@ -1986,19 +1072,6 @@ "E741": 3, "F821": 1 }, - "tests/integration/defs/triton_server/local_venv.py": { - "D205": 4, - "D212": 4 - }, - "tests/integration/defs/triton_server/rcca/bug_4323566/inflight_batcher_llm_client_with_end_id.py": { - "E721": 2 - }, - "tests/integration/defs/triton_server/runner_interface.py": { - "D205": 3, - "D208": 4, - "D212": 3, - "D415": 2 - }, "tests/integration/defs/triton_server/test_list_parser.py": { "D202": 5, "D205": 5, @@ -2013,22 +1086,6 @@ "F811": 1, "F821": 2 }, - "tests/integration/defs/triton_server/test_triton_llm.py": { - "F403": 2, - "F405": 174 - }, - "tests/integration/defs/triton_server/test_triton_memleak.py": { - "F403": 2, - "F405": 7 - }, - "tests/integration/defs/triton_server/test_triton_multi_node.py": { - "F403": 2, - "F405": 3 - }, - "tests/integration/defs/triton_server/test_triton_rcca.py": { - "F403": 2, - "F405": 29 - }, "tests/integration/defs/triton_server/trt_test_alternative.py": { "D200": 1, "D212": 1, @@ -2068,12 +1125,6 @@ "tests/unittest/_torch/misc/test_autotuner.py": { "E731": 1 }, - "tests/unittest/_torch/modeling/test_modeling_exaone4.py": { - "E402": 12 - }, - "tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py": { - "E731": 4 - }, "tests/unittest/api_stability/api_stability_core.py": { "D202": 1, "D205": 1, @@ -2085,13 +1136,13 @@ }, "tests/unittest/bindings/test_executor_bindings.py": { "D202": 1, - "E712": 56, + "E712": 54, "F403": 2, "F405": 3, "F811": 1 }, "tests/unittest/conftest.py": { - "D200": 4, + "D200": 3, "D202": 2, "D212": 5 }, @@ -2150,30 +1201,17 @@ "D209": 2 }, "tests/unittest/llmapi/test_llm.py": { - "D200": 1, "D205": 1, "D209": 1, "D210": 1, - "D212": 1, - "D300": 2, - "E712": 1, - "E722": 1 + "D300": 1, + "E712": 1 }, "tests/unittest/llmapi/test_llm_args.py": { "D200": 1, - "D202": 2, - "D212": 2, - "E712": 20, - "F811": 1 - }, - "tests/unittest/llmapi/test_llm_multi_gpu.py": { - "D210": 2, - "D300": 2 + "E712": 19 }, "tests/unittest/llmapi/test_llm_pytorch.py": { - "D202": 1, - "D210": 2, - "D212": 1, "F811": 1, "F821": 3 }, @@ -2181,10 +1219,6 @@ "D205": 1, "D212": 1 }, - "tests/unittest/llmapi/test_llm_utils.py": { - "F403": 2, - "F405": 14 - }, "tests/unittest/llmapi/test_mpi_session.py": { "D200": 1, "D212": 1, @@ -2194,20 +1228,8 @@ "tests/unittest/llmapi/test_serialization.py": { "E721": 1 }, - "tests/unittest/others/test_graph_rewriter.py": { - "D205": 1, - "D212": 1, - "D300": 1, - "D415": 1 - }, "tests/unittest/others/test_kv_cache_transceiver.py": { - "D205": 1, - "D212": 2 - }, - "tests/unittest/others/test_layer.py": { - "D205": 1, - "D212": 1, - "E711": 1 + "D212": 1 }, "tests/unittest/others/test_leak.py": { "D200": 1, @@ -2215,15 +1237,6 @@ "D300": 1, "D415": 1 }, - "tests/unittest/others/test_model_dtype.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, - "tests/unittest/others/test_module.py": { - "E741": 1 - }, "tests/unittest/others/test_time_breakdown.py": { "D212": 1, "D415": 1 @@ -2234,27 +1247,6 @@ "tests/unittest/scaffolding/test_scaffolding.py": { "F811": 10 }, - "tests/unittest/test_model_runner_cpp.py": { - "F403": 2, - "F405": 2 - }, - "tests/unittest/tools/plugin_gen/kernel_config.py": { - "F403": 1, - "F405": 22 - }, - "tests/unittest/tools/plugin_gen/test_core.py": { - "D200": 1, - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "F403": 1, - "F405": 7 - }, - "tests/unittest/tools/plugin_gen/test_shape_infer.py": { - "F403": 1, - "F405": 4 - }, "tests/unittest/tools/test_prepare_dataset.py": { "D205": 2, "D212": 7 @@ -2263,76 +1255,6 @@ "D202": 1, "E402": 1 }, - "tests/unittest/trt/attention/test_gpt_attention_IFB.py": { - "E711": 1, - "E712": 1, - "E731": 2 - }, - "tests/unittest/trt/functional/test_moe.py": { - "D202": 1, - "D210": 4, - "D415": 5 - }, - "tests/unittest/trt/functional/test_unbind.py": { - "F811": 1 - }, - "tests/unittest/trt/model/test_gpt_e2e.py": { - "E402": 1 - }, - "tests/unittest/trt/model/test_llama.py": { - "E741": 1 - }, - "tests/unittest/trt/model/test_mamba.py": { - "E741": 1 - }, - "tests/unittest/trt/model/test_mistral.py": { - "E741": 1 - }, - "tests/unittest/trt/model/test_nemotron_nas.py": { - "D415": 1 - }, - "tests/unittest/trt/model/test_phi.py": { - "F601": 1 - }, - "tests/unittest/trt/model/test_unet.py": { - "F403": 1, - "F405": 13 - }, - "tests/unittest/trt/model_api/test_model_api_multi_gpu.py": { - "D200": 1, - "D300": 1, - "D415": 1 - }, - "tests/unittest/trt/model_api/test_model_level_api.py": { - "D200": 1, - "D205": 1, - "D208": 3, - "D300": 2, - "D403": 1, - "D415": 1 - }, - "tests/unittest/trt/quantization/test_moe_weight_only_quant_matmul.py": { - "E741": 1 - }, - "tests/unittest/trt/quantization/test_qserve_gemm.py": { - "E402": 3 - }, - "tests/unittest/trt/quantization/test_quant_layer.py": { - "D200": 2, - "D205": 1, - "D208": 1, - "D210": 1, - "D212": 2, - "D415": 3 - }, - "tests/unittest/utils/test_medusa_utils.py": { - "D202": 2, - "D205": 2, - "D212": 3, - "D405": 1, - "D411": 1, - "D415": 2 - }, "tests/unittest/utils/test_prebuilt_whl_cpp_extensions.py": { "D200": 1, "D212": 1 @@ -2348,92 +1270,14 @@ }, "tests/unittest/utils/util.py": { "D200": 1, - "D202": 1, - "D205": 2, "D210": 2, - "D212": 3, + "D212": 1, "D300": 3, "D403": 3, - "D405": 1, - "D410": 1, - "D411": 3, - "D415": 4, - "E731": 1 - }, - "triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/1/model.py": { - "D205": 3, - "D416": 1, - "E722": 1 - }, - "triton_backend/all_models/gpt/postprocessing/1/model.py": { - "D202": 1, - "D205": 4, - "D411": 1, - "D415": 2, - "D416": 1 - }, - "triton_backend/all_models/gpt/preprocessing/1/model.py": { - "D200": 1, - "D202": 1, - "D205": 5, - "D208": 1, - "D212": 2, - "D300": 1, - "D403": 1, - "D411": 1, - "D415": 3, - "D416": 1, - "E711": 1 - }, - "triton_backend/all_models/gpt/tensorrt_llm/1/model.py": { - "D205": 4, - "D416": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/postprocessing/1/model.py": { - "D202": 1, - "D205": 4, - "D411": 1, - "D415": 2, - "D416": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/preprocessing/1/model.py": { - "D200": 1, - "D202": 2, - "D205": 6, - "D208": 1, - "D210": 1, - "D212": 5, - "D300": 1, - "D403": 1, - "D411": 1, "D415": 3, - "D416": 1, - "E711": 3, - "F821": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py": { - "D205": 5, - "D212": 1, - "D415": 1, - "D416": 1, - "E402": 1, - "E712": 2, - "E722": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py": { - "D415": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py": { - "F403": 1, - "F405": 43 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/model.py": { - "D205": 1 + "E731": 1 }, "triton_backend/all_models/llmapi/tensorrt_llm/1/helpers.py": { - "D205": 3, - "D212": 3, - "D415": 3, "E722": 1 }, "triton_backend/all_models/llmapi/tensorrt_llm/1/model.py": { @@ -2444,75 +1288,12 @@ "D411": 1, "D415": 2 }, - "triton_backend/all_models/multimodal/multimodal_encoders/1/model.py": { - "D202": 3, - "D205": 6, - "D212": 2, - "D411": 1, - "D415": 2, - "D416": 1, - "E711": 2 - }, - "triton_backend/all_models/multimodal/multimodal_encoders/1/multimodal_utils.py": { - "D212": 1, - "D214": 2, - "D410": 1, - "D411": 2 - }, - "triton_backend/all_models/tests/test_decode.py": { - "F403": 1, - "F405": 34 - }, "triton_backend/all_models/tests/test_llmapi_python_backend.py": { "E402": 2, "E712": 3, "F403": 1, "F405": 2 }, - "triton_backend/all_models/tests/test_multi_image_preprocess.py": { - "E402": 1 - }, - "triton_backend/all_models/tests/test_multimodal_encoders.py": { - "E402": 1 - }, - "triton_backend/all_models/tests/test_python_backend.py": { - "E402": 2, - "E712": 40, - "F403": 1, - "F405": 72 - }, - "triton_backend/all_models/tests/test_triton_decoder.py": { - "E402": 4, - "F601": 1 - }, - "triton_backend/all_models/whisper/whisper_bls/1/fbank.py": { - "D205": 2, - "D212": 2, - "D403": 1, - "D415": 2, - "D416": 1 - }, - "triton_backend/all_models/whisper/whisper_bls/1/model.py": { - "D205": 4, - "D212": 2, - "D415": 1, - "D416": 2 - }, - "triton_backend/ci/L0_backend_trtllm/base_metrics_verification_tests.py": { - "D205": 1, - "D212": 3, - "E402": 3 - }, - "triton_backend/inflight_batcher_llm/client/e2e_grpc_speculative_decoding_client.py": { - "E721": 1 - }, - "triton_backend/inflight_batcher_llm/client/end_to_end_grpc_client.py": { - "E721": 1, - "E722": 1 - }, - "triton_backend/inflight_batcher_llm/client/inflight_batcher_llm_client.py": { - "E721": 2 - }, "triton_backend/tools/inflight_batcher_llm/benchmark_core_model.py": { "E722": 1 }, @@ -2529,10 +1310,5 @@ "D205": 1, "D212": 1, "F811": 3 - }, - "triton_backend/tools/whisper/client.py": { - "D205": 1, - "D212": 1, - "E721": 3 } } diff --git a/setup.py b/setup.py index be0f840caaa5..7a3b7d2af8f7 100644 --- a/setup.py +++ b/setup.py @@ -438,9 +438,6 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], license_files=get_license(), entry_points={ 'console_scripts': [ - 'trtllm-build=tensorrt_llm.commands.build:main', - 'trtllm-prune=tensorrt_llm.commands.prune:main', - 'trtllm-refit=tensorrt_llm.commands.refit:main', 'trtllm-bench=tensorrt_llm.commands.bench:main', 'trtllm-serve=tensorrt_llm.commands.serve:main', 'trtllm-eval=tensorrt_llm.commands.eval:main' diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 0b269405274f..91e5f98a500d 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -104,32 +104,42 @@ def _setup_vendored_triton_kernels(): # ImportError: libc10.so: cannot open shared object file: No such file or directory import torch # noqa + +def _preload_tensorrt_libs(): + """Preload the TensorRT libraries needed by the bindings extension. + + The C++ runtime still links against the TensorRT libraries until it is + decoupled from TensorRT. Importing the tensorrt package loads libnvinfer + from the tensorrt_libs wheel for environments where it is not on the + system loader path; without it, importing tensorrt_llm.bindings raises + ImportError: libnvinfer.so.10: cannot open shared object file. + """ + try: + import tensorrt # noqa: F401 + except ImportError: + pass + + +_preload_tensorrt_libs() + import tensorrt_llm._torch.models as torch_models -import tensorrt_llm.functional as functional import tensorrt_llm.math_utils as math_utils import tensorrt_llm.models as models import tensorrt_llm.quantization as quantization import tensorrt_llm.runtime as runtime import tensorrt_llm.tools as tools -from ._common import _init, default_net, default_trtnet, precision +from ._common import _init from ._mnnvl_utils import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo from ._utils import (default_gpus_per_node, local_mpi_rank, local_mpi_size, mpi_barrier, mpi_comm, mpi_rank, mpi_world_size, - set_mpi_comm, str_dtype_to_torch, str_dtype_to_trt, - torch_dtype_to_trt) -from .builder import BuildConfig, Builder, BuilderConfig, build + set_mpi_comm, str_dtype_to_torch) from .disaggregated_params import DisaggregatedParams -from .functional import Tensor, constant from .llmapi import LLM, AsyncLLM, MultimodalEncoder -from .llmapi.llm_args import LlmArgs, TorchLlmArgs, TrtLlmArgs +from .llmapi.llm_args import LlmArgs, TorchLlmArgs from .logger import logger from .mapping import Mapping from .models.automodel import AutoConfig, AutoModelForCausalLM -from .module import Module -from .network import Network, net_guard -from .parameter import Parameter -from .python_plugin import PluginBase from .sampling_params import SamplingParams from .version import __version__ from .visual_gen import (ExtraParamSchema, VisualGen, VisualGenArgs, @@ -140,8 +150,6 @@ def _setup_vendored_triton_kernels(): 'AutoConfig', 'AutoModelForCausalLM', 'logger', - 'str_dtype_to_trt', - 'torch_dtype_to_trt', 'str_dtype_to_torch', 'default_gpus_per_node', 'local_mpi_rank', @@ -151,27 +159,12 @@ def _setup_vendored_triton_kernels(): 'mpi_rank', 'set_mpi_comm', 'mpi_world_size', - 'constant', - 'default_net', - 'default_trtnet', - 'precision', - 'net_guard', 'torch_models', - 'Network', 'Mapping', 'MnnvlMemory', 'MnnvlMoe', 'MoEAlltoallInfo', - 'PluginBase', - 'Builder', - 'BuilderConfig', - 'build', - 'BuildConfig', - 'Tensor', - 'Parameter', 'runtime', - 'Module', - 'functional', 'models', 'quantization', 'tools', @@ -180,7 +173,6 @@ def _setup_vendored_triton_kernels(): 'MultimodalEncoder', 'LlmArgs', 'TorchLlmArgs', - 'TrtLlmArgs', 'SamplingParams', 'VisualGenArgs', 'ExtraParamSchema', diff --git a/tensorrt_llm/_common.py b/tensorrt_llm/_common.py index 871120aabf13..00bac35cc11b 100644 --- a/tensorrt_llm/_common.py +++ b/tensorrt_llm/_common.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -12,7 +12,7 @@ # 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 contextlib + import ctypes import os import platform @@ -20,27 +20,12 @@ import time from functools import wraps from pathlib import Path -from typing import TYPE_CHECKING - -import numpy as np -# isort: off import torch -import tensorrt as trt - -# isort: on -if TYPE_CHECKING: - from .network import Network -else: - Network = None - -from ._utils import print_all_stacks, str_dtype_to_trt +from ._utils import print_all_stacks from .bindings import MpiComm from .logger import logger -from .plugin import _load_plugin_lib - -net = None _inited = False @@ -50,7 +35,6 @@ def _init(log_level: object = None) -> None: if _inited: return _inited = True - # Move to __init__ if log_level is not None: logger.set_level(log_level) @@ -60,11 +44,19 @@ def _init(log_level: object = None) -> None: logger.info("Starting TensorRT LLM init.") - # load plugin lib - _load_plugin_lib() - - # load FT decoder layer and torch custom ops project_dir = str(Path(__file__).parent.absolute()) + + # Promote libtensorrt_llm.so symbols to the process's global scope. The KV + # cache transfer agent wrappers (libtensorrt_llm_nixl_wrapper.so, + # libtensorrt_llm_ucx_wrapper.so) are dlopen'ed at runtime without linking + # against libtensorrt_llm.so and resolve its symbols from the global symbol + # table, while Python extension modules and their dependencies load with + # RTLD_LOCAL. This promotion was previously a side effect of loading the + # TensorRT plugin library with RTLD_GLOBAL. + if platform.system() != "Windows": + ctypes.CDLL(project_dir + "/libs/libtensorrt_llm.so", mode=ctypes.RTLD_GLOBAL) + + # Load FT decoder layer and torch custom ops. if platform.system() == "Windows": ft_decoder_lib = project_dir + "/libs/th_common.dll" else: @@ -99,112 +91,9 @@ def _print_stacks(): logger.info("TensorRT LLM inited.") -def default_net() -> Network: - assert net, ( - "Use builder to create network first, and use `set_network` or `net_guard` to set it to default" - ) - return net - - -def default_trtnet(): - return default_net().trt_network - - -def set_network(network): - global net - net = network - - -def switch_net_dtype(cur_dtype): - prev_dtype = default_net().dtype - default_net().dtype = cur_dtype - return prev_dtype - - -@contextlib.contextmanager -def precision(dtype): - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - prev_dtype = switch_net_dtype(dtype) - yield - switch_net_dtype(prev_dtype) - - -def serialize_engine(engine, path): - logger.info(f"Serializing engine to {path}...") - tik = time.time() - if isinstance(engine, trt.ICudaEngine): - engine = engine.serialize() - with open(path, "wb") as f: - f.write(engine) - tok = time.time() - t = time.strftime("%H:%M:%S", time.gmtime(tok - tik)) - logger.info(f"Engine serialized. Total time: {t}") - - -def deserialize_engine(path): - runtime = trt.Runtime(logger.trt_logger) - with open(path, "rb") as f: - logger.info(f"Loading engine from {path}...") - tik = time.time() - - engine = runtime.deserialize_cuda_engine(f.read()) - assert engine is not None - - tok = time.time() - t = time.strftime("%H:%M:%S", time.gmtime(tok - tik)) - logger.info(f"Engine loaded. Total time: {t}") - return engine - - -_field_dtype_to_np_dtype_dict = { - trt.PluginFieldType.FLOAT16: np.float16, - trt.PluginFieldType.FLOAT32: np.float32, - trt.PluginFieldType.FLOAT64: np.float64, - trt.PluginFieldType.INT8: np.int8, - trt.PluginFieldType.INT16: np.int16, - trt.PluginFieldType.INT32: np.int32, -} - - -def field_dtype_to_np_dtype(dtype): - ret = _field_dtype_to_np_dtype_dict.get(dtype) - assert ret is not None, f"Unsupported dtype: {dtype}" - return ret - - -def convert_capsule_to_void_p(capsule): - ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p - ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] - return ctypes.pythonapi.PyCapsule_GetPointer(capsule, None) - - -def get_nparray_from_void_p(void_pointer, elem_size, field_dtype): - ctypes.pythonapi.PyMemoryView_FromMemory.restype = ctypes.py_object - ctypes.pythonapi.PyMemoryView_FromMemory.argtypes = [ - ctypes.c_char_p, - ctypes.c_ssize_t, - ctypes.c_int, - ] - logger.info(f"get_nparray: pointer = {void_pointer}, elem_size = {elem_size}") - char_pointer = ctypes.cast(void_pointer, ctypes.POINTER(ctypes.c_char)) - np_dtype = field_dtype_to_np_dtype(field_dtype) - buf_bytes = elem_size * np.dtype(np_dtype).itemsize - logger.info(f"get_nparray: buf_bytes = {buf_bytes}") - mem_view = ctypes.pythonapi.PyMemoryView_FromMemory( - char_pointer, buf_bytes, 0 - ) # number 0 represents PyBUF_READ - logger.info(f"get_nparray: mem_view = {mem_view}, field_dtype = {field_dtype}") - buf = np.frombuffer(mem_view, np_dtype) - return buf - - -def get_scalar_from_field(field): - void_p = convert_capsule_to_void_p(field.data) - np_array = get_nparray_from_void_p(void_p, 1, field.type) - return np_array[0] - - +# TODO: dead on the Python side (no remaining @_is_building users); the IS_BUILDING +# env var is only read by the C++ isBuilding() in the TensorRT plugins. Remove this +# together with that C++ half in the C++ decouple step. class _BuildingFlag: def __enter__(self): os.environ["IS_BUILDING"] = "1" @@ -224,83 +113,3 @@ def decorated(*args, **kwargs): return f(*args, **kwargs) return decorated - - -def check_max_num_tokens( - max_num_tokens, - opt_num_tokens, - max_batch_size, - max_input_len, - max_seq_len, - max_beam_width, - remove_input_padding, - enable_context_fmha, - tokens_per_block, - multiple_profiles, -): - if not remove_input_padding: - if max_num_tokens is not None or opt_num_tokens is not None: - max_num_tokens = max_batch_size * max_seq_len - logger.warning( - "remove_input_padding is not enabled, the specified " - "max_num_tokens/opt_num_tokens will be ignored." - ) - return max_num_tokens, opt_num_tokens - else: - if max_num_tokens is None: - max_num_tokens = max_seq_len * max_batch_size - logger.warning( - "remove_input_padding is enabled, while max_num_tokens " - "is not set, setting to max_batch_size*max_seq_len. \n" - "It may not be optimal to set max_num_tokens=max_batch_size*max_seq_len " - "when remove_input_padding is enabled, because the number " - "of packed input tokens are very likely to be smaller, " - "we strongly recommend to set max_num_tokens according " - "to your workloads." - ) - if opt_num_tokens is None and not multiple_profiles: - opt_num_tokens = min(max_batch_size * max_beam_width, max_num_tokens) - logger.warning( - "remove_input_padding is enabled, while opt_num_tokens " - "is not set, setting to max_batch_size*max_beam_width. \n" - ) - if max_num_tokens > 16384: - logger.warning( - "Specifying a `max_num_tokens` larger than 16384 is usually " - "not recommended, we do not expect perf gain with that and too " - "large `max_num_tokens` could possibly exceed the TensorRT " - "tensor volume, causing runtime errors. " - f"Got `max_num_tokens` = {max_num_tokens}" - ) - if max_num_tokens > max_seq_len * max_batch_size: - logger.warning( - f"max_num_tokens ({max_num_tokens}) shouldn't be greater than " - f"max_seq_len * max_batch_size ({max_seq_len * max_batch_size}), " - f"specifying to max_seq_len * max_batch_size ({max_seq_len * max_batch_size})." - ) - max_num_tokens = max_seq_len * max_batch_size - if max_num_tokens < max_input_len and not enable_context_fmha: - logger.warning( - f"When enable_context_fmha is not turned on, max_num_tokens ({max_num_tokens}) " - f"should be at least max_input_len ({max_input_len}), specifying to " - f"max_input_len ({max_input_len})." - ) - max_num_tokens = max_input_len - elif max_num_tokens < tokens_per_block and enable_context_fmha: - logger.warning( - f"When enable_context_fmha is turned on, max_num_tokens ({max_num_tokens}) " - f"should be at least tokens_per_block ({tokens_per_block}), specifying to " - f"tokens_per_block ({tokens_per_block}). At this time, you also need to enable " - f"context chunking at runtime, otherwise you may encounter errors." - ) - max_num_tokens = tokens_per_block - - if opt_num_tokens is not None and opt_num_tokens > max_num_tokens: - logger.warning( - f"opt_num_tokens ({opt_num_tokens}) shouldn't be greater than " - f"max_num_tokens ({max_num_tokens}), " - f"specifying to max_num_tokens ({max_num_tokens})." - ) - opt_num_tokens = max_num_tokens - - return max_num_tokens, opt_num_tokens diff --git a/tensorrt_llm/_tensorrt_engine/__init__.py b/tensorrt_llm/_tensorrt_engine/__init__.py deleted file mode 100644 index 39669a168fd3..000000000000 --- a/tensorrt_llm/_tensorrt_engine/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from tensorrt_llm.llmapi.llm import _TrtLLM as LLM - -__all__ = ['LLM'] diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index 283e12ab8366..258b3048d491 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -21,7 +21,6 @@ from pydantic_settings import BaseSettings, SettingsConfigDict from tensorrt_llm.llmapi.llm_args import ( - BuildConfig, EagleDecodingConfig, MTPDecodingConfig, TorchLlmArgs, @@ -75,13 +74,6 @@ class LlmArgs(DynamicYamlMixInForSettings, TorchLlmArgs, BaseSettings): model_config = _get_config_dict() - build_config: Optional[BuildConfig] = Field( - default_factory=BuildConfig, - description="!!! DO NOT USE !!! Internal only; needed for BaseLlmArgs compatibility.", - exclude_from_json=True, - frozen=True, - repr=False, - ) backend: Literal["_autodeploy"] = Field( default="_autodeploy", description="The backend to use for this LLM instance.", @@ -101,12 +93,6 @@ def ensure_no_beam_search(cls, value: Any) -> Any: raise ValueError("AutoDeploy does not support beam search (max_beam_width > 1).") return value - @field_validator("build_config", mode="before") - @classmethod - def ensure_no_build_config(cls, value: Any, info: ValidationInfo) -> Any: - msg = "build_config is not in use by AutoDeploy's LlmArgs" - return _check_for_default_value_only(cls, value, info, msg) - @field_validator( "tensor_parallel_size", "pipeline_parallel_size", diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 3d2695075986..1d24db4fcc64 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -25,10 +25,11 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm import deep_gemm +from tensorrt_llm._torch.distributed.allreduce_helper import \ + CustomAllReduceHelper from tensorrt_llm._utils import get_sm_version from tensorrt_llm.functional import AllReduceFusionOp, AllReduceStrategy from tensorrt_llm.logger import logger -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper from tensorrt_llm.quantization.utils import fp8_quantize from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py b/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py index 9f9707808f69..21b81b250ff5 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py @@ -26,8 +26,7 @@ except ImportError: from cuda import cudart -from tensorrt_llm._utils import prefer_pinned -from tensorrt_llm.runtime.generation import CUASSERT +from tensorrt_llm._utils import CUASSERT, prefer_pinned try: import torch diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index d285a548c37a..8fefc2c761e4 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -34,7 +34,7 @@ TransferOp, TransferRequest, ) -from tensorrt_llm.runtime.generation import CUASSERT +from tensorrt_llm._utils import CUASSERT from .buffer import SlotAllocator from .config import SizingContext, fit_within_free diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index b808e6c360b3..fe5105bc8652 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -51,9 +51,8 @@ from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm._utils import nvtx_range +from tensorrt_llm._utils import CUASSERT, nvtx_range from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle -from tensorrt_llm.runtime.generation import CUASSERT if TYPE_CHECKING: from .bounce import Config diff --git a/tensorrt_llm/_torch/distributed/allreduce_helper.py b/tensorrt_llm/_torch/distributed/allreduce_helper.py new file mode 100644 index 000000000000..ef555be727f6 --- /dev/null +++ b/tensorrt_llm/_torch/distributed/allreduce_helper.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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 os +from typing import List, Tuple + +import torch + +from tensorrt_llm._ipc_utils import IpcMemory, can_access_peer +from tensorrt_llm.bindings.internal.runtime import ( + lamport_initialize, + max_workspace_size_lowprecision, +) +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + + +def force_all_reduce_deterministic(): + return ( + os.getenv("FORCE_DETERMINISTIC", "0") == "1" + or os.getenv("FORCE_ALL_REDUCE_DETERMINISTIC", "0") == "1" + ) + + +class CustomAllReduceHelper: + """ + Helper for custom all-reduce workspace/IPC buffer allocation and sizing, + used by the PyTorch backend. + """ + + POINTERS_PER_RANK = 7 + POINTERS_OF_COUNTER = 3 + + @staticmethod + def max_workspace_size_auto(tp_size: int, support_deterministic=True) -> int: + """Calculate workspace size for allreduce fusion kernel. + + The workspace is used for lamport buffers in the fusion kernel. + Required size calculation: + - Each GPU needs 3 sub-buffers (for triple buffering) + - Each sub-buffer stores: max_num_tokens * hidden_size * dtype_size (bf16=2) + - The lamport allocation multiplies by tp_size, so: + lamport_size = 3 * size * tp_size (per GPU) + + Example: Llama 8B (hidden=4096), max_tokens=8192, bf16, TP=4 + - Data per sub-buffer: 8192 * 4096 * 2 = 64 MiB + - Total lamport: 3 * 64MB * 4 = 768 MiB per GPU + - Required 'size' parameter: 64 MiB (gets multiplied by tp_size in allocation) + + Default (67,108,864 = 64 MiB) supports: + - Models up to hidden_size=4096 with max_num_tokens=8192 + - Or hidden_size=8192 with max_num_tokens=4096 + + Override with TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE env var if needed for larger models. + """ + if force_all_reduce_deterministic() and support_deterministic: + workspace_size = os.getenv("FORCE_ALLREDUCE_KERNEL_WORKSPACE_SIZE", "1000000000") + return int(workspace_size) + + # Allow override via environment variable for edge cases + workspace_size_env = os.getenv("TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE") + if workspace_size_env: + size = int(workspace_size_env) + logger.info( + f"Using custom allreduce fusion workspace size: {size} bytes ({size / (1024**2):.1f} MiB)" + ) + return size + + # Default: 64 MiB - supports most common model configurations + # Increase via env var if you see CUDA illegal memory access errors with large models + default_size = 67_108_864 # Exactly 64 MiB + return default_size + + @staticmethod + def max_workspace_size_lowprecision(tp_size: int) -> int: + return max_workspace_size_lowprecision(tp_size) + + @staticmethod + def initialize_lowprecision_buffers(workspace: "torch.tensor", tp_size: int) -> None: + return torch.ops.trtllm.initialize_static_lowprecision_buffers(workspace, tp_size) + + @staticmethod + def allocate_lowprecision_workspace( + mapping: Mapping, size: int + ) -> Tuple[List[IpcMemory], "torch.tensor"]: + # Force pull mode and disable lamport when force deterministic is enabled, for reducing device memory usage. + is_p2p_supported = can_access_peer(mapping) + ipc_buffers_size = size + ipc_buffers_ping = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) + ipc_buffers_pong = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) + ipc_barriers_in = IpcMemory( + mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, is_p2p_supported + ) + ipc_barriers_out = IpcMemory( + mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, is_p2p_supported + ) + buffers = [ipc_buffers_ping, ipc_buffers_pong, ipc_barriers_in, ipc_barriers_out] + + return buffers, torch.tensor( + ipc_buffers_ping.serialize() + + ipc_buffers_pong.serialize() + + ipc_barriers_in.serialize() + + ipc_barriers_out.serialize() + + [0] + + [0], + dtype=torch.int64, + device="cpu", + ) + + @staticmethod + def allocate_allreduce_fusion_workspace( + mapping: Mapping, size: int + ) -> Tuple[List[IpcMemory], "torch.tensor"]: + is_p2p_supported = can_access_peer(mapping) + ipc_buffers_size = size * mapping.tp_size + ipc_buffers = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) + ipc_barriers = IpcMemory(mapping, 256 * mapping.tp_size, is_p2p_supported) + lamport_buffers_size = size * mapping.tp_size + lamport_buffers = IpcMemory(mapping, 3 * lamport_buffers_size, is_p2p_supported) + if is_p2p_supported: + lamport_initialize( + lamport_buffers.local_ptr, + 3 * lamport_buffers_size, + ) + # flag_buffer[0], atomic flag read counter + # flag_buffer[1], non-lamport flag + # flag_buffer[2], lamport flag + flag_buffer = torch.tensor([0, 0, 0], dtype=torch.int, device="cuda") + # layout_buffer[0], clear size for next lamport kernel + # layout_buffer[1], triple buffer offset for lamport kernel + layout_buffer = torch.tensor([0, lamport_buffers_size], dtype=torch.int64, device="cuda") + + buffers = [ipc_buffers, ipc_barriers, lamport_buffers, flag_buffer, layout_buffer] + + return buffers, torch.tensor( + ipc_buffers.serialize() + + ipc_barriers.serialize() + + lamport_buffers.serialize() + + [flag_buffer.data_ptr()] + + [layout_buffer.data_ptr()], + dtype=torch.int64, + device="cuda", + ) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 63bf9d76c2c0..95570ccdce4d 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -8,6 +8,8 @@ from torch import nn from tensorrt_llm._mnnvl_utils import HelixCpMnnvlMemory, MnnvlMemory +from tensorrt_llm._torch.distributed.allreduce_helper import \ + CustomAllReduceHelper from tensorrt_llm._torch.distributed.symm_mem_allreduce import \ SymmetricMemoryAllReduce from tensorrt_llm._torch.utils import get_model_extra_attrs @@ -19,7 +21,6 @@ AllReduceStrategy, MoEAllReduceParams) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper # Feature flag: GEMM→NCCL-window zero-copy (writes GEMM output directly into # the window buffer so the allreduce needs no extra copy). On by default diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py index 2ad24e129be7..57efbbec82bc 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py @@ -1,20 +1,15 @@ -from typing import Union - from torch import nn from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.qwen2_moe_weight_mapper import \ Qwen2MoeHfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper -from tensorrt_llm.models.modeling_utils import DecoderModelForCausalLM @register_mapper("HF", "Qwen3MoeForCausalLM") class Qwen3MoeHfWeightMapper(Qwen2MoeHfWeightMapper): - def init_model_and_config(self, model: Union[nn.Module, - DecoderModelForCausalLM], - config: ModelConfig): + def init_model_and_config(self, model: nn.Module, config: ModelConfig): super().init_model_and_config(model, config) def should_skip_module(self, module_name: str) -> bool: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 0fbf170aaae6..dba700626476 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -22,9 +22,10 @@ except ImportError: from cuda import cudart -from tensorrt_llm._utils import (customized_gc_thresholds, is_trace_enabled, - mpi_comm, mpi_disabled, nvtx_range, - set_thread_local_mpi_comm, trace_func) +from tensorrt_llm._utils import (CUASSERT, customized_gc_thresholds, + is_trace_enabled, mpi_comm, mpi_disabled, + nvtx_range, set_thread_local_mpi_comm, + trace_func) from tensorrt_llm.bindings.executor import (DisServingRequestStats, FinishReason, InflightBatchingStats, IterationStats, KvCacheStats, @@ -38,7 +39,6 @@ from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType -from tensorrt_llm.runtime.generation import CUASSERT from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from tensorrt_llm.tools.profiler.host_profile_tools.host_profiler import ( diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index cf242cd205c9..d70190b3293f 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -38,12 +38,15 @@ import nvtx from mpi4py import MPI from mpi4py.util import pkl5 -from packaging import version from typing_extensions import ParamSpec # isort: off import torch -import tensorrt as trt + +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart try: from pynvml import ( @@ -107,6 +110,17 @@ def numpy_to_torch(x): return torch.from_numpy(x) +def CUASSERT(cuda_ret): + err = cuda_ret[0] + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" + ) + if len(cuda_ret) > 1: + return cuda_ret[1:] + return None + + def numpy_to_dtype(x, dtype: str): if str_dtype_to_np(dtype) == x.dtype: return x @@ -124,28 +138,12 @@ def numpy_to_dtype(x, dtype: str): bool_array = partial(np.array, dtype=np.bool_) -def dims_array(x): - is_int64_dims = True - try: - trt.Dims([np.iinfo(np.int64).max]) - except TypeError: - is_int64_dims = False - return int64_array(x) if is_int64_dims else int32_array(x) - - def bf16_array(x): x = torch.tensor(x, dtype=torch.bfloat16) x = torch_to_numpy(x) return x -def numpy_array(data, trt_dtype): - # convenient wrapper due to numpy not support bf16 yet - if trt_dtype == trt.bfloat16: - return bf16_array(data) - return np.array(data, trt_dtype_to_np(trt_dtype)) - - def copy_torch_to_numpy(x: torch.Tensor, ndarray: np.array): if x.dtype == torch.bfloat16: torch.from_numpy(ndarray.view(np.int16)).copy_(x.view(torch.int16)) @@ -156,18 +154,6 @@ def copy_torch_to_numpy(x: torch.Tensor, ndarray: np.array): return ndarray -def trt_version(): - return trt.__version__ - - -def trt_gte(major: int, minor: int = 0): - """ - Check if TRT version is greater than or equal to major.minor - """ - trt_ver = version.parse(trt_version()) - return trt_ver.major >= major and trt_ver.minor >= minor - - def torch_version(): return torch.__version__ @@ -276,82 +262,6 @@ def torch_dtype_to_str(dtype): return _torch_dtype_to_str_dict[dtype] -_str_to_trt_dtype_dict = dict(float16=trt.float16, - float32=trt.float32, - int64=trt.int64, - int32=trt.int32, - int8=trt.int8, - bool=trt.bool, - bfloat16=trt.bfloat16, - fp8=trt.fp8, - nvfp4=trt.fp4) - - -def str_dtype_to_trt(dtype): - if dtype == "fp4": - # Special handling for FP4 since CI's trt version is not recent enough. - if not hasattr(trt, 'fp4'): - raise ValueError( - "fp4 unsupported, trt version needs to be upgraded.") - return trt.fp4 - - ret = _str_to_trt_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - -_trt_to_str_dtype_dict = {v: k for k, v in _str_to_trt_dtype_dict.items()} - - -def trt_dtype_to_str(dtype: trt.DataType) -> str: - assert isinstance(dtype, trt.DataType) - return _trt_to_str_dtype_dict[dtype] - - -_np_to_trt_dtype_dict = { - np.int8: trt.int8, - np.int32: trt.int32, - np.int64: trt.int64, - np.float16: trt.float16, - np.float32: trt.float32, - np.bool_: trt.bool, - - # hash of np.dtype('int32') != np.int32 - np.dtype('int8'): trt.int8, - np.dtype('int32'): trt.int32, - np.dtype('int64'): trt.int64, - np.dtype('float16'): trt.float16, - np.dtype('float32'): trt.float32, - np.dtype('bool'): trt.bool, - np_bfloat16: trt.bfloat16, - np_float8: trt.fp8, -} - - -def np_dtype_to_trt(dtype): - ret = _np_to_trt_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - -_trt_to_np_dtype_dict = { - trt.int8: np.int8, - trt.int32: np.int32, - trt.int64: np.int64, - trt.float16: np.float16, - trt.float32: np.float32, - trt.bool: np.bool_, - trt.bfloat16: np_bfloat16, - trt.fp8: np_float8, -} - - -def trt_dtype_to_np(dtype): - ret = _trt_to_np_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - _torch_to_np_dtype_dict = { torch.bool: np.bool_, torch.uint8: np.uint8, @@ -398,54 +308,6 @@ def np_dtype_to_torch(dtype): return ret -_trt_to_torch_dtype_dict = { - trt.float16: torch.float16, - trt.float32: torch.float32, - trt.int64: torch.int64, - trt.int32: torch.int32, - trt.int8: torch.int8, - trt.bool: torch.bool, - trt.bfloat16: torch.bfloat16, - trt.fp8: torch.float8_e4m3fn, -} - - -def trt_dtype_to_torch(dtype): - ret = _trt_to_torch_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - -def is_same_dtype(type_a: Union[str, trt.DataType], - type_b: Union[str, trt.DataType]) -> bool: - if isinstance(type_a, str): - type_a = str_dtype_to_trt(type_a) - - if isinstance(type_b, str): - type_b = str_dtype_to_trt(type_b) - - return type_a == type_b - - -_torch_to_trt_dtype_dict = { - torch.float16: trt.float16, - torch.float32: trt.float32, - torch.int64: trt.int64, - torch.int32: trt.int32, - torch.int8: trt.int8, - torch.float8_e4m3fn: trt.fp8, - torch.qint8: trt.int8, - torch.bool: trt.bool, - torch.bfloat16: trt.bfloat16 -} - - -def torch_dtype_to_trt(dtype): - ret = _torch_to_trt_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - _torch_to_binding_dtype_dict = { torch.float16: DataType.HALF, torch.float32: DataType.FLOAT, @@ -1062,7 +924,7 @@ class TensorWrapper: def __init__( self, data_ptr: int, - dtype: Union[torch.dtype, str, np.dtype, trt.DataType, DataType], + dtype: Union[torch.dtype, str, np.dtype, DataType], shape: Sequence[int], strides: Optional[Sequence[int]] = None, ): @@ -1084,16 +946,13 @@ def shape(self): return getattr(self, "_shape", None) @dtype.setter - def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType, - DataType]): + def dtype(self, dtype: Union[torch.dtype, str, np.dtype, DataType]): if isinstance(dtype, torch.dtype): self._dtype = dtype elif isinstance(dtype, str): self._dtype = str_dtype_to_torch(dtype) elif isinstance(dtype, np.dtype): self._dtype = np_dtype_to_torch(dtype) - elif isinstance(dtype, trt.DataType): - self._dtype = trt_dtype_to_torch(dtype) elif isinstance(dtype, DataType): self._dtype = binding_to_torch_dtype(dtype) else: @@ -1122,10 +981,6 @@ def __cuda_array_interface__(self): 3, } - @staticmethod - def from_trt_desc(desc: trt.PluginTensorDesc, pointer: int): - return TensorWrapper(pointer, trt_dtype_to_torch(desc.type), desc.dims) - def convert_to_torch_tensor( tensor: Union[TensorWrapper, torch.Tensor]) -> torch.Tensor: diff --git a/tensorrt_llm/bench/benchmark/__init__.py b/tensorrt_llm/bench/benchmark/__init__.py index 53fa8b337d66..8386e115f505 100644 --- a/tensorrt_llm/bench/benchmark/__init__.py +++ b/tensorrt_llm/bench/benchmark/__init__.py @@ -5,7 +5,6 @@ from pydantic import AliasChoices, BaseModel, Field from tensorrt_llm import LLM as PyTorchLLM -from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm.bench.benchmark.utils.processes import IterationWriter from tensorrt_llm.bench.build.build import get_model_config from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig @@ -128,7 +127,7 @@ def get_llm(runtime_config: RuntimeConfig, kwargs: dict): Returns: An instance of the appropriate LLM class for the specified backend. """ - llm_cls = LLM + llm_cls = PyTorchLLM if runtime_config.backend != None: ignore_trt_only_args(kwargs, runtime_config.backend) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 18782061bccf..6fbf8f1f724b 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -12,8 +12,7 @@ from zmq import PUSH from zmq.asyncio import Context -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._utils import EnergyMonitor from tensorrt_llm.bench.dataclasses.general import InferenceRequest from tensorrt_llm.bench.dataclasses.reporting import PerfItemTuple, StatsKeeper diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index e35708df1bd2..e28121b18aba 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -23,7 +23,7 @@ QuantAlgo.NVFP4.value: "fp8", } -ALL_SUPPORTED_BACKENDS = ["pytorch", "_autodeploy", "tensorrt"] +ALL_SUPPORTED_BACKENDS = ["pytorch", "_autodeploy"] def get_settings_from_engine( diff --git a/tensorrt_llm/bench/build/build.py b/tensorrt_llm/bench/build/build.py index 3290828107b4..16bf6d23d9aa 100644 --- a/tensorrt_llm/bench/build/build.py +++ b/tensorrt_llm/bench/build/build.py @@ -1,28 +1,28 @@ from __future__ import annotations from pathlib import Path -from typing import Tuple, get_args -import click -from click_option_group import AllOptionGroup, optgroup - -from tensorrt_llm._torch.pyexecutor.config_utils import is_nemotron_hybrid, is_qwen3_hybrid, load_pretrained_config -from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment -from tensorrt_llm.bench.utils.data import create_dataset_from_stream, initialize_tokenizer -from tensorrt_llm.bench.utils import VALID_QUANT_ALGOS -from tensorrt_llm.builder import BuildConfig -from tensorrt_llm._tensorrt_engine import LLM +from typing import Tuple + +from tensorrt_llm._torch.pyexecutor.config_utils import (is_nemotron_hybrid, + is_qwen3_hybrid, + load_pretrained_config) +from tensorrt_llm.bench.build.dataclasses import (ModelConfig, + NemotronHybridConfig, + Qwen3HybridConfig) +from tensorrt_llm.bench.build.tuning import calc_engine_setting +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.llmapi.llm_utils import QuantConfig from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo -from tensorrt_llm.bench.build.dataclasses import ModelConfig, NemotronHybridConfig, Qwen3HybridConfig -from tensorrt_llm.bench.build.tuning import calc_engine_setting TUNED_QUANTS = { QuantAlgo.NVFP4, QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES, QuantAlgo.NO_QUANT, None } -DEFAULT_MAX_BATCH_SIZE = BuildConfig.model_fields["max_batch_size"].default -DEFAULT_MAX_NUM_TOKENS = BuildConfig.model_fields["max_num_tokens"].default +# Sourced from TorchLlmArgs so the bench defaults track the args-class field +# defaults and can't drift. +DEFAULT_MAX_BATCH_SIZE = TorchLlmArgs.model_fields["max_batch_size"].default +DEFAULT_MAX_NUM_TOKENS = TorchLlmArgs.model_fields["max_num_tokens"].default def get_benchmark_engine_settings( @@ -98,273 +98,3 @@ def get_model_config(model_name: str, model_path: Path = None) -> ModelConfig: if is_qwen3_hybrid(pretrained_config): return Qwen3HybridConfig.from_hf(model_name, model_path) return ModelConfig.from_hf(model_name, model_path) - - -def apply_build_mode_settings(params): - """ Validate engine build options and update the necessary values for engine - build settings. - """ - dataset_path = params.get("dataset") - max_batch_size = params.get("max_batch_size") - target_input_len = params.get("target_input_len") - max_seq_len = params.get("max_seq_len") - tp_size = params.get("tp_size") - pp_size = params.get("pp_size") - - # Check of engine build method. User must choose one engine build option. - build_options = [dataset_path, max_batch_size, target_input_len] - # If no engine build option is provided, fall back to build engine with - # TRT-LLM's default max_batch_size and max_num_tokens. - if sum([bool(opt) for opt in build_options]) == 0: - logger.warning( - "No engine build option is selected, use TRT-LLM default " - "max_batch_size and max_num_tokens to build the engine.") - params['max_batch_size'] = DEFAULT_MAX_BATCH_SIZE - params['max_num_tokens'] = DEFAULT_MAX_NUM_TOKENS - elif sum([bool(opt) for opt in build_options]) > 1: - raise ValueError("Multiple engine build options detected, please " - "choose only one engine build option. Exiting.") - - # Check for supported parallelism mappings: only world size <= 8 for now. - if tp_size * pp_size > 8: - raise ValueError( - f"Parallelism mapping of TP{tp_size}-PP{pp_size} is " - "currently unsupported. Please try with a mapping with <=8 GPUs.") - - # If dataset is not specified, max_seq_len must be provided. - if not dataset_path and not max_seq_len: - raise ValueError("Unspecified max_seq_len for engine build. Exiting.") - - -@click.command(name="build") -@optgroup.group("Engine Configuration", - help="Configuration of the TensorRT LLM engine.") -@optgroup.option( - "--tp_size", - "-tp", - type=int, - default=1, - required=False, - help="Number of tensor parallel shards to run the benchmark with.", -) -@optgroup.option( - "--pp_size", - "-pp", - type=int, - default=1, - required=False, - help="Number of pipeline parallel shards to run the benchmark with.", -) -@optgroup.option( - "--quantization", - "-q", - type=click.Choice(tuple(get_args(VALID_QUANT_ALGOS))), - default=None, - help= - ("The quantization algorithm to be used when benchmarking. See the " - "documentations for more information.\n" - " - https://nvidia.github.io/TensorRT-LLM/precision.html" - " - https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/quantization-in-TRT-LLM.md" - ), -) -@optgroup.option( - "--max_seq_len", - default=None, - type=click.IntRange(min=1), - help="Maximum total length of one request, including prompt and outputs.", -) -@optgroup.option( - "--no_weights_loading", - type=bool, - default=False, - help= - "Do not load the weights from the checkpoint. Use dummy weights instead.") -@optgroup.option( - "--trust_remote_code", - type=bool, - default=False, - help= - "Trust remote code for the HF models that are not natively implemented in the transformers library. " - "This is needed when using LLM API when loading the HF config to build the engine." -) -@optgroup.group( - "Build Engine with Dataset Information", - cls=AllOptionGroup, - help="Optimize engine build parameters with user-specified dataset " - "statistics, e.g., average input/output length, max sequence length.", -) -@optgroup.option( - "--dataset", - type=click.Path(exists=True, - readable=True, - path_type=Path, - resolve_path=True), - default=None, - help="Dataset file to extract the sequence statistics for engine build.", -) -@optgroup.group( - "Build Engine with IFB Scheduler Limits", - cls=AllOptionGroup, - help="Optimize engine build parameters with user-specified inflight " - "batching scheduler settings.", -) -@optgroup.option( - "--max_batch_size", - default=None, - type=click.IntRange(min=1), - help="Maximum number of requests that the engine can schedule.", -) -@optgroup.option( - "--max_num_tokens", - default=None, - type=click.IntRange(min=1), - help="Maximum number of batched tokens the engine can schedule.", -) -@optgroup.group( - "[Experimental Feature] Build Engine with Tuning Heuristics Hints", - cls=AllOptionGroup, - help="Optimize engine build parameters with user-specified target " - "sequence length information.", -) -@optgroup.option( - "--target_input_len", - default=None, - type=click.IntRange(min=1), - help="Target (average) input length for tuning heuristics.", -) -@optgroup.option( - "--target_output_len", - default=None, - type=click.IntRange(min=1), - help="Target (average) sequence length for tuning heuristics.", -) -@click.pass_obj -def build_command( - bench_env: BenchmarkEnvironment, - **params, -) -> None: - """Build engines for benchmarking.""" - - apply_build_mode_settings(params) - # Collect configuration parameters from CLI parameters. - tp_size = params.get("tp_size") - pp_size = params.get("pp_size") - quantization = params.get("quantization") - max_seq_len: int = params.get("max_seq_len") - # Dataset options - dataset_path: Path = params.get("dataset") - # IFB scheduler options - max_batch_size = params.get("max_batch_size") - max_num_tokens = params.get("max_num_tokens") - # Tuning heuristics options - target_input_len: int = params.get("target_input_len") - target_output_len: int = params.get("target_output_len") - - load_format = "dummy" if params.get("no_weights_loading") else "auto" - trust_remote_code: bool = params.get("trust_remote_code") - model_name = bench_env.model - checkpoint_path = bench_env.checkpoint_path or model_name - model_config = get_model_config(model_name, bench_env.checkpoint_path) - engine_dir = Path(bench_env.workspace, model_name, - f"tp_{tp_size}_pp_{pp_size}") - - # Set the compute quantization. - quant_algo = QuantAlgo(quantization) if quantization is not None else None - quant_config = QuantConfig(quant_algo=quant_algo) - # If the quantization is NVFP4 or FP8, force the KV cache dtype to FP8. - if quant_algo in [QuantAlgo.NVFP4, QuantAlgo.FP8]: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - - # Initialize the HF tokenizer for the specified model. - tokenizer = initialize_tokenizer(checkpoint_path) - # If we receive dataset from a path or stdin, parse and gather metadata. - if dataset_path: - logger.info("Found dataset.") - # Dataset Loading and Preparation - with open(dataset_path, "r") as dataset: - metadata, _ = create_dataset_from_stream( - tokenizer, - dataset, - ) - max_seq_len = metadata.max_sequence_length - target_input_len = metadata.avg_isl - target_output_len = metadata.avg_osl - logger.info(metadata.get_summary_for_print()) - - # Use user-specified engine settings if provided. - if max_batch_size and max_num_tokens: - logger.info("Use user-provided max batch size and max num tokens for " - "engine build and benchmark.") - # If not provided, use the engine setting provided by trtllm-bench. - else: - logger.info( - "Max batch size and max num tokens are not provided, " - "use tuning heuristics or pre-defined setting from trtllm-bench.") - max_batch_size, max_num_tokens = get_benchmark_engine_settings( - model_config, - quant_config, - tp_size, - pp_size, - target_input_len, - target_output_len, - ) - - # Construct a TRT-LLM build config. - build_config = BuildConfig(max_batch_size=max_batch_size, - max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens) - - build_config.plugin_config.dtype = model_config.dtype - # Enable multiple profiles and paged context FMHA. - build_config.plugin_config.multiple_profiles = True - # build_config.plugin_config._reduce_fusion = True - - # Enable FHMA, and FP8 FMHA if NVFP4 or FP8 quantization is enabled. - # TODO: Revisit, there is an issue with enabling FHMA. If only - # paged FMHA is enabled with NVFP4 or FP8 quantization, the Builder - # will not enable the FP8 FMHA. - build_config.plugin_config.use_paged_context_fmha = True - if quant_algo in [QuantAlgo.NVFP4, QuantAlgo.FP8]: - build_config.plugin_config.use_fp8_context_fmha = True - # Enable nvfp4 gemm_plugin explicitly for Blackwell - if quant_algo == QuantAlgo.NVFP4: - build_config.plugin_config.gemm_plugin = "nvfp4" - - # Build the LLM engine with the LLMAPI. - llm = LLM(checkpoint_path, - tokenizer, - dtype=model_config.dtype, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - build_config=build_config, - quant_config=quant_config, - workspace=str(bench_env.workspace), - load_format=load_format, - trust_remote_code=trust_remote_code, - telemetry_config=bench_env.telemetry_config) - # Save the engine. - llm.save(engine_dir) - llm.shutdown() - - logger.info( - "\n===========================================================\n" - "= ENGINE BUILD INFO\n" - "===========================================================\n" - f"Model Name:\t\t{bench_env.model}\n" - f"Model Path:\t\t{bench_env.checkpoint_path}\n" - f"Workspace Directory:\t{bench_env.workspace}\n" - f"Engine Directory:\t{engine_dir}\n\n" - "===========================================================\n" - "= ENGINE CONFIGURATION DETAILS\n" - "===========================================================\n" - f"Max Sequence Length:\t\t{max_seq_len}\n" - f"Max Batch Size:\t\t\t{max_batch_size}\n" - f"Max Num Tokens:\t\t\t{max_num_tokens}\n" - f"Quantization:\t\t\t{quant_config.quant_algo}\n" - f"KV Cache Dtype:\t\t\t{quant_config.kv_cache_quant_algo}\n" - "===========================================================\n") - - logger.info( - "\n\n===========================================================\n" - f"ENGINE SAVED: {engine_dir}\n" - "===========================================================\n") diff --git a/tensorrt_llm/builder.py b/tensorrt_llm/builder.py deleted file mode 100644 index 764c83db3668..000000000000 --- a/tensorrt_llm/builder.py +++ /dev/null @@ -1,1292 +0,0 @@ -# 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 json -import math -import os -import shutil -import time -from pathlib import Path -from typing import Dict, Optional, Union - -import numpy as np -import tensorrt as trt -from pydantic import Field - -from ._common import _is_building, check_max_num_tokens, serialize_engine -from ._deprecation import emit_engine_arch_deprecation -from ._utils import (get_sm_version, np_bfloat16, np_float8, str_dtype_to_trt, - to_json_file, trt_gte) -from .functional import PositionEmbeddingType -from .graph_rewriting import optimize -from .llmapi.kv_cache_type import KVCacheType -from .llmapi.utils import StrictBaseModel -from .logger import logger -from .lora_helper import LoraConfig -from .models import PretrainedConfig, PretrainedModel -from .models.modeling_utils import SpeculativeDecodingMode, optimize_model -from .network import Network, net_guard -from .plugin import PluginConfig -from .quantization import QuantAlgo, QuantMode -from .version import __version__ - - -class ConfigEncoder(json.JSONEncoder): - - def default(self, obj): - if hasattr(obj, 'model_dump'): - # Handle Pydantic models (including DecodingBaseConfig and subclasses) - return obj.model_dump(mode='json') - else: - return super().default(obj) - - -class BuilderConfig(object): - - def __init__(self, **kwargs): - # intentionally use **kwargs, user should never call this ctor directly, - # use Builder.create_builder_config() instead - pass - - def _init(self, trt_builder_config, **kwargs): - self._trt_builder_config = trt_builder_config - for key, value in kwargs.items(): - setattr(self, key, value) - return self - - @property - def trt_builder_config(self) -> trt.IBuilderConfig: - return self._trt_builder_config - - def to_dict(self) -> Dict: - '''return a dict with keys - { - "builder_config": { - # all key values set by the _init function - }, - "plugin_config": { - # the network plugin_config (if any) attached to this BuilderConfig object - # inside the Builder.build_engine - } - } - ''' - config = {'builder_config': {}} - for k in self.__dict__.keys(): - if k not in ['_trt_builder_config', 'plugin_config']: - config['builder_config'][k] = self.__getattribute__(k) - if hasattr(self, 'plugin_config'): - assert isinstance(self.plugin_config, PluginConfig), \ - f"Found unexpected plugin_config object with type: {type(self.plugin_config)}" - config['plugin_config'] = self.plugin_config.model_dump(mode="json") - return config - - -class Builder(): - - _ALLOWED_PRECISIONS = [ - 'float32', 'float16', 'bfloat16', trt.DataType.HALF, trt.DataType.FLOAT, - trt.DataType.BF16 - ] - - def __init__(self): - super().__init__() - self._trt_builder = trt.Builder(logger.trt_logger) - self.strongly_typed = True - - @property - def trt_builder(self) -> trt.Builder: - return self._trt_builder - - def create_network(self) -> Network: - explicit_batch_flag = 0 - # Explicit batch flag will be deprecated in TRT 10 - if "EXPLICIT_BATCH" in trt.NetworkDefinitionCreationFlag.__members__.keys( - ): - explicit_batch_flag = 1 << int( - trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) - - if self.strongly_typed: - return Network()._init( - self.trt_builder.create_network( - explicit_batch_flag - | (1 << int( - trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)))) - else: - return Network()._init( - self.trt_builder.create_network(explicit_batch_flag)) - - def create_builder_config(self, - precision: Union[str, trt.DataType], - timing_cache: Union[str, Path, - trt.ITimingCache] = None, - tensor_parallel: int = 1, - use_refit: bool = False, - int8: bool = False, - strongly_typed: bool = True, - force_num_profiles: Optional[int] = None, - profiling_verbosity: str = "layer_names_only", - use_strip_plan: bool = False, - weight_streaming: bool = False, - precision_constraints: Optional[str] = "obey", - **kwargs) -> BuilderConfig: - ''' @brief Create a builder config with given precisions and timing cache - @param precision: one of allowed precisions, defined in Builder._ALLOWED_PRECISIONS - @param timing_cache: a timing cache object or a path to a timing cache file - @param tensor_parallel: number of GPUs used for tensor parallel - @param kwargs: any other arguments users would like to attach to the config object as attributes - @param refit: set to accelerate multi-gpu building, build engine for 1 gpu and refit for the others - @param int8: whether to build with int8 enabled or not. Can't be used together with refit option - @return: A BuilderConfig object, return None if failed - ''' - self.strongly_typed = self.strongly_typed and strongly_typed - - quant_mode = kwargs.get("quant_mode", QuantMode(0)) - if not strongly_typed and precision not in self._ALLOWED_PRECISIONS: - logger.error( - f"precision should be one of {self._ALLOWED_PRECISIONS}") - - config = self.trt_builder.create_builder_config() - if weight_streaming: - config.set_flag(trt.BuilderFlag.WEIGHT_STREAMING) - if not self.strongly_typed: - fp8 = quant_mode.has_fp8_qdq() or quant_mode.has_fp8_kv_cache() - if precision == 'float16' or precision == trt.DataType.HALF: - config.set_flag(trt.BuilderFlag.FP16) - if precision_constraints == 'obey': - config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) - elif precision == 'bfloat16' or precision == trt.DataType.BF16: - config.set_flag(trt.BuilderFlag.BF16) - if precision_constraints == 'obey': - config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) - if int8: - config.set_flag(trt.BuilderFlag.INT8) - if fp8: - config.set_flag(trt.BuilderFlag.FP8) - if precision_constraints == 'obey': - config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) - - if use_refit: - config.set_flag(trt.BuilderFlag.REFIT) - - # Use fine-grained refit when strip plan is enabled in TRT10.2+. - if use_strip_plan: - config.set_flag(trt.BuilderFlag.REFIT_INDIVIDUAL) - - if use_strip_plan: - config.set_flag(trt.BuilderFlag.STRIP_PLAN) - - # Set TRT Engine profiling verbosity - if profiling_verbosity == "detailed": - config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED - elif profiling_verbosity == "none": - config.profiling_verbosity = trt.ProfilingVerbosity.NONE - else: - config.profiling_verbosity = trt.ProfilingVerbosity.LAYER_NAMES_ONLY - - # set timing cache - cache = None - if timing_cache is not None: - # use given cache - if isinstance(timing_cache, trt.ITimingCache): - cache = timing_cache - # read cache from file - elif isinstance(timing_cache, - (str, Path)) and os.path.exists(timing_cache): - with open(timing_cache, "rb") as f: - cache = config.create_timing_cache(f.read()) - else: - logger.warning( - "Invalid timing cache, using freshly created one") - if cache is None: - cache = config.create_timing_cache(b"") - # When user does not given any existing cache, internally always created one - # so the cache should never None here - assert cache is not None and isinstance(cache, trt.ITimingCache) - config.set_timing_cache(cache, ignore_mismatch=False) - - # set weight sparsity - weight_sparsity = kwargs.get("weight_sparsity", False) - if weight_sparsity: - config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS) - - # TODO: remove this constraint after trt 10.6 is integrated - if trt_gte(10, 6): - # set monitor memory - monitor_memory = kwargs.get("monitor_memory", False) - if monitor_memory: - config.set_flag(trt.BuilderFlag.MONITOR_MEMORY) - - return BuilderConfig()._init(config, - precision=precision, - tensor_parallel=tensor_parallel, - use_refit=use_refit, - int8=int8, - force_num_profiles=force_num_profiles, - strongly_typed=self.strongly_typed, - use_strip_plan=use_strip_plan, - **kwargs) - - def _add_optimization_profile(self, network: Network, - builder_config: BuilderConfig): - assert isinstance(builder_config, BuilderConfig) - assert isinstance(network, Network) - input_tensors = network._inputs - if len(input_tensors) == 0: - logger.warning("There are no inputs in the network!") - return - num_profiles = len(list(input_tensors.values())[0].profiles) - force_num_profiles = getattr(builder_config, "force_num_profiles", None) - for i in range(num_profiles): - logger.debug(f'Adding optimization profile {i+1}/{num_profiles}') - profile = self.trt_builder.create_optimization_profile() - for input_name in input_tensors.keys(): - if len(input_tensors[input_name].profiles) == 0: - continue - shape_profile = input_tensors[input_name].profiles[i] - min_shape = [*shape_profile.min] - opt_shape = [*shape_profile.opt] - max_shape = [*shape_profile.max] - profile.set_shape(input_name, min_shape, opt_shape, max_shape) - logger.debug( - f'{input_name}, min: {min_shape}, opt: {opt_shape}, max: {max_shape}, dimension names: {shape_profile.dimension_names}' - ) - ret = builder_config.trt_builder_config.add_optimization_profile( - profile) - logger.debug(f"Added optimization profile: #{ret}") - if force_num_profiles is not None and ( - i + 1 - ) == force_num_profiles and force_num_profiles < num_profiles: - logger.warning( - f"Only adding {force_num_profiles} profiles instead of {num_profiles}." - ) - break - assert self._validate_named_dimensions( - network, builder_config - ), "Validation of the tensor dimension ranges failed, please check the dimension ranges, find the offensive tensor and dimension name in above the error log" - - def _validate_named_dimensions(self, network: Network, - builder_config) -> bool: - ''' - For each profile, validate that the named dimensions of different input tensors in this profile all have same range. - TRT will validate the same condition, validate it earlier to make sure the modeling in TensorRT LLM are correct and - makes the error msg more user friendly. - ''' - valid = True - for profile_idx in range( - builder_config.trt_builder_config.num_optimization_profiles): - dimension_to_range = {} - for input_name, input_tensor in network._inputs.items(): - # it's legal that a Tensor does not have dim_range? - if len(input_tensor.profiles) != 0: - profile = input_tensor.profiles[profile_idx] - for dim_idx, dim_name in enumerate(profile.dimension_names): - if dim_name not in dimension_to_range: - dimension_to_range[dim_name] = [] - min, opt, max = profile.min[dim_idx], profile.opt[ - dim_idx], profile.max[dim_idx] - dimension_to_range[dim_name].append( - (input_name, (min, opt, max))) - for dim, ranges in dimension_to_range.items(): - unique_ranges = set([r[1] for r in ranges]) - logger.debug( - f"Validating dimension:{dim}, ranges for this dim are:{unique_ranges}" - ) - if len(unique_ranges) != 1: - logger.error( - f"Found illegal dimension setting for profile {profile_idx}, dimension name is: {dim}" - ) - logger.error( - "Offensive tensors which have this dimension are:\n" + - "\n".join([f"{r[1]} {dim} {r[0]}" for r in ranges])) - valid = False - return valid - - @_is_building - def refit_engine(self, network: Network, engine_buffer) -> trt.IHostMemory: - ''' - @brief: Refit one TensorRT engine using weights from the network, - user should guarantee that the engine is built with REFIT flag, and the network has the same structure with the engine. - @param engine_buffer: A serialized TensorRT engine. - @param network: Network object. - @return: A serialized TRT engine if refit successfully, None otherwise - ''' - assert isinstance(network, Network) - logger.info('Refit TRT engine') - runtime = trt.Runtime(logger.trt_logger) - engine = runtime.deserialize_cuda_engine(engine_buffer) - - tik = time.time() - - # Refit engine - refitter = trt.Refitter(engine, logger.trt_logger) - if network.named_parameters is not None: - for name, param in network.named_parameters: - if param._get_weights( - ) is None or not refitter.set_named_weights( - name, param._get_weights()): - logger.error(f'Failed to refit weight: {name}') - return None - else: - logger.error( - 'Please set named parameters before building multiple engines.') - return None - - if not refitter.refit_cuda_engine(): - logger.error('Failed to refit engine.') - return None - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of refitting {engine.name}: {t}') - serialized_engine = engine.serialize() - return serialized_engine - - @_is_building - def build_engine(self, - network: Network, - builder_config: BuilderConfig, - managed_weights: dict = None) -> trt.IHostMemory: - ''' - @brief: Build one TensorRT engine from the network. - @param network: Network object. - @param builder_config: BuilderConfig object. - @return: A serialized TRT engine. - ''' - assert isinstance(network, Network) - builder_config.plugin_config = network.plugin_config - if builder_config.trt_builder_config.num_optimization_profiles == 0: - self._add_optimization_profile(network, builder_config) - logger.info( - f"Total optimization profiles added: {builder_config.trt_builder_config.num_optimization_profiles}" - ) - engine = None - - tik = time.time() - # Rename weights - if network.named_parameters is not None: - managed_parameters = [] - for name, param in network.named_parameters: - if param.is_managed(network): - assert managed_weights is not None, "managed_weights should be provided when enabled" - managed_parameters.append(param) - param.set_name(name, network) - continue - if param._get_weights(network) is None: - if not param.is_buffer: - logger.debug( - f"Parameter {name} {param.raw_value.shape} {param.raw_value.dtype} was created" - " but unused in forward method, so not materialized to TRT network" - ) - continue - if not param.set_name(name, network): - raise RuntimeError(f'Failed to set weight: {name}') - # This mark_weights_refittable has no side effect when refit_individual is not enabled. - network.trt_network.mark_weights_refittable(name) - - network._fill_weights() - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info( - f'Total time to initialize the weights in network {network.trt_network.name}: {t}' - ) - - # Build engine - logger.info(f'Build TensorRT engine {network.trt_network.name}') - tik = time.time() - engine = self.trt_builder.build_serialized_network( - network.trt_network, builder_config.trt_builder_config) - assert engine is not None, 'Engine building failed, please check the error log.' - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of building {network.trt_network.name}: {t}') - - if managed_weights is not None and network.named_parameters is not None: - for param in managed_parameters: - name = param.name - value: np.ndarray = param._value - if value is None: - logger.error(f'Failed to get weight: {name}') - continue - if param.need_transpose: - # MOE has ndim=3 and uses plugin, no need to transpose - value = value.transpose(1, 0) # WAR for bug 4641821 - managed_weights[name] = value - - return engine - - @staticmethod - def save_timing_cache(builder_config: BuilderConfig, out_path: str) -> bool: - '''Serialize timing cache of given builder config to file specified by out_path - return True if the cache is successfully serialized, False otherwise - ''' - cache = builder_config.trt_builder_config.get_timing_cache() - if cache is None: - logger.warning( - 'No timing cache found in the given builder config, skip saving.' - ) - return False - with cache.serialize() as buffer: - with open(out_path, "wb") as f: - f.write(buffer) - f.flush() - os.fsync(f) - logger.info(f'Timing cache serialized to {out_path}') - return True - - @staticmethod - def save_config(builder_config: BuilderConfig, config_path: str): - config = builder_config.to_dict() - to_json_file(config, config_path) - logger.info(f'Config saved to {config_path}.') - - -class BuildConfig(StrictBaseModel): - """Configuration class for TensorRT LLM engine building parameters. - - This class contains all the configuration parameters needed to build a TensorRT LLM engine, - including sequence length limits, batch sizes, optimization settings, and various features. - """ - max_input_len: int = Field(default=1024, - description="Maximum length of input sequences.") - max_seq_len: Optional[int] = Field( - default=None, - description= - "The maximum possible sequence length for a single request, including both input and generated " - "output tokens.") - opt_batch_size: int = Field( - default=8, description="Optimal batch size for engine optimization.") - max_batch_size: int = Field( - default=2048, description="Maximum batch size the engine can handle.") - max_beam_width: int = Field( - default=1, description="Maximum beam width for beam search decoding.") - max_num_tokens: int = Field( - default=8192, - description="Maximum number of batched input tokens after padding is " - "removed in each batch.") - opt_num_tokens: Optional[int] = Field( - default=None, - description= - "Optimal number of batched input tokens for engine optimization.") - max_prompt_embedding_table_size: int = Field( - default=0, - description="Maximum size of prompt embedding table for prompt tuning.") - kv_cache_type: Optional[KVCacheType] = Field( - default=None, - description= - "Type of KV cache to use (CONTINUOUS or PAGED). If None, defaults to PAGED." - ) - gather_context_logits: bool = Field( - default=False, - description="Whether to gather logits during context phase.") - gather_generation_logits: bool = Field( - default=False, - description="Whether to gather logits during generation phase.") - strongly_typed: bool = Field(default=True, - description="Whether to use strongly_typed.") - force_num_profiles: Optional[int] = Field( - default=None, - description= - "Force a specific number of optimization profiles. If None, auto-determined." - ) - profiling_verbosity: str = Field( - default='layer_names_only', - description= - "Verbosity level for TensorRT profiling ('layer_names_only', 'detailed', 'none')." - ) - enable_debug_output: bool = Field( - default=False, - description="Whether to enable debug output during building.") - max_draft_len: int = Field( - default=0, - description="Maximum length of draft tokens for speculative decoding.") - speculative_decoding_mode: SpeculativeDecodingMode = Field( - default=SpeculativeDecodingMode.NONE, - description="Mode for speculative decoding (NONE, MEDUSA, EAGLE, etc.)." - ) - use_refit: bool = Field( - default=False, - description="Whether to enable engine refitting capabilities.") - input_timing_cache: Optional[str] = Field( - default=None, - description= - "Path to input timing cache file. If None, no input cache used.") - output_timing_cache: str = Field( - default='model.cache', description="Path to output timing cache file.") - lora_config: LoraConfig = Field( - default_factory=LoraConfig, - description="Configuration for LoRA (Low-Rank Adaptation) fine-tuning.") - weight_sparsity: bool = Field( - default=False, - description="Whether to enable weight sparsity optimization.") - weight_streaming: bool = Field( - default=False, - description="Whether to enable weight streaming for large models.") - plugin_config: PluginConfig = Field( - default_factory=PluginConfig, - description="Configuration for TensorRT LLM plugins.") - use_strip_plan: bool = Field( - default=False, - description="Whether to use stripped plan for engine building.") - max_encoder_input_len: int = Field( - default=1024, - description="Maximum encoder input length for encoder-decoder models.") - dry_run: bool = Field( - default=False, - description= - "Whether to perform a dry run without actually building the engine.") - visualize_network: Optional[str] = Field( - default=None, - description= - "Path to save network visualization. If None, no visualization generated." - ) - monitor_memory: bool = Field( - default=False, - description="Whether to monitor memory usage during building.") - use_mrope: bool = Field( - default=False, - description= - "Whether to use Multi-RoPE (Rotary Position Embedding) optimization.") - - # Since we have some overlapping between kv_cache_type, paged_kv_cache, and paged_state (later two will be deprecated in the future), - # we need to handle it given model architecture. - def update_kv_cache_type(self, model_architecture: str): - paged_kv_cache_attr = 'paged_state' if model_architecture in [ - 'MambaForCausalLM', 'RecurrentGemmaForCausalLM' - ] else 'paged_kv_cache' - assert self.plugin_config is not None - paged_kv_cache_val = getattr(self.plugin_config, paged_kv_cache_attr) - - if self.kv_cache_type is not None: - if paged_kv_cache_val is not None: - assert (paged_kv_cache_val == True - and self.kv_cache_type == KVCacheType.PAGED) or ( - paged_kv_cache_val == False - and self.kv_cache_type != KVCacheType.PAGED) - else: - setattr(self.plugin_config, paged_kv_cache_attr, - self.kv_cache_type == KVCacheType.PAGED) - else: - if paged_kv_cache_val is not None: - self.kv_cache_type = KVCacheType.PAGED if paged_kv_cache_val else KVCacheType.CONTINUOUS - else: - self.kv_cache_type = KVCacheType.PAGED - setattr(self.plugin_config, paged_kv_cache_attr, - self.kv_cache_type == KVCacheType.PAGED) - - assert self.kv_cache_type is not None and getattr( - self.plugin_config, paged_kv_cache_attr) is not None - - def override_attri(attr_name, value): - val = getattr(self.plugin_config, attr_name) - if val is not None and val != value: - logger.warning(f'Overriding {attr_name} to {value}') - setattr(self.plugin_config, attr_name, value) - - # Init other paged kvcache attri to false. For RecurrentGemma, we only support paged_state and paged_kv_cache have - # the same values. All other models should only consume either of the value and set other to False. - is_recurrent_gemma = model_architecture == 'RecurrentGemmaForCausalLM' - - if paged_kv_cache_attr == 'paged_state': - override_attri( - 'paged_kv_cache', - getattr(self.plugin_config, paged_kv_cache_attr) - if is_recurrent_gemma else False) - else: - override_attri('paged_state', False) - - @classmethod - def from_json_file(cls, config_file): - with open(config_file) as f: - config = json.load(f) - return BuildConfig(**config) - - -class EngineConfig: - - def __init__(self, pretrained_config: 'PretrainedConfig', - build_config: 'BuildConfig', version: str): - self.pretrained_config = pretrained_config - self.build_config = build_config - self.version = version - - @classmethod - def from_json_file(cls, config_file): - with open(config_file) as f: - return cls.from_json_str(f.read()) - - @classmethod - def from_json_str(cls, config_str): - config = json.loads(config_str) - return cls(PretrainedConfig.from_dict(config['pretrained_config']), - BuildConfig(**config['build_config']), config['version']) - - def to_dict(self): - build_config = self.build_config.model_dump(mode="json") - build_config.pop('dry_run', None) # Not an Engine Characteristic - build_config.pop('visualize_network', - None) # Not an Engine Characteristic - return { - 'version': self.version, - 'pretrained_config': self.pretrained_config.to_dict(), - 'build_config': build_config, - } - - -class Engine: - - def __init__( - self, - config: EngineConfig, - engine: Union[trt.IHostMemory, None], - managed_weights: dict[str, np.ndarray] = {}, - ): - self.config = config - self.engine = engine - self.managed_weights = managed_weights - if self.managed_weights is None: - self.managed_weights = {} - for name, value in self.managed_weights.items(): - if not value.flags['C_CONTIGUOUS']: - self.managed_weights[name] = np.ascontiguousarray(value) - - def save(self, engine_dir: str): - os.makedirs(engine_dir, exist_ok=True) - lora_config = self.config.build_config.lora_config - lora_dirs = lora_config.lora_dir - root_lora_dir = os.path.join(engine_dir, 'lora') - if len(lora_dirs) > 0: - os.makedirs(root_lora_dir, exist_ok=True) - for index, lora_dir in enumerate(lora_dirs): - if lora_config.lora_ckpt_source == 'hf': - target_lora_dir = f"{root_lora_dir}/{index}" - os.makedirs(target_lora_dir, exist_ok=True) - shutil.copy2(os.path.join(lora_dir, 'adapter_config.json'), - target_lora_dir) - weight_file = os.path.join(lora_dir, 'adapter_model.bin') - if os.path.exists(weight_file): - shutil.copy2(weight_file, target_lora_dir) - weight_file = os.path.join(lora_dir, - 'adapter_model.safetensors') - if os.path.exists(weight_file): - shutil.copy2(weight_file, target_lora_dir) - lora_config.lora_dir[index] = f"lora/{index}" - elif lora_config.lora_ckpt_source == 'nemo': - target_lora_file = f"{root_lora_dir}/{index}.nemo" - shutil.copyfile(lora_dir, target_lora_file) - lora_config.lora_dir[index] = f"lora/{index}.nemo" - else: - if os.path.exists(root_lora_dir) and os.path.isdir(root_lora_dir): - shutil.rmtree(root_lora_dir) - if self.config.pretrained_config.mapping.rank == 0: - config_dict = self.config.to_dict() - if self.config.pretrained_config.quant_algo == QuantAlgo.MIXED_PRECISION: - quant_dict = { - 'version': self.config.version, - } - quant_dict.update( - config_dict['pretrained_config']['quantization']) - config_dict['pretrained_config']['quantization'].pop( - 'quantized_layers', None) - with open(os.path.join(engine_dir, 'quant_cfg.json'), - "w", - encoding="utf-8") as f: - json.dump(quant_dict, f, indent=4, cls=ConfigEncoder) - - with open(os.path.join(engine_dir, 'config.json'), - "w", - encoding="utf-8") as f: - json.dump(config_dict, f, indent=4, cls=ConfigEncoder) - if self.engine is not None: - serialize_engine( - self.engine, - os.path.join( - engine_dir, - f'rank{self.config.pretrained_config.mapping.rank}.engine')) - if self.managed_weights is not None and len(self.managed_weights) > 0: - fn = os.path.join( - engine_dir, - f'rank{self.config.pretrained_config.mapping.rank}_managed_weights.safetensors' - ) - serialize_managed_weights(self.managed_weights, fn) - - @classmethod - def from_dir(cls, engine_dir: str, rank: int = 0): - with open(os.path.join(engine_dir, f'rank{rank}.engine'), 'rb') as f: - engine_buffer = f.read() - - mw_path = os.path.join(engine_dir, - f'rank{rank}_managed_weights.safetensors') - managed_weights = deserialize_managed_weights( - mw_path) if os.path.exists(mw_path) else None - - config = EngineConfig.from_json_file( - os.path.join(engine_dir, 'config.json')) - config.pretrained_config.set_rank(rank) - - return cls(config, engine_buffer, managed_weights) - - @classmethod - def from_buffer(cls, - engine_buffer: Union[trt.IHostMemory, bytes], - json_config_str: str, - rank: int = 0): - config = EngineConfig.from_json_str(json_config_str) - config.pretrained_config.set_rank(rank) - return cls(config, engine_buffer) - - -def get_engine_version(engine_dir: str) -> Union[None, str]: - engine_dir = Path(engine_dir) - config_path = engine_dir / "config.json" - with open(config_path, 'r') as f: - config = json.load(f) - - if 'version' not in config: - return None - - return config['version'] - - -def optimize_model_with_config(model: PretrainedModel, - build_config: BuildConfig): - gemm_swiglu_plugin = build_config.plugin_config.gemm_swiglu_plugin - low_latency_gemm_swiglu_plugin = build_config.plugin_config.low_latency_gemm_swiglu_plugin - if gemm_swiglu_plugin or low_latency_gemm_swiglu_plugin: - if not build_config.plugin_config.use_fused_mlp: - raise RuntimeError( - "GemmSwiGLU plugin requires --use_fused_mlp flag") - if gemm_swiglu_plugin not in [ - "fp8" - ] and low_latency_gemm_swiglu_plugin not in ["fp8"]: - raise RuntimeError( - f"GemmSwiGLU plugin currently has limited support: fp8 only, " - f"got: {gemm_swiglu_plugin}" - f"got: {low_latency_gemm_swiglu_plugin}") - - if build_config.plugin_config.lora_plugin is not None: - model.use_lora(build_config.lora_config) - - is_enc_dec = model.config.architecture in ["EncoderModel", "DecoderModel"] - # FusedMLP does not support RecurrentGemma FP8 currently. - is_recurrent_gemma = model.config.architecture in [ - "RecurrentGemmaForCausalLM" - ] - is_fp8 = model.config.quantization.quant_algo == QuantAlgo.FP8 - model = optimize_model( - model, - share_embedding_table=True, - use_ootb_moe=build_config.plugin_config.moe_plugin is None, - use_fused_mlp=(build_config.plugin_config.use_fused_mlp - and not is_enc_dec - and not (is_recurrent_gemma and is_fp8)), - gemm_swiglu_plugin_dtype=gemm_swiglu_plugin, - low_latency_gemm_swiglu_plugin_dtype=low_latency_gemm_swiglu_plugin, - use_fused_rg_lru=is_recurrent_gemma, - use_unfused_qkv_gemm=False, - use_prompt_tuning=(build_config.max_prompt_embedding_table_size > 0), - use_lora=build_config.plugin_config.lora_plugin is not None, - max_lora_rank=build_config.lora_config.max_lora_rank, - use_fp8_context_fmha=(model.config.quantization.quant_algo in [ - QuantAlgo.FP8, QuantAlgo.W4A8_AWQ, QuantAlgo.NVFP4 - ] and build_config.plugin_config.use_fp8_context_fmha), - fuse_fp4_quant=build_config.plugin_config.fuse_fp4_quant, - use_optimize_cross_qkv=True, - use_dora=build_config.plugin_config.dora_plugin) - - if is_enc_dec: - model.precompute_relative_attention_bias(build_config) - return model - - -def _init_max_seq_len(model_config, build_config): - """ - If max_seq_len is not specified, set it to max_position_embeddings * rotary_factor - Additional checks to ensure max_seq_len, max_input_len, and max_num_tokens have valid values. - """ - # Extract rotary scaling which will be used for checks and default value of max_seq_len - rotary_scaling = getattr(model_config, "rotary_scaling", None) - if rotary_scaling is not None: - rotary_type = rotary_scaling.get('type', - rotary_scaling.get('rope_type')) - rotary_factor = rotary_scaling.get( - 'factor', 1.0) if rotary_type not in ("su", "longrope", - "llama3") else 1 - else: - rotary_factor = 1 - - if model_config.architecture == "EncoderModel": - if build_config.max_seq_len is None: - build_config.max_seq_len = build_config.max_input_len - logger.info( - f'max_seq_len is not specified for EncoderModel, using --max_input_len.' - ) - assert build_config.max_input_len == build_config.max_seq_len, f"EncoderModel should have same --max_input_len ({build_config.max_input_len}) and --max_seq_len ({build_config.max_seq_len})." - - if build_config.max_seq_len is None: - # Step 1: Find the upper bound of max_seq_len - deduced_max_seq_len = 2048 - if model_config.max_position_embeddings is not None: - deduced_max_seq_len = model_config.max_position_embeddings - - # Step 2: Scale max_seq_len with rotary scaling - if rotary_factor != 1: - deduced_max_seq_len = math.ceil(deduced_max_seq_len * rotary_factor) - logger.warning( - f'max_seq_len is scaled to {deduced_max_seq_len} by rotary scaling {rotary_factor}' - ) - - # Step 3: Assign the new max_seq_len - build_config.max_seq_len = int(deduced_max_seq_len) - logger.info( - f'max_seq_len is not specified, using deduced value {deduced_max_seq_len}' - ) - else: - if not build_config.plugin_config.streamingllm and model_config.max_position_embeddings is not None \ - and model_config.position_embedding_type != PositionEmbeddingType.relative: - if build_config.max_seq_len > model_config.max_position_embeddings * rotary_factor: - logger.warning( - f'max_seq_len {build_config.max_seq_len} is larger than max_position_embeddings {model_config.max_position_embeddings} * rotary scaling {rotary_factor}, ' - 'the model accuracy might be affected') - - if build_config.max_input_len > build_config.max_seq_len: - logger.warning( - f'max_input_len is {build_config.max_input_len} is larger than max_seq_len {build_config.max_seq_len}, clipping it to max_seq_len' - ) - build_config.max_input_len = build_config.max_seq_len - - # Check and may modify max_num_tokens and opt_num_tokens (need to happen after max_seq_len is deduced) - max_num_tokens, opt_num_tokens = check_max_num_tokens( - max_num_tokens=build_config.max_num_tokens, - opt_num_tokens=build_config.opt_num_tokens, - max_batch_size=build_config.max_batch_size, - max_input_len=build_config.max_input_len, - max_seq_len=build_config.max_seq_len, - max_beam_width=build_config.max_beam_width, - remove_input_padding=build_config.plugin_config.remove_input_padding, - enable_context_fmha=build_config.plugin_config.context_fmha, - tokens_per_block=build_config.plugin_config.tokens_per_block, - multiple_profiles=build_config.plugin_config.multiple_profiles, - ) - build_config.max_num_tokens, build_config.opt_num_tokens = max_num_tokens, opt_num_tokens - - if build_config.plugin_config.remove_input_padding and build_config.plugin_config.context_fmha: - if build_config.max_input_len: - logger.warning( - 'padding removal and fMHA are both enabled, max_input_len is not required and will be ignored' - ) - else: - assert build_config.max_input_len is not None, 'padding removal and fMHA aren\'t both enabled, max_input_len is required' - if build_config.max_seq_len: - assert build_config.max_input_len <= build_config.max_seq_len, 'max_input_len should not be larger than max_seq_len' - - -def serialize_managed_weights(managed_weights: dict[str, np.ndarray], - path: str | Path, - metadata=None) -> None: - header = {} - if metadata is not None: - header["__metadata__"] = metadata - begin = 0 - for name, value in managed_weights.items(): - size = value.size * value.itemsize - if value.dtype == np.float32: - dtype = "F32" - elif value.dtype == np.float16: - dtype = "F16" - elif value.dtype == np_bfloat16: - dtype = "BF16" - elif value.dtype == np_float8: - dtype = "F8_E4M3" - elif value.dtype == np.int64: - dtype = "I64" - elif value.dtype == np.int32: - dtype = "I32" - elif value.dtype == np.int8: - dtype = "I8" - else: - raise RuntimeError(f"Unsupported dtype: {value.dtype}") - header[name] = { - "dtype": dtype, - "shape": value.shape, - "data_offsets": [begin, begin + size], - } - begin += size - - header_json = json.dumps(header) - header_json_len = len(header_json) - with open(path, "wb") as f: - logger.info( - f"Serializing {len(managed_weights)} managed weights to {path}...") - f.write(header_json_len.to_bytes(8, byteorder="little")) - f.write(header_json.encode()) - for name, value in managed_weights.items(): - logger.debug(f"Serializing managed weight: {name}") - buf = value.data - f.write(buf) - - -def deserialize_managed_weights(path: str | Path) -> dict[str, np.ndarray]: - with open(path, "rb") as f: - header_json_len = int.from_bytes(f.read(8), byteorder="little") - header_json = f.read(header_json_len).decode() - header = json.loads(header_json) - - managed_weights = {} - for name, info in header.items(): - dtype = info["dtype"] - shape = info["shape"] - data_offsets = info["data_offsets"] - if dtype == "F32": - dtype = np.float32 - elif dtype == "F16": - dtype = np.float16 - elif dtype == "BF16": - dtype = np_bfloat16 - elif dtype == "F8_E4M3": - dtype = np_float8 - elif dtype == "I64": - dtype = np.int64 - elif dtype == "I32": - dtype = np.int32 - else: - raise RuntimeError(f"Unsupported dtype: {dtype}") - - f.seek(data_offsets[0] + header_json_len + 8) - buf = f.read(data_offsets[1] - data_offsets[0]) - value = np.frombuffer(buf, dtype=dtype).reshape(shape) - managed_weights[name] = value - - return managed_weights - - -def build(model: PretrainedModel, build_config: BuildConfig) -> Engine: - '''Build engine from given model and optimization options specified in the build_config - WARNING: this function may change the given model object state in some optimization passes - to avoid cloning a model since normally the LLM models consumes large memory. - Create a new fresh model object if you need to build with different options. - ''' - emit_engine_arch_deprecation("builder.build()") - tic = time.time() - # avoid changing the input config - build_config = build_config.model_copy(deep=True) - build_config.plugin_config.dtype = model.config.dtype - build_config.update_kv_cache_type(model.config.architecture) - - _init_max_seq_len(model.config, build_config) - - if build_config.plugin_config.streamingllm: - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "Paged Context FMHA is disabled because StreamingLLM is not supported when enabling paged KV context FMHA." - ) - if build_config.plugin_config.reduce_fusion and ( - model.config.mapping.tp_size == 1 or - (model.config.architecture != "LlamaForCausalLM" - and model.config.architecture != "Gemma2ForCausalLM" - and model.config.architecture != "MedusaForCausalLM")): - logger.warning('Overriding reduce_fusion to False') - build_config.plugin_config.reduce_fusion = False - if build_config.plugin_config.user_buffer and not build_config.plugin_config.reduce_fusion: - logger.warning('Overriding user_buffer to False') - build_config.plugin_config.user_buffer = False - if build_config.plugin_config.norm_quant_fusion and ( - build_config.plugin_config.reduce_fusion - or model.config.architecture != "LlamaForCausalLM" - or model.config.quantization.quant_algo != QuantAlgo.NVFP4): - logger.warning('Overriding norm_quant_fusion to False') - build_config.plugin_config.norm_quant_fusion = False - - if model.config.quantization.quant_algo == QuantAlgo.FP8 or \ - model.config.quantization.kv_cache_quant_algo == QuantAlgo.FP8: - build_config.strongly_typed = True - - if hasattr(model.config, 'max_draft_len'): - # If model.config has 'max_draft_len' but build_config not specified, - # use the value of model.config.max_draft_len to set the value of build_config.max_draft_len - if build_config.max_draft_len == 0: - build_config.max_draft_len = model.config.max_draft_len - - if hasattr(model.config, 'redrafter_num_beams') and hasattr( - model.config, 'redrafter_draft_len_per_beam'): - build_config.max_draft_len = model.config.redrafter_num_beams * model.config.redrafter_draft_len_per_beam - if build_config.speculative_decoding_mode != SpeculativeDecodingMode.EXPLICIT_DRAFT_TOKENS: - logger.warning( - 'speculative_decoding_mode is not EXPLICIT_DRAFT_TOKENS for ReDrafter model. Overwriting speculative_decoding_mode' - ) - build_config.speculative_decoding_mode = SpeculativeDecodingMode.EXPLICIT_DRAFT_TOKENS - - if build_config.speculative_decoding_mode != SpeculativeDecodingMode.NONE: - logger.info( - f'Increasing max_seq_len ({build_config.max_seq_len}) ' - f'by max_draft_len ({build_config.max_draft_len}) ' - 'to account for speculative decoding implementation specifics. ' - 'Maximum number of generated tokens remains the same. ' - f'New max_seq_len is set to {build_config.max_seq_len + build_config.max_draft_len}' - ) - build_config.max_seq_len += build_config.max_draft_len - - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.EAGLE: - assert hasattr(model.config, 'num_eagle_layers') - num_eagle_layers = model.config.num_eagle_layers - logger.info( - f'Increasing max_seq_len ({build_config.max_seq_len}) ' - f'by num_eagle_layers ({num_eagle_layers}) ' - 'to account for EAGLE implementation specifics. ' - 'Maximum number of generated tokens remains the same. ' - f'New max_seq_len is set to {build_config.max_seq_len + num_eagle_layers}' - ) - build_config.max_seq_len += num_eagle_layers - - if build_config.speculative_decoding_mode != SpeculativeDecodingMode.NONE: - num_tokens = build_config.max_batch_size * (build_config.max_draft_len + - 1) - if build_config.max_num_tokens < num_tokens: - logger.info( - f'max_num_tokens ({build_config.max_num_tokens}) is smaller than ' - 'max_batch_size * (max_draft_len + 1) = ' - f'({build_config.max_batch_size} * ({build_config.max_draft_len} + 1)). ' - f'New max_num_tokens is set to {num_tokens}.') - build_config.max_num_tokens = num_tokens - - # Logics to control paged_context_fmha and fp8_context_fmha - if not build_config.plugin_config.context_fmha: - build_config.plugin_config.use_fp8_context_fmha = False - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "Context FMHA is disabled, FP8 Context FMHA and Paged Context FMHA are disabled." - ) - elif model.config.quantization.quant_algo not in [ - QuantAlgo.FP8, QuantAlgo.W4A8_AWQ, QuantAlgo.NVFP4 - ]: - if build_config.plugin_config.use_fp8_context_fmha: - build_config.plugin_config.use_fp8_context_fmha = False - logger.warning( - "FP8 Context FMHA is disabled because it must be used together with the fp8 quantization workflow." - ) - if build_config.plugin_config.use_paged_context_fmha and model.config.quant_mode.has_fp8_kv_cache( - ): - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "FP8 Paged Context FMHA is disabled because FP8 context FMHA is disabled." - ) - elif get_sm_version() < 89: - build_config.plugin_config.use_fp8_context_fmha = False - logger.warning( - "FP8 context FMHA is disabled because it is only supported on Ada and Hopper Arch." - ) - if build_config.plugin_config.use_paged_context_fmha and model.config.quant_mode.has_fp8_kv_cache( - ): - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "FP8 Paged Context FMHA is disabled because FP8 context FMHA is disabled." - ) - elif build_config.plugin_config.use_paged_context_fmha: - if not model.config.quant_mode.has_fp8_kv_cache( - ) and build_config.plugin_config.use_fp8_context_fmha: - build_config.plugin_config.use_fp8_context_fmha = False - logger.warning( - "FP8 Paged Context FMHA is disabled because it must be used together with fp8 KV Cache." - ) - elif model.config.quant_mode.has_fp8_kv_cache( - ) and not build_config.plugin_config.use_fp8_context_fmha: - build_config.plugin_config.use_fp8_context_fmha = True - logger.warning( - "FP8 Context FMHA is enabled to support FP8 Paged Context FMHA." - ) - - if build_config.plugin_config.use_paged_context_fmha and model.config.quant_mode.has_int8_kv_cache( - ): - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "Paged Context FMHA is disabled because it doesn't work with int8 kv cache currently." - ) - - if get_sm_version() >= 100 and get_sm_version() < 120: - if model.config.quant_mode.is_int8_weight_only( - ) or model.config.quant_mode.is_int4_weight_only( - ) or model.config.quant_mode.has_int8_kv_cache(): - raise RuntimeError( - "INT8/INT4 quantization is not supported on SM>=100.") - if model.config.quant_mode.has_act_and_weight_quant(): - raise RuntimeError("SmoothQuant is not supported on SM>=100.") - if model.config.quant_mode.has_per_channel_scaling( - ) or model.config.quant_mode.has_per_token_dynamic_scaling(): - raise RuntimeError( - "Per-channel or per-token scaling is not supported on SM>=100.") - - model = optimize_model_with_config(model, build_config) - - builder = Builder() - builder_config = builder.create_builder_config( - precision=model.config.dtype, - use_refit=build_config.use_refit, - timing_cache=build_config.input_timing_cache, - int8=(model.config.quant_mode.has_act_or_weight_quant() - and not model.config.quant_mode.has_per_group_scaling()) - or model.config.quant_mode.has_int8_kv_cache(), - strongly_typed=build_config.strongly_typed, - force_num_profiles=build_config.force_num_profiles, - profiling_verbosity=build_config.profiling_verbosity, - quant_mode=model.config.quant_mode, - use_strip_plan=build_config.use_strip_plan, - weight_sparsity=build_config.weight_sparsity, - weight_streaming=build_config.weight_streaming, - monitor_memory=build_config.monitor_memory, - ) - - network = builder.create_network() - network.plugin_config = build_config.plugin_config - - use_weight_only = model.config.quant_mode.is_weight_only() - per_group = model.config.quant_mode.has_per_group_scaling() - use_smooth_quant = model.config.quant_mode.has_act_and_weight_quant() - use_qserve = model.config.quant_mode.is_qserve_w4a8() - use_fp8_rowwise = model.config.quant_mode.has_fp8_rowwise() - disable_weight_only_quant_plugin = model.config.disable_weight_only_quant_plugin if hasattr( - model.config, 'disable_weight_only_quant_plugin') else False - use_fp8_rowwise = model.config.quant_mode.has_fp8_rowwise() - use_fp4_gemm = model.config.quant_mode.has_nvfp4() - if use_fp4_gemm and network.plugin_config._explicitly_disable_gemm_plugin is False: - logger.info( - 'NVFP4 quantization detected, by default enabling NVFP4 GEMM plugin. To use OOTB GEMM, please explicitly set gemm_plugin to "disable"' - ) - network.plugin_config.gemm_plugin = "nvfp4" - - if build_config.plugin_config.manage_weights: - if use_weight_only and disable_weight_only_quant_plugin: - raise RuntimeError( - "Manage weights of weight only quant works only with plugin currently." - ) - - if use_weight_only and not disable_weight_only_quant_plugin: - if per_group: - network.plugin_config.weight_only_groupwise_quant_matmul_plugin = model.config.dtype - else: - network.plugin_config.weight_only_quant_matmul_plugin = model.config.dtype - if use_smooth_quant and model.config.quantization._use_plugin_sq and build_config.plugin_config.smooth_quant_plugins: - network.plugin_config.set_smooth_quant_plugins(model.config.dtype) - if use_qserve: - network.plugin_config.set_qserve_plugins(model.config.dtype) - if use_fp8_rowwise: - network.plugin_config.set_fp8_rowwise_quant_plugins(model.config.dtype) - nccl_plugin = model.config.dtype if model.config.mapping.world_size > 1 else None - network.plugin_config.set_nccl_plugin(nccl_plugin) - - with net_guard(network): - # Prepare - network.set_named_parameters(model.named_parameters()) - - # Forward - prepare_input_args = { - "max_batch_size": - build_config.max_batch_size, - "max_input_len": - build_config.max_input_len, - "max_seq_len": - build_config.max_seq_len, - "use_cache": - build_config.kv_cache_type != KVCacheType.DISABLED, - "max_beam_width": - build_config.max_beam_width, - "max_num_tokens": - build_config.max_num_tokens, - "opt_num_tokens": - build_config.opt_num_tokens, - "prompt_embedding_table_size": - build_config.max_prompt_embedding_table_size, - "max_draft_len": - build_config.max_draft_len, - "speculative_decoding_draft_tokens_external": - build_config.speculative_decoding_mode == - SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL, - "gather_context_logits": - build_config.gather_context_logits, - "lora_target_modules": - build_config.lora_config.lora_target_modules - } - - if model.config.architecture == "DecoderModel" or "mllama" in model.config.architecture.lower( - ): - prepare_input_args["max_seq_len"] = build_config.max_seq_len - prepare_input_args[ - "max_decoder_input_len"] = build_config.max_input_len - prepare_input_args[ - "max_encoder_input_len"] = build_config.max_encoder_input_len - - if model.config.architecture == "WhisperEncoder": - - prepare_input_args = { - "max_batch_size": build_config.max_batch_size, - } - - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.EAGLE: - prepare_input_args[ - "spec_decoding_is_generation_length_variable"] = True - assert build_config.max_batch_size <= 512, "Max batch size > 512 is not supported for EAGLE" - assert build_config.max_draft_len <= 256, "Max draft len > 256 is not supported for EAGLE" - - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.LOOKAHEAD_DECODING: - prepare_input_args[ - "spec_decoding_is_generation_length_variable"] = True - if model.config.architecture == "Qwen2VLForConditionalGeneration" or model.config.architecture == "Qwen2VLModel": - prepare_input_args[ - 'mrope_rotary_cos_sin_size'] = model.config.max_position_embeddings * model.config.rotary_embedding_dim - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.EAGLE and not build_config.plugin_config.use_paged_context_fmha: - logger.warning( - "Paged Context FMHA is required for EAGLE. Turning it on") - build_config.plugin_config.use_paged_context_fmha = True - - inputs = model.prepare_inputs(**prepare_input_args) - model(**inputs) - - if build_config.enable_debug_output: - for k, v in model.named_network_outputs(): - network._mark_output(v, k, str_dtype_to_trt(model.config.dtype)) - - if model.config.architecture != "DecoderModel": - optimize(network) - - if build_config.visualize_network is not None: - with net_guard(network): - network.to_onnx(build_config.visualize_network) - - # Network -> Engine - logger.info( - f"Total time of constructing network from module object {time.time()-tic} seconds" - ) - managed_weights = {} if network.plugin_config.manage_weights else None - engine = None if build_config.dry_run else builder.build_engine( - network, builder_config, managed_weights) - engine_config = EngineConfig(model.config, build_config, __version__) - - if build_config.output_timing_cache is not None and model.config.mapping.rank == 0: - ok = builder.save_timing_cache(builder_config, - build_config.output_timing_cache) - assert ok, "Failed to save timing cache." - - import psutil - - # Get the current process - current_process = psutil.Process() - # Get resource usage for the current process (self) - rusage_s = current_process.memory_info() - # Get resource usage for all child processes - children = current_process.children(recursive=True) - rusage_c = [child.memory_info() for child in children] - logger.info( - f"Build phase peak memory: {rusage_s.rss / 1024 / 1024:.2f} MB, children: {sum([ru.rss for ru in rusage_c]) / 1024 / 1024:.2f} MB" - ) - - return Engine(engine_config, engine, managed_weights) diff --git a/tensorrt_llm/commands/bench.py b/tensorrt_llm/commands/bench.py index 269d4d1d3f2e..63873b5701a1 100644 --- a/tensorrt_llm/commands/bench.py +++ b/tensorrt_llm/commands/bench.py @@ -6,7 +6,6 @@ from tensorrt_llm.bench.benchmark.low_latency import latency_command from tensorrt_llm.bench.benchmark.throughput import throughput_command from tensorrt_llm.bench.benchmark.visual_gen import visual_gen_command -from tensorrt_llm.bench.build.build import build_command from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataset.prepare_dataset import prepare_dataset from tensorrt_llm.logger import logger, severity_map @@ -88,7 +87,6 @@ def main( ctx.obj.workspace.mkdir(parents=True, exist_ok=True) -main.add_command(build_command) main.add_command(throughput_command) main.add_command(latency_command) main.add_command(prepare_dataset) diff --git a/tensorrt_llm/commands/build.py b/tensorrt_llm/commands/build.py deleted file mode 100644 index 0f68382607bb..000000000000 --- a/tensorrt_llm/commands/build.py +++ /dev/null @@ -1,553 +0,0 @@ -# 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 copy -import os -import time -import traceback -from concurrent.futures import ProcessPoolExecutor, as_completed -from importlib.machinery import SourceFileLoader -from multiprocessing import get_context -from typing import Optional, Union - -import torch - -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import (local_mpi_rank, local_mpi_size, mpi_barrier, - mpi_comm, mpi_rank, mpi_world_size) -from tensorrt_llm.builder import BuildConfig, Engine, build -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.logger import logger, severity_map -from tensorrt_llm.lora_helper import LoraConfig -from tensorrt_llm.lora_manager import LoraManager -from tensorrt_llm.models import MODEL_MAP, PretrainedConfig -from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode -from tensorrt_llm.plugin import PluginConfig, add_plugin_argument -from tensorrt_llm.quantization.mode import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument( - '--checkpoint_dir', - type=str, - default=None, - help="The directory path that contains TensorRT LLM checkpoint.") - parser.add_argument( - '--model_config', - type=str, - default=None, - help="The file path that saves TensorRT LLM checkpoint config.") - parser.add_argument( - '--build_config', - type=str, - default=None, - help="The file path that saves TensorRT LLM build config.") - parser.add_argument( - '--model_cls_file', - type=str, - default=None, - help="The file path that defines customized TensorRT LLM model.") - parser.add_argument('--model_cls_name', - type=str, - default=None, - help="The customized TensorRT LLM model class name.") - parser.add_argument( - '--output_dir', - type=str, - default='engine_outputs', - help= - "The directory path to save the serialized engine files and engine config file." - ) - - parser.add_argument( - '--max_batch_size', - type=int, - default=BuildConfig.model_fields["max_batch_size"].default, - help="Maximum number of requests that the engine can schedule.") - parser.add_argument( - '--max_input_len', - type=int, - default=BuildConfig.model_fields["max_input_len"].default, - help="Maximum input length of one request.") - parser.add_argument( - '--max_seq_len', - '--max_decoder_seq_len', - dest='max_seq_len', - type=int, - default=BuildConfig.model_fields["max_seq_len"].default, - help="Maximum total length of one request, including prompt and outputs. " - "If unspecified, the value is deduced from the model config.") - parser.add_argument( - '--max_beam_width', - type=int, - default=BuildConfig.model_fields["max_beam_width"].default, - help="Maximum number of beams for beam search decoding.") - parser.add_argument( - '--max_num_tokens', - type=int, - default=BuildConfig.model_fields["max_num_tokens"].default, - help= - "Maximum number of batched input tokens after padding is removed in each batch. " - "Currently, the input padding is removed by default; " - "you may explicitly disable it by specifying ``--remove_input_padding disable``." - ) - parser.add_argument( - '--opt_num_tokens', - type=int, - default=BuildConfig.model_fields["opt_num_tokens"].default, - help= - "Optimal number of batched input tokens after padding is removed in each batch " - "It equals to ``max_batch_size * max_beam_width`` by default, set this " - "value as close as possible to the actual number of tokens on your workload. " - "Note that this argument might be removed in the future.") - parser.add_argument( - '--max_encoder_input_len', - type=int, - default=BuildConfig.model_fields["max_encoder_input_len"].default, - help="Maximum encoder input length for enc-dec models. " - "Set ``max_input_len`` to 1 to start generation from decoder_start_token_id of length 1." - ) - parser.add_argument( - '--max_prompt_embedding_table_size', - '--max_multimodal_len', - type=int, - default=BuildConfig.model_fields["max_prompt_embedding_table_size"]. - default, - help= - "Maximum prompt embedding table size for prompt tuning, or maximum multimodal input size for multimodal models. " - "Setting a value > 0 enables prompt tuning or multimodal input.") - parser.add_argument( - '--kv_cache_type', - default=argparse.SUPPRESS, - type=KVCacheType, - help= - "Set KV cache type (continuous, paged, or disabled). For disabled case, KV cache is disabled and only context phase is allowed." - ) - parser.add_argument( - '--paged_kv_cache', - type=str, - default=argparse.SUPPRESS, - help= - "Deprecated. Enabling this option is equivalent to ``--kv_cache_type paged`` for transformer based models." - ) - - parser.add_argument( - '--input_timing_cache', - type=str, - default=BuildConfig.model_fields["input_timing_cache"].default, - help= - "The file path to read the timing cache. This option is ignored if the file does not exist." - ) - parser.add_argument( - '--output_timing_cache', - type=str, - default=BuildConfig.model_fields["output_timing_cache"].default, - help="The file path to write the timing cache.") - parser.add_argument( - '--profiling_verbosity', - type=str, - default=BuildConfig.model_fields["profiling_verbosity"].default, - choices=['layer_names_only', 'detailed', 'none'], - help= - "The profiling verbosity for the generated TensorRT engine. Setting to detailed allows inspecting tactic choices and kernel parameters." - ) - parser.add_argument( - '--strip_plan', - default=BuildConfig.model_fields["use_strip_plan"].default, - action='store_true', - help= - "Enable stripping weights from the final TensorRT engine under the assumption that the refit weights are identical to those provided at build time." - ) - parser.add_argument( - '--weight_sparsity', - default=BuildConfig.model_fields["weight_sparsity"].default, - action='store_true', - help="Enable weight sparsity.") - parser.add_argument( - '--weight_streaming', - default=BuildConfig.model_fields["weight_streaming"].default, - action='store_true', - help= - "Enable offloading weights to CPU and streaming loading at runtime.", - ) - parser.add_argument( - '--fast_build', - default=False, - action='store_true', - help= - "Enable features for faster engine building. This may cause some performance degradation and is currently incompatible with int8/int4 quantization without plugin.", - ) - - parser.add_argument('--workers', - type=int, - default=1, - help="The number of workers for building in parallel.") - parser.add_argument('--log_level', - type=str, - default='info', - choices=severity_map.keys(), - help="The logging level.") - parser.add_argument( - '--enable_debug_output', - default=BuildConfig.model_fields["enable_debug_output"].default, - action='store_true', - help="Enable debug output.") - parser.add_argument( - '--visualize_network', - type=str, - default=None, - help= - "The directory path to export TensorRT Network as ONNX prior to Engine build for debugging." - ) - parser.add_argument( - '--dry_run', - default=BuildConfig.model_fields["dry_run"].default, - action='store_true', - help= - "Run through the build process except the actual Engine build for debugging." - ) - parser.add_argument('--monitor_memory', - default=False, - action='store_true', - help="Enable memory monitor during Engine build.") - - logits_parser = parser.add_argument_group("Logits arguments") - logits_parser.add_argument('--logits_dtype', - type=str, - default=None, - choices=['float16', 'float32'], - help="The data type of logits.") - logits_parser.add_argument('--gather_context_logits', - action='store_true', - default=False, - help="Enable gathering context logits.") - logits_parser.add_argument('--gather_generation_logits', - action='store_true', - default=False, - help="Enable gathering generation logits.") - logits_parser.add_argument( - '--gather_all_token_logits', - action='store_true', - default=False, - help= - "Enable both ``gather_context_logits`` and ``gather_generation_logits``." - ) - - lora_parser = parser.add_argument_group("LoRA arguments") - lora_parser.add_argument( - '--lora_dir', - type=str, - default=None, - nargs="+", - help="The directory of LoRA weights. " - "If multiple directories are provided, the first one is used for configuration." - ) - lora_parser.add_argument('--lora_ckpt_source', - type=str, - default="hf", - choices=["hf", "nemo"], - help="The source type of LoRA checkpoint.") - lora_parser.add_argument( - '--lora_target_modules', - nargs='+', - default=None, - choices=LoraManager.LORA_MODULE_IDS.keys(), - help= - "The target module names that LoRA is applied. Only effective when ``lora_plugin`` is enabled." - ) - lora_parser.add_argument( - '--max_lora_rank', - type=int, - default=64, - help="Maximum LoRA rank for different LoRA modules. " - "It is used to compute the workspace size of LoRA plugin.") - - spec_parser = parser.add_argument_group("Speculative decoding arguments") - spec_parser.add_argument('--speculative_decoding_mode', - default=None, - choices=[ - "draft_tokens_external", "lookahead_decoding", - "medusa", "explicit_draft_tokens", "eagle" - ], - help="Mode of speculative decoding.") - spec_parser.add_argument( - '--max_draft_len', - type=int, - default=0, - help= - "Maximum lengths of draft tokens for speculative decoding target model." - ) - - plugin_config_parser = parser.add_argument_group("Plugin config arguments") - add_plugin_argument(plugin_config_parser) - return parser - - -def build_model( - build_config: BuildConfig, - rank: int = 0, - ckpt_dir: str = None, - model_config: Union[str, PretrainedConfig] = None, - model_cls=None, - dry_run: - bool = False, # return the modified BuildConfig without actually building the engine - **kwargs -) -> Union[Engine, BuildConfig]: - model_config = copy.deepcopy(model_config) - - logits_dtype = kwargs.get('logits_dtype') - if logits_dtype is not None: - model_config.logits_dtype = logits_dtype - - architecture = model_config.architecture - assert not build_config.plugin_config.streamingllm, \ - "StreamingLLM is no longer supported because attention sink cannot work with the non-cyclic kv cache kernel & runtime changes." - assert not build_config.plugin_config.pp_reduce_scatter or architecture == "MixtralForCausalLM", \ - "PP reduce scatter is only supported in the mixtral model." - - assert rank < model_config.mapping.world_size - - rank_config = copy.deepcopy(model_config) - rank_config.set_rank(rank) - - if model_cls is None: - assert architecture in MODEL_MAP, \ - f"Unsupported model architecture: {architecture}" - model_cls = MODEL_MAP[architecture] - if ckpt_dir is None: - model = model_cls(rank_config) - else: - model = model_cls.from_checkpoint(ckpt_dir, config=rank_config) - is_checkpoint_pruned = getattr(rank_config, 'is_pruned', False) - - if build_config.plugin_config.lora_plugin is not None: - lora_config = LoraConfig(lora_dir=kwargs['lora_dir'] or [], - lora_ckpt_source=kwargs['lora_ckpt_source'], - max_lora_rank=kwargs['max_lora_rank']) - if kwargs['lora_target_modules'] is not None: - # command line options is preferred over the modules in the lora dir - lora_config.lora_target_modules = kwargs['lora_target_modules'] - build_config.lora_config = lora_config - - if is_checkpoint_pruned or kwargs.pop('strip_plan', False): - build_config.use_strip_plan = True - build_config.use_refit = kwargs.get('refit', False) - - return build(model, build_config) - - -def build_and_save(rank, gpu_id, ckpt_dir, build_config, output_dir, log_level, - model_config, model_cls, **kwargs): - torch.cuda.set_device(gpu_id) - logger.set_level(log_level) - engine = build_model(build_config, - rank, - ckpt_dir, - model_config, - model_cls=model_cls, - **kwargs) - assert engine is not None - engine.save(output_dir) - return True - - -def parallel_build(model_config: PretrainedConfig, - ckpt_dir: Optional[str], - build_config: BuildConfig, - output_dir: str, - workers: int = 1, - log_level: str = 'info', - model_cls=None, - **kwargs): - - world_size = model_config.mapping.world_size - use_mpi = mpi_world_size() > 1 - - if not use_mpi and workers == 1: - for rank in range(world_size): - passed = build_and_save(rank, rank % workers, ckpt_dir, - build_config, output_dir, log_level, - model_config, model_cls, **kwargs) - assert passed, "Engine building failed, please check error log." - elif not use_mpi: - with ProcessPoolExecutor(mp_context=get_context('spawn'), - max_workers=workers) as p: - futures = [ - p.submit(build_and_save, rank, rank % workers, ckpt_dir, - build_config, output_dir, log_level, model_config, - model_cls, **kwargs) 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, "Engine building failed, please check error log." - else: - mpi_local_rank = local_mpi_rank() - node_gpu_count = local_mpi_size() - exceptions = [] - for engine_rank in range(world_size): - if engine_rank % mpi_world_size() != mpi_rank(): - continue - try: - build_and_save(engine_rank, mpi_local_rank % node_gpu_count, - ckpt_dir, build_config, output_dir, log_level, - model_config, model_cls, **kwargs) - except Exception as e: - traceback.print_exc() - exceptions.append(e) - mpi_barrier() - if len(exceptions) != 0: - print("Engine building failed, please check error log.", flush=True) - mpi_comm().Abort() - - -def main(): - emit_engine_arch_deprecation("trtllm-build") - - parser = parse_arguments() - args = parser.parse_args() - - if hasattr(args, 'gather_generation_logits'): - logger.warning( - 'Option --gather_generation_logits is deprecated, a build flag is not required anymore. Use --output_generation_logits at runtime instead.' - ) - - if args.gather_all_token_logits: - args.gather_context_logits = True - args.gather_generation_logits = True - if args.gather_context_logits and args.max_draft_len > 0: - raise RuntimeError( - "Gather context logits is not support with draft len > 0. " - "If want to get the accepted tokens' logits from target model, please just enable gather_generation_logits" - ) - - logger.set_level(args.log_level) - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir, exist_ok=True) - - model_cls = None - if args.model_cls_file is not None: - assert args.model_cls_name is not None - loader = SourceFileLoader('models', args.model_cls_file) - mod = loader.load_module() - model_cls = getattr(mod, args.model_cls_name) - - workers = min(torch.cuda.device_count(), args.workers) - - if hasattr(args, 'paged_kv_cache'): - logger.warning( - 'Option --paged_kv_cache is deprecated, use --kv_cache_type=paged/disabled instead.' - ) - - plugin_config = PluginConfig.from_arguments(args) - if args.fast_build: - plugin_config.manage_weights = True - - kwargs = { - 'logits_dtype': args.logits_dtype, - 'use_fused_mlp': args.use_fused_mlp, - 'lora_dir': args.lora_dir, - 'lora_ckpt_source': args.lora_ckpt_source, - 'max_lora_rank': args.max_lora_rank, - 'lora_target_modules': args.lora_target_modules, - 'strip_plan': args.strip_plan, - 'refit': False, - } - speculative_decoding_mode = SpeculativeDecodingMode.from_arguments(args) - - ckpt_dir_or_model_config = args.checkpoint_dir if args.checkpoint_dir is not None else args.model_config - if ckpt_dir_or_model_config.lower().endswith('.json'): - config_path = ckpt_dir_or_model_config - ckpt_dir = None - else: - config_path = os.path.join(ckpt_dir_or_model_config, 'config.json') - ckpt_dir = ckpt_dir_or_model_config - model_config = PretrainedConfig.from_json_file(config_path) - - # avoid ValueError if not supported quantization is chosen with use_fused_mlp - quant_algo = model_config.quantization.quant_algo - if quant_algo and quant_algo not in (QuantAlgo.FP8, - QuantAlgo.MIXED_PRECISION): - kwargs['use_fused_mlp'] = False - - if args.build_config is None: - if args.multiple_profiles == "enable" and args.opt_num_tokens is not None: - raise RuntimeError( - "multiple_profiles is enabled, while opt_num_tokens is set. " - "They are not supposed to be working in the same time for now.") - - # This should only be used for debugging. - # The env var BUILDER_FORCE_NUM_PROFILES should override the number of - # optimization profiles during TRT build. - # BUILDER_FORCE_NUM_PROFILES must be less than or equal to the number of - # optimization profiles set by model's prepare_inputs(). - force_num_profiles_from_env = os.environ.get( - "BUILDER_FORCE_NUM_PROFILES", None) - if force_num_profiles_from_env is not None: - logger.warning( - f"Overriding # of builder profiles <= {force_num_profiles_from_env}." - ) - - build_config = BuildConfig( - max_input_len=args.max_input_len, - max_seq_len=args.max_seq_len, - max_batch_size=args.max_batch_size, - max_beam_width=args.max_beam_width, - max_num_tokens=args.max_num_tokens, - opt_num_tokens=args.opt_num_tokens, - max_prompt_embedding_table_size=args. - max_prompt_embedding_table_size, - kv_cache_type=getattr(args, "kv_cache_type", None), - gather_context_logits=args.gather_context_logits, - gather_generation_logits=args.gather_generation_logits, - strongly_typed=True, - force_num_profiles=force_num_profiles_from_env, - weight_sparsity=args.weight_sparsity, - profiling_verbosity=args.profiling_verbosity, - enable_debug_output=args.enable_debug_output, - max_draft_len=args.max_draft_len, - speculative_decoding_mode=speculative_decoding_mode, - input_timing_cache=args.input_timing_cache, - output_timing_cache=args.output_timing_cache, - dry_run=args.dry_run, - visualize_network=args.visualize_network, - max_encoder_input_len=args.max_encoder_input_len, - weight_streaming=args.weight_streaming, - monitor_memory=args.monitor_memory, - use_mrope=getattr(model_config, "qwen_type", None) == "qwen2_vl", - plugin_config=plugin_config) - else: - build_config = BuildConfig.from_json_file(args.build_config) - build_config.plugin_config = plugin_config - - parallel_build(model_config, ckpt_dir, build_config, args.output_dir, - workers, args.log_level, model_cls, **kwargs) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of building all engines: {t}') - - -if __name__ == '__main__': - main() diff --git a/tensorrt_llm/commands/eval.py b/tensorrt_llm/commands/eval.py index 2f0e4b9cd6b0..0e25bb80736f 100644 --- a/tensorrt_llm/commands/eval.py +++ b/tensorrt_llm/commands/eval.py @@ -19,18 +19,22 @@ import tensorrt_llm.profiler as profiler from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..evaluate import (AALCR, AIME2025, AIME2026, GSM8K, HLE, MMLU, MMMU, ArenaHard, CnnDailymail, CoVoST2, GPQADiamond, GPQAExtended, GPQAMain, GPQANemoSkills, IFBench, JsonModeEval, LongBenchV1, LongBenchV2, MMMUPro, SciCode) -from ..llmapi import BuildConfig, KvCacheConfig +from ..llmapi import KvCacheConfig +from ..llmapi.llm_args import TorchLlmArgs from ..llmapi.llm_utils import update_llm_args_with_extra_options from ..logger import logger, severity_map from ..usage import config as _telemetry_config from .utils import collect_explicit_cli_keys +# CLI defaults are sourced from the TorchLlmArgs field defaults so they stay in +# lock-step with the args class and can't drift. +_LLM_ARGS_FIELDS = TorchLlmArgs.model_fields + # Map Click parameter names to the LlmArgs field name (or merge-function CLI # scalar name) used by `update_llm_args_with_extra_options`. _CLICK_TO_LLM_ARG = { @@ -64,7 +68,7 @@ ) @click.option( "--backend", - type=click.Choice(["pytorch", "tensorrt"]), + type=click.Choice(["pytorch"]), default="pytorch", help="The backend to use for evaluation. Default is pytorch backend.") @click.option('--log_level', @@ -73,23 +77,23 @@ help="The logging level.") @click.option("--max_beam_width", type=int, - default=BuildConfig.model_fields["max_beam_width"].default, + default=_LLM_ARGS_FIELDS["max_beam_width"].default, help="Maximum number of beams for beam search decoding.") @click.option("--max_batch_size", type=int, - default=BuildConfig.model_fields["max_batch_size"].default, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, help="Maximum number of requests that the engine can schedule.") @click.option( "--max_num_tokens", type=int, - default=BuildConfig.model_fields["max_num_tokens"].default, + default=_LLM_ARGS_FIELDS["max_num_tokens"].default, help= "Maximum number of batched input tokens after padding is removed in each batch." ) @click.option( "--max_seq_len", type=int, - default=BuildConfig.model_fields["max_seq_len"].default, + default=None, help="Maximum total length of one request, including prompt and outputs. " "If unspecified, the value is deduced from the model config.") @click.option("--tp_size", type=int, default=1, help='Tensor parallelism size.') @@ -187,13 +191,6 @@ def main(ctx, model: str, tokenizer: Optional[str], max_num_tokens=max_num_tokens, max_beam_width=max_beam_width, max_seq_len=max_seq_len) - elif backend == 'tensorrt': - llm_cls = LLM - build_config = BuildConfig(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) - llm_args.update(build_config=build_config) else: raise click.BadParameter( f"{backend} is not a known backend, check help for available options.", diff --git a/tensorrt_llm/commands/prune.py b/tensorrt_llm/commands/prune.py deleted file mode 100644 index 5e7994e9b4da..000000000000 --- a/tensorrt_llm/commands/prune.py +++ /dev/null @@ -1,122 +0,0 @@ -# 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. -''' -Script that prunes TRT-LLM checkpoints. -''' -import argparse -import json -import os -from pathlib import Path -from typing import Dict - -import safetensors -import torch -from safetensors.torch import save_file - -from tensorrt_llm.logger import logger -from tensorrt_llm.models import MODEL_MAP, PretrainedConfig - -SUPPORTED_MODELS = list(MODEL_MAP.keys()) -PRUNABLE_WEIGHTS = [ - 'attention.qkv.weight', - 'attention.proj.weight', - 'mlp.fc.weight', - 'mlp.proj.weight', - 'mlp.gate.weight', -] - - -def can_prune(key: str) -> bool: - for w in PRUNABLE_WEIGHTS: - if w in key: - return True - return False - - -def load_config(config_path: Path) -> Dict[str, any]: - if not config_path.exists(): - return {} - - with open(str(config_path), 'r') as f: - return json.load(f) - - -def prune_and_save(ckpt_dir: str, out_dir: str, prune_all: bool): - logger.info(f'Checkpoint Dir: {ckpt_dir}, Out Dir: {out_dir}') - model_config = PretrainedConfig.from_json_file( - os.path.join(ckpt_dir, 'config.json')) - - architecture = model_config.architecture - if architecture not in MODEL_MAP: - raise RuntimeError(f'Unsupported model architecture: {architecture}') - - if not os.path.exists(out_dir): - os.makedirs(out_dir) - - for rank in range(model_config.mapping.world_size): - pruned_weights = {} - with safetensors.safe_open(os.path.join(ckpt_dir, - f'rank{rank}.safetensors'), - framework='pt', - device='cpu') as f: - for key in f.keys(): - tensor = f.get_tensor(key) - if prune_all or can_prune(key): - pruned_weights[key] = torch.tensor([], dtype=tensor.dtype) - else: - pruned_weights[key] = tensor - - save_file(pruned_weights, - os.path.join(out_dir, f'rank{rank}.safetensors')) - - config_path = Path(ckpt_dir, 'config.json') - with open(str(Path(out_dir, 'config.json')), 'w') as f: - config = load_config(config_path) - config['is_pruned'] = True - json.dump(config, f) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--checkpoint_dir', type=str, default=None) - parser.add_argument('--prune_all', - default=False, - action='store_true', - help='Remove all weights in the checkpoint') - parser.add_argument( - '--out_dir', - type=str, - default=None, - help= - 'Path to write pruned checkpoint. Defaults to the same directory append with `.pruned`' - ) - args = parser.parse_args() - - if args.checkpoint_dir is None: - raise RuntimeError( - "No `--checkpoint_dir` supplied to checkpoint pruner.") - - if args.out_dir is None: - ckpt_path = Path(args.checkpoint_dir) - ckpt_name = ckpt_path.name - args.out_dir = str( - Path(args.checkpoint_dir).with_name(ckpt_name + '.pruned')) - - prune_and_save(os.path.abspath(args.checkpoint_dir), - os.path.abspath(args.out_dir), args.prune_all) - - -if __name__ == '__main__': - main() diff --git a/tensorrt_llm/commands/refit.py b/tensorrt_llm/commands/refit.py deleted file mode 100644 index b77296758f61..000000000000 --- a/tensorrt_llm/commands/refit.py +++ /dev/null @@ -1,183 +0,0 @@ -''' -Script that refits TRT-LLM engine(s) with weights in a TRT-LLM checkpoint. -''' -import argparse -import json -import os -import re -import shutil -import time -from pathlib import Path - -import tensorrt as trt - -from tensorrt_llm._common import _is_building -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import np_dtype_to_trt -from tensorrt_llm.builder import EngineConfig, optimize_model_with_config -from tensorrt_llm.models import MODEL_MAP, PretrainedConfig - -from ..logger import logger - -ENGINE_RE = re.compile(r'rank(\d+).engine') - - -@_is_building -def refit_engine(engine_path: str, refit_engine_dir: str, checkpoint_dir: str, - engine_config: EngineConfig, fixed_weights_names: list): - # This function loops through all weights in the model and does a textual match between - # checkpoint weight names and engine weight names. - rank = int(ENGINE_RE.fullmatch(os.path.basename(engine_path)).group(1)) - tik = time.time() - with open(engine_path, - "rb") as f, trt.Runtime(logger.trt_logger) as runtime: - engine = runtime.deserialize_cuda_engine(f.read()) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Load TRT engine time: {t}') - - refitter = trt.Refitter(engine, logger.trt_logger) - refittable_weights = set(refitter.get_all_weights()) - - # Load model. - tik = time.time() - rank_config = PretrainedConfig.from_dict( - engine_config.pretrained_config.to_dict()) - rank_config.set_rank(rank) - - architecture = rank_config.architecture - assert architecture in MODEL_MAP, \ - f"Unsupported model architecture: {architecture}" - model_cls = MODEL_MAP[architecture] - model = model_cls.from_checkpoint(checkpoint_dir, config=rank_config) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Load checkpoint(s) time: {t}') - - # There are weights preprocess during optimize model. - tik = time.time() - build_config = engine_config.build_config.model_copy(deep=True) - optimize_model_with_config(model, build_config) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Preprocess weights time: {t}') - - # Refit engine. - tik = time.time() - refitted_weights = [] - for name, buf in model.named_parameters(): - if name in refittable_weights: - assert buf.is_inited, f"Failed because weight `{name}` is not initialized in model." - weight = buf._value - weights_value = trt.Weights(np_dtype_to_trt(weight.dtype), - weight.ctypes.data, weight.size) - assert refitter.set_named_weights( - name, weights_value), f'Failed to refit weight: `{name}`' - refitted_weights.append(name) - else: - if name not in fixed_weights_names: - logger.warning( - f"model weights `{name}` (shape: {buf._value.shape}) is not refittable, this means that we might not be able to update the engine using fine-tuned checkpoint!" - ) - - # Validate all refittable weights are provided. - if len(refitted_weights) != len(refittable_weights): - raise RuntimeError( - f'Missing refittable weights {refittable_weights.difference(refitted_weights)} from {checkpoint_dir}' - ) - - assert refitter.refit_cuda_engine(), f'Failed to refit engine.' - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Execute GPU refit graph time: {t}') - - tik = time.time() - refit_engine_path = os.path.join(refit_engine_dir, - os.path.basename(engine_path)) - with open(refit_engine_path, 'wb') as f: - logger.info(f'\nWriting refitted engine to `{refit_engine_path}`') - s_config = engine.create_serialization_config() - s_config.flags &= ~(1 << int(trt.SerializationFlag.EXCLUDE_WEIGHTS)) - f.write(engine.serialize_with_config(s_config)) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Write TRT engine to disk time: {t}') - - del refitter - - -def refit(engine_dir: str, checkpoint_dir: str, engine_config: EngineConfig, - output_dir: str, fixed_weights_names: list): - refit_engine_dir = output_dir - os.makedirs(refit_engine_dir, exist_ok=True) - shutil.copyfile(os.path.join(engine_dir, 'config.json'), - os.path.join(refit_engine_dir, 'config.json')) - engine_paths = list(Path(engine_dir).glob('*.engine')) - for path in engine_paths: - refit_engine(path, refit_engine_dir, checkpoint_dir, engine_config, - fixed_weights_names) - - -def main(): - emit_engine_arch_deprecation("trtllm-refit") - - parser = argparse.ArgumentParser() - parser.add_argument( - '--engine_dir', - type=str, - default=None, - help= - 'Path to trt-llm engines. These engines must have been built from a pruned checkpoint, or otherwise be refittable.' - ) - parser.add_argument('--checkpoint_dir', - type=str, - default=None, - help='Path to checkpoint containing desired weights') - parser.add_argument('--output_dir', - type=str, - default=None, - help="Output path of the refit model") - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - logger.set_level(args.log_level) - if args.engine_dir is None or not Path(args.engine_dir).exists(): - raise RuntimeError( - f'Please supply a valid --engine_dir (found `{args.engine_dir}`)') - if args.checkpoint_dir is None or not Path(args.checkpoint_dir).exists(): - raise RuntimeError( - f'Please supply a valid --checkpoint_dir (found `{args.checkpoint_dir}`)' - ) - - engine_config = EngineConfig.from_json_file( - os.path.join(args.engine_dir, 'config.json')) - - with open(os.path.join(args.checkpoint_dir, 'config.json'), 'r') as f: - checkpoint_config = json.load(f) - - engine_arch = engine_config.pretrained_config.architecture - checkpoint_arch = checkpoint_config['architecture'] - if engine_arch != checkpoint_arch: - raise RuntimeError( - f'Engine Architecture and Checkpoint Architecture do not match. ' + - f'Engine Architecture: `{engine_arch}`, Checkpoint Architecture: `{checkpoint_arch}`' - ) - - # The fixed weights are not read from checkpoint, they are hardcoded buffer from the model itself. These values remain constant across different fine-tuned checkpoints. - fixed_wts_in_model = [] - model_cls = MODEL_MAP[engine_arch] - model = model_cls.from_config(engine_config.pretrained_config) - for name, param in model.named_parameters(): - if param.is_inited(): - fixed_wts_in_model.append(name) - - refit(engine_dir=os.path.normpath(args.engine_dir), - checkpoint_dir=os.path.normpath(args.checkpoint_dir), - engine_config=engine_config, - output_dir=os.path.normpath(args.output_dir), - fixed_weights_names=fixed_wts_in_model) - - -if __name__ == '__main__': - main() diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index ab7dffce6e3d..079e668917d4 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -11,7 +11,7 @@ import sys import uuid from pathlib import Path -from typing import Any, Dict, Mapping, Optional, Sequence, Set +from typing import Any, Dict, Optional, Sequence, Set import click import torch @@ -22,24 +22,20 @@ from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm import MultimodalEncoder -from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._utils import mpi_rank from tensorrt_llm.commands._serve_stability import stability_option from tensorrt_llm.commands.utils import (collect_explicit_cli_keys, get_is_diffusion_only_model) from tensorrt_llm.executor.utils import LlmLauncherEnvs from tensorrt_llm.inputs.multimodal import MultimodalServerConfig -from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy, - DynamicBatchConfig, KvCacheConfig, - SchedulerConfig) +from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole, extract_disagg_cluster_config, parse_disagg_config_file, parse_metadata_server_config_file, validate_config_bool) -from tensorrt_llm.llmapi.llm_args import (MultimodalConfig, TorchLlmArgs, - TrtLlmArgs) +from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -157,9 +153,7 @@ def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): for s in cli_derived_fields.get(param_name, ())): return True - if backend == "tensorrt": - llm_args_class = TrtLlmArgs - elif backend == "_autodeploy": + if backend == "_autodeploy": from tensorrt_llm._torch.auto_deploy.llm_args import \ LlmArgs as AutoDeployLlmArgs llm_args_class = AutoDeployLlmArgs @@ -179,19 +173,21 @@ def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): return value != default +# CLI/API defaults are sourced from the TorchLlmArgs field defaults so they stay +# in lock-step with the args class and can't drift. +_LLM_ARGS_FIELDS = TorchLlmArgs.model_fields + + def get_llm_args( model: str, tokenizer: Optional[str] = None, custom_tokenizer: Optional[str] = None, post_processor_hook: Optional[str] = None, backend: str = "pytorch", - max_beam_width: int = BuildConfig.model_fields["max_beam_width"]. - default, - max_batch_size: int = BuildConfig.model_fields["max_batch_size"]. - default, - max_num_tokens: int = BuildConfig.model_fields["max_num_tokens"]. - default, - max_seq_len: int = BuildConfig.model_fields["max_seq_len"].default, + max_beam_width: int = _LLM_ARGS_FIELDS["max_beam_width"].default, + max_batch_size: int = _LLM_ARGS_FIELDS["max_batch_size"].default, + max_num_tokens: int = _LLM_ARGS_FIELDS["max_num_tokens"].default, + max_seq_len: int = None, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, context_parallel_size: int = 1, @@ -250,19 +246,8 @@ def get_llm_args( dtype=kv_cache_dtype), "cp_config": cp_config, - "build_config": - BuildConfig(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) if backend == "tensorrt" else None, "scheduler_config": - SchedulerConfig(capacity_scheduler_policy=CapacitySchedulerPolicy. - GUARANTEED_NO_EVICT, - dynamic_batch_config=DynamicBatchConfig( - enable_batch_size_tuning=True, - enable_max_num_tokens_tuning=False, - dynamic_batch_moving_average_window=128)) - if backend == "tensorrt" else None, + None, "max_batch_size": max_batch_size, "max_beam_width": @@ -419,9 +404,6 @@ def launch_server( # AutoDeploy does not support build_config llm_args.pop("build_config", None) llm = AutoDeployLLM(**llm_args) - elif backend == 'tensorrt' or backend == 'trt': - llm_args.pop("backend") - llm = LLM(**llm_args) else: raise click.BadParameter( f"{backend} is not a known backend, check help for available options.", @@ -489,9 +471,6 @@ async def serve_grpc_async(): from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM llm_args.pop("build_config", None) llm = AutoDeployLLM(**llm_args) - elif backend == "tensorrt" or backend == "trt": - llm_args.pop("backend") - llm = LLM(**llm_args) else: raise click.BadParameter( f"{backend} is not a known backend, check help for available options.", @@ -653,27 +632,6 @@ def launch_visual_gen_server( uvloop.run(server(host, port, sockets=[s])) -class ChoiceWithAlias(click.Choice): - - def __init__(self, - choices: Sequence[str], - aliases: Mapping[str, str], - case_sensitive: bool = True) -> None: - super().__init__(choices, case_sensitive) - self.aliases = aliases - - def to_info_dict(self) -> Dict[str, Any]: - info_dict = super().to_info_dict() - info_dict["aliases"] = self.aliases - return info_dict - - def convert(self, value: Any, param: Optional["click.Parameter"], - ctx: Optional["click.Context"]) -> Any: - if value in self.aliases: - value = self.aliases[value] - return super().convert(value, param, ctx) - - @click.command("serve") @click.argument("model", type=str) @stability_option( @@ -715,8 +673,7 @@ def convert(self, value: Any, param: Optional["click.Parameter"], status="beta") @stability_option( "--backend", - type=ChoiceWithAlias(["pytorch", "tensorrt", "_autodeploy"], - {"trt": "tensorrt"}), + type=click.Choice(["pytorch", "_autodeploy"]), default="pytorch", help="The backend to use to serve the model. Default is pytorch backend.", status="beta") @@ -736,26 +693,26 @@ def convert(self, value: Any, param: Optional["click.Parameter"], status="beta") @stability_option("--max_beam_width", type=int, - default=BuildConfig.model_fields["max_beam_width"].default, + default=_LLM_ARGS_FIELDS["max_beam_width"].default, help="Maximum number of beams for beam search decoding.", status="beta") @stability_option( "--max_batch_size", type=int, - default=BuildConfig.model_fields["max_batch_size"].default, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, help="Maximum number of requests that the engine can schedule.", status="beta") @stability_option( "--max_num_tokens", type=int, - default=BuildConfig.model_fields["max_num_tokens"].default, + default=_LLM_ARGS_FIELDS["max_num_tokens"].default, help= "Maximum number of batched input tokens after padding is removed in each batch.", status="beta") @stability_option( "--max_seq_len", type=int, - default=BuildConfig.model_fields["max_seq_len"].default, + default=None, help="Maximum total length of one request, including prompt and outputs. " "If unspecified, the value is deduced from the model config.", status="beta") @@ -1254,7 +1211,7 @@ def _serve_visual_gen(): @stability_option( "--max_batch_size", type=int, - default=BuildConfig.model_fields["max_batch_size"].default, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, help="Maximum number of requests that the engine can schedule.", status="beta") @stability_option( diff --git a/tensorrt_llm/disaggregated_params.py b/tensorrt_llm/disaggregated_params.py index 93b9a38f2215..5c91dae3f5de 100644 --- a/tensorrt_llm/disaggregated_params.py +++ b/tensorrt_llm/disaggregated_params.py @@ -4,11 +4,6 @@ import numpy as np -# isort: off -# needed before trying to import bindings to load tensorrt_libs -import tensorrt as trt # noqa -# isort: on - from tensorrt_llm.bindings import executor as tllme diff --git a/tensorrt_llm/evaluate/cnn_dailymail.py b/tensorrt_llm/evaluate/cnn_dailymail.py index f46cf65cc4c3..a676c0f5f69a 100644 --- a/tensorrt_llm/evaluate/cnn_dailymail.py +++ b/tensorrt_llm/evaluate/cnn_dailymail.py @@ -12,14 +12,13 @@ # 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 Iterable, List, Optional, Union +from typing import Iterable, List, Optional import click import datasets import evaluate from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -124,7 +123,7 @@ def command(ctx, dataset_path: Optional[str], num_samples: int, apply_chat_template: bool, system_prompt: Optional[str], max_input_length: int, max_output_length: int, output_dir: Optional[str]) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length) diff --git a/tensorrt_llm/evaluate/covost2.py b/tensorrt_llm/evaluate/covost2.py index 74e5dea4cd22..0e2051b9464a 100644 --- a/tensorrt_llm/evaluate/covost2.py +++ b/tensorrt_llm/evaluate/covost2.py @@ -22,14 +22,13 @@ task is added here for apples-to-apples comparison. """ -from typing import Iterable, List, Optional, Tuple, Union +from typing import Iterable, List, Optional, Tuple import click import datasets import numpy as np from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -350,7 +349,7 @@ def command( output_dir: Optional[str], dump_samples_path: Optional[str], ) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length, diff --git a/tensorrt_llm/evaluate/json_mode_eval.py b/tensorrt_llm/evaluate/json_mode_eval.py index 8e6b0b814e50..ce35d474b333 100644 --- a/tensorrt_llm/evaluate/json_mode_eval.py +++ b/tensorrt_llm/evaluate/json_mode_eval.py @@ -14,7 +14,7 @@ # limitations under the License. import json import os -from typing import Iterable, List, Optional, Union +from typing import Iterable, List, Optional import click import datasets @@ -22,7 +22,6 @@ import numpy as np from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import GuidedDecodingParams, SamplingParams @@ -132,7 +131,7 @@ def command(ctx, dataset_path: Optional[str], num_samples: int, random_seed: int, system_prompt: Optional[str], max_input_length: int, max_output_length: int, output_dir: Optional[str]) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index d41ede2e4180..b077747b6c03 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -17,7 +17,7 @@ import os from contextlib import contextmanager from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple import click import numpy as np @@ -33,7 +33,6 @@ TemplateLM = object from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..inputs import (ConversationMessage, MultimodalDataTracker, add_multimodal_placeholders, convert_image_mode) from ..inputs.content_format import ContentFormat @@ -55,7 +54,7 @@ class LmEvalWrapper(TemplateLM): def __init__(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False, chat_template_kwargs: Optional[dict[str, Any]] = None, @@ -199,7 +198,7 @@ class MultimodalLmEvalWrapper(LmEvalWrapper): """ def __init__(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False, max_images: int = 999, @@ -263,7 +262,7 @@ def __init__(self, self.interleave = MULTIMODAL_PLACEHOLDER_REGISTRY.get_interleave_placeholders( self.model_type) - def _get_model_type(self, llm: Union[LLM, PyTorchLLM]) -> str: + def _get_model_type(self, llm: PyTorchLLM) -> str: """Extract model type from the model configuration.""" config_path = os.path.join(llm._hf_model_dir, 'config.json') @@ -602,7 +601,7 @@ def save_results(self, results: dict) -> None: logger.info(f"Results saved to {result_path}") def evaluate(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False, scores_filter: str = None, @@ -663,7 +662,7 @@ def evaluate(self, @classmethod def command_harness(cls, ctx, **kwargs): - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj # Resolve the post-processor: accept a callable (already-bound) or the # string key "strip_thinking_mmmu" coming from CLI flags. @@ -1336,7 +1335,7 @@ def _get_group_score(metrics: Dict[str, Any], return None def evaluate(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False) -> float: import lm_eval @@ -1445,7 +1444,7 @@ def evaluate(self, @click.pass_context @staticmethod def command(ctx, **kwargs) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj evaluator = LongBenchV1( dataset_path=kwargs.pop("dataset_path", None), diff --git a/tensorrt_llm/evaluate/longbench_v2.py b/tensorrt_llm/evaluate/longbench_v2.py index 1b2ccaf221ba..40a7d236efa3 100644 --- a/tensorrt_llm/evaluate/longbench_v2.py +++ b/tensorrt_llm/evaluate/longbench_v2.py @@ -24,13 +24,12 @@ import os import re from datetime import datetime -from typing import Any, Dict, Iterable, List, Optional, Union +from typing import Any, Dict, Iterable, List, Optional import click from datasets import load_dataset from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -779,7 +778,7 @@ def command(ctx, dataset_path: str, prompts_dir: Optional[str], max_len: int, max_input_length: int, max_output_length: int, chat_template_kwargs: Optional[dict[str, Any]], temperature: float, top_p: float) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams(max_tokens=max_output_length, temperature=temperature, diff --git a/tensorrt_llm/evaluate/mmlu.py b/tensorrt_llm/evaluate/mmlu.py index 0b26e52a5af9..8463a639bd94 100644 --- a/tensorrt_llm/evaluate/mmlu.py +++ b/tensorrt_llm/evaluate/mmlu.py @@ -17,14 +17,13 @@ import json import math -from typing import Any, Iterable, List, Optional, Union +from typing import Any, Iterable, List, Optional import click import numpy as np import pandas as pd from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -312,7 +311,7 @@ def command(ctx, dataset_path: Optional[str], num_samples: int, system_prompt: Optional[str], max_input_length: int, max_output_length: int, check_accuracy: bool, accuracy_threshold: float, output_dir: Optional[str]) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length) diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index c4cb84d6bbf8..670d6f30508f 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -33,7 +33,6 @@ from .._utils import (global_mpi_rank, global_mpi_size, mpi_comm, mpi_rank, nvtx_range_debug) from ..bindings import executor as tllm -from ..builder import ConfigEncoder, Engine, EngineConfig from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType, PybindMirror from ..llmapi.tokenizer import TokenizerBase from ..llmapi.tracer import global_tracer @@ -41,7 +40,6 @@ from ..lora_manager import LoraManager from ..prompt_adapter_manager import PromptAdapterManager from ..runtime import ModelConfig -from ..runtime.model_runner import _engine_config_to_model_config from ..sampling_params import BatchedLogitsProcessor, SamplingParams from .executor import GenerationExecutor, IterationResultQueue from .ipc import FusedIpcQueue, IpcQueue @@ -91,7 +89,7 @@ class WorkerExit(GeneratorExit): def __init__( self, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, postproc_worker_config: Optional[PostprocWorkerConfig] = None, @@ -220,14 +218,6 @@ def _create_engine(executor_config): executor_config.parallel_config = tllm.ParallelConfig( participant_ids=comm_ranks, device_ids=device_ids) - if isinstance(engine, Engine): - return tllm.Executor(engine.engine, - json.dumps(engine.config.to_dict(), - cls=ConfigEncoder), - tllm.ModelType.DECODER_ONLY, - executor_config=executor_config, - managed_weights=engine.managed_weights) - assert not hasattr(executor_config, "backend") return tllm.Executor(engine, tllm.ModelType.DECODER_ONLY, executor_config) @@ -239,26 +229,6 @@ def _create_engine(executor_config): self._lora_manager: Optional[LoraManager] = None self._prompt_adapter_manager: Optional[PromptAdapterManager] = None self._runtime_model_config: Optional[ModelConfig] = None - if self.rank == 0 and isinstance(self.engine, tllm.Executor): - if isinstance(self.engine, Engine): - engine_config = self.engine.config - else: - engine_config = EngineConfig.from_json_file( - f"{self._engine}/config.json") - self._runtime_model_config = _engine_config_to_model_config( - engine_config) - if engine_config.build_config.plugin_config.lora_plugin: - # TODO(azuker): Passing peft cache manager to LoraManager is used for LoRA optimization - # (see LoraManager constructor docstring). Getting the peft cache manager from this - # point in the TRT flow is currently not supported (it's at the CPP - # Executor->ExecutorImpl->TrtGptModel->mPeftCacheManager) therefore for now this LoRA - # optimization is not available in TRT-python flow. - self._lora_manager = LoraManager( - mapping=engine_config.pretrained_config.mapping, - model_config=self._runtime_model_config, - cpp_peft_cache_manager=None) - if engine_config.build_config.max_prompt_embedding_table_size > 0: - self._prompt_adapter_manager = PromptAdapterManager() if self._backend == "pytorch" and self._lora_config is not None: from tensorrt_llm._torch.pyexecutor.resource_manager import \ diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 82f5aed52110..d9c8ab6c5468 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -20,7 +20,6 @@ from .._utils import mpi_world_size from ..bindings import executor as tllm -from ..builder import Engine from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..llmapi.llm_args import BaseLlmArgs, TorchLlmArgs @@ -529,7 +528,7 @@ def _create_ipc_executor( @staticmethod def create( - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, model_world_size: int = 1, diff --git a/tensorrt_llm/executor/ray_gpu_worker.py b/tensorrt_llm/executor/ray_gpu_worker.py index 0b86d5255c8a..7a312aa003f4 100644 --- a/tensorrt_llm/executor/ray_gpu_worker.py +++ b/tensorrt_llm/executor/ray_gpu_worker.py @@ -5,7 +5,7 @@ from functools import wraps from pathlib import Path from queue import Queue -from typing import Any, List, Optional, Type, Union +from typing import Any, List, Optional, Type import ray import torch @@ -17,7 +17,6 @@ from .. import TorchLlmArgs from ..bindings import executor as tllm -from ..builder import Engine from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ..llmapi.tokenizer import TokenizerBase from ..llmapi.utils import configure_cpu_affinity @@ -214,7 +213,7 @@ class RayGPUWorker(RpcWorkerMixin, BaseWorker): def __init__( self, device_id: int, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, postproc_worker_config: Optional[PostprocWorkerConfig] = None, diff --git a/tensorrt_llm/executor/rpc_worker.py b/tensorrt_llm/executor/rpc_worker.py index 5433cbd37b54..667fa985921f 100644 --- a/tensorrt_llm/executor/rpc_worker.py +++ b/tensorrt_llm/executor/rpc_worker.py @@ -1,7 +1,7 @@ from pathlib import Path from queue import Queue from threading import Event -from typing import Optional, Union +from typing import Optional import nvtx @@ -10,7 +10,6 @@ from .._utils import mpi_rank from ..bindings import executor as tllm -from ..builder import Engine from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.tokenizer import TokenizerBase from ..logger import set_level @@ -48,7 +47,7 @@ class RpcWorker(RpcWorkerMixin, BaseWorker): def __init__( self, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, is_llm_executor: Optional[bool] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, @@ -111,7 +110,7 @@ def start(self): @staticmethod def main_task( - engine: Union[Path, Engine], + engine: Path, rpc_addr: str, *, executor_config: Optional[tllm.ExecutorConfig] = None, diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index c05e44f02cd3..87841c1edd8e 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -5,7 +5,7 @@ import traceback from concurrent.futures import ProcessPoolExecutor from pathlib import Path -from typing import List, Optional, Union +from typing import List, Optional import zmq @@ -13,7 +13,6 @@ from .._utils import mpi_comm, mpi_rank, print_all_stacks from ..bindings import executor as tllm -from ..builder import Engine from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.mpi_session import set_mpi_session_cpp from ..llmapi.tokenizer import TokenizerBase @@ -38,7 +37,7 @@ class GenerationExecutorWorker(RpcWorkerMixin, BaseWorker): def __init__( self, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, postproc_worker_config: Optional[PostprocWorkerConfig] = None, @@ -165,7 +164,7 @@ def block_subordinates(self): @print_traceback_on_error def worker_main( - engine: Path | Engine, + engine: Path, worker_queues: WorkerCommIpcAddrs, log_level: str, executor_config: Optional[tllm.ExecutorConfig] = None, diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py old mode 100755 new mode 100644 index df80459359aa..f4b97746f615 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -12,661 +12,14 @@ # 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 math -import weakref -from collections import OrderedDict from enum import IntEnum -from functools import partial -from typing import List, Optional, Sequence, Tuple, Union +from typing import List, Optional import numpy as np import torch - -# isort: off -import tensorrt as trt -# isort: on - -from . import graph_rewriting as gw -from ._common import default_net, default_trtnet, precision -from ._utils import (QuantModeWrapper, bf16_array, bool_array, - dim_resolve_negative, dim_to_trt_axes, dims_array, - fp16_array, fp32_array, get_sm_version, int32_array, - int64_array, np_dtype_to_trt, str_dtype_to_trt, - trt_dtype_to_np, trt_dtype_to_str) -from .network import PluginInfo, get_np_weight, set_np_weight, set_plugin_info -from .plugin import TRT_LLM_PLUGIN_NAMESPACE, current_all_reduce_helper -from .quantization import QuantMode - - -class DimRange(object): - ''' - One DimRange object stores the ranges of all the dimensions of one tensor in one optimization profile. - For example, tensor has 2 dimensions. Then the data members are: - self.min = [dim 0 min, dim 1 min] - self.opt = [dim 0 opt, dim 1 opt] - self.max = [dim 0 max, dim 1 max] - For static dimension, it has min==opt==max, thus the shape param in the ctor can be an integer - ''' - - def __init__(self, shape: List[Union[int, List[int], Tuple[int, int, int]]], - names: List[str]): - ''' - Parameters: - shape: a list with length N, each element is an integer or a 3-elements tuple/list of int, - where N is the number of dimensions for a tensor. - When one element is an integer, it means that dimension is static. - Otherwise, when one element is a tuple/list, it means the dimension is dynamic. - The 3 elements in one tuple/list is ordered by (min, opt, max), and this function asserts - 0 <= min <= opt <= max. - - Example, for a 3 rank tensor, with 1st dimension being static and has value 16, and second dimension being dynamic with - min/opt/max values being 1/8/32, and 3rd dimension being static and has value 8. - The shape parameter could be: - [16, (1, 8, 32), 8] - It has same semantics of - [(16, 16, 16), (1, 8, 32), (8, 8, 8)] - ''' - self.min = [] - self.opt = [] - self.max = [] - self.dimension_names = names - assert len(names) == len( - shape - ), "Expecting shape list and name list must have same length, got {shape=}, {name=}" - for dim in shape: - if isinstance(dim, (list, tuple)): - assert len(dim) == 3 and 0 <= dim[0] <= dim[1] <= dim[2], \ - "Each dimension must specify a 3-elements tuple or list in the order of (min,opt,max), got {dim=}" - self.min.append(dim[0]) - self.opt.append(dim[1]) - self.max.append(dim[2]) - elif isinstance(dim, int): - self.min.append(dim) - self.opt.append(dim) - self.max.append(dim) - else: - raise AttributeError( - f'Dimension should be [min, opt, max] (dynamic shape) or int (specific value). Got {type(dim)}' - ) - - def __eq__(self, __value: object) -> bool: - return isinstance(__value, DimRange) and \ - self.dimension_names == __value.dimension_names and \ - self.min == __value.min and self.opt == __value.opt and self.max == __value.max - - def __repr__(self) -> str: - return str(self) - - def __str__(self) -> str: - return f"{self.dimension_names=} {self.min=}, {self.opt=}, {self.max=})" - - def __hash__(self) -> int: - return hash(str(self)) - - -class Tensor(object): - ''' - The class to represent dense tensors. - - A dense tensor is named, has a shape and contains typed elements. Each - dimension of a tensor can either be static or dynamic. Static dimensions - are known at engine compilation by TensorRT. Dynamic dimensions can take - values determined at runtime. The tensor can be located on the host (CPU) - or the device (GPU). - ''' - - def __init__(self, - name=None, - dtype=None, - shape=None, - dim_range=None, - is_network_input=True, - location=trt.TensorLocation.DEVICE, - network=None, - trt_tensor=None): - ''' - Parameters: - name : str - The name of the tensor. - - dtype : tensorrt.DataType - The type of the elements of the tensor. See the TensorRT - documentation for list of supported data types. - - shape : tensorrt.Dims - The dimensions of the tensor. In TensorRT-LLM, tensors can have - static or dynamic dimensions (it is possible to mix static and - dynamic dimensions). A static dimension is known when the - TensorRT engine is built. A dynamic dimension can be set when - the engine is executed. Use -1 for dynamic dimensions. - - dim_range : OrderedDict - An ordered dictionary (the positions of the elements matter) - that associates a name and a range of values to the dimensions. - For a static dimension, the range must be limited to a single - value. For a dynamic dimension, the range is defined by three - values [min, opt, max] where min and max are, respectively, the - smallest and largest possible values of that dimension. The - opt value is used by TensorRT to optimize the engine for the - most common case. - - Assume there is N optimization profiles, each item dim_range dict is ordered by: - (dynamic dimension name : [profile 0 (min, opt, max), profile 1 (min, opt, max), ... profile N(min, opt, max)]) - or it's following when the dimension is static (can think as min==opt==max): - (static dimension name : [profile 0 value, profile 1 value, ... profile N value]) - For static dimension the profile 0-N value must be same, (TODO: can it be simplified to be only 1 value?) - And number of keys is equal to number of optimization profiles. - - is_network_input : bool - A boolean indicating if that tensor is an input of the network. - Inputs must be provided by the user to run the engine. - - location : tensorrt.TensorLocation - A flag to indicate where the tensor will be located. It can be - on the host (CPU) or the device (GPU). - - network: Network - A parent Network instance, that helps to fine the users of this tensor. - - trt_tensor: trt.ITensor - Construct with the ITensor instance directly, and no shape profiles are required. - ''' - - # Layout of self.profiles - # Opt profile 0: dim 0 (min, opt, max), dim 1 (min, opt, max) ... dim M - # Opt profile 1: dim 0 (min, opt, max), dim 1 (min, opt, max) ... dim M - # ... - # Opt profile N: dim 0 ... dim M - - # So from the dim_range arg to self.profiles conversion, there is a layout transpose - # dim_range arg is: {M dimension x N profiles}, while self.profiles layout is {N profiles x M dimensions} - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - - self.profiles = [] - - self.is_tensor_wrapper = False # specially for the graph rewriter - - # work as a wrapper for a trt.ITensor, this is used specially in the graph rewriter - if trt_tensor is not None: - self.is_tensor_wrapper = True - assert network is not None - self.trt_tensor = trt_tensor - self._network = weakref.ref(network) - assert not is_network_input, "is_network_input should be False when trt_tensor is not None" - return - - # be cautious here, the weakref is critical to avoid circular referencing before Network and Tensor - # using strong reference will likely cause significant peak memory increase, since Network objects - # holds the weights data. - self._network = weakref.ref(default_net()) - self.is_network_input = is_network_input - if is_network_input: - if dim_range is not None: - assert isinstance(dim_range, OrderedDict) - assert len( - dim_range - ) >= 1, f"Each input tensor shall have at least one dimension, tensor '{name}' found {dim_range=}" - - found_profiles = [ - len(ranges) for _, ranges in dim_range.items() - ] - assert all( - [x == found_profiles[0] for x in found_profiles] - ), f"Expecting all the dimensions in the dim_range has same number of profiles, tensor '{name}' got {dim_range=}" - - num_opt_profile = len(list(dim_range.items())[0][1]) - assert num_opt_profile >= 1 - for i in range(num_opt_profile): - range_shape = [] - dimension_names = [] - for dim, ranges in dim_range.items(): - assert isinstance(ranges, (list, tuple)) - range_shape.append(ranges[i]) - dimension_names.append(dim) - self.profiles.append(DimRange(range_shape, dimension_names)) - - default_net()._add_input(self, name, dtype, shape, dim_range) - self.name = name - self.dtype = dtype - self.shape = shape - self.location = location - - @property - def network(self): - return self._network() - - @property - def name(self): - ''' - The name of the tensor. - ''' - return self.trt_tensor.name - - @name.setter - def name(self, name): - ''' - Set the name of the tensor. - ''' - if name is not None: - self.trt_tensor.name = name - - @property - def dtype(self): - ''' - The type of the elements in the tensor. - ''' - return self.trt_tensor.dtype - - @dtype.setter - def dtype(self, dtype): - ''' - Set the type of the elements in the tensor. - ''' - if dtype is not None: - self.trt_tensor.dtype = dtype - - @property - def shape(self): - ''' - The shape of the tensor. - ''' - return self.size() - - @shape.setter - def shape(self, shape): - ''' - Set the shape of the tensor. See __init__. - ''' - if shape is not None: - self.trt_tensor.shape = shape - - @property - def location(self): - ''' - The physical location of the tensor (on the host or the device). - ''' - return self.trt_tensor.location - - @location.setter - def location(self, location): - ''' - Set the physical location of the tensor (on the host or the device). See __init__. - ''' - if location is not None: - self.trt_tensor.location = location - - def mark_output(self, - name: Optional[str] = None, - dtype: Optional[Union[str, trt.DataType]] = None): - ''' - Mark a tensor as a network output. - - When a tensor is marked as an output, its content can be obtained after - the execution of the TensorRT engine. The user is responsible for - allocating buffers to store the output tensors when preparing the - execution of the TensorRT engine. - ''' - if name is None: - name = self.name - - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - - assert dtype is None or isinstance(dtype, trt.DataType) - default_net()._mark_output(self, name, dtype) - - def __add__(self, b): - ''' - See functional.add. - ''' - return add(self, b) - - def __radd__(self, b): - ''' - See functional.add. - ''' - return add(b, self) - - def __sub__(self, b): - ''' - See functional.sub. - ''' - return sub(self, b) - - def __rsub__(self, b): - ''' - See functional.sub. - ''' - return sub(b, self) - - def __mul__(self, b): - ''' - See functional.mul. - ''' - return mul(self, b) - - def __rmul__(self, b): - ''' - See functional.mul. - ''' - return mul(b, self) - - def __truediv__(self, b): - ''' - See functional.div. - ''' - return div(self, b) - - def __floordiv__(self, b): - ''' - See functional.floordiv. - ''' - return floordiv(self, b) - - def __mod__(self, b): - ''' - See functional.floordiv. - ''' - return modulo(self, b) - - def __lt__(self, b): - ''' - See functional.lt. - ''' - return lt(self, b) - - def __gt__(self, b): - ''' - See functional.gt. - ''' - return gt(self, b) - - def __eq__(self, b): - ''' - See functional.eq. - ''' - if self.is_tensor_wrapper: - # for graph rewriter - return hash(self) == hash(b) - else: - # for creating the network - return eq(self, b) - - def __ge__(self, b): - ''' - Maps to functional.gt or functional.eq. - ''' - return op_or(self.__gt__(b), self.__eq__(b)) - - def __le__(self, b): - ''' - Maps to functional.lt or functional.eq. - ''' - return op_or(self.__lt__(b), self.__eq__(b)) - - def view(self, shape, zero_is_placeholder=True): - ''' - See functional.view. - ''' - return view(self, shape, zero_is_placeholder) - - def flatten(self, start_dim=0, end_dim=-1): - ''' - See functional.flatten. - ''' - return flatten(self, start_dim, end_dim) - - def permute(self, dims): - ''' - See functional.permute. - ''' - return permute(self, dims) - - def transpose(self, dim0, dim1): - ''' - See functional.transpose. - ''' - return transpose(self, dim0, dim1) - - def mean(self, dim, keepdim=False): - ''' - See functional.mean. - ''' - return mean(self, dim, keepdim) - - def max(self, dim, keepdim=False): - ''' - See functional.max. - ''' - return max(self, dim, keepdim) - - def abs(self): - ''' - See functional.abs. - ''' - return abs(self) - - def sqrt(self): - ''' - See functional.sqrt. - ''' - return sqrt(self) - - def squeeze(self, dim, zero_is_placeholder): - ''' - See functional.squeeze. - ''' - return squeeze(self, dim, zero_is_placeholder) - - def unsqueeze(self, dim): - ''' - See functional.squeeze. - ''' - return unsqueeze(self, dim) - - def log(self): - ''' - See functional.log. - ''' - return log(self) - - def cast(self, dtype): - ''' - See functional.cast. - ''' - return cast(self, dtype) - - def size(self, dim=None): - ''' - Returns the shape of the tensor if the dim parameter is None. - Otherwise, returns a size of the dimension indicated by dim. The - behavior is undefined if dim is negative or exceeds the rank of the - tensor. - ''' - if dim is None: - return self.trt_tensor.shape - - return self.trt_tensor.shape[dim] - - def rank(self): - ''' - Returns the rank (i.e. the number of dimensions) of the tensor. - ''' - return len(self.trt_tensor.shape) - - def ndim(self): - ''' - Returns the rank (i.e. the number of dimensions) of the tensor. - ''' - return self.rank() - - def split(self, split_size_or_sections, dim=0): - ''' - See functional.split. - ''' - return split(self, split_size_or_sections, dim) - - def select(self, dim, index): - ''' - See functional.select. - ''' - return select(self, dim, index) - - def unbind(self, dim=0): - ''' - See functional.unbind. - ''' - return unbind(self, dim) - - def repeat(self, sizes): - ''' - See functional.repeat - ''' - return repeat(self, sizes) - - def is_dynamic(self, dim=None): - ''' - If the argument 'dim' is None, that function returns a boolean that - indicates if the tensor contains a dynamic dimension (True) or not - (False). In that case, the first dimension is excluded (as it usually - corresponds to the batch size). If the argument is an integer, that - functions returns a boolean that indicates if the dimension 'dim' is - dynamic (True) or not (False). - ''' - if dim is not None: - return self.trt_tensor.shape[dim] == -1 - - for i, s in enumerate(self.trt_tensor.shape): - if i != 0 and s == -1: - return True - - return False - - # graph writer related functions - - def get_parent(self): - ''' Get the layer that produces this tensor. ''' - return self.network.get_tensor_parent(self) - - def get_users(self): - ''' Get the layers that use this tensor as an input. ''' - return self.network.get_tensor_users(self) - - def replace_all_uses_with(self, new_tensor): - ''' - Replace all uses of this tensor as an input to consumer layers - ''' - - self.network.is_graph_altered = True - users = self.get_users() - for user in users: - inputs_changed = 0 - for i in range(user.num_inputs): - if user.get_inputs(i)[0].trt_tensor is self.trt_tensor: - inputs_changed += 1 - user.set_input(i, new_tensor.trt_tensor) - assert inputs_changed >= 1, "Tensor not found in layer inputs" - - # update the FLayerMetadata as well - flayer = gw.FLayerInfoMemo.instance().get(user.name) - flayer and flayer.replace_input_with(self, new_tensor) - - def is_trt_wrapper(self): - ''' - Check if there is a trt.ITensor member inside, which is required for - graph rewriter. In order to differentiate usages, it may be necessary - to have an inheritance hierarchy. - ''' - if hasattr(self, 'trt_tensor'): - return True - else: - return False - - def __hash__(self): - if self.is_trt_wrapper(): - return id(self.trt_tensor) - else: - return id(None) - - def __repr__(self): - return f"TensorRT LLM Tensor: {self.name=} {self.dtype=} {self.shape=}" - - def __xor__(self, b): - ''' - Maps to functional.gt or functional.eq. - ''' - print(f"self.shape: {self.shape}, b.shape: {b.shape}") - a, b = broadcast_helper(self, b) - print(f"a.shape: {a.shape}, b.shape: {b.shape}") - return op_xor(a, b) - - -def _create_tensor(trt_tensor: trt.ITensor, producer: trt.ILayer) -> Tensor: - ''' - A helper function to create a TensorRT LLM Tensor object that encapsulates - the connection between the TensorRT tensor (trt.ITensor) and the layer - (trt.ILayer) that produces it. - - That function is expected to be used as: - - # Insert a new layer in the network using the TensorRT API: - layer = default_trtnet().add_(...) - # Extract the first output of that layer and connect it to the layer. - return _create_tensor(layer.get_output(0), layer) - - That function also sets the precision of the layer/producer to the default - precision of the network. - - Parameters: - trt_tensor : trt.ITensor - The TensorRT tensor to connect to its producer (the layer). - - producer : trt.ILayer - The producer. - - Returns: - The TensorRT LLM tensor (functional.Tensor) that encapsulates the - TensorRT tensor and the layer that produces it. The former is - accessible through the attribute 'trt_tensor' and the latter using the - attribute 'producer'. - ''' - assert trt_tensor is not None - assert producer is not None - - # Set the layer name since this is the only - # centralized location to pass the name from - # module space to the TRT IR - default_net()._set_layer_name(producer) - - assert trt_tensor.shape.__len__( - ) >= 0, f"tensor {trt_tensor.name} has an invalid shape" - tensor = Tensor(name=trt_tensor.name, - dtype=trt_tensor.dtype, - shape=trt_tensor.shape, - is_network_input=False) - tensor.trt_tensor = trt_tensor - tensor.producer = producer - - # tb.print_stack(limit=10) # FOR DEBUGGING: filter producer.name if needed - if default_net().dtype is not None and not default_net().strongly_typed: - if producer.type not in [ - trt.LayerType.SHAPE, trt.LayerType.CONSTANT, - trt.LayerType.GATHER, trt.LayerType.CONCATENATION - ]: - producer.precision = default_net().dtype - assert tensor is not None - - if gw.FLayerInfoMemo.instance().cur_flayer is not None: - gw.FLayerInfoMemo.instance().cur_flayer.layer_name = producer.name - - return tensor - - -def _add_plugin_info(layer, plugin_creator: trt.IPluginCreator, - plugin_name: str, pfc: trt.PluginFieldCollection) -> None: - plugin_info = PluginInfo(plugin_creator, plugin_name, pfc) - set_plugin_info(default_net().trt_network, layer.name, plugin_info) +from torch import Tensor class RotaryScalingType(IntEnum): @@ -691,7 +44,7 @@ def from_string(s): try: return RotaryScalingType[key] except KeyError: - raise ValueError(f'Unsupported rotary scaling type: {s}') + raise ValueError(f"Unsupported rotary scaling type: {s}") class PositionEmbeddingType(IntEnum): @@ -705,7 +58,9 @@ class PositionEmbeddingType(IntEnum): chatglm = 7 yarn = 8 mrope = 9 - deferred = 10 # Apply customized positional embedding by using an external embedder. K will be cached before embedding. + # Apply customized positional embedding by using an external embedder. + # K will be cached before embedding. + deferred = 10 def is_rope(self) -> bool: return self in [ @@ -736,7 +91,7 @@ def from_string(s): try: return PositionEmbeddingType[s] except KeyError: - raise ValueError(f'Unsupported position embedding type: {s}') + raise ValueError(f"Unsupported position embedding type: {s}") class AttentionMaskType(IntEnum): @@ -749,6943 +104,418 @@ class AttentionMaskType(IntEnum): custom_mask = 6 -class LayerNormType(IntEnum): - LayerNorm = 0 - RmsNorm = 1 - GroupNorm = 2 - - -class LayerNormPositionType(IntEnum): - pre_layernorm = 0 - post_layernorm = 1 - - -class MLPType(IntEnum): - MLP = 0 - GatedMLP = 1 - FusedGatedMLP = 2 - - -class SliceInputType(IntEnum): - data = 0 - start = 1 - size = 2 - stride = 3 - fill_value = 4 - axes = 5 - - -def activation(input: Tensor, act_type: trt.ActivationType) -> Tensor: - ''' - Add an activation function. - - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - act_type : trt.ActivationType - The type of the activation (RELU, TANH, SIGMOID, ...). - - The following closures are defined in functional.*: - - relu for op=trt.ActivationType.RELU - tanh for op=trt.ActivationType.TANH - sigmoid for op=trt.ActivationType.SIGMOID - - Returns: - The tensor produced by the activation layer. - ''' - layer = default_trtnet().add_activation(input.trt_tensor, act_type) - return _create_tensor(layer.get_output(0), layer) - - -def int_clip(input: Tensor, lower: int, upper: int) -> Tensor: - assert lower <= upper, f"Lower bound must be less than or equal to upper bound i.e. {lower} <= {upper}" - res = minimum(input, upper) - res = maximum(res, lower) - return res - - -def clip(input: Tensor, alpha: float, beta: float) -> Tensor: - ''' - Add a CLIP operation that sets the range to [alpha, beta]. +class AllReduceStrategy(IntEnum): + NCCL = 0 + MIN_LATENCY = 1 + UB = 2 + AUTO = 3 + ONESHOT = 4 + TWOSHOT = 5 + LOWPRECISION = 6 + MNNVL = 7 + NCCL_SYMMETRIC = 8 + SYMM_MEM = 9 # PyTorch symmetric memory with MULTIMEM - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - alpha : float - The lower bound of the CLIP function. +class AllReduceFusionOp(IntEnum): + NONE = 0 + RESIDUAL_RMS_NORM = 1 + LAST_PROCESS_FOR_UB = 2 + RESIDUAL_RMS_PREPOST_NORM = 3 + RESIDUAL_RMS_NORM_QUANT_FP8 = 4 + RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5 + RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6 + RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7 + MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 8 + RMS_NORM = 9 - beta : float - The upper bound of the CLIP function. - Returns: - The tensor produced by the activation layer. - ''' - layer = default_trtnet().add_activation(input.trt_tensor, - trt.ActivationType.CLIP) - layer.alpha = alpha - layer.beta = beta - return _create_tensor(layer.get_output(0), layer) +class AllReduceParams: + + def __init__( + self, + strategy: AllReduceStrategy = AllReduceStrategy.AUTO, + fusion_op: AllReduceFusionOp = AllReduceFusionOp.NONE, + bias: Optional[Tensor] = None, + residual: Optional[Tensor] = None, + norm_weight: Optional[Tensor] = None, + scale: Optional[Tensor] = None, + norm_pre_residual_weight: Optional[Tensor] = None, + eps: float = 1e-06, + enable_allreduce: bool = True, + trigger_completion_at_end: bool = True, + ): + self.strategy = strategy + self.fusion_op = fusion_op + self.bias = bias + self.residual = residual + self.norm_weight = norm_weight + self.scale = scale + self.norm_pre_residual_weight = norm_pre_residual_weight + self.eps = eps + # For torch path only, has no effect on TRT path + self.enable_allreduce = enable_allreduce + self.trigger_completion_at_end = trigger_completion_at_end + assert fusion_op in (AllReduceFusionOp.NONE.value, + AllReduceFusionOp.RMS_NORM.value) or (residual + is not None) + def has_affine(self): + return 1 if self.norm_weight is not None else 0 -relu = partial(activation, act_type=trt.ActivationType.RELU) -tanh = partial(activation, act_type=trt.ActivationType.TANH) -sigmoid = partial(activation, act_type=trt.ActivationType.SIGMOID) + def has_bias(self): + return 1 if self.bias is not None else 0 + def has_scale(self): + return 1 if self.scale is not None else 0 -def silu(input: Tensor) -> Tensor: - ''' - Add a SiLU (`x * sigmoid(x)`) operation. - Parameters: - input : Tensor - The input tensor on which the activation function is applied. +class MoEAllReduceParams(AllReduceParams): - Returns: - The tensor produced by the activation layer. - ''' - return input * sigmoid(input) + def __init__( + self, + device_num_experts: Optional[Tensor] = None, + expert_scale_factor: Optional[Tensor] = None, + expanded_idx_to_permuted_idx: Optional[Tensor] = None, + shared_expert_output: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + residual: Optional[Tensor] = None, + norm_weight: Optional[Tensor] = None, + scale: Optional[Tensor] = None, + norm_pre_residual_weight: Optional[Tensor] = None, + eps: float = 1e-06, + enable_allreduce: bool = True, + is_cutlass_min_latency: bool = False, + ): + super().__init__( + bias=bias, + residual=residual, + norm_weight=norm_weight, + scale=scale, + norm_pre_residual_weight=norm_pre_residual_weight, + eps=eps, + enable_allreduce=enable_allreduce, + ) + self.device_num_experts = device_num_experts + self.expert_scale_factor = expert_scale_factor + self.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx + self.shared_expert_output = shared_expert_output + self.is_cutlass_min_latency = is_cutlass_min_latency + def is_valid(self): + if self.is_cutlass_min_latency: + return (self.device_num_experts is not None + and self.expert_scale_factor is not None + and self.shared_expert_output is not None) + else: + return self.expanded_idx_to_permuted_idx is not None -def swiglu(input: Tensor) -> Tensor: - ''' - Add a SwiGLU (`x * SiLU(gate)`) operation. - That function takes a tensor, splits it into two halves along the last - dimension, applies SiLU to the second half and multiply the results. The - behavior is undefined if the last dimension is not even. +class RopeEmbeddingUtils: - Parameters: - input : Tensor - The input tensor on which the activation function is applied. + @staticmethod + # ref: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py#L298 + def apply_llama3_scaling(inv_freqs: np.ndarray, rope_scaling_config: dict): + scale_factor = rope_scaling_config.get("factor", 8.0) + low_freq_factor = rope_scaling_config.get("low_freq_factor", 1.0) + high_freq_factor = rope_scaling_config.get("high_freq_factor", 4.0) + old_context_len = rope_scaling_config.get( + "original_max_position_embeddings", 8192) - Returns: - The tensor produced by the activation layer. - ''' - x, gate = chunk(input, 2, dim=-1) - return silu(gate) * x + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + new_inv_freqs = [] + for inv_freq in inv_freqs: + wavelen = 2 * math.pi / inv_freq + if wavelen < high_freq_wavelen: + new_inv_freqs.append(inv_freq) + elif wavelen > low_freq_wavelen: + new_inv_freqs.append(inv_freq / scale_factor) + else: + assert low_freq_wavelen != high_freq_wavelen + smooth = (old_context_len / wavelen - low_freq_factor) / ( + high_freq_factor - low_freq_factor) + new_inv_freqs.append((1 - smooth) * inv_freq / scale_factor + + smooth * inv_freq) + return np.array(new_inv_freqs, dtype=inv_freqs.dtype) + @staticmethod + def create_sinusoidal_positions(num_pos: int, + dim: int, + theta: float = 10000.0, + dtype=np.float32): + inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) + sinusoid_inp = np.einsum("i , j -> i j", + np.arange(num_pos, dtype=dtype), + inv_freq, + dtype=dtype) + concat = np.concatenate((np.sin(sinusoid_inp), np.cos(sinusoid_inp)), + axis=1) + return np.expand_dims(concat, axis=0).astype(dtype) -def squared_relu(x: Tensor) -> Tensor: - ''' - Add a Squared ReLU operation. + @staticmethod + def create_sinusoidal_positions_for_attention_plugin( + num_pos: int, + dim: int, + theta: float = 10000.0, + scale: float = 1.0, + scale_type: RotaryScalingType = RotaryScalingType.none, + # Other scaling configs that only used by certain scaling types. + rope_scaling_config: dict = None, + duplicate_data: bool = False, + dtype=np.float32, + ): + if scale_type == RotaryScalingType.linear: + scale = 1.0 / scale + if scale_type == RotaryScalingType.llama3: + assert rope_scaling_config is not None, "rotary_scaling config must be provided." + inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) + inv_freq = RopeEmbeddingUtils.apply_llama3_scaling( + inv_freq, rope_scaling_config) + elif scale_type == RotaryScalingType.dynamic: + # Make sure scaling_alpha exists in rope_scaling + # Ref: https://huggingface.co/tencent/Hunyuan-A13B-Instruct-FP8/blob/main/modeling_hunyuan.py#L346 + assert rope_scaling_config["alpha"] is not None, ( + "rope_scaling_config.alpha must be provided.") + scaling_alpha = rope_scaling_config["alpha"] + adjusted_base = theta * (scaling_alpha**(dim / (dim - 2))) + inv_freq = 1.0 / (adjusted_base**( + np.arange(0, dim, 2, dtype=dtype) / dim)).astype(dtype) + else: + inv_freq = scale / (theta + **(np.arange(0, dim, 2) / dim)).astype(dtype) + sinusoid_inp = np.expand_dims( + np.einsum("i , j -> i j", + np.arange(num_pos, dtype=dtype), + inv_freq, + dtype=dtype), + axis=-1, + ) + if duplicate_data: + sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) + # fuse cos/sin into float2 (cos, sin). + concat = np.concatenate( + (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), + axis=-1) # np.cos(sinusoid_inp).shape = (32768, 64, 1) - This function applies ReLU and squares the output. + return inv_freq, concat.reshape(1, -1).astype(dtype) - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - Returns: - The tensor produced by the activation layer. - ''' - return pow(relu(x), 2.0) - - -def cast(input: Tensor, dtype: Union[str, trt.DataType]): - ''' - Add a cast operation. - - For an input tensor of type INT8, this function sets the dynamic range of - the input to [-127, 127] for automatic dequantization. For a cast into - INT8, that function sets the dynamic range of the output to [-127, 127] for - automatic quantization. - - Parameters: - input : Tensor - The input tensor on which the cast is applied. - - dtype : str or trt.DataType - The data type of the output tensor after the cast. When 'dtype' is - provided as a string, it must be a name amongst the valid names. - See _str_to_trt_dtype_dict in _utils.py for a list of supported - types and type names. - - Returns: - The tensor produced by the inserted layer. - ''' - if isinstance(dtype, str): - cvt_dtype = str_dtype_to_trt(dtype) - elif isinstance(dtype, trt.DataType): - cvt_dtype = dtype - else: - raise TypeError("%s is not supported" % type(dtype)) - - if input.dtype == cvt_dtype: - # If input type and cast dtype are the same, do nothing - return input - - layer = default_trtnet().add_cast(input.trt_tensor, cvt_dtype) - if not default_net().strongly_typed: - layer.set_output_type(0, cvt_dtype) - output = _create_tensor(layer.get_output(0), layer) - if not default_net().strongly_typed: - if input.dtype == str_dtype_to_trt('int8'): - layer.get_input(0).set_dynamic_range(-127, 127) - if cvt_dtype == str_dtype_to_trt('int8'): - layer.get_output(0).set_dynamic_range(-127, 127) - - return output - - -def flip(input: Tensor, dims: Sequence[int]) -> Tensor: - ''' - Reverses the order of an n-D tensor along given axis in dims. - - That flip operation maps to a TensorRT ISliceLayer. For the dimensions - listed in dims it copies the elements from the last one to the first one - (from (N-1) down to 0 with a step of -1). For the dimensions not in 'dims', - it copies the elements from the first one to the last one (from 0 to N-1 - with a step of 1). - - Parameters: - input : Tensor - The input tensor on which the cast is applied. - - dims : list or tuple - The axes to flip. Negative indices are supported. - - Returns: - The tensor produced by the inserted layer. - ''' - assert not input.is_dynamic() - - ndim = input.ndim() - - for index, value in enumerate(dims): - assert -ndim <= value < ndim - if -ndim <= value < 0: - dims[index] += ndim - - assert len(dims) == len(set(dims)) - - start_values = [ - input.size()[i] - 1 if i in dims else 0 for i in range(ndim) - ] - stride_values = [-1 if i in dims else 1 for i in range(ndim)] - - layer = default_trtnet().add_slice(input.trt_tensor, - start=start_values, - shape=input.size(), - stride=stride_values) - - return _create_tensor(layer.get_output(0), layer) - - -def interpolate(input: Tensor, - size: Union[int, List[int]] = None, - scale_factor: Union[float, List[float]] = None, - mode: str = 'nearest', - align_corners: bool = False, - recompute_scale_factor: bool = False, - antialias: bool = False) -> Tensor: - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic() - - input_ndim = input.ndim() - - assert 2 < input_ndim < 6, "Only 3D, 4D and 5D input Tensors supported" - assert (size is not None) ^ ( - scale_factor - is not None), "Only one of out_shape or scales should be defined" - - assert mode in ('nearest', 'linear', 'bilinear', 'bicubic', 'trilinear', - 'nearest-exact') - - if mode == 'trilinear' and input_ndim != 5: - raise ValueError("trilinear only supports 5D tensor") - - if mode == "bilinear" and input_ndim != 4: - raise ValueError("bilinear only supports 4D tensor") - - if mode == "linear" and input_ndim != 3: - raise ValueError("linear only supports 3D tensor") - - layer = default_trtnet().add_resize(input.trt_tensor) - - input_shape = input.size() - - updated_shape = [] - if scale_factor: - scale_len = 1 if isinstance(scale_factor, - (float, int)) else len(scale_factor) - if scale_len == 1 and isinstance(scale_factor, (float, int)): - updated_scale = [scale_factor for _ in range(input_ndim - 2)] - - else: - updated_scale = scale_factor - updated_shape = [ - int(math.floor(updated_scale[i - 2] * - input_shape[i])) if i > 1 else input_shape[i] - for i in range(input_ndim) - ] - - else: - size_len = 1 if isinstance(size, int) else len(size) - assert size_len == input_ndim - 2 - if size_len == 1 and isinstance(size, int): - updated_size = [size for _ in range(input_ndim - 2)] - else: - updated_size = size - - updated_shape = [ - input_shape[i] if i < 2 else updated_size[i - 2] - for i in range(input_ndim) - ] - layer.shape = updated_shape - - if mode in ['nearest', 'nearest-exact'] or mode is None: - layer.resize_mode = trt.InterpolationMode.NEAREST - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.ASYMMETRIC - elif mode in ['linear', 'bilinear', 'trilinear']: - layer.resize_mode = trt.InterpolationMode.LINEAR - if align_corners: - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.ALIGN_CORNERS - else: - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.HALF_PIXEL - # TODO, need to confirm the align_corners effect on bilinear mode. - if mode == 'bilinear': - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.HALF_PIXEL - - elif mode in ['bicubic']: - layer.resize_mode = trt.InterpolationMode.CUBIC - - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.HALF_PIXEL - - else: - layer.resize_mode = trt.InterpolationMode.NEAREST - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.ASYMMETRIC - - return _create_tensor(layer.get_output(0), layer) - - -def matmul(input: Tensor, - mat2: Tensor, - transa: bool = False, - transb: bool = False, - use_fp32_acc: bool = True) -> Tensor: - ''' - Add a matrix multiplication. - - That operation maps to a tensorrt.IMatrixMultiplyLayer layer. As explained - in the TensorRT documentation, it computes the inner product between the - two inputs after applying an optional transposition on the inputs. - - Parameters: - input : Tensor - The first tensor (often called A). - - mat2 : Tensor - The second tensor (often called B). - - transa : bool - Is the first input transposed? Set to 'True' if you want the first - input to be transposed, 'False' otherwise. - - transb : bool - Is the second input transposed? Set to 'True' if you want the - second input to be transposed, 'False' otherwise. - - use_fp32_acc: bool - Set to 'True' if for accuracy reason, this fp16 matmul needs to use - fp32 accumulation. This can be a per model and per matmul decision. - Returns: - The tensor produced by the inserted layer. - ''' - # This option is only supported for fp16, but not bf16 or any other precisions. - use_fp32_acc = use_fp32_acc and input.dtype == trt.DataType.HALF and mat2.dtype == trt.DataType.HALF - - if use_fp32_acc: - input = cast(input, 'float32') - mat2 = cast(mat2, 'float32') - - input, mat2 = broadcast_helper(input, mat2) - op0 = trt.MatrixOperation.TRANSPOSE if transa \ - else trt.MatrixOperation.NONE - op1 = trt.MatrixOperation.TRANSPOSE if transb \ - else trt.MatrixOperation.NONE - layer = default_trtnet().add_matrix_multiply(input.trt_tensor, op0, - mat2.trt_tensor, op1) - output = _create_tensor(layer.get_output(0), layer) - if use_fp32_acc: - output = cast(output, "float16") - - return output - - -def gemm_swiglu(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - scale_d0: float = 1.0, - scale_d1: float = 1.0, - scale_output: float = 1.0) -> Tensor: - ''' - Add a matrix multiplication, followed by SwiGLU (`x * SiLU(gate)`) operation. - - The second SwiGLU operation takes the preceding tensor, splits it into two halves - along the last dimension, applies SiLU to the second half and multiply the results. The - behaviour is undefined if the last dimension is not even. - - Parameters: - input : Tensor - The first tensor (often called A). - - weight : Tensor - The second tensor (often called B). - - bias : Optional[Tensor] - The per-channel bias. The plugin with fp8 dtype does not support bias yet. - - scale_d0 : float - The scale for dequantizing x, used for fp8 - - scale_d1 : float - The scale for dequantizing gate, used for fp8 - - scale_output : float - The scale for quantizing output, used for fp8 - - Returns: - The tensor produced by the inserted layer. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'GemmSwiglu', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.gemm_swiglu_plugin - if p_dtype == "fp8": - assert bias is None, "fp8 gemm_swiglu does not support bias yet" - - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_has_bias = trt.PluginField( - "has_bias", np.array(np.int8(0 if bias is None else 1), np.int8), - trt.PluginFieldType.INT8) - pf_scale_d0 = trt.PluginField("scale_d0", - np.array(scale_d0, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_d1 = trt.PluginField("scale_d1", - np.array(scale_d1, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_output = trt.PluginField("scale_output", - np.array(scale_output, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - pfc = trt.PluginFieldCollection( - [pf_type, pf_has_bias, pf_scale_d0, pf_scale_d1, pf_scale_output]) - gemm_swiglu_plug = plg_creator.create_plugin("gemm_swiglu", pfc) - - # TODO(anchengc) pass nullptr when no bias - if bias is None: - bias = constant( - np.zeros([weight.shape[0]], dtype=trt_dtype_to_np(input.dtype))) - plug_inputs = [input.trt_tensor, weight.trt_tensor, bias.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_swiglu_plug) - - return _create_tensor(layer.get_output(0), layer) - - -def constant(ndarray: np.ndarray, - as_dtype: trt.DataType | None = None, - as_shape=None) -> Tensor: - ''' - Add a constant layer. - - TensorRT graphs encapsulate constant values in the form of constant layers - (tensorrt.IConstantLayer). That function creates such a layer from a Numpy - array of values. After compilation of the network by TensorRT, those - weights are stored in the serialized TensorRT engine. - - Parameters: - ndarray : numpy.ndarray - The array of values (weights) encapsulated by this constant layer. - - Returns: - The tensor produced by the inserted layer. - ''' - trt_dtype = np_dtype_to_trt(ndarray.dtype) if as_dtype is None else as_dtype - trt_shape = trt.Dims( - ndarray.shape) if as_shape is None else trt.Dims(as_shape) - trt_count = 1 - for i in range(len(trt_shape)): - trt_count *= trt_shape[i] - weights = trt.Weights(trt_dtype, ndarray.ctypes.data, trt_count) - # Prevent underlying numpy array from going out of scope - default_net().register_ndarray(ndarray) - layer = default_trtnet().add_constant(trt_shape, weights) - if not default_net().strongly_typed: - layer.set_output_type(0, trt_dtype) - tensor = _create_tensor(layer.get_output(0), layer) - # TODO: remove this WAR after https://nvbugs/4359151 fixed. - set_np_weight(default_trtnet(), layer.name, ndarray) - return tensor - - -# TODO: TensorRT uses sizes of the output dimensions. -# DL framework uses ends usually. Will change it to ends. -def slice(input: Tensor, - starts: Union[Tensor, Sequence[int]], - sizes: Union[Tensor, Sequence[int]], - strides: Union[Tensor, Sequence[int]] = None, - mode: trt.SampleMode = None, - fill_value: Union[float, Tensor] = None) -> Tensor: - ''' - Add an operation to extract a slice from a tensor. - - As described in the TensorRT documentation of the ISliceLayer, the slice - layer has two variants: Static and dynamic. - - For static slicing, this function takes the starts and sizes values in the - different dimensions to slice at layer creation time via a sequence of - integers. For dynamic slicing, it accepts starts and sizes as - tensorrt.ITensor`s. - - The slice layer selects for each dimension a start location from within the - input tensor, and copies elements to the output tensor using a stride of 1 - across the input tensor. Start and size tensors must be 1-D int32 shape - tensors if not specified as a sequence of integers. - - As an example, on input = [[0, 2, 4], [1, 3, 5]], the call to - - slice(input, start=[1, 0], size=[1, 2]) - - will produce the tensor [[1, 3]] as output. The slice operator when - executed by TensorRT will copy one row (because size[0] == 1) starting from - the 2nd row (because start[0] == 1) and two columns (size[1] == 2) starting - from the 1st column (because start[1] == 0). - - In pseudo-code the behavior of that operation can be described as follows - for a 2D tensor (and easily be extended to more dimensions): - - output = Tensor(shape=sizes) - for ii in range(sizes[0]): - for jj in range(sizes[1]): - output[ii][jj] = input[starts[0]+ii][starts[1]+jj] - - Note that it is common in deep-learning frameworks to use ranges - [start:end] for similar operations. It can be emulated by setting the sizes - argument such that in each dimension [start:start+size] == [start:end] i.e. - size = end-start. - - TensorRT supports different slice modes but that function restricts that - choice to `mode == tensorrt.SampleMode.STRICT_BOUNDS`. - - Parameters: - input : Tensor - The input tensor on which the slicing is performed. - - starts : Union[Tensor, Sequence[int]] - The starting points, in the input tensor, and each dimension. - - sizes : Union[Tensor, Sequence[int]] - The number of elements in each dimension of the sliced tensor (output). - - strides : Union[Tensor, Sequence[int]] - The step be taken from start, in input tensor. - - mode : trt.SampleMode - The mode that controls how the slice operation handles out of bounds coordinates. - - Returns: - The tensor produced by the slice layer. - ''' - input_ndim = input.ndim() - - trt_starts = starts - if isinstance(starts, Tensor): - trt_starts = [0 for _ in range(input_ndim)] # unused dummy value - - trt_sizes = sizes - if isinstance(sizes, Tensor): - trt_sizes = [1 for _ in range(input_ndim)] # unused dummy value - - trt_strides = strides - if isinstance(strides, Tensor) or strides is None: - trt_strides = [1 for _ in range(input_ndim)] - - if fill_value is not None and isinstance(fill_value, float): - fill_value = constant(fp32_array(fill_value)) - - layer = default_trtnet().add_slice(input.trt_tensor, - start=trt_starts, - shape=trt_sizes, - stride=trt_strides) - if mode is not None: - layer.mode = mode - - if isinstance(starts, Tensor): - layer.set_input(1, starts.trt_tensor) - - if isinstance(sizes, Tensor): - layer.set_input(2, sizes.trt_tensor) - - if isinstance(strides, Tensor): - layer.set_input(3, strides.trt_tensor) - - if mode is trt.SampleMode.FILL and isinstance(fill_value, Tensor): - layer.set_input(4, fill_value.trt_tensor) - - return _create_tensor(layer.get_output(0), layer) - - -def pad(input: Tensor, - pad: Union[Sequence[int], Tensor], - mode: str = 'constant', - value: Optional[float] = None) -> Tensor: - ''' - Add a pad layer. - - The padding layer adds zero-padding at the start and end of the input tensor. And the - padding size by which to pad some dimensions of input are described starting from the - last dimension and moving forward. - - `[len(pad) / 2]` dimensions of input will be padded. For example, to pad only the last - dimension of the input tensor, then pad has the form [padding_left, padding_right]; to - pad the last 2 dimensions of the input tensor, then use [padding_left, padding_right, - padding_top, padding_bottom]; to pad the last 3 dimensions, use [padding_left, - padding_right, padding_top, padding_bottom, padding_front, padding_back]. - - Parameters: - input : Tensor - The input tensor on which the padding_2d is performed. - pad : sequence of int - An m-elements tuple for padding, where its length m meets the requirement that - m <= 2*input dimensions, and m is even. - mode : str - Only \'constant\' is supported. - value : float - Fill value for 'constant' padding. Default: 0. - - Returns: - The tensor produced by the inserted layer. - ''' - assert mode == "constant", "Only `'constant'` is supported now." - - if isinstance(pad, list) or isinstance(pad, tuple): - assert ( - len(pad) % 2 == 0 and len(pad) <= 2 * input.ndim() - ), "The length of `pad` should be even and less than 2*input.ndim" - pad = constant(np.array(pad).astype(np.int32)).view([-1, 2]) - elif isinstance(pad, Tensor): - pad = pad.flatten() - assert ( - pad.size(0) % 2 == 0 and pad.size(0) <= 2 * input.ndim() - ), "The length of `pad` should be even and less than 2*input.ndim" - pad = pad.cast("int32").view([-1, 2]) - else: - raise NotImplementedError(f"pad type {type(pad)} not supported") - if value is None: - value = 0 - - pad = concat([constant(np.zeros((1, 2), dtype=np.int32)), - pad]) # pre-padding the indices - padding_index = [0] * input.ndim() - padding_index[-(pad.size(0) - 1):] = list(range(pad.size(0) - 1, 0, - -1)) # reverse the indices - pad = index_select(pad, - dim=0, - index=constant(np.array(padding_index, dtype=np.int32))) - pre_padding, post_padding = chunk(pad, chunks=2, dim=1) - start = (pre_padding.flatten() * (-1)).cast('int32') - extend_size = (pre_padding + post_padding).flatten() - size = (extend_size + shape(input)).cast('int32') - layer = default_trtnet().add_slice(input.trt_tensor, - start=[0] * input.ndim(), - shape=[0] * input.ndim(), - stride=[1] * input.ndim()) - layer.mode = trt.SampleMode.FILL - layer.set_input(SliceInputType.start, start.trt_tensor) - layer.set_input(SliceInputType.size, size.trt_tensor) - layer.set_input(SliceInputType.fill_value, - constant_to_tensor_(value, dtype=input.dtype).trt_tensor) - output = _create_tensor(layer.get_output(0), layer) - return output - - -def rand(shape: Tensor, - low: float = 0, - high: float = 1, - dtype: Union[str, trt.DataType] = 'float32') -> Tensor: - ''' - This operation adds a fill layer that generates a random (uniform) tensor with the specified shape and data type. - - Parameters: - shape: Tensor - The shape of the tensor needed to be generated. - low: float - The minimum value (inclusive) of the range used for random. - high: float - The maximum value (inclusive) of the range used for random. - dtype: Union[str, trt.DataType] - The desired data type for the output tensor. - Returns: - The generated random tensor produced by the fill layer. - ''' - # NOTE: DISABLED FOR NOW UNTIL THE FILL LAYER (RANDOM_UNIFORM) in TRT IS FIXED - assert False, "The rand() op is temporarily disabled." - low = constant(fp32_array(low)) - high = constant(fp32_array(high)) - trt_dtype = dtype if isinstance(dtype, - trt.DataType) else str_dtype_to_trt(dtype) - - layer = default_trtnet().add_fill([0], trt.FillOperation.RANDOM_UNIFORM, - trt_dtype) - - layer.set_input(0, shape.trt_tensor) - layer.set_input(1, low.trt_tensor) - layer.set_input(2, high.trt_tensor) - return _create_tensor(layer.get_output(0), layer) - - -def categorical_sample(probs: Tensor, rand_data: Tensor = None) -> Tensor: - ''' - This is a sampling operation and an equivalent of torch.distributions.Categorical.sample() - i.e. given a probability distribution tensor, it samples an index of that tensor. - See: https://pytorch.org/docs/stable/distributions.html#torch.distributions.categorical.Categorical.sample - NOTE: This assumes that the given probabilities are **not** normalized. - - Parameters: - probs: Tensor - A 1-D floating point tensor representing the probability distributions. - rand_data: Tensor (optional) - A random tensor of same shape as `probs` tensor. - If not provided, this function will add a rand() op to generate it and use for sampling. - Returns: - A tensor containing a single index of the `probs` tensor representing the sample. - ''' - probs = probs / sum(probs, dim=-1, keepdim=True) - rand_shape = [] - assert probs.ndim() > 0 - for i in range(probs.ndim() - 1): - rand_shape.append(shape(probs, i)) - rand_shape = concat(rand_shape) - if rand_data is None: - rand_data = rand(rand_shape, low=0, high=1, dtype=probs.dtype) - assert rand_shape == shape(rand_data) - rand_data = expand(unsqueeze(rand_data, -1), shape(probs)) - cum_probs = cumsum(probs, dim=-1) - cmp = cast(cum_probs >= rand_data, probs.dtype) - samples = argmax(cmp, dim=-1) - return samples - - -class Conditional: - ''' - Add an operation to conditionally execute two code paths/subgraphs. - - Usage: - 1. conditional = Conditional(condition) - 2. input_1_ = conditional.add_input(input_1) - ... - input_n_ = conditional.add_input(input_n) - 3. Construct the graph to get true_output_value and false_output_value using input_1_, ..., input_n_ - 4. output = conditional.add_output(true_output_value, false_output_value) - ''' - - def __init__(self, condition: Tensor): - self.layer = default_trtnet().add_if_conditional() - if condition.ndim() > 0: - condition = view(condition, []) - self.layer.set_condition(condition.trt_tensor) - - def add_input(self, input: Tensor) -> Tensor: - in_node = self.layer.add_input(input.trt_tensor) - return _create_tensor(in_node.get_output(0), in_node) - - def add_output(self, true_value: Tensor, false_value: Tensor) -> Tensor: - out_node = self.layer.add_output(true_value.trt_tensor, - false_value.trt_tensor) - return _create_tensor(out_node.get_output(0), out_node) - - -# TODO: support step. -def arange(start: Union[Tensor, int], end: Union[Tensor, int], - dtype: str) -> Tensor: - ''' - Add an operation to fill a 1D tensor. - - The tensor is filled with the values between start and end with a step of 1 - between the different elements. In pseudo-code, it corresponds to a tensor - populated with the values: - - output = Tensor([dtype(ii) for ii in range(start, end, 1)]) - - For example, a call to arange(3, 6, 'int32') will add an operation to the - TensorRT graph that will produce [3, 4, 5] when executed. The call to - arange(2, 5, 'float32') will add a layer to generate [2.0, 3.0, 4.0]. - - This operation is implemented using a tensorrt.IFillLayer in - trt.FillOperation.LINSPACE mode. - - Parameters: - start : Union[Tensor, int] - The starting point of the range. - - end : Union[Tensor, int] - The end point of the range. - - dtype : str - The type of the elements. See _str_to_trt_dtype_dict in _utils.py - for a list of supported types and type names. - - Returns: - The tensor produced by the fill layer. It is a 1D tensor containing - `end-start` elements of type `dtype`. - ''' - res_dtype = str_dtype_to_trt(dtype) - if isinstance(start, int): - assert isinstance(end, int) - array_func = int32_array if res_dtype == trt.int32 else int64_array - start = constant(array_func(start)) - end = constant(array_func(end)) - elif isinstance(start, Tensor): - assert isinstance(end, Tensor) - assert start.dtype == trt.int32 or start.dtype == trt.int64 - assert end.dtype == trt.int32 or end.dtype == trt.int64 - if start.dtype != end.dtype: - if start.dtype == trt.int32: # end == trt.int64 - if res_dtype == trt.int32: - end = cast(end, "int32") - else: - start = cast(start, "int64") - else: # start == trt.int64 and end == trt.int32 - if res_dtype == trt.int32: - start = cast(start, "int32") - else: - end = cast(end, "int64") - else: - raise TypeError("%s is not supported" % type(start)) - - assert start.dtype == end.dtype, f"start type ({start.dtype}) != end type ({end.dtype})" - step = constant_to_tensor_(1, dtype=start.dtype, to_array=True) - - num = end - start - num = num.view([1]).cast(trt.int64) - - layer = default_trtnet().add_fill([0], trt.FillOperation.LINSPACE, - start.dtype) - layer.set_input(0, num.trt_tensor) # rank = 1 - layer.set_input(1, start.trt_tensor) # rank = 0 - layer.set_input(2, step.trt_tensor) # rank = 1 - tensor = _create_tensor(layer.get_output(0), layer) - if tensor.dtype != res_dtype: - tensor = tensor.cast(dtype) - return tensor - - -def expand(input: Tensor, expand_shape: Tensor) -> Tensor: - ''' - Add an operation to expand a tensor. - - The operation expands the input tensor in the singleton dimensions to the - size indicated by the corresponding dimension in the `expand_shape` tensor. - In other words, given an input tensor with dimensions of size 1, those - dimensions will be expanded to the size in `expand_shape`. - - For example, a tensor of shape [4, 3, 1, 3] will be expanded to a tensor of - shape [4, 3, 2, 3] by the layer created using expand(input, [4, 3, 2, 3]). - - The expansion may either replicate the values or be mapped to a view with a - stride of 0 in the expanded dimensions. For example, for a tensor [[3, 2]] of - shape [1, 2], - - expand([[3, 2]], [2, 2]) - - can be used to expand the input to [[3, 2], [3, 2]]. - - This operation is implemented using a tensorrt.ISliceLayer. The current - implementation does not verify that non singleton dimensions are not - shrunk. In other words, for an input of shape [4, 1, 2], - - expand(input, [3, 2, 2]) - - will produce a tensor of shape [3, 2, 2]. That behavior is subject to - change in the future. - - Parameters: - input : Tensor - The input tensor. - - expand_shape : Tensor - The new shape of the expanded tensor. - - Returns: - The tensor produced by the expand layer. - ''' - ndim = input.rank() - layer = default_trtnet().add_slice( - input.trt_tensor, - start=[0 for _ in range(ndim)], - shape=[1 for _ in range(ndim)], # unused dummy value - stride=[1 for _ in range(ndim)] # unused dummy value - ) - - # The stride is either: - # 0 for dimensions of size 1 (i.e. shape(input, i) - 1 == 1 - 1 == 0) or, - # 1 for dimensions of size > 1 since minimum(value >= 1, 1) == 1. - stride_tensor = concat( - [minimum((shape(input, i) - 1), 1) for i in range(ndim)]) - - layer.set_input(2, expand_shape.trt_tensor) - layer.set_input(3, stride_tensor.trt_tensor) - return _create_tensor(layer.get_output(0), layer) - - -def einsum(einsum_eq: str, inputs: Sequence[Tensor]) -> Tensor: - ''' - Add an Einsum operation. - - That operation maps to tensorrt.IEinsumLayer. As explained in the TensorRT - documentation, this layer implements a summation over the elements of the - inputs along dimensions specified by the equation parameter, based on the - Einstein summation convention. The layer can have one or more inputs of - rank >= 0. All the inputs must be of same data type. This layer supports - all TensorRT data types except bool. There is one output tensor of the same - type as the input tensors. The shape of output tensor is determined by the - equation. - - The equation specifies ASCII lower-case letters for each dimension in the - inputs in the same order as the dimensions, separated by comma for each - input. The dimensions labeled with the same subscript must match or be - able to be broadcasted. Repeated subscript labels in one input take the diagonal. - Repeating a label across multiple inputs means that those axes will be - multiplied. Omitting a label from the output means values along those axes - will be summed. In implicit mode, the indices which appear once in the - expression will be part of the output in increasing alphabetical order. In - explicit mode, the output can be controlled by specifying output subscript - labels by adding an arrow (‘->’) followed by subscripts for the output. For - example, “ij,jk->ik” is equivalent to “ij,jk”. Ellipsis (‘…’) can be used - in place of subscripts to broadcast the dimensions. See the TensorRT - Developer Guide for more details on equation syntax. - - Many common operations can be expressed using the Einsum equation. For - example: - Matrix Transpose: ij->ji - Sum: ij-> Matrix-Matrix - Multiplication: ik,kj->ij - Dot Product: i,i-> - Matrix-Vector Multiplication: ik,k->i - Batch Matrix Multiplication: ijk,ikl->ijl - Batch Diagonal: …ii->…i - - Note that TensorRT does not support ellipsis or diagonal operations so, - neither, does TensorRT-LLM. - - Parameters: - einsum_eq : str - The Einsum equation. - - inputs: Sequence[Tensor] - The sequence of inputs consumed by the Einsum operation. - - Returns: - The tensor produced by the Einsum operation. - ''' - layer = default_trtnet().add_einsum([i.trt_tensor for i in inputs], - einsum_eq) - return _create_tensor(layer.get_output(0), layer) - - -def permute(input: Tensor, dims: Sequence[int]) -> Tensor: - ''' - Add an operation to permute the dimensions of a tensor. - - The dimensions of the input tensor are permuted according to the sequence - of dimensions in 'dims'. That operation maps to tensorrt.IShuffleLayer where - the second transposition is described by the indices in 'dims'. - - Given a tensor of rank N, the result of the permutation is a tensor of rank - N in which the i-th input dimension maps to the dims[i]-th dimension. - - For example, permute(input, [1, 0]) will transpose a 2D tensor by permuting - the rows and columns. - - Parameters: - input : Tensor - The input tensor to permute. - - dims : Sequence[int] - The description of the permutation. - - Returns: - The tensor produced by the permutation layer. - ''' - dims = dim_resolve_negative(tuple(dims), input.ndim()) - layer = default_trtnet().add_shuffle(input.trt_tensor) - layer.second_transpose = dims - return _create_tensor(layer.get_output(0), layer) - - -def transpose(input: Tensor, dim0: int, dim1: int) -> Tensor: - ''' - Add an operation to transpose two dimensions of a tensor. - - That operation produces a tensor in which the dimensions 'dim0' and 'dim1' - are permuted. The other dimensions, if the rank of the tensor is greater - than 2, remain untouched. - - That function is a helper built on the 'functional.permute' function. - - Parameters: - input : Tensor - The input tensor to transpose. - - dim0 : int - The first dimension to transpose. - - dim1 : int - The second dimension to transpose. - - Returns: - The tensor produced by the permutation layer. - ''' - permutation = list(range(input.ndim())) - permutation[dim0] = dim1 - permutation[dim1] = dim0 - - return permute(input, permutation) - - -def view(input: Tensor, - shape: Union[Tensor, Sequence[int]], - zero_is_placeholder: bool = True) -> Tensor: - ''' - Add an operation to create a view of a tensor. - - That operation adds a tensorrt.IShuffleLayer to the network. If the 'shape' - parameter is a Tensor, that view is dynamic. Otherwise, it is a static - view. - - Note that TensorRT limits the number of inferred dimensions to 1. It means - that the shape sequence or tensor cannot contain more than one -1. This - function enforces that constraint and will assert if it is not respected. - - Parameters: - input : Tensor - The input tensor to transpose. - - shape : Union[Tensor, Sequence[int]] - The shape of the new tensor. - - zero_is_placeholder : bool - When that parameter is True, the 0s in 'shape' are replaced by the - sizes of the corresponding dimensions from the 'input'. Otherwise, - the dimensions corresponding to 0s are shrunk. - - Returns: - The tensor produced by the view/shuffle layer. - ''' - - # TensorRT demands that at most one dimension is permitted to be specified as -1 - def assert_no_more_than_one_inferred_dim(list): - inferred_dim_list = [i for i in list if i == -1] - assert len(inferred_dim_list) <= 1 - - layer = default_trtnet().add_shuffle(input.trt_tensor) - layer.zero_is_placeholder = zero_is_placeholder - if isinstance(shape, Tensor): - assert_no_more_than_one_inferred_dim(shape.shape) - layer.set_input(1, shape.trt_tensor) - elif isinstance(shape, (list, tuple)): - assert_no_more_than_one_inferred_dim(shape) - layer.reshape_dims = tuple(shape) - else: - raise TypeError("%s is not supported" % type(shape)) - return _create_tensor(layer.get_output(0), layer) - - -def flatten(input: Tensor, start_dim: int = 0, end_dim: int = -1): - ''' - Flattens input by reshaping it into a one-dimensional tensor. - - If start_dim or end_dim are passed, only dimensions starting with start_dim and - ending with end_dim are flattened. The order of elements in input is unchanged. - - Parameters: - input : Tensor - The input tensor to flatten. - - start_dim : int - The first dim to flatten. - - end_dim : int - The last dim to flatten. - - Returns: - The tensor produced by the flatten layer. - - ''' - shape = input.shape - ndim = input.ndim() - if start_dim < 0: start_dim += ndim - if end_dim < 0: end_dim += ndim - new_shape = list() - for i in range(start_dim): - new_shape.append(shape[i]) - if end_dim - start_dim >= 0: - flat_dim = 1 - for i in range(start_dim, end_dim + 1): - flat_dim *= shape[i] - new_shape.append(flat_dim) - for i in range(end_dim + 1, ndim): - new_shape.append(shape[i]) - return view(input, new_shape) - - -def expand_dims(input: Tensor, - dim: Union[int, Sequence[int]], - shape_cast_dtype=None) -> Tensor: - ''' - Add an operation to expand the tensor shape with singleton dimensions. - - That function adds a tensorrt.IShuffleLayer to the network. Given an 'input' - of rank N and a sequence of M dimensions, the output tensor produced by - this operation (when executed by TensorRT) will have a rank of N+M. Singleton - dimensions will be inserted at the different positions in 'dim'. - - The pseudo-code for that operation is: - - new_shape, ii = [], 0 - for jj in range(input.rank() + len(dim)): - new_shape.append(1 if jj in dims else input.shape[ii++]) - - For example, for a tensor of shape [3, 4, 1, 5] - - expand_dims(input, [0, 2]) - - will produce a tensor of shape [1, 3, 1, 4, 1, 5]. - - Parameters: - input : Tensor - The input tensor to expand. - - dim : Union[int, Sequence[int]] - The positions in the output tensor where to insert singleton - dimensions. - - Returns: - The tensor produced by the shuffle layer. - ''' - if isinstance(dim, int): - dim = (dim, ) - - out_ndim = len(dim) + input.ndim() - - input_shape = shape(input, cast_to_dtype=shape_cast_dtype) - out_shapes = [] - j = 0 - for i in range(out_ndim): - if i in dim: - out_shapes.append(1) - else: - out_shapes.append(gather(input_shape, 0, j)) - j = j + 1 - - out_shape = concat(out_shapes) - - return view(input, out_shape, zero_is_placeholder=False) - - -# NOTE: Jointly added with Apple -def squeeze(input: Tensor, - dim: Optional[Union[int, Sequence[int]]] = None, - zero_is_placeholder: bool = False): - ''' - Add an operation to remove singleton dimensions of a tensor. - - This functions creates an operation that removes singleton dimension - (dimension of size 1) at positions 'dim' in the input tensor. It works with - negative values for the 'dim'. - - For example, for a tensor 'input' of shape [1, 4, 1, 4]: - - squeeze(input, 0) will produce an output of shape [4, 1, 4], - squeeze(input, 2) will produce an output of shape [1, 4, 4], - squeeze(input, [0, 2]) will produce an output of shape [4, 4], - squeeze(input, [-2]) will produce an output of shape [1, 4, 4], - - Parameters: - input : Tensor - The input tensor for which the singleton dimensions will be removed. - - dim : Union[int, Sequence[int]] - The index of the singleton dimensions in the input tensor. - - Returns: - The tensor produced by the layer. - ''' - if dim is None: - dim = list(range(input.ndim())) - if isinstance(dim, int): - dim = (dim, ) - dim = dim_resolve_negative(dim, input.ndim()) - - new_shape = [] - for i, s in enumerate(input.shape): - if s == 1 and i in dim: - continue - new_shape.append(shape(input, i)) - - new_shape = concat(new_shape) if len(new_shape) > 0 else [] - input = input.view(new_shape, zero_is_placeholder=zero_is_placeholder) - return input - - -def unsqueeze(input: Tensor, axis: int): - ''' - Add an operation to insert a singleton dimension to a tensor. - - That functions creates an operation that insert a singleton dimension - (dimension of size 1) at position 'axis' in the output tensor. It works with - negative values for the 'axis'. - - For example, for a tensor 'input' of shape [4, 4]: - - unsqueeze(input, 0) will produce an output of shape [1, 4, 4], - unsqueeze(input, 1) will produce an output of shape [4, 1, 4], - unsqueeze(input, -1) will produce an output of shape [4, 4, 1], - unsqueeze(input, -2) will produce an output of shape [4, 1, 4], - - Parameters: - input : Tensor - The input tensor to expand with a singleton dimension. - - axis : int - The index of the singleton dimension in the output tensor. - - Returns: - The tensor produced by the layer. - ''' - if axis < 0: - axis = axis + input.ndim() + 1 - - return expand_dims(input, axis) - - -def stack(inputs: Sequence[Tensor], dim: int = 0) -> Tensor: - ''' - Add an operation to contact input tensors along a new dimension. - - The function creates an operation that creates a new dim for all the - input tensors and then concatenates them along that new dim. -. - - All the tensors in 'inputs' must have the same shape. - - for ii in range(inputs[0].rank()): - assert all(inp.shape[ii] == inputs[0].shape[ii] for inp in inputs) - - The shape of the output tensor is defined as: - - output.rank() = inputs[0].rank() + 1 - - output.shape[dim] = len(inputs) - - for ii in range(inputs[0].rank()): - if ii < dim: - output.shape[ii] = inputs[0].shape[ii] - else: - output.shape[ii+1] = inputs[0].shape[ii] - - For example, given a sequence of two 2D tensors [[0, 1], [2, 3]] and - [[4, 5], [6, 7]] both of shape [2, 2], - - stack(inputs, 0) - - will produce [[[0, 1], [2, 3]], [[4, 5], [6, 7]]] of shape [2, 2, 2] and - - stack(inputs, 1) - - will produce [[[0, 1], [4, 5]], [[2, 3], [6, 7]]] of shape [2, 2, 2]. - - Parameters: - inputs : Sequence[Tensor] - The sequence of tensors to stack. - - dim : int - The dimension in which the stack is performed. - - Returns: - A tensor that contains the input tensors stacked along a new dimension. - ''' - return concat([unsqueeze(inp, axis=dim) for inp in inputs], dim=dim) - - -def expand_dims_like(left: Union[Tensor, int, float], right: Tensor) -> Tensor: - ''' - Add an operation to expand the first tensor to the same rank as the second - tensor. - - That function takes a first tensor. It also accepts an integer or a float, - in which case it creates a constant tensor from it. In both cases, the rank - of that first tensor is compared to the rank of the second tensor. If they - are of the same rank, the first tensor is returned. Otherwise, the first - tensor is expanded on the left to match the rank of the second tensor. - - Note that the shapes do not have to match, only the rank is considered in - that function. - - For example, for a pair of tensors of shapes [3, 4] and [4, 3, 2], the - first tensor will be expanded to a tensor of rank 3 and shape [1, 3, 4]. - - Parameters: - left : Union[Tensor, int, float] - The first tensor to expand. When a scalar value is provided as a - parameter, that function first creates a tensor before expanding it - (if needed). - - right : Tensor - The reference tensor to match. - - Returns: - The tensor produced by the shuffle layer. - ''' - if isinstance(left, int): - left = constant(dims_array([left])) - elif isinstance(left, float): - if isinstance(right, Tensor) and right.dtype == trt.DataType.HALF: - left = constant(fp16_array([left])) - else: - left = constant(fp32_array([left])) - left_ndim = left.ndim() - right_ndim = right.ndim() - if right_ndim > left_ndim: - new_ndim = list(range(right_ndim - left_ndim)) - return expand_dims(left, new_ndim) - return left - - -# If dim is None, return a 1-D TensorRT LLM tensor of the size -# If dim is not None, return a 0-D TensorRT LLM tensor of the dimension size -def shape(input: Tensor, - dim: Optional[int] = None, - cast_to_dtype: Optional[Union[str, trt.DataType]] = None, - clip_before_cast: Sequence[int] = None) -> Tensor: - ''' - Add an operation to create a shape tensor. - - The shape tensor can either be the shape of the input tensor when the - parameter dim is None or a scalar (tensor of rank 0) that corresponds to - the size of dim-th dimension. - - Parameters: - input : Tensor - The input tensor from which we want to extract the shape or the - size in one dimension. - - dim : Optional[int] - The dimension from which to extract the size. If it is None, the - entire shape of the input tensor is returned. - - Returns: - A tensor that contains the shape of the input tensor (if 'dim' is None) - or the size in the dimension 'dim' of the input tensor. If 'dim' is - 'None', that tensor has the same rank as the input tensor, otherwise - its rank is 0. - ''' - layer = default_trtnet().add_shape(input.trt_tensor) - res = _create_tensor(layer.get_output(0), layer) - if cast_to_dtype is not None: - if clip_before_cast is not None and (cast_to_dtype == 'int32' - or cast_to_dtype == trt.int32): - assert len( - clip_before_cast - ) == 2, f"This parameter only expects a tuple of 2 integers (lower, upper) but got {clip_before_cast}" - res = int_clip(res, clip_before_cast[0], clip_before_cast[1]) - res = cast(res, cast_to_dtype) - - if dim is None: - return res - - return gather(res, dim=0, indices=dim).view([]) - - -def gather(input: Tensor, dim: int, indices: Union[Tensor, int]) -> Tensor: - ''' - Add an operation to gather elements from a tensor. - - That function implements the GatherElements operator from the ONNX - specification as described in - - https://github.com/onnx/onnx/blob/main/docs/Operators.md#GatherElements - - The input and indices arguments must have the same rank >= 1. The operation - will produce a tensor with the same shape as the indices tensor. The axis - is the dimension to gather on. - - As shown in the ONNX description, for a 3D tensor, the output is: - - out[i][j][k] = input[indices[i][j][k]][j][k] if axis = 0, - out[i][j][k] = input[i][indices[i][j][k]][k] if axis = 1, - out[i][j][k] = input[i][j][indices[i][j][k]] if axis = 2. - - For example, - - gather([[4, 2], [5, 3]], 0, [[1, 0], [0, 1]]) - - will produce [[5, 2], [4, 3]]. - - gather([[1, 2, 3], [4, 5, 6], 1, [[1], [0]]) - - will produce [[2], [4]]. See the ONNX documentation for more examples. - - That operation maps to the TensorRT IGatherLayer. - - Parameters: - input : Tensor - The input tensor to gather elements from. - - dim : int - The dimension to gather on. - - indices : Union[Tensor, int] - The positions in the 'dim' dimension to gather from. - - Returns: - The tensor containing the gathered elements. It has the same shape as - the indices tensor. - ''' - if isinstance(indices, int): - indices = constant(int32_array([indices])) - - # The input and indices tensors must have the same rank. - assert input.rank() == indices.rank() - - layer = default_trtnet().add_gather_v2(input.trt_tensor, - indices.trt_tensor, - mode=trt.GatherMode.ELEMENT) - - if dim < 0: - dim = input.ndim() + dim - layer.axis = dim - return _create_tensor(layer.get_output(0), layer) - - -def select(input: Tensor, dim: int, index: Union[Tensor, int]) -> Tensor: - ''' - Add an operation to select a slice of elements from a tensor. - - Given an input tensor, that function creates an operation that selects the - index-th slice of elements in the dimension 'dim' to create a new tensor. - The output tensor has a shape in which the input dimension 'dim' is - removed. - - The 'index' can either be an integer or a 1D tensor containing a single - element. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - select(input, 0, 1) - - will create a tensor of shape [3] that contains the [2, 1, 2]. - - Regarding the shape of the output tensor, the dimension 'dim' is removed. - It means that for a tensor of shape [4, 2, 6, 3], - - select(input, 2, 4) - - will select the 5th slice (index == 4) from the 3rd dimension (dim == 2) - and return a tensor of shape [4, 2, 3] (i.e. the 3rd dimension is removed). - - That operation maps to the TensorRT IGatherLayer. - - Parameters: - input : Tensor - The input tensor to select from. - - dim : int - The dimension to select from. - - index : Union[Tensor, int] - The index of the slice in the 'dim' dimension to select. - - Returns: - The tensor containing the selected slice. - ''' - if isinstance(index, int): - index = constant(int32_array([index])) - assert index.rank() == 1 and index.size( - 0) == 1, f"index should have rank 1, got {index.rank()}" - - new_shape = [] - for i in range(input.rank()): - if i != dim: - new_shape.append(shape(input, i)) - - layer = default_trtnet().add_gather(input.trt_tensor, index.trt_tensor, dim) - return _create_tensor(layer.get_output(0), layer).view(concat(new_shape)) - - -def index_select(input: Tensor, dim: int, index: Tensor) -> Tensor: - ''' - Add an operation to select slices of elements from a tensor. - - Given an input tensor, that function creates an operation that selects the - slices of elements in the dimension 'dim' at the indices listed in 'index' - to create a new tensor. The output tensor has the same rank as the input - tensor. - - The 'index' is a tensor of rank 1. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - index_select(input, 0, [0, 1]) - - will create a tensor of shape [2, 3] that contains the [[4, 2, 5], [2, 1, 2]]. - - Regarding the shape of the output tensor, the dimension 'dim' has the same - size as the 'index' tensor. It means that for a input tensor of shape [4, 2, 6, 3], - - index_select(input, 2, [1, 4]) - - will select the 2nd and 5th slices (index == 1 or 4) from the 3rd dimension - (dim == 2) and return a tensor of shape [4, 2, 2, 3] (i.e. the 3rd - dimension is shrunk to 2). - - Note that this operation can also be used to expand a tensor in the 'dim' - dimension, for example, on input [[0, 1], [2, 3]], - - index_select(input, 1, [0, 0, 0]) - - will produce a tensor of shape [2, 3] containing [[0, 0, 0], [2, 2, 2]]. - - That operation maps to the TensorRT IGatherLayer. - - Parameters: - input : Tensor - The input tensor to select from. - - dim : int - The dimension to select from. - - index : Tensor - The indices of the slices in the 'dim' dimension to select. - - Returns: - The tensor containing the selected slices. - ''' - assert index.rank() == 1, f"index should have rank 1, got {index.rank()}" - - new_shape = [] - for i in range(input.rank()): - if i != dim: - new_shape.append(shape(input, i)) - else: - new_shape.append(shape(index, 0)) - - layer = default_trtnet().add_gather(input.trt_tensor, index.trt_tensor, dim) - return _create_tensor(layer.get_output(0), layer).view(concat(new_shape)) - - -# NOTE: Jointly added with Apple -def scatter(input: Tensor, dim: int, indices: Tensor, - updates: Tensor) -> Tensor: - ''' - This operation adds a layer that creates an output tensor by element-wise - copying values from the input tensor and then updating values by the given - `indices` and `updates` tensors. - For a 2D input tensor, it first copies the input to output, - then updates the output tensor like the following for each entry in `updates`: - output[indices[i][j]][j] = updates[i][j] if dim=0 - output[i][indices[i][j]] = updates[i][j] if dim=1 - If the `input` tensor is [[1, 2, 3], [4, 5, 6]], - the indices tensor is [[1, 2], [0, 1]], - the updates tensor is [[-1, -2], [-3, -4]], and dim=1 - the output tensor will be [[1, -1, -2], [-3, -4, 6]]. - Parameters: - input: Tensor - The input data that needs to be updated. - dim: int - The axis on which the scatter is to be performed. - indices: Tensor - An integer tensor of the same rank as input that indicates the positions to be updated. - updates: Tensor - A data tensor of same shape as the `indices` tensor that contains the update values. - Returns: - A tensor created by the element-wise scatter layer. - ''' - layer = default_trtnet().add_scatter(input.trt_tensor, - indices.trt_tensor, - updates.trt_tensor, - mode=trt.ScatterMode.ELEMENT) - layer.axis = dim - return _create_tensor(layer.get_output(0), layer) - - -def gather_nd(input: Tensor, indices: Tensor, batch_dims: int = 1) -> Tensor: - ''' - Adds a layer that performs a gather with some element-wise dimensions. - See: https://onnx.ai/onnx/operators/onnx__GatherND.html - The gather is performed on dim=batch_dims. - - Parameters: - input: Tensor - The tensor on which the gather operation is performed. - indices: Tensor - The tensor that indicates which entries to be gathered. - batch_dims: int - The number of first dimensions that should be skipped before gather starts. - Returns: - A tensor created by the gather layer with GatherMode.ND. - ''' - gather_layer = default_trtnet().add_gather_v2(input.trt_tensor, - indices.trt_tensor, - mode=trt.GatherMode.ND) - gather_layer.num_elementwise_dims = batch_dims - return _create_tensor(gather_layer.get_output(0), gather_layer) - - -def nonzero(input: Tensor) -> Tensor: - ''' - Adds a layer that finds the indices of non-zero values of the input tensor. - - Parameters: - input: Tensor - The input tensor for which we need to find the indices of non-zero values. - Returns: - A tensor of shape [D, C] where D is the number of dimensions of `input` and - C is the number of non-zero values in it. - Each column of this 2D tensor represents the index tuple for each non-zero value. - ''' - non_zero_layer = default_trtnet().add_non_zero(input.trt_tensor) - return _create_tensor(non_zero_layer.get_output(0), non_zero_layer) - - -def masked_select(input: Tensor, mask: Tensor) -> Tensor: - ''' - Add an operation to select elements from a tensor according to a boolean - mask tensor. - - Given an input tensor, that function creates an operation that selects - elements at the indices indicated by the boolean mask tensor to create - a new tensor. The output tensor is a 1-D tensor. - - The input tensor must have rank >= 1. The shapes of the input tensor and - the mask tensor don’t need to match, but they must be able to be broadcasted. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - masked_select(input, [[True, False, True], [False, True, False], [True, False, True]]) - - will create a tensor of shape [5] that contains the [4, 5, 1, 4, 1]. - - masked_select(input, [[True], [False], [True]]) - - will create a tensor of shape [6] that contains the [4, 2, 5, 4, 7, 1]. - - masked_select(input, [[False, False, True]]) - - will create a tensor of shape [3] that contains the [5, 2, 1]. - - masked_select(input, [False]) - - will create a tensor of shape [0] which is empty. - - That operation is implemented by NonZero, Shuffle and GatherV2 layers - in TensorRT. - - Parameters: - input : Tensor - The input tensor to select from. - - mask : Tensor - The boolean mask tensor that indicates elements to select. - - Returns: - The 1-D tensor containing the selected elements. - ''' - assert input.rank() >= 1, "input should have rank >= 1" - input, mask = broadcast_helper(input, mask) - expanded_mask = expand(mask, shape(input)) - - non_zero_layer = default_trtnet().add_non_zero(expanded_mask.trt_tensor) - - shuffle_layer = default_trtnet().add_shuffle(non_zero_layer.get_output(0)) - shuffle_layer.second_transpose = (1, 0) - - gather_layer = default_trtnet().add_gather_v2(input.trt_tensor, - shuffle_layer.get_output(0), - mode=trt.GatherMode.ND) - return _create_tensor(gather_layer.get_output(0), gather_layer) - - -def cumsum(input: Tensor, dim: int, prefer_plugin: bool = True) -> Tensor: - ''' - Add an operation to calculate inclusive cumulative sum of elements of - a tensor in a given dimension. - - Given an input tensor, that function creates an operation that calculates - inclusive cumulative sum of elements in the dimension 'dim' to create - a new tensor. The output tensor has the same shape as the input tensor. - - The input tensor must have rank >= 1. The 'dim' must be valid, and negative - value is supported. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - cumsum(input, 0) - - will produce [[4, 2, 5], [6, 3, 7], [10, 10, 8]]. - - cumsum(input, 1) - - will produce [[4, 6, 11], [2, 3, 5], [4, 11, 12]]. - - That operation is implemented by TensorRT ILoopLayer. - - Parameters: - input : Tensor - The input tensor to calculate the inclusive cumulative sum. - - dim : int - The dimension to calculate the inclusive cumulative sum. Negative - value is supported. - - prefer_plugin : bool - Whether to use the cumsumLastDim plugin if dim is last dim. - - Returns: - The tensor containing the inclusive cumulative sum of input. - ''' - assert input.rank() >= 1, "input should have rank >= 1" - assert dim < input.rank() and dim >= -input.rank( - ), f"dim should be in [{-input.rank()}, {input.rank()}) when input have rank {input.rank()}" - - dim = dim_resolve_negative(dim, input.ndim())[0] - - if dim == input.ndim() - 1: - if prefer_plugin: - last_dim = input.size(-1) - if last_dim == -1: # dynamic? - last_dim = shape(input, -1) - old_shape = shape(input) - if input.ndim() == 1: - input_2d = unsqueeze( - input, 0) # special handling of rank-1 dynamic tensor - elif input.ndim() != 2: - input_2d = input.view(concat([-1, last_dim]), - zero_is_placeholder=False) - else: - input_2d = input - cumsum_last_dim_plg_creator = trt.get_plugin_registry( - ).get_plugin_creator('CumsumLastDim', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert cumsum_last_dim_plg_creator is not None - input_length = trt.PluginField( - "input_length", np.array(input_2d.size(-1), dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField("type_id", - np.array([int(input_2d.dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([input_length, pf_type]) - cumsum_last_dim_plug = cumsum_last_dim_plg_creator.create_plugin( - "cumsum_last_dim", pfc) - plug_inputs = [input_2d] - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, - cumsum_last_dim_plug) - _add_plugin_info(layer, cumsum_last_dim_plg_creator, - "cumsum_last_dim", pfc) - output = _create_tensor(layer.get_output(0), layer) - output = output.view(old_shape, zero_is_placeholder=False) - return output - else: - # credit to Apple - reduction_length = shape(input, -1) - reduction_range = arange(constant_to_tensor_(0, - dtype='int64', - to_array=False), - reduction_length, - dtype='int64') - lower_triangle = cast(unsqueeze(reduction_range, 0) - <= unsqueeze(reduction_range, 1), - dtype=input.dtype) - output = sum(unsqueeze(input, -2) * lower_triangle, dim=-1) - return output - else: - slice_shape = [] - for i in range(input.ndim()): - if i != dim: - slice_shape.append(shape(input, i)) - - zero_tensor = constant_to_tensor_(0, input.dtype, False) - if len(slice_shape) > 0: - zero_tensor = expand_dims(zero_tensor, - [i for i in range(len(slice_shape))]) - slice_shape = concat(slice_shape) - zero_tensor = expand(zero_tensor, slice_shape) - - loop_layer = default_trtnet().add_loop() - trip_limit = shape(input, dim).trt_tensor - loop_layer.add_trip_limit(trip_limit, trt.TripLimit.COUNT) - - iterator_layer = loop_layer.add_iterator(input.trt_tensor, dim) - cur_slice = iterator_layer.get_output(0) - - running_sum_layer = loop_layer.add_recurrence(zero_tensor.trt_tensor) - running_sum = running_sum_layer.get_output(0) - - cur_sum_layer = default_trtnet().add_elementwise( - cur_slice, running_sum, trt.ElementWiseOperation.SUM) - cur_sum = cur_sum_layer.get_output(0) - running_sum_layer.set_input(1, cur_sum) - - loop_output_layer = loop_layer.add_loop_output( - cur_sum, trt.LoopOutput.CONCATENATE, dim) - loop_output_layer.set_input(1, trip_limit) - return _create_tensor(loop_output_layer.get_output(0), - loop_output_layer) - - -def masked_scatter(input: Tensor, mask: Tensor, source: Tensor) -> Tensor: - ''' - Add the masked_scatter base on PyTorch definition. - - See https://pytorch.org/docs/stable/generated/torch.Tensor.masked_scatter_.html#torch-tensor-masked-scatter for a - description of that function. - - Parameters: - input : Tensor - The input tensor. - - mask : Tensor - The boolean mask tensor that indicates elements to select. - - source: Tensor - The tensor to copy from - Returns: - The tensor containing the source tensor selected by mask. - - ''' - assert input.rank() >= 1, "input should have rank >= 1" - input, mask = broadcast_helper(input, mask) - expanded_mask = expand(mask, shape(input)) - - non_zero_layer = default_trtnet().add_non_zero(expanded_mask.trt_tensor) - - shuffle_layer = default_trtnet().add_shuffle(non_zero_layer.get_output(0)) - shuffle_layer.second_transpose = (1, 0) - source = source.view([-1]) - - scatter_layer = default_trtnet().add_scatter(input.trt_tensor, - shuffle_layer.get_output(0), - source.trt_tensor, - mode=trt.ScatterMode.ND) - - return _create_tensor(scatter_layer.get_output(0), scatter_layer) - - -def concat(inputs: Sequence[Union[Tensor, int]], dim: int = 0) -> Tensor: - ''' - Add an operation to concatenate tensors. - - The function creates an operation that concatenates the tensors from the - sequence 'inputs'. The concatenation is done along the dimension 'dim'. - - All the tensors in 'inputs' must have the same shape expect for the - dimension 'dim'. - - for ii in range(inputs[0].rank()): - assert (ii == dim) or all(inp.shape[ii] == inputs[0].shape[ii] for inp in inputs) - - The shape of the output tensor is defined as: - - for ii in range(inputs[0].rank()): - # Same size as all the inputs in dimension ii != dim. - output.shape[ii] = inputs[0].shape[ii] - - # Sum of the sizes in the different inputs in dimension 'dim'. - if ii == dim: - for jj in range(1, len(inputs)): - output.shape[ii] += inputs[jj].shape[ii] - - For example, given a sequence of two 2D tensors [[0, 1], [2, 3]] and - [[4, 5], [6, 7]] both of shape [2, 2], - - concat(inputs, 0) - - will produce [[0, 1], [2, 3], [4, 5], [6, 7]] of shape [4, 2] and - - concat(inputs, 1) - - will produce [[0, 1, 4, 5], [2, 3, 6, 7]] of shape [2, 4]. - - Parameters: - inputs : Sequence[Union[Tensor, int]] - The sequence of tensors to concatenate. For integers, that function - creates constant tensors. - - dim : int - The dimension in which the concatenation is performed. - - Returns: - A tensor that contains the concatenation of the tensors. - ''' - assert len( - inputs - ) > 0, f"Number of inputs ({len(inputs)}) to the concatenation layer must be > 0." - tmp = [] - inputs = constants_to_tensors_(*inputs) - for i in inputs: - if i.rank() == 0: - tmp.append(i.view([1])) - else: - tmp.append(i) - - layer = default_trtnet().add_concatenation([i.trt_tensor for i in tmp]) - layer.axis = dim_resolve_negative(dim, tmp[0].ndim())[0] - return _create_tensor(layer.get_output(0), layer) - - -def softmax(input: Tensor, dim: Optional[int] = None) -> Tensor: - ''' - Add an operation to compute softmax on a tensor. - - That operation computes the softmax on the input tensor in the dimension - 'dim' if specified. Otherwise, it is applied on the last dimension. - - It inserts a ISoftmaxLayer to the TensorRT graph. - - Parameters: - input : Tensor - The input tensor on which to apply softmax. - - dim : Optional[int] - The dimension used to apply softmax. - - Returns: - The output tensor of the softmax layer. - ''' - if dim is None: - dim = input.ndim() - 1 - if dim < 0: - dim = input.ndim() + dim - axes = dim_to_trt_axes(dim) - - layer = default_trtnet().add_softmax(input.trt_tensor) - layer.axes = axes - - return _create_tensor(layer.get_output(0), layer) - - -def _lookup_plugin(input: Tensor, weight: Tensor, rank: int, - per_token_scale: Tensor) -> Tensor: - ''' - Add an operation to perform lookup in a tensor. - - That operation performs the lookup needed by embedding layers. Given a - 'weight' tensor of shape [rows, cols], it produces a tensor of shape - [inputs.size(0), cols] where the ith row corresponds to the input[i] row in - the weight tensor. - - It inserts a IPluginV2Layer. - - Parameters: - input : Tensor - The input tensor contains the indices to perform the lookup. - - weight : Tensor - The table to gather from. - - rank : int - The mpi rank. - - Returns: - The output tensor of the lookup layer. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Lookup', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = per_token_scale.dtype - pf_type = trt.PluginField("type_id", np.array([int(p_dtype)], np.int32), - trt.PluginFieldType.INT32) - - rank = trt.PluginField("rank", np.array([int(rank)], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, rank]) - lookup_plug = plg_creator.create_plugin("lookup", pfc) - plug_inputs = [input.trt_tensor, weight.trt_tensor] - if per_token_scale is not None: - plug_inputs.append(per_token_scale.trt_tensor) - weight.trt_tensor.set_dynamic_range(-127, 127) - layer = default_trtnet().add_plugin_v2(plug_inputs, lookup_plug) - _add_plugin_info(layer, plg_creator, "lookup", pfc) - return _create_tensor(layer.get_output(0), layer) - - -def embedding(input: Tensor, - weight: Tensor, - tp_size=1, - tp_group=None, - sharding_dim=0, - tp_rank=None, - per_token_scale=None, - padding=None) -> Tensor: - ''' - Add an operation to perform embedding lookup. - - That operation performs the embedding lookup. The 'input' tensor contains - the identifiers of the rows of 'weight' to gather. - - 1. Distribute the embedding lookup table over multiple GPU - When 'tp_size' is greater than 1 and the 'tp_group' is defined, this - embedding lookup is distributed among multiple GPUs. - - When 'sharding_dim==0', each GPU stores a subset of the rows of the embedding - table rows(that number of rows per GPU is given by weights.shape[0] and the offset to - the 1st row stored on the GPU is given by rank * weights.shape[0]). Each - parallel rank will query all the indices and set 0s for the weights that - are not stored on the associated GPU. To compute the final result, a - parallel all-reduce operation is added to the TensorRT graph. That lookup - can be performed using either the plugin or the operators TensorRT support. - - When'sharding_dim==1', each GPU stores a subset of the embedding table's columns. - Each rank can obtain a portion of the embedding results. - Then the embedding is collected using the all-gather operation. - Related transposition operations are also used to obtain the final results. - - 2. Store embedding lookup table as a whole - When 'tp_size' is not greater than 1, the embedding lookup table will not - be divided. In this case, when the default_net().plugin_config.lookup_plugin is set, - the operation is implemented using a plugin (without the all-reduce operation). - Otherwise, this operation is implemented using the standard IGatherLayer in TensorRT. - - Parameters: - input : Tensor - The input tensor the contains the indices to perform the lookup. - - weight : Tensor - The table to gather from. - - tp_size : int - The number of GPUs collaborating to perform that embedding. - - tg_group : Optional[List[int]] - The group of world ranks participating in the all-reduce when - tp_size > 1. - - sharding_dim : int - sharding_dim = 0 means that we shard the embedding table in vocab dim; - sharding_dim = 1 means that we shard the embedding table in embedding dim. - - tp_rank : int - The tensor parallelism rank. Used to calculate offset in TP on vocab dim. - - padding: Tensor - Additional padding added to the end of the embedding table before feeding into gather op. - - Returns: - The tensor produced by the embedding lookup layer. - ''' - - # Per token scale is only supported by lookup plugin so if per_token_scale is not None, we must use lookup plugin - # Otherwise, we prefer to use ootb - use_lookup_plugin = per_token_scale is not None - - if padding is not None: - padded_weight = concat([weight, padding], dim=0) - else: - padded_weight = weight - - # Distribute embedding lookup table across multiple GPU - if tp_size > 1 and tp_group is not None: - if sharding_dim == 0: # TP on vocab_size dimension - if tp_rank is None: - raise ValueError( - "Rank cannot be none for tensor parallelism on vocab dim") - - if use_lookup_plugin: - x = _lookup_plugin(input, weight, tp_rank, per_token_scale) - x = allreduce(x, tp_group) - else: - shape_weight = shape(weight) - vocab_size = slice(shape_weight, starts=[0], sizes=[1]) - tmp_input = input - vocab_size * tp_rank - - # Identify the valid indices - is_qualified = op_and(tmp_input >= 0, tmp_input < vocab_size) - is_qualified_expand = expand_dims(is_qualified, - [is_qualified.ndim()]) - - # Replace the invalid ones to zero - placeholder_input = where(is_qualified, tmp_input, 0) - - # Get the temporal results - layer = default_trtnet().add_gather( - padded_weight.trt_tensor, placeholder_input.trt_tensor, 0) - tmp_output = _create_tensor(layer.get_output(0), layer) - - # Set zero for invalid results - placeholder_tmp = cast(is_qualified_expand, tmp_output.dtype) - placeholder = placeholder_tmp - placeholder_tmp - x = where(is_qualified_expand, tmp_output, placeholder) - - # Use all reduce to collect the results - x = allreduce(x, tp_group) - - elif sharding_dim == 1: # TP on hidden dimension - layer = default_trtnet().add_gather(padded_weight.trt_tensor, - input.trt_tensor, 0) - x = _create_tensor(layer.get_output(0), layer) - - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, tp_group, gather_dim=-1) - - else: - raise ValueError( - 'Tensor Parallelism only support splitting Embedding lookup along hidden (sharding_dim==1) and vocab (sharding_dim==0) dimensionis' - ) - - # Store embedding lookup table as a whole - else: - if use_lookup_plugin: - x = _lookup_plugin(input, - padded_weight, - rank=0, - per_token_scale=per_token_scale) - else: - layer = default_trtnet().add_gather(padded_weight.trt_tensor, - input.trt_tensor, 0) - x = _create_tensor(layer.get_output(0), layer) - return x - - -def constant_to_tensor_(input: Union[Tensor, int, float, bool], - dtype: Union[trt.DataType, str] = None, - to_array=True) -> Tensor: - if dtype is None: - # deduce the type from the given value - # NOTE: bool is a subtype of int, so bool needs to be checked first - if isinstance(input, bool): - dtype = trt.bool - elif isinstance(input, int): - dtype = trt.int32 - else: - dtype = trt.float32 - - if not isinstance(input, Tensor): - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - array_fn_dict = { - trt.int64: int64_array, - trt.int32: int32_array, - trt.float32: fp32_array, - trt.float16: fp16_array, - trt.bfloat16: bf16_array, - trt.bool: bool_array, - } - assert dtype in array_fn_dict - return constant(array_fn_dict[dtype]([input] if to_array else input)) - - return input - - -def constants_to_tensors_( - *inputs: Union[Tensor, int, float]) -> Tuple[Tensor, ...]: - ''' - Helper function to create tensors from multiple inputs. - - For each inputs, that function first creates a constant tensor if the input - is an integer or a float. Then, if any input is int64, it upcasts other - integer inputs to int64. - - Parameters: - inputs : Tuple[Union[Tensor, int, float], ...] - The inputs to create tensors from. - - Returns: - A tuple of tensors. - ''' - has_int64: bool = False - for i in inputs: - if isinstance(i, int) and (i >= 2**31 or i < -2**31)\ - or isinstance(i, Tensor) and i.dtype == trt.int64: - has_int64 = True - break - - if not has_int64: - return tuple(constant_to_tensor_(i) for i in inputs) - - result = [] - for i in inputs: - if isinstance(i, int) or isinstance(i, Tensor) and i.dtype == trt.int32: - result.append( - constant_to_tensor_(i, trt.int64 if has_int64 else trt.int32)) - else: - result.append(constant_to_tensor_(i)) - return tuple(result) - - -def broadcast_helper(left: Union[Tensor, int, float], - right: Union[Tensor, int, float]) -> Tuple[Tensor, Tensor]: - ''' - Helper function to perform a broadcast. - - For each input, that function first creates a constant tensor if the input - is an integer or a float. Then, if needed, it expands the smaller tensor to - make sure its rank is the same as the larger one. - - Parameters: - left : Union[Tensor, int, float] - The first input. If that input is an integer or a float, the - function creates a constant tensor. - - right : Union[Tensor, int, float] - The second input. If that input is an integer or a float, the - function creates a constant tensor. - - Returns: - A pair of tensors of same rank. - ''' - if not default_net().strongly_typed: - left = constant_to_tensor_(left) - right = constant_to_tensor_(right) - else: - left = constant_to_tensor_( - left, right.dtype if isinstance(right, Tensor) else None) - right = constant_to_tensor_(right, left.dtype) - - if left.rank() == right.rank(): - return (left, right) - - if left.rank() < right.rank(): - left = expand_dims_like(left, right) - return (left, right) - - if left.rank() > right.rank(): - right = expand_dims_like(right, left) - return (left, right) - - -def elementwise_binary(left: Union[Tensor, int, - float], right: Union[Tensor, int, float], - op: trt.ElementWiseOperation) -> Tensor: - ''' - Add an elementwise operation with two inputs. - - For each input, that function first creates a constant tensor if the input - is an integer or a float. Then, if needed, it expands the smaller tensor to - make sure its rank is the same as the larger one. Then, it performs the - elementwise operation 'op'. - - The following closures are defined in functional.*: - - add for op=trt.ElementWiseOperation.SUM - sub for op=trt.ElementWiseOperation.SUB - mul for op=trt.ElementWiseOperation.PROD - div for op=trt.ElementWiseOperation.DIV - floordiv for op=trt.ElementWiseOperation.FLOOR_DIV - gt for op=trt.ElementWiseOperation.GREATER - lt for op=trt.ElementWiseOperation.LESS - op_and for op=trt.ElementWiseOperation.AND - op_or for op=trt.ElementWiseOperation.OR - eq for op=trt.ElementWiseOperation.EQUAL - minimum for op=trt.ElementWiseOperation.MIN - maximum for op=trt.ElementWiseOperation.MAX - pow for op=trt.ElementWiseOperation.POW - - It is implemented using the IElementWiseLayer from TensorRT. - - Parameters: - left : Union[Tensor, int, float] - The first input. If that input is an integer or a float, the - function creates a constant tensor. - - right : Union[Tensor, int, float] - The second input. If that input is an integer or a float, the - function creates a constant tensor. - - op : trt.ElementWiseOperation - The binary operation to perform. - - Returns: - The tensor produced by this elementwise operation. - ''' - left, right = broadcast_helper(left, right) - if left.dtype == trt.int32 and right.dtype == trt.int64: - left = cast(left, trt.int64) - if left.dtype == trt.int64 and right.dtype == trt.int32: - right = cast(right, trt.int64) - layer = default_trtnet().add_elementwise(left.trt_tensor, right.trt_tensor, - op) - return _create_tensor(layer.get_output(0), layer) - - -add = partial(elementwise_binary, op=trt.ElementWiseOperation.SUM) -sub = partial(elementwise_binary, op=trt.ElementWiseOperation.SUB) -mul = partial(elementwise_binary, op=trt.ElementWiseOperation.PROD) -div = partial(elementwise_binary, op=trt.ElementWiseOperation.DIV) -floordiv = partial(elementwise_binary, op=trt.ElementWiseOperation.FLOOR_DIV) -gt = partial(elementwise_binary, op=trt.ElementWiseOperation.GREATER) -lt = partial(elementwise_binary, op=trt.ElementWiseOperation.LESS) -op_and = partial(elementwise_binary, op=trt.ElementWiseOperation.AND) -op_or = partial(elementwise_binary, op=trt.ElementWiseOperation.OR) -eq = partial(elementwise_binary, op=trt.ElementWiseOperation.EQUAL) -minimum = partial(elementwise_binary, op=trt.ElementWiseOperation.MIN) -maximum = partial(elementwise_binary, op=trt.ElementWiseOperation.MAX) -pow = partial(elementwise_binary, op=trt.ElementWiseOperation.POW) -op_xor = partial(elementwise_binary, op=trt.ElementWiseOperation.XOR) - - -def modulo(x: Tensor, y: Union[Tensor, int]) -> Tensor: - ''' - This function adds an element-wise modulo (x % y) operation for a given tensor. - Since there is no TensorRT layer that can directly perform this, - this function implements it using some of the basic operations. - - Returns: - A tensor that represents (x % y) modulo operation. - ''' - return x - (x // y) * y - - -def where(condition: Union[Tensor, bool], left: Union[Tensor, int, float], - right: Union[Tensor, int, float]) -> Tensor: - ''' - Add a where (aka select or if-then-else) operation. - - Assuming the three input parameters have the same shape, that function creates - the operation to compute a tensor of the same shape such that: - - for ii in range(mul(condition.shape)): - output[ii] = left[ii] if condition[ii] else right[ii] - - For each input, that function first creates a constant tensor if the - condition is boolean or the left/right input is an integer or a float. - Then, if needed, it expands the smaller tensor to make sure its - rank is the same as the larger one. Then, it performs the selection. - - It is implemented using the ISelectLayer from TensorRT. - - Parameters: - condition : Union[Tensor, bool] - The condition. If that input is a boolean, the function - creates a constant tensor. - - left : Union[Tensor, int, float] - The first input. If that input is an integer or a float, the - function creates a constant tensor. - - right : Union[Tensor, int, float] - The second input. If that input is an integer or a float, the - function creates a constant tensor. - - Returns: - The tensor produced by this where operation. - ''' - # Convert to tensors. - condition = constant_to_tensor_(condition) - left, right = constants_to_tensors_(left, right) - - # Find the tensor with the largest rank of the three. - largest = condition - if largest.rank() < left.rank(): - largest = left - if largest.rank() < right.rank(): - largest = right - - # Expand the tensors to match the largest one. - if condition is not largest: - condition = expand_dims_like(condition, largest) - if left is not largest: - left = expand_dims_like(left, largest) - if right is not largest: - right = expand_dims_like(right, largest) - - # Insert the operation. - layer = default_trtnet().add_select(condition.trt_tensor, left.trt_tensor, - right.trt_tensor) - return _create_tensor(layer.get_output(0), layer) - - -def unary(input: Tensor, op: trt.UnaryOperation) -> Tensor: - ''' - Add an elementwise operation on a single input. - - The following closures are defined in functional.*: - - round for op=trt.UnaryOperation.ROUND - sqrt for op=trt.UnaryOperation.SQRT - exp for op=trt.UnaryOperation.EXP - sin for op=trt.UnaryOperation.SIN - cos for op=trt.UnaryOperation.COS - abs for op=trt.UnaryOperation.ABS - log for op=trt.UnaryOperation.LOG - - It is implemented using the IUnaryLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - op : trt.UnaryOperation - The unary operation to perform. - - Returns: - The tensor produced by this elementwise operation. - ''' - layer = default_trtnet().add_unary(input.trt_tensor, op) - return _create_tensor(layer.get_output(0), layer) - - -round = partial(unary, op=trt.UnaryOperation.ROUND) -sqrt = partial(unary, op=trt.UnaryOperation.SQRT) -exp = partial(unary, op=trt.UnaryOperation.EXP) -sin = partial(unary, op=trt.UnaryOperation.SIN) -cos = partial(unary, op=trt.UnaryOperation.COS) -abs = partial(unary, op=trt.UnaryOperation.ABS) -log = partial(unary, op=trt.UnaryOperation.LOG) -not_op = partial(unary, op=trt.UnaryOperation.NOT) - - -def log_softmax(input: Tensor, dim: int) -> Tensor: - ''' - This function is equivalent of torch.nn.functional.log_softmax() i.e. - it performs log(softmax(input)) in a safer and faster way. - - Parameters: - input: Tensor - The data tensor on which log_softmax to be computed. - dim: int - The dimension of the input tensor along which log_softmax will be computed. - Returns: - A tensor of same shape as input with log_softmax computed on the specified dim. - ''' - x_max = max(input, dim=dim, keepdim=True) - x = input - x_max - return x - log(sum(exp(x), dim=dim, keepdim=True)) - - -def reduce(input: Tensor, - op: trt.ReduceOperation, - dim: Union[int, Tuple[int]], - keepdim: bool = False) -> Tensor: - ''' - Add an reduction operation to do along a dimension. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - op : trt.ReduceOperation - The reduction operation to perform. - Options: SUM, PROD, MAX, MIN, AVG - - dim : int - The dimension along which the reduction is performed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - dim = dim_resolve_negative(dim, input.ndim()) - axes = dim_to_trt_axes(dim) - - layer = default_trtnet().add_reduce(input.trt_tensor, - op, - axes, - keep_dims=keepdim) - return _create_tensor(layer.get_output(0), layer) - - -prod = partial(reduce, op=trt.ReduceOperation.PROD) -min = partial(reduce, op=trt.ReduceOperation.MIN) - - -def mean(input: Tensor, - dim: Union[int, Tuple[int]], - keepdim: bool = False) -> Tensor: - ''' - Add an operation to compute the mean along a dimension. - - Computes the mean along the dimension 'dim' of the input tensor. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension along which the mean is computed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - return reduce(input, op=trt.ReduceOperation.AVG, dim=dim, keepdim=keepdim) - - -def max(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: - ''' - Add an operation to compute the max along a dimension. - - Computes the max along the dimension 'dim' of the input tensor. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension along which the mean is computed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - return reduce(input, op=trt.ReduceOperation.MAX, dim=dim, keepdim=keepdim) - - -def sum(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: - ''' - Add an operation to compute the sum along a dimension. - - Computes the sum along the dimension 'dim' of the input tensor. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension along which the mean is computed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - return reduce(input, op=trt.ReduceOperation.SUM, dim=dim, keepdim=keepdim) - - -def identity(input: Tensor) -> Tensor: - ''' - Add an identity operation. - - Parameters: - input : Tensor - The input tensor. - - Returns: - The tensor produced by this identity operation. - ''' - if not default_net().plugin_config.identity_plugin: - layer = default_trtnet().add_identity(input.trt_tensor) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Identity', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - pfc = trt.PluginFieldCollection() - id_plug = plg_creator.create_plugin("identity", pfc) - plug_inputs = [input.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, id_plug) - _add_plugin_info(layer, plg_creator, "identity", pfc) - return _create_tensor(layer.get_output(0), layer) - - -def argmax(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: - ''' - Add an argmax operation. - - As explained in the ONNX documentation, - - https://github.com/onnx/onnx/blob/main/docs/Operators.md#argmax - - that function creates a layer computing the indices of the max elements of - the input tensor's element along the provided dim. The resulting tensor - has the same rank as the input if keepdims is True. If keepdims is False, - then the resulting tensor has the reduced dimension pruned. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension in which to compute the argmax indices. - - keepdim : bool - Do we keep the dimension along which the reduction is performed? - Yes, if set to True, no otherwise. - - Returns: - The tensor produced by this argmax operation. - ''' - dim = dim_resolve_negative(dim, input.ndim()) - axes = dim_to_trt_axes(dim) - - layer = default_trtnet().add_topk(input.trt_tensor, trt.TopKOperation.MAX, - 1, axes) - output = layer.get_output(1) - - if keepdim: - return _create_tensor(output, layer) - - output = _create_tensor(output, layer) - a = list(range(input.ndim())) - for d in dim: - a.pop(d) - indices = constant(int32_array(a)) - output_shape = shape(output) - new_shape = gather(output_shape, 0, indices) - return view(output, new_shape) - - -def gelu(x: Tensor) -> Tensor: - ''' - Add a GELU operation. - - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - Returns: - The tensor produced by the activation layer. - ''' - return 0.5 * x * ( - tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * pow(x, 3.0))) + 1.0) - - -def geglu(x: Tensor) -> Tensor: - ''' - Add a Gated-GELU operation. - - That function takes a tensor, splits it into two halves along the last - dimension, applies GELU to the second half and multiply the results. The - behavior is undefined if the last dimension is not even. - - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - Returns: - The tensor produced by the activation layer. - ''' - a, b = chunk(x, 2, dim=-1) - return a * gelu(b) - - -def quick_gelu(x: Tensor) -> Tensor: - return x * sigmoid(1.702 * x) - - -def gegelu(x: Tensor, limit: Optional[float] = None) -> Tensor: - # a, b = x[..., ::2], x[..., 1::2] - ndim = x.ndim() - a_starts = [0 for i in range(ndim)] - b_starts = [1 if i == (ndim - 1) else 0 for i in range(ndim)] - shapes = concat([ - shape(x, i) / 2 if i == (ndim - 1) else shape(x, i) for i in range(ndim) - ]) - strides = [2 if i == (ndim - 1) else 1 for i in range(ndim)] - - a = slice(x, a_starts, shapes, strides) - b = slice(x, b_starts, shapes, strides) - - if limit is not None: - a = clip(a, alpha=float(-1e20), beta=limit) - b = clip(b, alpha=-limit, beta=limit) - - # C = B + 1 - const1 = arange(constant(int32_array(1)), constant(int32_array(2)), - trt_dtype_to_str(b.dtype)) - for _ in range(ndim - 1): - const1 = expand_dims(const1, 0) - - b_shape = concat([shape(b, i) for i in range(ndim)]) - const1_arr = expand(const1, b_shape) - - return quick_gelu(a) * (b + const1_arr) - - -def group_norm(input: Tensor, - num_groups: int, - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - eps: float = 1e-05): - - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic(1) - num_channels = input.size()[1] - - ndim = input.ndim() - old_shape = shape(input) - new_shape = concat([ - input.size(0), - num_groups, - num_channels // num_groups, - ] + [input.size(i) for i in range(2, ndim)]) - x = input.view(new_shape) - - # instance norm - w_shape = [1, num_groups] + [1 for i in range(ndim - 1)] - instance_weight = constant(np.ones(w_shape, dtype=trt_dtype_to_np(x.dtype))) - instance_bias = constant(np.zeros(w_shape, dtype=trt_dtype_to_np(x.dtype))) - axes_mask = 0 - for i in range(2, x.ndim()): - axes_mask |= 1 << i - layer = default_trtnet().add_normalization(x.trt_tensor, - instance_weight.trt_tensor, - instance_bias.trt_tensor, - axes_mask) - layer.epsilon = eps - y = _create_tensor(layer.get_output(0), layer) - y = y.view(old_shape) - - new_shape = concat([num_channels] + [1 for _ in range(2, ndim)]) - if weight is not None: - y = y * weight.view(new_shape) - if bias is not None: - y = y + bias.view(new_shape) - - return y - - -def softplus(input: Tensor, beta: float, threshold: float) -> Tensor: - ''' - Add the softplus activation base on PyTorch definition. - - See https://pytorch.org/docs/stable/generated/torch.nn.functional.softplus.html#torch-nn-functional-softplus for a - description of that function. - - Parameters: - input : Tensor - Input TensorRT LLM Tensor. - beta : float - The parameter for softplus computation. - threshold : float - The threshold for reverting to the linear function when input * beta > threshold - - Returns: - The output tensor created by that layer. - ''' - sf_layer = default_trtnet().add_activation(input.trt_tensor, - trt.ActivationType.SOFTPLUS) - sf_layer.alpha = 1 / beta - sf_layer.beta = beta - - prod_tensor = input * beta - result = prod_tensor > threshold - - return where(result, input, _create_tensor(sf_layer.get_output(0), - sf_layer)) - - -def outer(input: Tensor, vec2: Tensor) -> Tensor: - ''' - Add an operation to compute the outer product between two tensors. - - That operation creates an Einsum node. - - Parameters: - input : Tensor - The first input tensor. - - vec2 : Tensor - The second input tensor. - - Returns: - The output tensor produced by this layer. - ''' - return einsum('i,j->ij', [input, vec2]) - - -def avg_pool2d(input: Tensor, - kernel_size: Tuple[int], - stride: Optional[Tuple[int]] = None, - padding: Optional[Tuple[int]] = (0, 0), - ceil_mode: bool = False, - count_include_pad: bool = True) -> Tensor: - - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic() - ndim = input.ndim() - if ndim == 3: - input = expand_dims(input, 0) - - layer = default_trtnet().add_pooling_nd(input.trt_tensor, - trt.PoolingType.AVERAGE, - kernel_size) - if stride is None: - stride = kernel_size - layer.stride_nd = stride - - output = _create_tensor(layer.get_output(0), layer) - - if ndim == 3: - return output.view( - concat([output.size(1), - output.size(2), - output.size(3)])) - - return output - - -def _get_trt_weight(weight: Tensor) -> Tuple[trt.Weights, bool]: - is_weight_constant = (weight.producer is not None - and weight.producer.type == trt.LayerType.CONSTANT) - if is_weight_constant: - ndarray = get_np_weight(default_trtnet(), weight.producer.name) - if ndarray is not None: - trt_weight = trt.Weights(np_dtype_to_trt(ndarray.dtype), - ndarray.ctypes.data, - int(np.prod(ndarray.shape))) - else: - weight.producer.__class__ = trt.IConstantLayer - trt_weight = weight.producer.weights - else: - trt_weight = trt.Weights() - - return trt_weight, is_weight_constant - - -def conv1d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: int = 1, - padding: int = 0, - dilation: int = 1, - groups: int = 1) -> Tensor: - - noutput = weight.size()[0] - kernel_size = weight.size()[-2] - kernel_shape = trt.Dims([kernel_size, 1]) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - input_shuffled = stack([input], dim=input.ndim()) - - layer = default_trtnet().add_convolution_nd(input_shuffled.trt_tensor, - noutput, kernel_shape, - trt_weight, trt_bias) - layer.stride_nd = (stride, 2) - layer.padding_nd = (padding, 0) - layer.dilation_nd = (dilation, 2) - layer.num_groups = groups - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output_2d = _create_tensor(layer.get_output(0), layer) - output_1d = squeeze(output_2d, dim=-1) - return output_1d - - -def conv2d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1, - pre_padding: Optional[Tuple[int, int]] = None, - post_padding: Optional[Tuple[int, int]] = None) -> Tensor: - ## - ## TODO: Document that function! - ## - - ndim = input.ndim() - if ndim == 3: - input = expand_dims(input, 0) - - noutput = weight.size()[0] - kernel_size = (weight.size()[-2], weight.size()[-1]) - kernel_shape = trt.Dims(list(kernel_size)) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - layer = default_trtnet().add_convolution_nd(input.trt_tensor, noutput, - kernel_shape, trt_weight, - trt_bias) - layer.stride_nd = stride - layer.padding_nd = padding - layer.dilation_nd = dilation - layer.num_groups = groups - layer.dilation_nd = dilation - if pre_padding: - layer.pre_padding = pre_padding - if post_padding: - layer.post_padding = post_padding - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output = _create_tensor(layer.get_output(0), layer) - - if ndim == 3: - return output.view( - concat([output.size(1), - output.size(2), - output.size(3)])) - - return output - - -def conv3d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: Union[int, Tuple[int, int]] = (1, 1, 1), - padding: Union[int, Tuple[int, int]] = (0, 0, 0), - dilation: Union[int, Tuple[int, int]] = (1, 1, 1), - groups: int = 1) -> Tensor: - ## - ## TODO: Document this function! - ## - - ndim = input.ndim() - # TRT requires the input of Conv3D layer to be 5-dimentional tensor. - if ndim == 4: - input = expand_dims(input, 0) - assert input.ndim() == 5 - - if isinstance(stride, int): - stride = tuple([stride] * 3) - if isinstance(padding, int): - padding = tuple([padding] * 3) - if isinstance(dilation, int): - dilation = tuple([dilation] * 3) - - noutput = weight.size()[0] - kernel_size = (weight.size()[-3], weight.size()[-2], weight.size()[-1]) - kernel_shape = trt.Dims(list(kernel_size)) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - layer = default_trtnet().add_convolution_nd(input.trt_tensor, noutput, - kernel_shape, trt_weight, - trt_bias) - layer.stride_nd = stride - layer.padding_nd = padding - layer.dilation_nd = dilation - layer.num_groups = groups - layer.dilation_nd = dilation - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output = _create_tensor(layer.get_output(0), layer) - return output - - -def conv_transpose2d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - output_padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1) -> Tensor: - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic() - - ndim = input.ndim() - if ndim == 3: - input = expand_dims(input, 0) - - noutput = weight.size()[1] - kernel_size = (weight.size()[-2], weight.size()[-1]) - kernel_shape = trt.Dims(list(kernel_size)) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - layer = default_trtnet().add_deconvolution_nd(input.trt_tensor, noutput, - kernel_shape, trt_weight, - trt_bias) - layer.stride_nd = stride - layer.padding_nd = padding - layer.num_groups = groups - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output = _create_tensor(layer.get_output(0), layer) - - if ndim == 3: - return output.view( - concat([output.size(1), - output.size(2), - output.size(3)])) - - return output - - -def split(tensor: Tensor, - split_size_or_sections: Union[int, Sequence[int]], - dim: int = 0) -> Sequence[Tensor]: - ''' - Add an operation that splits a tensor into sub-tensors. - - This operation creates a list of tensors that are obtained from the input - tensor by slicing it along the dimension 'dim'. If 'split_size_or_sections' - is an integer, the tensor is split into 'input.shape[dim] / - split_size_or_sections' slices. If 'split_size_or_sections' is a list of - sizes, the tensor is split into 'len(split_size_or_sections)' slices and - the size of the ith slice is given by 'split_size_or_sections[i]'. - - There are several constraints with the current implementation: - - - The input tensor must be static (no dynamic dimension), - - If 'split_size_or_sections' is an integer, the number of elements in - the 'dim' dimension of the input must be a multiple of - 'split_size_or_sections': 'input.shape[dim] % split_size_or_sections == 0'. - - If 'split_size_or_sections' is a sequence, the sum of the elements in - 'split_size_or_sections' must be equal to the size in the dimension - 'dim': 'input.shape[dim] == sum(ii for ii in split_size_or_sections)'. - - That operation is implemented using a 'slice' operation for each output - slice. - - Parameters: - tensor : Tensor - The input tensor to slice. - - split_size_or_sections : Union[int, Sequence[int]] - If it is an integer, it encodes the size of each slice. Otherwise, - if it is a sequence, it is the size of each slice. - - dim : int - The dimension of the tensor to slice. - - Returns: - The list of tensors produced by the different operations. - ''' - assert not tensor.is_dynamic(dim) - - ndim = tensor.ndim() - if dim < 0: - dim += ndim - dim_value = tensor.size()[dim] - starts = [constant(dims_array([0])) for _ in range(ndim)] - sizes = [shape(tensor, i) for i in range(ndim)] - - if isinstance(split_size_or_sections, int): - # TODO: support non-divisible cases - assert dim_value % split_size_or_sections == 0 - num_sections = dim_value // split_size_or_sections - sizes[dim] = constant(dims_array([split_size_or_sections])) - - outputs = [] - for i in range(num_sections): - starts[dim] = constant(dims_array([split_size_or_sections * i])) - outputs.append(slice(tensor, concat(starts), concat(sizes))) - return outputs - else: - total_size = 0 - for i in split_size_or_sections: - total_size += i - assert dim_value == total_size - num_sections = len(split_size_or_sections) - - outputs = [] - for i in range(num_sections): - if i > 0: - starts[dim] = starts[dim] + sizes[dim] - sizes[dim] = constant(dims_array([split_size_or_sections[i]])) - outputs.append(slice(tensor, concat(starts), concat(sizes))) - return outputs - - -def chunk(tensor: Tensor, chunks: int, dim: int = 0) -> Tensor: - ''' - Add an operation that splits a tensor into sub-tensors. - - This operation creates a list of tensors that are obtained from the input - tensor by chunking it along the dimension 'dim'. It produces 'chunks' - sub-tensors. - - That operation is only defined for static tensors (no dynamic dimension) - and the size of the tensor in the dimension 'dim' must be a multiple of - 'chunks': 'input.shape[dim] % chunks == 0'. - - It maps to 'split' with 'split_size = input.shape[dim] / chunks'. - - Parameters: - tensor : Tensor - The input tensor to slice. - - chunks : int - The number of slices to split the input tensor into. - - dim : int - The dimension of the tensor to slice. - - Returns: - The list of tensors produced by the different operations. - ''' - assert not tensor.is_dynamic(dim) - - ndim = tensor.ndim() - if dim < 0: - dim += ndim - dim_value = tensor.size()[dim] - assert dim_value % chunks == 0 - - return split(tensor, dim_value // chunks, dim) - - -def unbind(input: Tensor, dim: int = 0): - ''' - Removes a tensor dimension. - - Returns a tuple of all slices along a given dimension, already without it. - ''' - ndim = input.ndim() - outputs = split(input, 1, dim) - output_shape = [input.shape[i] for i in range(ndim) if i != dim] - return [output.view(output_shape) for output in outputs] - - -class AllReduceStrategy(IntEnum): - NCCL = 0 - MIN_LATENCY = 1 - UB = 2 - AUTO = 3 - ONESHOT = 4 - TWOSHOT = 5 - LOWPRECISION = 6 - MNNVL = 7 - NCCL_SYMMETRIC = 8 - SYMM_MEM = 9 # PyTorch symmetric memory with MULTIMEM - - -class AllReduceFusionOp(IntEnum): - NONE = 0 - RESIDUAL_RMS_NORM = 1 - LAST_PROCESS_FOR_UB = 2 - RESIDUAL_RMS_PREPOST_NORM = 3 - RESIDUAL_RMS_NORM_QUANT_FP8 = 4 - RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5 - RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6 - RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7 - MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 8 - RMS_NORM = 9 - - -class AllReduceParams(): - - def __init__(self, - strategy: AllReduceStrategy = AllReduceStrategy.AUTO, - fusion_op: AllReduceFusionOp = AllReduceFusionOp.NONE, - bias: Optional[Tensor] = None, - residual: Optional[Tensor] = None, - norm_weight: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - norm_pre_residual_weight: Optional[Tensor] = None, - eps: float = 1e-06, - enable_allreduce: bool = True, - trigger_completion_at_end: bool = True): - self.strategy = strategy - self.fusion_op = fusion_op - self.bias = bias - self.residual = residual - self.norm_weight = norm_weight - self.scale = scale - self.norm_pre_residual_weight = norm_pre_residual_weight - self.eps = eps - # For torch path only, has no effect on TRT path - self.enable_allreduce = enable_allreduce - self.trigger_completion_at_end = trigger_completion_at_end - assert fusion_op in (AllReduceFusionOp.NONE.value, - AllReduceFusionOp.RMS_NORM.value) or (residual - is not None) - - def has_affine(self): - return 1 if self.norm_weight is not None else 0 - - def has_bias(self): - return 1 if self.bias is not None else 0 - - def has_scale(self): - return 1 if self.scale is not None else 0 - - def update_strategy(self): - if self.strategy == AllReduceStrategy.AUTO and default_net( - ).plugin_config.user_buffer: - self.strategy = AllReduceStrategy.UB - - -class MoEAllReduceParams(AllReduceParams): - - def __init__(self, - device_num_experts: Optional[Tensor] = None, - expert_scale_factor: Optional[Tensor] = None, - expanded_idx_to_permuted_idx: Optional[Tensor] = None, - shared_expert_output: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - residual: Optional[Tensor] = None, - norm_weight: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - norm_pre_residual_weight: Optional[Tensor] = None, - eps: float = 1e-06, - enable_allreduce: bool = True, - is_cutlass_min_latency: bool = False): - super().__init__( - bias=bias, - residual=residual, - norm_weight=norm_weight, - scale=scale, - norm_pre_residual_weight=norm_pre_residual_weight, - eps=eps, - enable_allreduce=enable_allreduce, - ) - self.device_num_experts = device_num_experts - self.expert_scale_factor = expert_scale_factor - self.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx - self.shared_expert_output = shared_expert_output - self.is_cutlass_min_latency = is_cutlass_min_latency - - def is_valid(self): - if self.is_cutlass_min_latency: - return (self.device_num_experts is not None - and self.expert_scale_factor is not None - and self.shared_expert_output is not None) - else: - return (self.expanded_idx_to_permuted_idx is not None) - - -def create_allreduce_plugin( - network: trt.INetworkDefinition, - tensor: trt.ITensor, - workspace: Optional[trt.ITensor], - group: np.array, - dtype: trt.DataType, - all_reduce_params: AllReduceParams, -): - allreduce_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'AllReduce', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert allreduce_plg_creator is not None - - pf_group = trt.PluginField("group", group, trt.PluginFieldType.INT32) - pf_dtype = trt.PluginField("type_id", np.array([int(dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = [pf_group, pf_dtype] - p_strategy = trt.PluginField( - "strategy", np.array([int(all_reduce_params.strategy)], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_strategy) - p_fusion_op = trt.PluginField( - "fusion_op", np.array([int(all_reduce_params.fusion_op)], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_fusion_op) - p_eps = trt.PluginField( - "eps", np.array([float(all_reduce_params.eps)], np.float32), - trt.PluginFieldType.FLOAT32) - pfc.append(p_eps) - p_affine = trt.PluginField( - "affine", np.array([int(all_reduce_params.has_affine())], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_affine) - p_bias = trt.PluginField( - "bias", np.array([int(all_reduce_params.has_bias())], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_bias) - p_scale = trt.PluginField( - "scale", np.array([int(all_reduce_params.has_scale())], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_scale) - - pfc = trt.PluginFieldCollection(pfc) - ar_plug = allreduce_plg_creator.create_plugin("allreduce", pfc) - plug_inputs = [tensor] - if all_reduce_params.strategy not in { - AllReduceStrategy.NCCL, AllReduceStrategy.UB, - AllReduceStrategy.NCCL_SYMMETRIC - }: - plug_inputs.append(workspace) - if all_reduce_params.fusion_op != AllReduceFusionOp.NONE: - if all_reduce_params.has_bias() == 1: - plug_inputs.append(all_reduce_params.bias.trt_tensor) - if all_reduce_params.residual is not None: - plug_inputs.append(all_reduce_params.residual.trt_tensor) - if all_reduce_params.has_affine() == 1: - plug_inputs.append(all_reduce_params.norm_weight.trt_tensor) - if all_reduce_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_PREPOST_NORM: - plug_inputs.append( - all_reduce_params.norm_pre_residual_weight.trt_tensor) - if all_reduce_params.has_scale() == 1: - plug_inputs.append(all_reduce_params.scale.trt_tensor) - - layer = network.add_plugin_v2(plug_inputs, ar_plug) - return layer, allreduce_plg_creator, pfc - - -allreduce_ub_counter = 0 - - -def allreduce( - tensor: Tensor, - group: List[int], - all_reduce_params: Optional[AllReduceParams] = AllReduceParams() -) -> Tensor: - ''' - Add an operation that performs a collective all-reduce. - - Let's define 'world_size' as the length of the 'group' list. That functions - creates a layer to compute the sum of 'world_size' tensors distributed - amongst the 'world_size' participating ranks (one GPU per rank). - - The list 'group' contains the identifiers of the ranks participating into - the collective operation. - - The tensors in the different ranks must be 1D tensors (or views) and the output - tensor will have that same shape. The output tensor will be replicated on - the 'world_size' ranks. - - That operation is implemented using a plugin that wraps the NCCL all-reduce - collective operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce - for details. - - Parameters: - tensor : Tensor - The input tensor. - - group : List[int] - The ranks participating into the all-reduce operation. - - strategy: AllReduceStrategy - NCCL delegates all-reduce to NCCL while ONESHOT and TWOSHOT are custom latency-optimal algorithms. - AUTO chooses amongst the three based on a message-size heuristic. - - Returns: - The tensor produced by that layer. - ''' - - global allreduce_ub_counter - allreduce_ub_counter += 1 - - if all_reduce_params is None: - all_reduce_params = AllReduceParams() - all_reduce_params.update_strategy() - - # TODO(TRTLLM-996): remove this WAR when custom allreduce is supported - # for encoder models in C++ runtime. - workspace = None - if all_reduce_params.strategy != AllReduceStrategy.NCCL and all_reduce_params.strategy != AllReduceStrategy.UB: - if current_all_reduce_helper().workspace is None: - all_reduce_params.strategy = AllReduceStrategy.NCCL_SYMMETRIC - else: - workspace = current_all_reduce_helper().workspace.trt_tensor - if all_reduce_params.strategy == AllReduceStrategy.UB: - tensor.mark_output("allreduce_ub_0_" + str(allreduce_ub_counter)) - dtype = default_net().plugin_config.nccl_plugin - layer, allreduce_plg_creator, pfc = create_allreduce_plugin( - network=default_trtnet(), - tensor=tensor.cast(dtype).trt_tensor, - workspace=workspace, - group=np.array(group, dtype=np.int32), - dtype=str_dtype_to_trt(dtype), - all_reduce_params=all_reduce_params, - ) - _add_plugin_info(layer, allreduce_plg_creator, "allreduce", pfc) - if all_reduce_params.fusion_op != AllReduceFusionOp.NONE: - inter_output = _create_tensor(layer.get_output(1), - layer).cast(tensor.dtype) - if all_reduce_params.strategy == AllReduceStrategy.UB and all_reduce_params.has_scale( - ) == 1: - final_output = _create_tensor(layer.get_output(0), layer) - if all_reduce_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4: - scale_factor = _create_tensor(layer.get_output(2), layer) - else: - final_output = _create_tensor(layer.get_output(0), - layer).cast(tensor.dtype) - if all_reduce_params.strategy == AllReduceStrategy.UB: - if all_reduce_params.has_scale() == 1: - final_output.mark_output("allreduce_ub_1_" + - str(allreduce_ub_counter)) - if all_reduce_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4: - scale_factor.mark_output("allreduce_ub_2_" + - str(allreduce_ub_counter)) - return (final_output, scale_factor), inter_output - else: - assert all_reduce_params.fusion_op == AllReduceFusionOp.LAST_PROCESS_FOR_UB - inter_output.mark_output("allreduce_ub_1_" + - str(allreduce_ub_counter)) - return final_output, inter_output - else: - final_output = _create_tensor(layer.get_output(0), - layer).cast(tensor.dtype) - return final_output - - -def allgather(tensor: Tensor, group: List[int], gather_dim: int = 0) -> Tensor: - ''' - Add an operation that performs a collective all-gather. - - Let's define 'group_size' as the length of the 'group' list. That functions - creates a layer to gather 'group_size' tensors distributed - amongst the 'group_size' participating ranks (one GPU per rank). - - The list 'group' contains the identifiers of the ranks participating into - the collective operation. - - Note that 'group' here can be either TP group or PP group, because allgather communication is not limited to a specific split pattern. Therefore 'group_size' does not need to equal MPI 'world_size'. - - The tensors in the different ranks must be 1D tensors (or views) and the - output tensor will have that same shape. - - Given the 'section_size = input.shape[0] / group_size', each rank - contributes a section of its input tensor that correspond to - 'rank*section_size:(rank+1)*section_size'. - - That operation is implemented using a plugin that wraps the NCCL all-gather - collective operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather - for details. - - Parameters: - tensor : Tensor - The input tensor. - - group : List[int] - The ranks participating into the all-gather operation. - - gather_dim: int = 0 - Gather along given dimension. By default 0, i.e. treated as 1D tensor. - - Returns: - The tensor produced by that layer. - ''' - allgather_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'AllGather', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert allgather_plg_creator is not None - - group_size = len(group) - group = trt.PluginField("group", np.array(group, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([group, pf_type]) - allgather = allgather_plg_creator.create_plugin("allgather", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, allgather) - _add_plugin_info(layer, allgather_plg_creator, "allgather", pfc) - - x = _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - # gather along a given dimension other than dim0 - if gather_dim != 0: - # also support -1 type of dim representation - if gather_dim < 0: - gather_dim = x.ndim() + gather_dim - - # plugin above gathers as 1D flattened tensor - # 1. [dim0, ...dimi, ...dimN] -> [group_size * dim0, ...dimi, ...dimN] - - # now we need to gather-by-dim via split-concat - # 2. [group_size * dim0, ...dimi, ...dimN] -> [dim0, ...group_size * dimi, ...dimN] - # 2.1 split - split_size = shape(x, dim=0) / group_size - ndim = x.ndim() - starts = [constant(dims_array([0])) for _ in range(ndim)] - sizes = [shape(x, dim=d) for d in range(ndim)] - sizes[0] = split_size - sections = [] - for i in range(group_size): - starts[0] = split_size * i - sections.append(slice(x, concat(starts), concat(sizes))) - # 2.2 concat - x = concat(sections, dim=gather_dim) - - return x - - -def reduce_scatter(tensor: Tensor, group: List[int]) -> Tensor: - - plg_creater = trt.get_plugin_registry().get_plugin_creator( - 'ReduceScatter', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creater is not None - - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - group = trt.PluginField("group", np.array(group, dtype=np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([group, pf_type]) - - reduce_scatter_plug = plg_creater.create_plugin("reduce_scatter", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, reduce_scatter_plug) - _add_plugin_info(layer, plg_creater, "reduce_scatter", pfc) - - return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - -def send(tensor: Tensor, tgt: int) -> Tensor: - ''' - Add an operation that performs a send from a rank to another. - - The send operation sends a tensor from one rank to another. If a rank 'i' - sends a tensor to a rank 'j', the rank 'j' must have a corresponding 'recv' - operation from rank 'i'. See 'recv'. - - That operation is implemented using a plugin that wraps the NCCL send - point-to-point operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclsend - for details. - - Parameters: - tensor : Tensor - The input tensor. - - tgt : int - The rank that receives the tensor. - - Returns: - The tensor produced by that layer. - ''' - send_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Send', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert send_plg_creator is not None - - tgt = trt.PluginField("tgt_rank", np.array(tgt, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([tgt, pf_type]) - send_plug = send_plg_creator.create_plugin("send", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, send_plug) - _add_plugin_info(layer, send_plg_creator, "send", pfc) - return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - -def recv(tensor: Tensor, src: int) -> Tensor: - ''' - Add an operation that performs a recv to a rank from another. - - The recv operation receives a tensor from on a rank from another. If a rank 'i' - receives a tensor from a rank 'j', the rank 'j' must have a corresponding 'send' - operation to rank 'j'. See 'send'. - - That operation is implemented using a plugin that wraps the NCCL recv - point-to-point operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclrecv - for details. - - Parameters: - tensor : Tensor - The input tensor. - - src : int - The rank that sends the tensor to. - - Returns: - The tensor produced by that layer. - ''' - recv_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Recv', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert recv_plg_creator is not None - - src = trt.PluginField("src_rank", np.array(src, dtype=np.int32), - trt.PluginFieldType.INT32) - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([src, pf_type]) - recv_plug = recv_plg_creator.create_plugin("recv", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, recv_plug) - _add_plugin_info(layer, recv_plg_creator, "recv", pfc) - return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - -def gemm_allreduce(a: Tensor, - b: Tensor, - group: List[int], - transa: bool = False, - transb: bool = False, - alpha: Optional[Union[np.ndarray, Tensor]] = None, - output_dtype: Optional[trt.DataType] = None, - fp8_inputs_override: bool = False, - a_sf: Optional[Tensor] = None, - b_sf: Optional[Tensor] = None): - ''' - Add an operation that performs fused GEMM+AllReduce. - - Parameters: - a: Tensor - Input tensor A - b: Tensor - Input tensor B - a_sf: Optional[Tensor] - Input tensor for scaling input A - b_sf: Optional[Tensor] - Input tensor for scaling input B - group: List[int] - Ranks participating in collective - transa: bool - Whether or not input tensor A is transposed - transb: bool - Whether or not input tensor B is transposed - alpha: float - Alpha for GEMM -> beta * C + (alpha * acc) - output_dtype: trt.DataType - Output type for plugin. If it is None, we - will use type set in plugin_config. - fp8_inputs_override: bool - TRT graph does not detect FP8 inputs correctly. This - flag is used to override the derived input tensor - types so that our plugin knows to issue FP8 MMAs. - - Returns: - Returns GEMM output tensor which has been reduced across ranks. - ''' - - # Output tensor needs to be bound to externally managed - # memory so keep track of layer index so we can assign - # output tensor unique label. - if not hasattr(gemm_allreduce, 'layer_idx'): - gemm_allreduce.layer_idx = 0 - - # Check inputs - assert isinstance(a.dtype, trt.DataType) - assert isinstance(b.dtype, trt.DataType) - - if fp8_inputs_override: - assert ( - isinstance(alpha, np.ndarray) and alpha.dtype == np.float32 - and alpha.size == 1 - ), "`alpha` must be passed as a float32 ndarray if `fp8_inputs_override` is enabled for gemm_allreduce_plugin" - assert a.dtype == trt.fp8 - assert b.dtype == trt.fp8 - - if output_dtype is None: - output_dtype = str_dtype_to_trt( - default_net().plugin_config.gemm_allreduce_plugin) - assert output_dtype in [trt.float16, trt.bfloat16] - - alpha_is_tensor = isinstance(alpha, Tensor) - if alpha is None or alpha_is_tensor: - alpha_value = np.array(1.0, dtype=np.float32) - else: - alpha_value = alpha - - plugin_creator = trt.get_plugin_registry().get_plugin_creator( - 'GemmAllReduce', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plugin_creator is not None - - trt_type_a = trt.fp8 if fp8_inputs_override else a.dtype - trt_type_b = trt.fp8 if fp8_inputs_override else b.dtype - - # create plugin fields - field_list = [] - field_list.append( - trt.PluginField('type_a', np.array([int(trt_type_a)], np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('type_b', np.array([int(trt_type_b)], np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('type_d', np.array([int(output_dtype)], np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('transa', np.array(transa, dtype=np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('transb', np.array(transb, dtype=np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('group', np.array(group, dtype=np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('has_sfa', np.array([int(a_sf is not None)], np.int8), - trt.PluginFieldType.INT8)) - field_list.append( - trt.PluginField('has_sfb', np.array([int(b_sf is not None)], np.int8), - trt.PluginFieldType.INT8)) - field_list.append( - trt.PluginField('alpha_is_ptr', np.array([int(alpha_is_tensor)], - np.int8), - trt.PluginFieldType.INT8)) - field_list.append( - trt.PluginField('alpha', alpha_value.flatten(), - trt.PluginFieldType.FLOAT32)) - - # create plugin - fields = trt.PluginFieldCollection(field_list) - plugin = plugin_creator.create_plugin("gemm_allreduce", fields) - # define symbolic input tensors. - # note this does NOT allocate memory. - inputs = [a.trt_tensor, b.trt_tensor] - if a_sf is not None: - inputs += [a_sf.trt_tensor] - if b_sf is not None: - inputs += [b_sf.trt_tensor] - if alpha_is_tensor: - inputs += [alpha.trt_tensor] - - layer = default_trtnet().add_plugin_v2(inputs, plugin) - _add_plugin_info(layer, plugin_creator, "gemm_allreduce", fields) - # define symbolic output tensors - # both output tensors point to same physical memory but - # one has unicast address and other has multicast address - uc_output = _create_tensor(layer.get_output(0), layer) - mc_output = _create_tensor(layer.get_output(1), layer) - ipc_output = _create_tensor(layer.get_output(2), layer) - assert uc_output is not None - assert mc_output is not None - assert ipc_output is not None - # mark outputs so that we can bind our own allocated memory in runtime - # (see generation.py) - uc_output.mark_output(f'gemm_allreduce_uc_out_{gemm_allreduce.layer_idx}') - mc_output.mark_output(f'gemm_allreduce_mc_out_{gemm_allreduce.layer_idx}') - ipc_output.mark_output(f'gemm_allreduce_ipc_out_{gemm_allreduce.layer_idx}') - gemm_allreduce.layer_idx += 1 - - return uc_output - - -def bert_attention(tensor: Tensor, - input_lengths: Tensor, - num_heads: int, - head_size: int, - q_scaling: float, - relative_attention: bool = False, - relative_attention_bias: Tensor = None, - max_distance: int = 0, - max_input_length: Tensor = None, - sage_attn: bool = False, - sage_attn_q_block_size: int = 0, - sage_attn_k_block_size: int = 0, - sage_attn_v_block_size: int = 0, - cp_group: list[int] = None, - cp_size: int = 1, - cp_rank: int = 0) -> Tuple[Tensor]: - ''' - Add an operation that performs the multi-head attention in BERT. - - The multi-head attention (MHA) is the sequence of a batched matmul, a - softmax and a batched matmul as described in - https://arxiv.org/abs/1706.03762. That function adds an operation that - performs those computations using a single GPU kernel. - - The input tensor contains the Q, K and V elements. It is a 2D tensor and - its shape is '[sum_of_tokens, 3*hidden_dim]' where the 'sum_of_tokens' is - the sum of the sequence lengths in the batch. - - In MHA, the output of the Q*K^T product is scaled by a constant value that - is computed as: - - 1.f / (q_scaling * sqrt(head_size)). - - That 'q_scaling' constant is the last argument of that function. - - That layer is implemented using a plugin (see bertAttentionPlugin). - - Parameters: - tensor : Tensor - The QKV input tensor. - - input_lengths : Tensor - The length of each sequence. It is a 1D tensor of size 'batch_size'. - - num_heads : int - The number of heads. - - head_size : int - The size of each head. - - q_scaling : float - The factor to compute the scaling factor to scale the output of the - 'Q*K^T' product. - - relative_attention: bool = False - If enable relative attention. - - relative_attention_bias: Tensor = None - The relative attention bias [num_heads, max_seq_len, max_seq_len], or The relative attention embedding table for implicit mode, [num_heads, num_buckets]. - - max_distance: int = 0 - The maximum distance of relative position in attention, for implicit mode. - Default value is 0, meaning to use the regular mode of relative attention bias. - Implicit mode is only enabled when passing in non-zero positive max_distance value. - See relative attention bias in docs/source/advanced/gpt-attention.md - - max_input_length: Tensor = None - The maximum input sequence length represented by Tensor shape. Requires for remove_input_padding to pre-define plugin workspace size. - - sage_attn: bool = False - SageAttention is a 8-bit implementation of attention kernel. It's input q, k, v and output datatypes are 16-bit. It performance dynamic quantization for q, k, v - tensor every time before attention. https://github.com/thu-ml/SageAttention - - sage_attn_q_quant_size: int = 0 - dynamic quant block size along sequence dimension of q tensor. Each quant block will share one scale. - - sage_attn_k_quant_size: int = 0 - dynamic quant block size along sequence dimension of k tensor. Each quant block will share one scale. - - sage_attn_v_quant_size: int = 0 - dynamic quant block size along sequence dimension of v tensor. Each quant block will share one scale. - - cp_group: list[int] = None - The communication group for context parallel - - cp_size: int = 1 - The communication size for context parallel - - cp_rank: int = 0 - The communication rank for context parallel - - Returns: - The tensor produced by that layer. - ''' - attn_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'BertAttention', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert attn_plg_creator is not None - - nheads = trt.PluginField("num_heads", np.array(num_heads, dtype=np.int32), - trt.PluginFieldType.INT32) - head_size = trt.PluginField("head_size", np.array(head_size, - dtype=np.int32), - trt.PluginFieldType.INT32) - q_scaling = trt.PluginField("q_scaling", - np.array(q_scaling, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - context_fmha_type = trt.PluginField( - "context_fmha_type", - np.array(np.int8(default_net().plugin_config.context_fmha_type), - dtype=np.int8), trt.PluginFieldType.INT8) - p_dtype = default_net().plugin_config.bert_attention_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - do_relative_attention = trt.PluginField( - "do_relative_attention", - np.array(np.int8(relative_attention), dtype=np.int8), - trt.PluginFieldType.INT8) - max_distance = trt.PluginField("max_distance", - np.array(max_distance, dtype=np.int32), - trt.PluginFieldType.INT32) - remove_padding = trt.PluginField( - "remove_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - - sage_attn = trt.PluginField("sage_attn", - np.array(np.int8(sage_attn), dtype=np.int8), - trt.PluginFieldType.INT8) - - sage_attn_q_block_size = trt.PluginField( - "sage_attn_q_block_size", - np.array(sage_attn_q_block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - sage_attn_k_block_size = trt.PluginField( - "sage_attn_k_block_size", - np.array(sage_attn_k_block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - sage_attn_v_block_size = trt.PluginField( - "sage_attn_v_block_size", - np.array(sage_attn_v_block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - if cp_size > 1: - # transpose q,k,v inside qkv to make kv contiguous, which is required by ring attention - # (b, s, 3d) - query, key, value = chunk(tensor, 3, dim=-1) - bs = shape(query, 0) - seq_len = shape(query, 1) - # (b, s, d) -> (b, s, 2d) -> (2b, s, d) - kv = concat([key, value], - dim=-1).view(concat((2 * bs, seq_len, query.shape[-1]))) - tensor = concat((query, kv), - dim=0).view(concat((bs, seq_len, query.shape[-1] * 3))) - - cp_size = trt.PluginField("cp_size", np.array(cp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_rank = trt.PluginField("cp_rank", np.array(cp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_group = cp_group or [0] - cp_group = np.array(cp_group, dtype=np.int32) - cp_group = trt.PluginField("cp_group", cp_group, trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - nheads, head_size, q_scaling, context_fmha_type, pf_type, - do_relative_attention, max_distance, remove_padding, sage_attn, - sage_attn_q_block_size, sage_attn_k_block_size, sage_attn_v_block_size, - cp_size, cp_rank, cp_group - ]) - - attn_plug = attn_plg_creator.create_plugin("padding_attn", pfc) - plug_inputs = [tensor, input_lengths] - if max_input_length is not None: - # for remove padding mode - plug_inputs += [max_input_length] - if relative_attention_bias is not None: - # for relative attention mode - plug_inputs += [relative_attention_bias] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, attn_plug) - _add_plugin_info(layer, attn_plg_creator, "padding_attn", pfc) - assert layer.num_outputs == 1, \ - f"Plugin outputs number mismatch with expected, got {layer.num_outputs}, expected 1" - output = _create_tensor(layer.get_output(0), layer) - assert output is not None - return output - - -class RopeEmbeddingUtils: - - @staticmethod - # ref: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py#L298 - def apply_llama3_scaling(inv_freqs: np.ndarray, rope_scaling_config: dict): - - scale_factor = rope_scaling_config.get("factor", 8.0) - low_freq_factor = rope_scaling_config.get("low_freq_factor", 1.0) - high_freq_factor = rope_scaling_config.get("high_freq_factor", 4.0) - old_context_len = rope_scaling_config.get( - "original_max_position_embeddings", 8192) - - low_freq_wavelen = old_context_len / low_freq_factor - high_freq_wavelen = old_context_len / high_freq_factor - new_inv_freqs = [] - for inv_freq in inv_freqs: - wavelen = 2 * math.pi / inv_freq - if wavelen < high_freq_wavelen: - new_inv_freqs.append(inv_freq) - elif wavelen > low_freq_wavelen: - new_inv_freqs.append(inv_freq / scale_factor) - else: - assert low_freq_wavelen != high_freq_wavelen - smooth = (old_context_len / wavelen - low_freq_factor) / ( - high_freq_factor - low_freq_factor) - new_inv_freqs.append((1 - smooth) * inv_freq / scale_factor + - smooth * inv_freq) - return np.array(new_inv_freqs, dtype=inv_freqs.dtype) - - @staticmethod - def create_sinusoidal_positions(num_pos: int, - dim: int, - theta: float = 10000.0, - dtype=np.float32): - inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) - sinusoid_inp = np.einsum("i , j -> i j", - np.arange(num_pos, dtype=dtype), - inv_freq, - dtype=dtype) - concat = np.concatenate((np.sin(sinusoid_inp), np.cos(sinusoid_inp)), - axis=1) - return np.expand_dims(concat, axis=0).astype(dtype) - - @staticmethod - def create_sinusoidal_positions_for_attention_plugin( - num_pos: int, - dim: int, - theta: float = 10000.0, - scale: float = 1.0, - scale_type: RotaryScalingType = RotaryScalingType.none, - # Other scaling configs that only used by certain scaling types. - rope_scaling_config: dict = None, - duplicate_data: bool = False, - dtype=np.float32): - if scale_type == RotaryScalingType.linear: - scale = 1.0 / scale - if scale_type == RotaryScalingType.llama3: - assert rope_scaling_config is not None, "rotary_scaling config must be provided." - inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) - inv_freq = RopeEmbeddingUtils.apply_llama3_scaling( - inv_freq, rope_scaling_config) - elif scale_type == RotaryScalingType.dynamic: - # Make sure scaling_alpha exists in rope_scaling - # Ref: https://huggingface.co/tencent/Hunyuan-A13B-Instruct-FP8/blob/main/modeling_hunyuan.py#L346 - assert rope_scaling_config[ - "alpha"] is not None, "rope_scaling_config.alpha must be provided." - scaling_alpha = rope_scaling_config["alpha"] - adjusted_base = theta * (scaling_alpha**(dim / (dim - 2))) - inv_freq = 1.0 / (adjusted_base**( - np.arange(0, dim, 2, dtype=dtype) / dim)).astype(dtype) - else: - inv_freq = scale / (theta - **(np.arange(0, dim, 2) / dim)).astype(dtype) - sinusoid_inp = np.expand_dims(np.einsum("i , j -> i j", - np.arange(num_pos, dtype=dtype), - inv_freq, - dtype=dtype), - axis=-1) - if duplicate_data: - sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) - # fuse cos/sin into float2 (cos, sin). - concat = np.concatenate( - (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), - axis=-1) #np.cos(sinusoid_inp).shape = (32768, 64, 1) - - return inv_freq, concat.reshape(1, -1).astype(dtype) - - @staticmethod - def create_sinusoidal_positions_for_cogvlm_attention_plugin( - num_pos: int, - dim: int, - theta: float = 10000.0, - scale: float = 1.0, - scale_type: RotaryScalingType = RotaryScalingType.none, - vision_start: int = 1, - vision_length: int = 1225, - dtype=np.float32): - if scale_type == RotaryScalingType.linear: - scale = 1.0 / scale - inv_freq = scale / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) - position_id = np.hstack([ - np.arange(0, vision_start + 1, dtype=dtype), - np.full(vision_length, vision_start + 1, dtype=dtype), - np.arange(vision_start + 2, - num_pos - (vision_length - 1), - dtype=dtype) - ]) - sinusoid_inp = np.expand_dims(np.einsum("i , j -> i j", - position_id, - inv_freq, - dtype=dtype), - axis=-1) - # fuse cos/sin into float2 (cos, sin). - concat = np.concatenate((np.cos(sinusoid_inp), np.sin(sinusoid_inp)), - axis=-1) - - return inv_freq, concat.reshape(1, -1).astype(dtype) - - def create_sinusoidal_positions_long_rope_for_attention_plugin( - num_pos: int, - num_orig_pos: int, - dim: int, - theta: float = 10000.0, - scaling_short_factors: Tensor = 1.0, - scaling_long_factors: Tensor = 1.0, - short_mscale=None, - long_mscale=None, - dtype=np.float32): - - def _calc_mscale(scale): - if scale <= 1.0: - return 1.0 - return math.sqrt(1 + math.log(scale) / math.log(num_orig_pos)) - - if short_mscale is None: - short_mscale = _calc_mscale(num_pos / num_orig_pos) - long_mscale = short_mscale - - def _compute_sinusoidal_positions(scale_factors, is_short, - for_attention_plugin): - inv_freq = 1 / (scale_factors * - (theta**(np.arange(0, dim, 2) / dim)).astype(dtype)) - sinusoid_inp = np.einsum("i , j -> i j", - np.arange(num_pos, dtype=dtype), - inv_freq, - dtype=dtype) - - if for_attention_plugin: - sinusoid_inp = np.expand_dims(sinusoid_inp, axis=-1) - concat = np.concatenate( - (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), axis=-1) - else: - concat = np.concatenate( - (np.sin(sinusoid_inp), np.cos(sinusoid_inp)), axis=1) - concat = np.expand_dims(concat, axis=0) - - mscale = short_mscale if is_short else long_mscale - concat = concat.astype(dtype) * mscale - - # gpt attention plugins also need inv_freq. - if for_attention_plugin: - return inv_freq.reshape(1, -1), concat.reshape(1, -1) - else: - return concat - - return _compute_sinusoidal_positions( - scaling_short_factors, True, False), _compute_sinusoidal_positions( - scaling_long_factors, - False, False), _compute_sinusoidal_positions( - scaling_short_factors, True, - True), _compute_sinusoidal_positions( - scaling_long_factors, False, True), short_mscale - - @staticmethod - def create_sinusoidal_positions_long_rope(num_pos: int, - dim: int, - theta: float, - original_max_pos: int, - short_factor: List[float], - long_factor: List[float], - dtype=np.float32, - max_seq_len: Optional[int] = None, - duplicate_data: bool = False): - short_factor = np.array(short_factor, dtype=np.float32) - long_factor = np.array(long_factor, dtype=np.float32) - - inv_freq = 1.0 / (theta**(np.arange(0, dim, 2, dtype=np.float32) / dim)) - t_pos = np.arange(np.max([num_pos, original_max_pos]), dtype=np.float32) - - # Choose proper freqs based on max_seq_len. - factor = long_factor if max_seq_len is None or max_seq_len > original_max_pos else short_factor - inv_freq = inv_freq / factor - freqs = np.einsum("i,j->ij", t_pos, inv_freq) - sinusoid_inp = freqs.astype(np.float32)[..., np.newaxis] - - # Apply scaling - scale = num_pos / original_max_pos - if scale <= 1.0: - scaling_factor = 1.0 - else: - scaling_factor = np.sqrt(1.0 + - np.log(scale) / np.log(original_max_pos)) - - if duplicate_data: - sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) - - # fuse cos/sin into float2 (cos, sin). - concat = np.concatenate( - (np.cos(sinusoid_inp) * scaling_factor, - np.sin(sinusoid_inp) * scaling_factor), - axis=-1, - ) - - return None, concat.reshape(1, -1).astype(dtype) - - @staticmethod - def create_fake_weight(dim: int, dtype=np.half): - return np.random.rand(dim).astype(dtype) - - # Note: When not using deepseek_yarn, make sure to set mscale_all_dim to 0.0. - @staticmethod - def create_sinusoidal_positions_yarn( - num_pos: int, - dim: int, - base: int = 10000, - scaling_factor: float = 1.0, - original_max_position_embeddings: int = 4096, - beta_fast: int = 32, - beta_slow: int = 1, - mscale: float = 1.0, - mscale_all_dim: float = 1.0, - duplicate_data: bool = True, - dtype=torch.float32): - - # Copy from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/main/modeling_deepseek.py - # Inverse dim formula to find dim based on number of rotations - def yarn_find_correction_dim(num_rotations, dim, base, - max_position_embeddings): - return (dim * math.log(max_position_embeddings / - (num_rotations * 2 * math.pi))) / ( - 2 * math.log(base)) - - # Find dim range bounds based on rotations - def yarn_find_correction_range(low_rot, high_rot, dim, base, - max_position_embeddings): - low = math.floor( - yarn_find_correction_dim(low_rot, dim, base, - max_position_embeddings)) - high = math.ceil( - yarn_find_correction_dim(high_rot, dim, base, - max_position_embeddings)) - if low < 0: - low = 0 - if high > dim - 1: - high = dim - 1 - return low, high # Clamp values just in case - - def yarn_get_mscale(scale, mscale): - if scale <= 1: - return 1.0 - return 0.1 * mscale * math.log(scale) + 1.0 - - def yarn_linear_ramp_mask(min, max, dim): - if min == max: - max += 0.001 # Prevent singularity - - linear_func = (torch.arange(dim, dtype=dtype) - min) / (max - min) - ramp_func = torch.clamp(linear_func, 0, 1) - return ramp_func - - pos_freqs = base**(torch.arange(0, dim, 2, dtype=dtype) / dim) - freq_extra = 1.0 / pos_freqs - freq_inter = 1.0 / (scaling_factor * pos_freqs) - - low, high = yarn_find_correction_range( - beta_fast, - beta_slow, - dim, - base, - original_max_position_embeddings, - ) - inv_freq_mask = (1 - yarn_linear_ramp_mask(low, high, dim // 2)) - inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask - t = torch.arange(num_pos, dtype=dtype) - sinusoid_inp = torch.einsum("i,j -> ij", t, inv_freq).unsqueeze(-1) - - _mscale = float( - yarn_get_mscale(scaling_factor, mscale) / - yarn_get_mscale(scaling_factor, mscale_all_dim)) - - if duplicate_data: - emb = torch.cat((sinusoid_inp, sinusoid_inp), dim=-2) - else: - emb = sinusoid_inp - - concat = torch.cat((torch.cos(emb) * _mscale, torch.sin(emb) * _mscale), - dim=-1) - return inv_freq.numpy(), concat.reshape((1, -1)).to(dtype).numpy() - - @staticmethod - def rotate_every_two(tensor: Tensor) -> Tensor: - assert tensor.ndim() == 4 - - shape_tensor = concat([ - shape(tensor, i) / 2 if i == (tensor.ndim() - - 1) else shape(tensor, i) - for i in range(tensor.ndim()) - ]) - x1 = slice(tensor, [0, 0, 0, 0], shape_tensor, [1, 1, 1, 2]) - x2 = slice(tensor, [0, 0, 0, 1], shape_tensor, [1, 1, 1, 2]) - x1 = expand_dims(x1, 4) - x2 = expand_dims(x2, 4) - zero = constant( - np.ascontiguousarray( - np.zeros([1], dtype=trt_dtype_to_np(tensor.dtype)))) - x2 = zero - x2 - x = concat([x2, x1], 4) - return view( - x, concat([shape(x, 0), - shape(x, 1), - shape(x, 2), - shape(x, 3) * 2])) - - @staticmethod - def rotate_half(tensor: Tensor) -> Tensor: - # [bs, num_attention_kv_heads, seqlen, attention_head_size] - assert tensor.ndim() == 4 - shape_tensor = concat([ - shape(tensor, i) / 2 if i == (tensor.ndim() - - 1) else shape(tensor, i) - for i in range(tensor.ndim()) - ]) - last_dim = shape(tensor, tensor.ndim() - 1) / 2 - x1 = slice(tensor, [0, 0, 0, 0], shape_tensor, [1, 1, 1, 1]) - x2 = slice(tensor, concat([0, 0, 0, last_dim]), shape_tensor, - [1, 1, 1, 1]) - zero = constant( - np.ascontiguousarray( - np.zeros([1], dtype=trt_dtype_to_np(tensor.dtype)))) - x2 = zero - x2 - x = concat([x2, x1], 3) - return x - - @staticmethod - def apply_rotary_pos_emb( - tensor: Tensor, - position_embedding: List[Tensor] = None, - pos_emb_type: PositionEmbeddingType = PositionEmbeddingType.rope_gptj - ) -> Tensor: - - rotate_func = None - if pos_emb_type == PositionEmbeddingType.rope_gpt_neox or pos_emb_type == PositionEmbeddingType.long_rope: - assert len(position_embedding) == 2 - cos, sin = position_embedding - sin = expand_dims(sin, 2) - cos = expand_dims(cos, 2) - sin = concat([sin, sin], 3) - cos = concat([cos, cos], 3) - rotate_func = RopeEmbeddingUtils.rotate_half - elif pos_emb_type == PositionEmbeddingType.rope_gptj: - assert len(position_embedding) == 2 - cos, sin = position_embedding - sin = expand_dims(sin, 2) - cos = expand_dims(cos, 2) - sin = repeat_interleave(sin, 2, 3) - cos = repeat_interleave(cos, 2, 3) - rotate_func = RopeEmbeddingUtils.rotate_every_two - elif pos_emb_type == PositionEmbeddingType.chatglm: - assert len(position_embedding) == 4 - cos0, cos1, sin0, sin1 = position_embedding - shape_tensor = concat([ - shape(tensor, i) / 2 if i == (tensor.ndim() - - 1) else shape(tensor, i) - for i in range(tensor.ndim()) - ]) - last_dim = shape(tensor, tensor.ndim() - 1) / 2 - x_part0 = slice(tensor, [0, 0, 0, 0], shape_tensor, [1, 1, 1, 1]) - x_part1 = slice(tensor, concat([0, 0, 0, last_dim]), shape_tensor, - [1, 1, 1, 1]) - - y_part0 = (x_part0 * - cos0) + (RopeEmbeddingUtils.rotate_half(x_part0) * sin0) - y_part1 = (x_part1 * - cos1) + (RopeEmbeddingUtils.rotate_half(x_part1) * sin1) - - result = concat([y_part0, y_part1], dim=3) - return result.view(shape(tensor)) - - else: - raise ValueError('The PositionEmbeddingType is not RoPE') - return (tensor * cos) + (rotate_func(tensor) * sin) - - @staticmethod - def apply_rotary_pos_emb_chatglm(qkv, position_embedding, - num_attention_heads, attention_head_size, - max_position_embeddings, - rotary_embedding_scale, - remove_input_padding) -> Tensor: - - half_head_size = attention_head_size // 2 - input = qkv[0] if isinstance(qkv, list) else qkv - input_shape = shape(input) - batch_size = 1 if remove_input_padding else shape(input, 0) - seqlen = shape(input, 0 if remove_input_padding else 1) - if isinstance(qkv, list): - query, key, value = qkv - else: - qkv = qkv.view( - concat([ - batch_size, - seqlen, - num_attention_heads, - 3, - attention_head_size, - ])) - query, key, value = split(qkv, 1, dim=3) - q_shape = concat([ - batch_size, - seqlen, - num_attention_heads, - attention_head_size, - ]) - query = query.view(q_shape) - key = key.view(q_shape) - value = value.view(q_shape) - - embedding_weight = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, half_head_size) - embedding_weight /= rotary_embedding_scale - embedding_weight = np.split(embedding_weight.squeeze(0), 2, axis=1) - embedding_weight = np.concatenate( - [ - embedding_weight[0], - embedding_weight[0], - embedding_weight[1], - embedding_weight[1], - ], - axis=1, - ) - - if remove_input_padding: - position_embedding = unsqueeze(position_embedding, 0) - - embedding_weight = embedding_weight.astype(trt_dtype_to_np(query.dtype)) - embedding_weight = constant(embedding_weight) - position_embedding = embedding(position_embedding, embedding_weight) - position_embedding, block_embedding = split( - position_embedding, - 1, - dim=1, - ) - sin0, cos0 = split(position_embedding, half_head_size, dim=3) - sin1, cos1 = split(block_embedding, half_head_size, dim=3) - - new_shape = concat([ - batch_size, - seqlen, - 1, - half_head_size, - ]) - position_embedding = [ - tensor.view(new_shape) for tensor in [cos0, cos1, sin0, sin1] - ] - - query = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=query, - position_embedding=position_embedding, - pos_emb_type=PositionEmbeddingType.chatglm) - key = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=key, - position_embedding=position_embedding, - pos_emb_type=PositionEmbeddingType.chatglm) - - if isinstance(qkv, list): - qkv = [ - query.view(input_shape), - key.view(input_shape), - value.view(input_shape), - ] - else: - qkv = concat([query, key, value], dim=2) - qkv = qkv.view(input_shape) - - return qkv - - @staticmethod - def apply_rotary_pos_emb_cogvlm(qkv, position_embedding, - num_attention_heads, attention_head_size, - max_position_embeddings, - rotary_embedding_scale, - remove_input_padding) -> Tensor: - input = qkv[0] if isinstance(qkv, list) else qkv - input_shape = shape(input) - batch_size = 1 if remove_input_padding else shape(input, 0) - seqlen = shape(input, 0 if remove_input_padding else 1) - if isinstance(qkv, list): - query, key, value = qkv - else: - qkv = qkv.view( - concat([ - batch_size, - seqlen, - 3, - num_attention_heads, - attention_head_size, - ])) - query, key, value = split(qkv, 1, dim=2) - q_shape = concat([ - batch_size, - seqlen, - num_attention_heads, - attention_head_size, - ]) - query = query.view(q_shape) - key = key.view(q_shape) - value = value.view(q_shape) - - embedding_weight = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, attention_head_size).squeeze(0) - embedding_weight /= rotary_embedding_scale # [max_position_embeddings, attention_head_size] - - if remove_input_padding: - position_embedding = unsqueeze(position_embedding, 0) # [1, seqlen] - - embedding_weight = constant(embedding_weight) # float32 - position_embedding = embedding( - position_embedding, - embedding_weight) # [1, seqlen, attention_head_size] - sin, cos = split(position_embedding, attention_head_size // 2, - dim=-1) # [1, seqlen, attention_head_size//2] - - input_dtype = query.dtype - fp32_query = cast(query, "float32") - fp32_key = cast(key, "float32") - fp32_query = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=fp32_query, - position_embedding=[cos, sin], - pos_emb_type=PositionEmbeddingType.rope_gpt_neox) - fp32_key = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=fp32_key, - position_embedding=[cos, sin], - pos_emb_type=PositionEmbeddingType.rope_gpt_neox) - - query = cast(fp32_query, input_dtype) - key = cast(fp32_key, input_dtype) - - if isinstance(qkv, list): - qkv = [ - query.view(input_shape), - key.view(input_shape), - value.view(input_shape), - ] - else: - qkv = concat([query, key, value], dim=2) - qkv = qkv.view(input_shape) - - return qkv - - -@gw.record_signature -def gpt_attention( - *, - qkv: Tensor, - past_key_value: Tensor, - attention_mask: Optional[Tensor] = None, - attention_packed_mask: Optional[Tensor] = None, - sequence_length: Tensor, - host_past_key_value_lengths: Optional[Tensor], - host_max_attention_window_sizes: Tensor, - host_sink_token_length: Tensor, - context_lengths: Optional[Tensor], - cache_indirection: Optional[Tensor], - host_request_types: Tensor, - layer_idx: int, - num_heads: int, - num_kv_heads: int, - hidden_size_per_head: int, - q_scaling: float, - attn_logit_softcapping_scale: float = 0.0, - rotary_embedding_dim: int = 0, - rotary_embedding_base: float = 10000.0, - rotary_embedding_scale_type: RotaryScalingType = RotaryScalingType.none, - rotary_embedding_short_m_scale: float = 1.0, - rotary_embedding_long_m_scale: float = 1.0, - rotary_embedding_scale: float = 1.0, - rotary_embedding_max_positions: int = 1024, - rotary_embedding_original_max_positions: int = 1024, - position_embedding_type: PositionEmbeddingType = PositionEmbeddingType. - learned_absolute, - rotary_inv_freq: Optional[Tensor] = None, - rotary_cos_sin: Optional[Tensor] = None, - kv_orig_quant_scale: Optional[Tensor] = None, - kv_quant_orig_scale: Optional[Tensor] = None, - attention_output_orig_quant_scale: Optional[Tensor] = None, - attention_output_sf_scale: Optional[Tensor] = None, - kv_cache_quant_mode: Union[QuantModeWrapper, QuantMode] = QuantMode(0), - max_context_length: Optional[int] = None, - mask_type: AttentionMaskType = AttentionMaskType.causal, - block_sparse_block_size: int = 64, - block_sparse_homo_head_pattern: bool = False, - block_sparse_num_local_blocks: int = 16, - block_sparse_vertical_stride: int = 8, - alibi_slopes: Optional[Tensor] = None, - tp_size: int = 1, - tp_rank: int = 0, - vision_start: int = -1, - vision_length: int = -1, - kv_cache_block_offsets: Optional[Tensor] = None, - host_kv_cache_block_offsets: Tensor = None, - host_kv_cache_pool_pointers: Tensor = None, - host_kv_cache_pool_mapping: Tensor = None, - do_cross_attention: bool = False, - cross_kv: Optional[Tensor] = None, # for cross attention - cross_kv_length: Optional[Tensor] = None, # for cross attention - encoder_input_lengths: Optional[Tensor] = None, # for cross attention - relative_attention_bias: Optional[Tensor] = None, # for relative attention - logn_scaling: Optional[Tensor] = None, # for logn scaling - max_distance: int = 0, # for relative attention - host_context_lengths: Optional[Tensor] = None, # for pad-free input mode - qkv_bias: Optional[Tensor] = None, - use_cache: bool = True, - spec_decoding_is_generation_length_variable: bool = False, - spec_decoding_max_generation_length: int = 0, - spec_decoding_generation_lengths: Tensor = None, - spec_decoding_position_offsets: Tensor = None, - spec_decoding_packed_mask: Tensor = None, - spec_decoding_use: Tensor = None, - long_rope_rotary_inv_freq: Optional[Tensor] = None, - long_rope_rotary_cos_sin: Optional[Tensor] = None, - mrope_rotary_cos_sin: Tensor = None, - mrope_position_deltas: Tensor = None, - host_runtime_perf_knobs: Optional[Tensor] = None, - host_context_progress: Tensor = None, - is_mla_enabled_flag: bool = False, - q_lora_rank: int = 0, - kv_lora_rank: int = 0, - qk_nope_head_dim: int = 0, - qk_rope_head_dim: int = 0, - v_head_dim: int = 0, - q_b_proj: Optional[Tensor] = None, - kv_b_proj: Optional[Tensor] = None, - k_b_proj_trans: Optional[Tensor] = None, - skip_attn=None, - cp_group: List[int] = [0], - cp_size: int = 1, - cp_rank: int = 0, - num_kv_heads_origin: int = -1, -) -> Tuple[Tensor, Optional[Tensor]]: - ''' - Add an operation that performs the multi-head attention in GPT-like models. - - The signature of the function will change in the future release - we are in - the process of simplifying the API. The current version is still - work-in-progress! The following API is provided with hints regarding the - arguments that are likely to be removed or merged with others in the future - release. - - See docs/source/advanced/gpt-attention.md for the documentation of that function. - - Parameters: - qkv: Tensor (On GPU) - The input QKV tensor. Its shape is [batch_beam_size, max_seqlen, qkv_dim] in padded mode and [num_tokens, qkv_dim] in - packed mode. Where qkv_dim depends on using MQA, GQA, or MHA. See QKV Input in docs/source/advanced/gpt-attention.md, - - past_key_value: Tensor (On GPU) - The tensor that stores KV cache data. Its shape is - [max_batch_size * max_beam_width, 2, num_kv_heads, max_seqlen, hidden_dim_per_head] - in contiguous mode and - [max_blocks, 2, num_kv_heads, num_tokens_per_block, hidden_dim_per_head] - in paged mode. See KV Cache in docs/source/advanced/gpt-attention.md, - - attention_mask: Tensor (On GPU) - The tensor that stores the attention mask for unfused MHA or MMHA. - Its shape is [num_tokens, max_kv_seqlen]. - - attention_packed_mask: Tensor (On GPU) - The tensor that stores the packed custom mask for fmha. - Its shape is [num_tokens, max_kv_seqlen / 32], where each bit represents one mask position. - - sequence_lengths: Tensor (On GPU) - The tensor that stores the length of each sequence. Its shape is - [batch_size]. See QKV Input in docs/source/advanced/gpt-attention.md, - - host_past_key_value_lengths: Tensor (On CPU) - An INT32 tensor of shape [batch_size], - - host_max_attention_window_sizes: Tensor (On CPU) - An INT32 tensor of shape [1]. - by default, the max_attention_window_size is determined by the shape of cache_indir_table. - And we support independent max_attention_window_size for each layer. - This controls the sliding-window-attention kv-cache features. - - context_lengths: Tensor (On GPU) - The tensor that stores the context-phase sequence length of each request. Its shape - is [batch_size]. See QKV Input in doc/functional.py, - - cache_indirection: Tensor (On GPU) - The tensor to reconstruct the paths when using beam-search. Its - shape is [batch_size, beam_width, max_seqlen]. See Beam-Search in - docs/source/advanced/gpt-attention.md, - - host_request_types: Tensor = None (On CPU) - The tensor on the host that indicates if a request is in context or - generation phase. Its shape is [batch_size]. See Inflight Batching - in docs/source/advanced/gpt-attention.md, - - layer_idx: int - The index of this attention layer, used to access kv_cache_block_offsets, - - num_heads: int - The number of heads, - - num_kv_heads: int - The number of KV heads, generic to handle MHA/MQA/GQA, - - hidden_size_per_head: int - The hidden size per head, - - q_scaling: float - The value used to compute the scaling factor applied to the output - of the Q*K^T product. See Scaling Factors in docs/source/advanced/gpt-attention.md, - - attn_logit_softcapping_scale: float - The scale * tanh(value / scale) used to compute the scaling factor applied to the output - of the Q*K^T product. - - rotary_embedding_dim: int - The dimension to compute RoPE. Use 0 when position_embedding_type is not RoPE. - - rotary_embedding_base: float - The theta value to use for RoPE. Ignored when position_embedding_type is not RoPE. - - rotary_embedding_scale_type: RotaryScalingType - The scaling type of RoPE. Ignored when position_embedding_type is not RoPE. - Possible rotary scaling type: - * RotaryScalingType.none - * RotaryScalingType.linear - * RotaryScalingType.dynamic - * RotaryScalingType.longrope - * RotaryScalingType.llama3 - - rotary_embedding_scale: float - The scale value to use for linear/dynamic scaling in RoPE. - Ignored when position_embedding_type is not RoPE. - Must be set to 1 (default) if rotary_embedding_scale_type is `none`. - - rotary_inv_freq: float Tensor - The rotary inv freq with shape [head_size / 2]. - - rotary_cos_sin: float2(cos/sin) Tensor - The rotary cos/sin cache, which will be reused among different requests. - It is taken as constant tensor. - - rotary_embedding_max_positions: int - Needed only for `dynamic` RoPE scaling. Ignored otherwise. - - position_embedding_type: PositionEmbeddingType - The position embedding type: - * PositionEmbeddingType.learned_absolute - * PositionEmbeddingType.relative - * PositionEmbeddingType.rope_gptj - * PositionEmbeddingType.rope_gpt_neox - * PositionEmbeddingType.alibi - * PositionEmbeddingType.alibi_with_scale - - kv_orig_quant_scale: Tensor - The tensor to store the scaling factor for quantization to INT8/FP8 - in the KV cache. Its shape is [1]. See INT8/FP8 KV Cache in - docs/source/advanced/gpt-attention.md, - - kv_quant_orig_scale: Tensor - The tensor to store the scaling factor for dequantization from - INT8/FP8 in the KV cache. Its shape is [1]. See INT8/FP8 KV Cache - in docs/source/advanced/gpt-attention.md, - - attention_output_orig_quant_scale: Tensor - The tensor to store the scaling factor for quantization to FP8 - in the KV cache. Its shape is [1]. - - kv_cache_quant_mode: QuantMode (int flags) - Do we enable the INT8 or FP8 KV cache? - - max_context_length: int32_t - The length of the longest input sequence. See QKV Input in - docs/source/advanced/gpt-attention.md, - - mask_type: int = 1 - The type of mask: - * tensorrt_llm.layers.AttentionMaskType.padding for BERT, - * tensorrt_llm.layers.AttentionMaskType.causal for GPT, - * tensorrt_llm.layers.AttentionMaskType.sliding_window_causal for GPT, - * tensorrt_llm.layers.AttentionMaskType.bidirectional for ChatGLM-6B, - * tensorrt_llm.layers.AttentionMaskType.bidirectionalglm for GLM-10B, - * tensorrt_llm.layers.AttentionMaskType.blocksparse for Phi-3-small, - * tensorrt_llm.layers.AttentionMaskType.custom_mask for any models. - - block_sparse_block_size: int - Block size in block sparse attention - - block_sparse_homo_head_pattern: bool - Do all attention heads share same vertical stride pattern? - - block_sparse_num_local_blocks: int - Number of active blocks near diagonal - - block_sparse_vertical_stride: int - Stride of active blocks in vertical dimension - - alibi_slopes: Tensor - The ALiBi slopes. The ALiBi bias is computed on-the-fly in the kernel - when possible, - - tp_size: int - The number of processes/GPUs when tensor parallelism is activated, - - tp_rank: int - The rank of that process (when running tensor parallelism), - - kv_cache_block_offsets: - The tensor of block offsets for the KV cache. Its shape is - [num_layers, max_batch_size, max_beam_width, 2, max_blocks_per_sequence * 2], - See KV cache section in docs/source/advanced/gpt-attention.md, on gpu, - - host_kv_cache_block_offsets: - The same as kv_cache_block_offsets, but on cpu, - - host_kv_cache_pool_pointers: - The tensor of pool pointers for the KV cache. Its shape is [num_layers, 2], - See KV cache section in docs/source/advanced/gpt-attention.md, on gpu, - - host_kv_cache_pool_mapping: - The tensor of pool mapping for the different memory pools. Its shape is [num_layers,2] - for each layer, the index of the pool, and the index of the layer within the pool, - - do_cross_attention: bool = False - Do we use this as cross attention instead of self attention, - - cross_kv: Tensor = None - The KV tensor of encoder output hidden states. Its shape is [batch_size, max_seqlen, 2 * kvHeadNum * headSize] in padded mode and [1, num_tokens, 2 * kvHeadNum * headSize] in - packed mode, - - cross_kv_length: Tensor = None - The length of the longest encoder output sequence, - - encoder_input_lengths: Tensor - The tensor that stores the length of each encoder input sequence. Its shape is [batch_size], - - logn_scaling: Tensor = None - The logn scaling tensor [max_position_embedding_len], which is applied to q in order to help extrapolation - - relative_attention_bias: Tensor = None - The relative attention bias [num_heads, max_seq_len, max_seq_len], or The relative attention embedding table for implicit mode, [num_heads, num_buckets]. - - max_distance: int = 0 - The maximum distance of relative position in attention, for implicit mode. - Default value is 0, meaning to use the regular mode of relative attention bias. - Implicit mode is only enabled when passing in non-zero positive max_distance value. - See relative attention bias in docs/source/advanced/gpt-attention.md - - host_context_lengths: Tensor = None (On CPU) - A host tensor that contains the lengths of the different inputs, - - qkv_bias: Tensor = None, - The qkv bias tensor. - - use_cache: bool = False - Do we need to store kv cache ? not needed if there is no generation phase. - - spec_decoding_is_generation_length_variable: bool = False, - Whether the generation lengths can be different for each sequence in a batch. - For Medusa, this should be set False. - For Redrafter, this should be set to True. - - spec_decoding_max_generation_length: int = 1, - The maximum number of tokens possible in the generation phase per sequence. - - spec_decoding_generation_lengths: Tensor = None, - The generation phase tokens' lengths for each sequence. - Shape: [batch_size] - - spec_decoding_position_offsets: Tensor = None, - The speculative decoding tokens's position offsets (shared by all sequences). - Shape: [batch_size, num_draft_tokens + 1]. - - spec_decoding_packed_mask: Tensor = None, - The speculative decoding tokens's attention mask (packed into uint32_t bits). - remove_input_padding is False: - Shape: [batch_size, num_draft_tokens + 1, divUp(num_draft_tokens + 1, 32)]. - remove_input_padding is True: - Shape: [sum(spec_decoding_generation_lengths), divUp(num_draft_tokens + 1, 32)]. - - long_rope_rotary_inv_freq: float Tensor - Additional rotary inv freq used for longer sequence lengths. Shape: [head_size / 2] - - long_rope_rotary_cos_sin: float2(cos/sin) Tensor - Additional rotary cos/sin cache used for longer sequence lengths. - - is_mla_enable: bool = False - Do we need to enable deepseekv2 mla? - - host_runtime_perf_knobs: Tensor = None, - The runtime perf knobs bit mask, controls whether to use certain perf knob in the runtime. - - host_context_progress: Tensor = None, - The structure used to track layer-wise progress in context phase. - - skip_attn: Tensor = None, - A bool tensor on CPU. If it is true, don't run attention plugin, returning directly. - - num_kv_heads_origin: int - The origin number of KV heads, without the process of TP - - Returns: - The tensor produced by that layer. - ''' - - assert host_request_types is not None - assert (alibi_slopes is not None) == (position_embedding_type.is_alibi()) - assert (mrope_rotary_cos_sin - is not None) == (position_embedding_type.is_mrope()) - attn_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'GPTAttention', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert attn_plg_creator is not None - assert host_context_lengths is not None or not default_net( - ).plugin_config.remove_input_padding - assert isinstance(max_context_length, int) - assert host_max_attention_window_sizes is not None - assert host_sink_token_length is not None - - paged_kv_cache_flag = default_net().plugin_config.paged_kv_cache - if isinstance(qkv, list): - is_unfuse_qkv_gemm = 1 - else: - is_unfuse_qkv_gemm = 0 - - default_net().plugin_config.context_fmha_type - if do_cross_attention and not paged_kv_cache_flag: - pass - if logn_scaling is not None: - use_logn_scaling = 1 - else: - use_logn_scaling = 0 - - if num_kv_heads_origin < 1: - num_kv_heads_origin = num_kv_heads - - unfuse_qkv_gemm = trt.PluginField( - "unfuse_qkv_gemm", np.array(np.int8(is_unfuse_qkv_gemm), dtype=np.int8), - trt.PluginFieldType.INT8) - - layer_idx = trt.PluginField("layer_idx", np.array(layer_idx, - dtype=np.int32), - trt.PluginFieldType.INT32) - nheads = trt.PluginField("num_heads", np.array(num_heads, dtype=np.int32), - trt.PluginFieldType.INT32) - vision_start = trt.PluginField("vision_start", - np.array(vision_start, dtype=np.int32), - trt.PluginFieldType.INT32) - vision_length = trt.PluginField("vision_length", - np.array(vision_length, dtype=np.int32), - trt.PluginFieldType.INT32) - num_kv_heads = trt.PluginField("num_kv_heads", - np.array(num_kv_heads, dtype=np.int32), - trt.PluginFieldType.INT32) - num_kv_heads_origin = trt.PluginField( - "num_kv_heads_origin", np.array(num_kv_heads_origin, dtype=np.int32), - trt.PluginFieldType.INT32) - head_size = trt.PluginField("head_size", - np.array(hidden_size_per_head, dtype=np.int32), - trt.PluginFieldType.INT32) - unidirectional = trt.PluginField("unidirectional", - np.array(1, dtype=np.int32), - trt.PluginFieldType.INT32) - q_scaling = trt.PluginField("q_scaling", - np.array(q_scaling, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - attn_logit_softcapping_scale = trt.PluginField( - "attn_logit_softcapping_scale", - np.array(attn_logit_softcapping_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_dim = trt.PluginField( - "rotary_embedding_dim", np.array(rotary_embedding_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - rotary_embedding_base = trt.PluginField( - "rotary_embedding_base", - np.array(rotary_embedding_base, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_scale_type = trt.PluginField( - "rotary_embedding_scale_type", - np.array(rotary_embedding_scale_type, dtype=np.int8), - trt.PluginFieldType.INT8) - rotary_embedding_scale = trt.PluginField( - "rotary_embedding_scale", - np.array(rotary_embedding_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_short_m_scale = trt.PluginField( - "rotary_embedding_short_m_scale", - np.array(rotary_embedding_short_m_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_long_m_scale = trt.PluginField( - "rotary_embedding_long_m_scale", - np.array(rotary_embedding_long_m_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_max_positions = trt.PluginField( - "rotary_embedding_max_positions", - np.array(rotary_embedding_max_positions, dtype=np.int32), - trt.PluginFieldType.INT32) - rotary_embedding_original_max_positions = trt.PluginField( - "rotary_embedding_original_max_positions", - np.array(rotary_embedding_original_max_positions, dtype=np.int32), - trt.PluginFieldType.INT32) - position_embedding_type = trt.PluginField( - "position_embedding_type", - np.array(int(position_embedding_type), dtype=np.int8), - trt.PluginFieldType.INT8) - context_fmha_type = trt.PluginField( - "context_fmha_type", - np.array(np.int8(default_net().plugin_config.context_fmha_type), - dtype=np.int8), trt.PluginFieldType.INT8) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - is_spec_decoding_enabled = trt.PluginField( - "is_spec_decoding_enabled", - np.array(np.int8(spec_decoding_packed_mask is not None), dtype=np.int8), - trt.PluginFieldType.INT8) - spec_decoding_is_generation_length_variable = trt.PluginField( - "spec_decoding_is_generation_length_variable", - np.array(np.int8(spec_decoding_is_generation_length_variable), - dtype=np.int8), trt.PluginFieldType.INT8) - spec_decoding_max_generation_length = trt.PluginField( - "spec_decoding_max_generation_length", - np.array(spec_decoding_max_generation_length, dtype=np.int32), - trt.PluginFieldType.INT32) - is_mla_enabled = trt.PluginField( - "is_mla_enabled", np.array(is_mla_enabled_flag, dtype=np.int8), - trt.PluginFieldType.INT8) - q_lora_rank = trt.PluginField("q_lora_rank", - np.array(q_lora_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - kv_lora_rank = trt.PluginField("kv_lora_rank", - np.array(kv_lora_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - qk_nope_head_dim = trt.PluginField( - "qk_nope_head_dim", np.array(qk_nope_head_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - qk_rope_head_dim = trt.PluginField( - "qk_rope_head_dim", np.array(qk_rope_head_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - v_head_dim = trt.PluginField("v_head_dim", - np.array(v_head_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - p_dtype = default_net().plugin_config.gpt_attention_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - # reset mask_type to custom_mask. - if (attention_mask is not None) or (attention_packed_mask is not None): - # context fmha needs packed mask. - assert attention_packed_mask is not None - if get_sm_version() < 100: - mask_type = AttentionMaskType.custom_mask - - mask_type_filed = trt.PluginField("mask_type", - np.array([int(mask_type)], np.int32), - trt.PluginFieldType.INT32) - block_sparse_block_size = trt.PluginField( - "block_sparse_block_size", np.array([block_sparse_block_size], - np.int32), - trt.PluginFieldType.INT32) - block_sparse_homo_head_pattern = trt.PluginField( - "block_sparse_homo_head_pattern", - np.array(np.int8(block_sparse_homo_head_pattern), np.int8), - trt.PluginFieldType.INT8) - block_sparse_num_local_blocks = trt.PluginField( - "block_sparse_num_local_blocks", - np.array([block_sparse_num_local_blocks], np.int32), - trt.PluginFieldType.INT32) - block_sparse_vertical_stride = trt.PluginField( - "block_sparse_vertical_stride", - np.array([block_sparse_vertical_stride], np.int32), - trt.PluginFieldType.INT32) - tp_size = trt.PluginField("tp_size", np.array(tp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - tp_rank = trt.PluginField("tp_rank", np.array(tp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - if isinstance(kv_cache_quant_mode, QuantModeWrapper): - # Now in TRT-LLM only use global kv_cache, so it's enough to get the first quant mode from list - kv_cache_quant_mode = kv_cache_quant_mode[0] - kv_cache_quant_mode_field = trt.PluginField( - "kv_cache_quant_mode", np.array(kv_cache_quant_mode, dtype=np.int32), - trt.PluginFieldType.INT32) - paged_kv_cache = trt.PluginField( - "paged_kv_cache", np.array(paged_kv_cache_flag, dtype=np.int32), - trt.PluginFieldType.INT32) - tokens_per_block = trt.PluginField( - "tokens_per_block", - np.array(default_net().plugin_config.tokens_per_block, dtype=np.int32), - trt.PluginFieldType.INT32) - max_context_length = trt.PluginField("max_context_length", - np.array(max_context_length, np.int32), - trt.PluginFieldType.INT32) - pos_shift_enabled = trt.PluginField( - "pos_shift_enabled", - np.array(np.int8(default_net().plugin_config.streamingllm), - dtype=np.int8), trt.PluginFieldType.INT8) - dense_context_fmha = trt.PluginField( - "dense_context_fmha", - np.array(np.int8(default_net().plugin_config.streamingllm), - dtype=np.int8), trt.PluginFieldType.INT8) - if qkv_bias is None: - qkv_bias_enabled = trt.PluginField("qkv_bias_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - qkv_bias_enabled = trt.PluginField("qkv_bias_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - do_cross_attention_field = trt.PluginField( - "do_cross_attention", - np.array(np.int8(do_cross_attention), dtype=np.int8), - trt.PluginFieldType.INT8) - max_distance = trt.PluginField("max_distance", - np.array(max_distance, dtype=np.int32), - trt.PluginFieldType.INT32) - use_paged_context_fmha_field = trt.PluginField( - "use_paged_context_fmha", - np.array(np.int8(default_net().plugin_config.use_paged_context_fmha), - dtype=np.int8), trt.PluginFieldType.INT8) - use_fp8_context_fmha_field = trt.PluginField( - "use_fp8_context_fmha", - np.array(np.int8(default_net().plugin_config.use_fp8_context_fmha), - dtype=np.int8), trt.PluginFieldType.INT8) - has_full_attention_mask_field = trt.PluginField( - "has_full_attention_mask", - np.array(np.int8(attention_mask is not None), dtype=np.int8), - trt.PluginFieldType.INT8) - use_cache_pf = trt.PluginField("use_cache", - np.array([use_cache], dtype=np.int32), - trt.PluginFieldType.INT32) - fuse_fp4_quant = default_net().plugin_config.fuse_fp4_quant - fuse_fp4_quant_pf = trt.PluginField( - "fuse_fp4_quant", np.array(np.int8(fuse_fp4_quant), dtype=np.int8), - trt.PluginFieldType.INT8) - skip_attn_pf = trt.PluginField( - "skip_attn", np.array([skip_attn is not None], dtype=np.int8), - trt.PluginFieldType.INT8) - cp_size = trt.PluginField("cp_size", np.array(cp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_rank = trt.PluginField("cp_rank", np.array(cp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_group = np.array(cp_group, dtype=np.int32) - cp_group = trt.PluginField("cp_group", cp_group, trt.PluginFieldType.INT32) - use_logn_scaling = trt.PluginField( - "use_logn_scaling", np.array(np.int8(use_logn_scaling), dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - layer_idx, nheads, vision_start, vision_length, num_kv_heads, - num_kv_heads_origin, head_size, unidirectional, q_scaling, - attn_logit_softcapping_scale, position_embedding_type, - rotary_embedding_dim, rotary_embedding_base, - rotary_embedding_scale_type, rotary_embedding_scale, - rotary_embedding_short_m_scale, rotary_embedding_long_m_scale, - rotary_embedding_max_positions, rotary_embedding_original_max_positions, - tp_size, tp_rank, unfuse_qkv_gemm, context_fmha_type, - kv_cache_quant_mode_field, remove_input_padding, mask_type_filed, - block_sparse_block_size, block_sparse_homo_head_pattern, - block_sparse_num_local_blocks, block_sparse_vertical_stride, - paged_kv_cache, tokens_per_block, pf_type, max_context_length, - qkv_bias_enabled, do_cross_attention_field, max_distance, - pos_shift_enabled, dense_context_fmha, use_paged_context_fmha_field, - use_fp8_context_fmha_field, has_full_attention_mask_field, use_cache_pf, - is_spec_decoding_enabled, spec_decoding_is_generation_length_variable, - spec_decoding_max_generation_length, is_mla_enabled, q_lora_rank, - kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, - fuse_fp4_quant_pf, skip_attn_pf, cp_size, cp_rank, cp_group, - use_logn_scaling - ]) - - attn_plug = attn_plg_creator.create_plugin("causal_attn", pfc) - assert attn_plug - plug_inputs = [*qkv] if is_unfuse_qkv_gemm else [qkv] - if attention_mask is not None and mask_type == AttentionMaskType.custom_mask: - # useFullCustomMask - plug_inputs += [attention_mask] - if attention_packed_mask is not None and get_sm_version() < 100: - # usePackedCustomMask - plug_inputs += [attention_packed_mask] - if use_cache: - plug_inputs += [ - sequence_length, - host_past_key_value_lengths, - host_max_attention_window_sizes, - host_sink_token_length, - context_lengths, - cache_indirection, - host_request_types, - ] - else: - plug_inputs += [ - host_max_attention_window_sizes, - host_sink_token_length, - context_lengths, - host_request_types, - ] - if use_cache: - if paged_kv_cache_flag: - assert kv_cache_block_offsets is not None, "Paged kv cache is enabled, the kv_cache_block_offsets tensor shall not be None" - assert host_kv_cache_block_offsets is not None, "Paged kv cache is enabled, the host_kv_cache_block_offsets tensor shall not be None" - assert host_kv_cache_pool_pointers is not None, "Paged kv cache is enabled, the host_kv_cache_pool_pointers tensor shall not be None" - assert host_kv_cache_pool_mapping is not None, "Paged kv cache is enabled, the host_kv_cache_pool_mapping tensor shall not be None" - plug_inputs += [ - kv_cache_block_offsets, host_kv_cache_block_offsets, - host_kv_cache_pool_pointers, host_kv_cache_pool_mapping - ] - else: - plug_inputs += [past_key_value] - - if use_cache and kv_cache_quant_mode.has_kv_cache_quant(): - plug_inputs += [kv_orig_quant_scale, kv_quant_orig_scale] - - if attention_output_orig_quant_scale is not None: - assert default_net( - ).plugin_config.use_fp8_context_fmha, "FP8 Context FMHA needs to be enabled" - plug_inputs += [attention_output_orig_quant_scale] - - if fuse_fp4_quant: - assert attention_output_sf_scale is not None, "attention_output_sf_scale must be provided when fuse_fp4_quant is enabled." - plug_inputs += [attention_output_sf_scale] - - if rotary_inv_freq is not None: - plug_inputs += [rotary_inv_freq] - if rotary_cos_sin is not None: - plug_inputs += [rotary_cos_sin] - - if alibi_slopes is not None: - plug_inputs += [alibi_slopes] - - if relative_attention_bias is not None: - plug_inputs += [relative_attention_bias] - - if do_cross_attention: - plug_inputs += [cross_kv, cross_kv_length, encoder_input_lengths] - - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - - if qkv_bias is not None: - plug_inputs += [qkv_bias] - - if spec_decoding_packed_mask is not None: - # add position_ids as well only if speculative decoding mode - assert spec_decoding_position_offsets is not None - assert spec_decoding_generation_lengths is not None - assert spec_decoding_use is not None - plug_inputs += [ - spec_decoding_generation_lengths, spec_decoding_packed_mask, - spec_decoding_position_offsets, spec_decoding_use - ] - - if long_rope_rotary_inv_freq is not None: - assert long_rope_rotary_cos_sin is not None - plug_inputs += [long_rope_rotary_inv_freq, long_rope_rotary_cos_sin] - - if mrope_rotary_cos_sin is not None: - assert mrope_position_deltas is not None - plug_inputs += [ - mrope_rotary_cos_sin, - mrope_position_deltas, - ] - if host_runtime_perf_knobs is not None: - plug_inputs += [host_runtime_perf_knobs] - - if host_context_progress is not None: - plug_inputs += [host_context_progress] - - if is_mla_enabled_flag: - assert q_b_proj is not None - assert kv_b_proj is not None - assert k_b_proj_trans is not None - plug_inputs += [q_b_proj, kv_b_proj, k_b_proj_trans] - - if skip_attn is not None: - plug_inputs += [skip_attn] - - if logn_scaling is not None: - plug_inputs += [logn_scaling] - - for idx, i in enumerate(plug_inputs): - assert i is not None, f"Found None input for {idx} th item in plugin inputs {plug_inputs}" - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, attn_plug) - _add_plugin_info(layer, attn_plg_creator, "causal_attn", pfc) - output = _create_tensor(layer.get_output(0), layer) - expected_outputs = 1 - - # The output scaling factor tensor. - output_sf = None - if fuse_fp4_quant: - output_sf = _create_tensor(layer.get_output(expected_outputs), layer) - expected_outputs += 1 - - present_key_value = None - if use_cache and not paged_kv_cache_flag: - present_key_value = _create_tensor(layer.get_output(expected_outputs), - layer) - assert present_key_value is not None - expected_outputs += 1 - - assert layer.num_outputs == expected_outputs, \ - f"Plugin outputs number mismatch with expected, got {layer.num_outputs}, expected {expected_outputs}" - - if kv_cache_quant_mode.has_int8_kv_cache( - ) and not default_net().strongly_typed: - if not paged_kv_cache_flag: - # past key value - layer.get_input(8).set_dynamic_range(-127, 127) - # present key value - layer.get_output(expected_outputs - 1).set_dynamic_range(-127, 127) - else: - layer.get_input(0).set_dynamic_range(-127, 127) - layer.get_input(1).set_dynamic_range(-127, 127) - layer.get_output(expected_outputs - 1).set_dynamic_range(-127, 127) - - assert output is not None - if fuse_fp4_quant: - assert output_sf is not None - return (output, output_sf), present_key_value - return output, present_key_value - - -def assertion(condition: Tensor, message: str = '') -> None: - default_trtnet().add_assertion(condition.trt_tensor, message) - - -def layer_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - eps: float = 1e-05, - use_diff_of_squares: bool = True) -> Tensor: - ''' - Add a layer-norm operation on a tensor. - - That operation applies the layer-normalization to its input tensor. In its - simplest form, for large language models, the 'normalized_shape' should be - set to the hidden dimension of the activation tensor. Otherwise, it is the - shape of the normalized fraction of the tensor (starting from the - right-most dimension). - - The 'weight' tensor corresponds to 'gamma' in the layer-norm formula and - 'bias' is 'beta'. The 'eps' value is added to the variance before computing - the squared-root. - - This implementation (when using the plugin) supports an additional flag to - enable/disable the use of a difference of squares ('Var = Mean(X^2) - - Mean(X)^2'). - - Parameters: - input : Tensor - The tensor to normalize. - - normalized_shape : Union[int, Tuple[int]] - The shape of the sub-tensor that is normalized. Use 'hidden_dim' to - normalize the inner-most dimension of an activation tensor in LLMs. - - weight : Optional[Tensor] = None - The 'gamma' term in layer-norm. Its shape must be - 'normalized_shape'. - - bias : Optional[Tensor] = None - The 'beta' term in layer-norm. Its shape must be - 'normalized_shape'. - - eps : float - The epsilon term to be added to the variance in the squared-root. - - use_diff_of_squares : bool - Does the plugin use the difference of squares to compute the - variance? - - Returns: - The output tensor of that operation. - ''' - input, weight = broadcast_helper(input, weight) - input, bias = broadcast_helper(input, bias) - if isinstance(normalized_shape, int): # FIXME: better way? - axis = input.ndim() - 1 - else: - axis = input.ndim() - len(normalized_shape) - axes_mask = 0 - for i in range(axis, input.ndim()): - axes_mask |= 1 << i - layer = default_trtnet().add_normalization(input.trt_tensor, - weight.trt_tensor, - bias.trt_tensor, axes_mask) - layer.epsilon = eps - return _create_tensor(layer.get_output(0), layer) - - -def rms_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - num_groups: int = 1, - weight: Optional[Tensor] = None, - eps: float = 1e-06) -> Tensor: - ''' - Add a RMS norm operation on a tensor. - - That operation applies the rms-normalization to its input tensor. In its - simplest form, for large language models, the 'normalized_shape' should be - set to the hidden dimension of the activation tensor. Otherwise, it is the - shape of the normalized fraction of the tensor (starting from the - right-most dimension). - - The 'weight' tensor corresponds to 'gamma' in the rms-norm formula. - The 'eps' value is added to the variance before computing the squared-root. - - Parameters: - input: Tensor - The tensor to normalize. - - normalized_shape : Union[int, Tuple[int]] - The shape of the sub-tensor that is normalized. Use 'hidden_dim' to - normalize the inner-most dimension of an activation tensor in LLMs. - - num_groups: int = 1 - The group size. - - weight : Optional[Tensor] = None - The 'gamma' term in layer-norm. Its shape must be - 'normalized_shape'. - - eps : float - The epsilon term to be added to the variance in the squared-root.weig - Returns: - The output tensor of that operation. - ''' - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - - dim = tuple([-i - 1 for i in range(len(normalized_shape))]) - - if num_groups > 1: - assert len(normalized_shape) == 1 - num_channels = input.size()[-1] - ndim = input.ndim() - old_shape = shape(input) - new_shape = concat([input.size(i) for i in range(ndim - 1)] + - [num_groups, num_channels // num_groups]) - input = input.view(new_shape) - - with precision("float32"): - input_dtype = input.dtype - fp32_input = cast(input, "float32") - varx = pow(fp32_input, 2.0) - - varx = varx.mean(dim=dim, keepdim=True) - denom = varx + eps - denom = denom.sqrt() - fp32_y = fp32_input / denom - y = cast(fp32_y, input_dtype) - - if num_groups > 1: - y = y.view(old_shape) - - if weight is not None: - y = y * weight - - return y - - -def rearrange(inputs: Union[Tensor, Sequence[Tensor]], expression: str, - **kwargs) -> Tensor: - ''' - Add a rearrange operation on a tensor. - - This operation is a reader-friendly smart element reordering for multidimensional tensors, - including functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze, - stack, concatenate and other operations. Please see: https://einops.rocks/api/rearrange/ - - For example, if the shape of input tensor is [32, 30, 40, 3], and run: - `rearrange(x, 'b (h h1) (w w1) c -> b h w 1 (c h1 w1) 1', h1=2, w1=2)` - it would produce a tensor with shape as [32, 15, 20, 1, 12, 1]. - - Parameters: - input: Union[Tensor, Sequence[Tensor]] - If it is a tensor, it will directly operate on it. - Otherwise, if it is a sequence, it will concat it to a tensor and then - operates on it. - - expression : str - The expression about how to reorder the tensor in a reader-friendly way. - - kwargs: - Keyword arguments to set some identifiers with specific values. - - Returns: - The output tensor of this operation. - ''' - import re - - def _init_expression(expr): - expr_items = expr.split(" ") - tmp_name_index = 0 - for idx, item in enumerate(expr_items): - values = re.findall(r'\b\d+\b', item) - if len(values) > 0: - prefix = "(" if "(" in item else "" - subfix = ")" if ")" in item else "" - expr_items[ - idx] = f"{prefix}NumericId{tmp_name_index}Val{values[0]}{subfix}" - tmp_name_index += 1 - return " ".join(expr_items) - - def _get_all_identifier(expr): - return re.findall(r'\b[a-zA-Z_]+\d*\b', expr) - - def _get_all_symbols(expr): - return re.findall(r'\b\w+\b', expr) - - def _get_dim_expr(expr): - return [ - _get_all_symbols(match.group()) - for match in re.finditer(r'\b\w+\b|\(.*?\)', expr) - ] - - src_shape_expr, _, dst_shape_expr = expression.partition("->") - unknown_identifiers = re.findall(r'[^a-zA-Z0-9_\(\)]', - src_shape_expr + dst_shape_expr) - assert len( - unknown_identifiers) > 0, f"Unknown identifiers: {unknown_identifiers}" - src_identifiers = _get_all_identifier(src_shape_expr) - dst_identifiers = _get_all_identifier(dst_shape_expr) - assert (len(src_identifiers) == len(set(src_identifiers)) - and len(dst_identifiers) == len(set(dst_identifiers)) - ), "Indexing expression contains duplicate dimension." - assert (set(src_identifiers) == set(dst_identifiers) - ), "Identifiers only on one side of expression (should be on both)." - - new_expression = _init_expression(expression) - src_shape_expr, _, dst_shape_expr = new_expression.partition("->") - - # concat if inputs are sequence of tensors - if isinstance(inputs, Sequence): - inputs = concat([unsqueeze(t, 0) for t in inputs], dim=0) - assert ( - inputs.ndim() == len(_get_dim_expr(src_shape_expr)) - ), f"inputs.ndim() is {inputs.ndim()} while indexing expression has {len(_get_dim_expr(src_shape_expr))}" - - src_symbols = _get_all_symbols(src_shape_expr) - dst_symbols = _get_all_symbols(dst_shape_expr) - - # find all the symbols-values mapping and store them in symbol_map - symbol_map = { - symbol: { - "updated": False, - "value": None - } - for symbol in set(src_symbols + dst_symbols) - } - for symbol in symbol_map: - if "NumericId" in symbol: - symbol_map[symbol]["value"] = int(symbol.partition("Val")[-1]) - symbol_map[symbol]["updated"] = True - for symbol, value in kwargs.items(): - symbol_map[symbol]["value"] = value - symbol_map[symbol]["updated"] = True - - for idx, dim_expr in enumerate(_get_dim_expr(src_shape_expr)): - if len(dim_expr) == 1: - symbol = dim_expr[0] - if not symbol_map[symbol]["updated"]: - symbol_map[symbol]["value"] = shape(inputs, idx) - symbol_map[symbol]["updated"] = True - else: - divisors = [] - unknown_symbol = None - for symbol in dim_expr: - if not symbol_map[symbol]["updated"]: - unknown_symbol = symbol - else: - divisors.append(symbol_map[symbol]["value"]) - if unknown_symbol is not None: - assert len(divisors) > 0 - divisor = prod(cast(concat(divisors), "int64"), dim=-1) - symbol_map[unknown_symbol]["value"] = shape(inputs, - idx) / divisor - symbol_map[unknown_symbol]["updated"] = True - - for symbol, item in symbol_map.items(): - assert (item["updated"] - ), f"{symbol} cannot be inferred, please set it manually" - - dst_dims = [] - for dim_expr in _get_dim_expr(dst_shape_expr): - if len(dim_expr) == 1: - dst_dims.append(symbol_map[dim_expr[0]]["value"]) - else: - accumulator = prod(cast( - concat([symbol_map[symbol]["value"] for symbol in dim_expr]), - "int64"), - dim=-1) - dst_dims.append(accumulator) - dst_dims = cast(concat(dst_dims, dim=-1), "int64") - - src_indices = {symbol: idx for idx, symbol in enumerate(src_identifiers)} - permute_dims = [src_indices[symbol] for symbol in dst_identifiers] - - symbol_shape = cast( - concat([symbol_map[symbol]["value"] for symbol in src_identifiers], - dim=-1), "int64") - tensor = inputs.view(symbol_shape) - tensor = permute(tensor, permute_dims) - tensor = tensor.view(dst_dims) - return tensor - - -def repeat(input: Tensor, sizes: Sequence[int]) -> Tensor: - ''' - Repeats the tensor along the specified dimensions. - - Parameters: - input : Tensor - The tensor to be repeated. - sizes : Sequence[int] - The number of times to repeat the tensor along each dimension. - - Returns: - A tensor except for repeated input tensors along specified dim. - - ''' - assert input.ndim() <= len(sizes), \ - "Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor" - repeated_tensor = input - for k in range(-1, -len(sizes) - 1, -1): - repeated_tensor = concat([repeated_tensor] * sizes[k], dim=k) - return repeated_tensor - - -def repeat_interleave(tensor: Tensor, repeats: int, dim: int) -> Tensor: - ''' - Repeats elements of a tensor along an axis. - - Parameters: - repeats : int - The number of repetitions along axis specified. - dim : int - The dimension along which repetitions are performed. - - Returns: - A tensor with the same shape as input except for repeated elements along specified dim. - - TODO: Allow repeats to be a list of integers and dim to be unspecified. - ''' - expanded_tensor = expand_dims(tensor, dim + 1) - tile_output_size = concat([ - repeats if i == (dim + 1) else shape(expanded_tensor, i) - for i in range(expanded_tensor.ndim()) - ]) - tile = expand(expanded_tensor, tile_output_size) - tile_reshape_size = [shape(tensor, i) for i in range(tensor.ndim())] - tile_reshape_size[dim] = tile_reshape_size[dim] * repeats - tensor = tile.view(concat(tile_reshape_size)) - return tensor - - -def meshgrid2d(x: Tensor, y: Tensor) -> Tuple[Tensor]: - ''' - Creates grids (2D) of coordinates specified by the 1D inputs (only supports `indexing=\'xy\'`). - - Parameters: - x : Tensor - The first input (1D) tensor. - y : Tensor - The second input (1D) tensor. - - Returns: - The tuple of two tensors produced. - - TODO: Add full support for torch.meshgrid. - See https://pytorch.org/docs/stable/generated/torch.meshgrid.html#torch-meshgrid - ''' - if x.ndim() == 1: - x = expand_dims(x, 0) - if y.ndim() == 1: - y = expand_dims(y, 0) - grid_x = repeat_interleave(x, shape(y, 1), - 1).view([x.shape[-1], y.shape[-1]]) - grid_y = repeat(y, (x.shape[-1], 1)) - return (grid_x, grid_y) - - -def generate_logn_scaling(seq_length: int = 8192, - max_position_embeddings: int = 32768) -> np.ndarray: - ''' - Compute the Log-N scaling vector for Qwen inference extrapolation - - Parameters: - seq_length : int - The max seq length in training (default to 8192 in Qwen-1) - max_position_embeddings : int - The max position embeddings. (default to 32768 in Qwen-1) - - Returns: - A constant np.ndarray that contains logn scaling vector - ''' - logn_list = [ - math.log(i, seq_length) if i > seq_length else 1 - for i in range(1, max_position_embeddings + 1) - ] - return np.asarray(logn_list, dtype=np.float32) - - -def generate_alibi_slopes(num_heads: int, - tp_size: int = 1, - tp_rank: int = 0, - alibi_scale: float = 1.0, - alibi_bias_max: int = 8) -> np.ndarray: - ''' - Compute the ALiBi slopes as described in https://arxiv.org/abs/2211.05100. - - Parameters: - num_heads : int - The number of heads. - dtype : trt.DataType - The data type of the returned slopes - tp_size : int - The tensor parallelism size - tp_rank : int - The tensor parallelism rank - - Returns: - A constant tensor that contains the ALiBi slopes. - ''' - start_head_id = 0 - end_head_id = num_heads - - if tp_size > 1: - rank_heads = num_heads // tp_size - start_head_id = rank_heads * tp_rank - end_head_id = start_head_id + rank_heads - - closest_power_of_2 = 2**np.floor(np.log2(num_heads)) - # FT's implementation - # https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/kernels/gen_relative_pos_bias.cu#L248 - slopes_ft = [] - for h_id in range(start_head_id, end_head_id): - if h_id < closest_power_of_2: - slopes_ft.append( - np.power( - 2**(-(2**-(np.log2(closest_power_of_2) - - np.log2(alibi_bias_max)))), h_id + 1)) - else: - slopes_ft.append( - np.power( - 2**(-(2**-(np.log2(closest_power_of_2 * 2) - - np.log2(alibi_bias_max)))), - (h_id - closest_power_of_2) * 2 + 1)) - slopes = np.asarray(slopes_ft, dtype=np.float32) - - slopes = alibi_scale * slopes - slopes = slopes.reshape(1, (end_head_id - start_head_id), 1, 1) - return slopes - - -def generate_alibi_biases(slopes: Tensor, key_length: Tensor) -> Tensor: - ''' - Compute the ALiBi biases as described in https://arxiv.org/abs/2211.05100. - - The ALiBi biases are added to the result of the Q*K^T product in the - multi-head attention block. - - Parameters: - slopes : Tensor - The slopes. - - key_length : Tensor - The size of the K vector per head. - - Returns: - A constant tensor that contains the ALiBi biases. - ''' - # We don't need to care about the batch size or query length since we can just broadcast - # across the batch and query dimensions - - trt_0 = constant(int32_array(0)) - arange_shape = concat([1, 1, 1, key_length]) - - arange_tensor = arange(trt_0, key_length, "float32").view(arange_shape) - return slopes * arange_tensor - - -def expand_mask(mask: Tensor, tgt_len: Optional[Tensor] = None) -> Tensor: - ''' - Expand an attention mask. - - That function adds the sequence of operations to expand from a tensor of - shape '[batch_size, src_seq_len]' to a tensor of shape - '[batch_size, 1, tgt_seq_len, src_seq_len]'. It can be used to create the - mask applied to the Q*K^T product before the softmax operation in the - multi-head attention block. - - Parameters: - mask : Tensor - The input mask - - tgt_len : Optional[Tensor] - The dimension of the 3rd dimension in the output tensor. If None, - the 2nd dimension of the input is used. - - Returns: - The tensor created by that sequence of operations. - ''' - bsz = shape(mask, 0) - src_len = shape(mask, 1) - tgt_len = tgt_len if tgt_len is not None else src_len - - mask = mask.view(concat([bsz, 1, 1, src_len])) - - mask = expand(mask, concat([bsz, 1, tgt_len, src_len])) - mask = where(mask == 0, float('-inf'), 0.0) - return mask - - -def gather_last_token_logits(hidden_states: Tensor, last_token_ids: Tensor, - remove_input_padding: bool) -> Tensor: - ''' - Extract the logits that correspond to the last token from the hidden states. - - That function adds the operations to extract the logits of the last tokens - in a batch of sequences. - - Depending on whether 'remove_input_padding' is 'True' or 'False', that - function assumes inputs of different shapes. - - When 'remove_input_padding' is 'True', the 'hidden_states' tensor is - assumed to be packed. It has a shape '[num_tokens, hidden_dim]' where - 'num_tokens' is the sum of the lengths of the sequences in the batch and - 'hidden_dim' is the hidden dimension. The 'last_tokens_ids' is a 1D tensor - that encodes the inclusive prefix-sums of the lengths of the sequences in - the batch. - - When 'remove_input_padding' is 'False', the 'hidden_states' tensor is - assumed to be padded. It has a shape '[batch_size, max_seqlen, hidden_dim]' - where 'max_seqlen' is the length of the longest sequence in the batch and - 'hidden_dim' is the hidden dimension. The 'last_token_ids' is a 1D tensor - that encodes the length of each sequence in the batch. - - In both cases, that function produces a tensor of shape '[batch_size, - hidden_size]' where the row at index 'i' corresponds to the logits of the - last token from the 'i'-th sequence. - - Parameters: - hidden_states : Tensor - The hidden states - - last_token_ids : Tensor - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - remove_input_padding : bool - Indicate if the hidden_states are packed ('True') or padded - ('False'). - - Returns: - The tensor created by that sequence of operations. - ''' - if last_token_ids is None: - return hidden_states - - if remove_input_padding: - hidden_states = index_select(hidden_states, 0, - last_token_ids - 1) # [seq_len, hidden] - - hidden_states = hidden_states.view( - concat([shape(last_token_ids, 0), - shape(hidden_states, 1)])) - else: - ndim = last_token_ids.ndim() - if ndim == 1: - # only calculate logits for the last token - # [batch_size, seqlen, hidden_size] -> [batch_size, hidden_size] - last_token_ids = last_token_ids.view( - concat([shape(last_token_ids, 0), 1, 1])) - last_token_ids = expand( - last_token_ids, - concat([shape(last_token_ids, 0), 1, - shape(hidden_states, 2)])) - last_token_ids = last_token_ids - 1 - hidden_states = gather( - hidden_states, dim=1, indices=last_token_ids).view( - concat([shape(hidden_states, 0), - shape(hidden_states, 2)])) - elif ndim == 2: # speculative decoding needs last few token's logits - # last_token_ids is of shape [batch_size, num_last_tokens] - # So [batch_size, seqlen, hidden_size] -> [batch_size, num_last_tokens, hidden_size] - last_token_ids = last_token_ids.view( - concat([shape(last_token_ids, 0), - shape(last_token_ids, 1), 1])) - last_token_ids = expand( - last_token_ids, - concat([ - shape(last_token_ids, 0), - shape(last_token_ids, 1), - shape(hidden_states, 2) - ])) - hidden_states = gather(hidden_states, dim=1, indices=last_token_ids) - return hidden_states - - -ACT2FN = { - 'relu': relu, - 'tanh': tanh, - 'gelu': gelu, - 'gelu_new': gelu, - 'gelu_fast': gelu, - 'gelu_pytorch_tanh': gelu, - 'openai-gelu': gelu, - 'geglu': geglu, - 'gegelu': gegelu, - 'identity': identity, - 'silu': silu, - 'softplus': softplus, - 'relu2': squared_relu, - 'squared-relu': squared_relu, - 'swiglu': swiglu, - 'fast-swiglu': swiglu, - 'sigmoid': sigmoid, - 'quick_gelu': quick_gelu, -} - -GATED_ACT_2_ACT = { - 'swiglu': 'silu', - 'fast-swiglu': 'silu', - 'geglu': 'gelu', -} - - -def is_gated_activation(activation): - ''' - Is a given activation function gated? - - Parameters: - activation : str - The name of the activation function. - - Returns: - True if the function is gated, False otherwise. - ''' - assert activation in ACT2FN - return activation in GATED_ACT_2_ACT - - -def non_gated_version(activation): - ''' - Given an activation function, get the non-gated version. - - If the activation function is non-gated, it returns the same activation - function name. - - For example, that function returns 'silu' for 'swiglu' and 'relu' for - 'relu'. - - Parameters: - activation : str - The name of the activation function. - - Returns: - The name of the non-gated activation function. - ''' - if is_gated_activation(activation): - return GATED_ACT_2_ACT[activation] - return activation - - -def lora_plugin( - input: Tensor = None, - in_hidden_size: int = 0, - out_hidden_sizes: List[int] = [0], - host_request_types: Tensor = None, - transa: bool = False, - transb: bool = False, - host_context_lengths: Tensor = None, # for pad-free input mode - max_low_rank: int = 0, - lora_ranks: List[Tensor] = None, - lora_weights_pointers: List[Tensor] = None, - weight_index: int = 0, -): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - in_hidden_size/out_hidden_size : int - the lora computation workflow is - [M, in_hidden_size] -> [M, low_rank] -> [M, out_hidden_size] - - host_request_types : Tensor = None - The tensor on the host that indicates if a request is in context or - generation phase. Its shape is [batch_size]. See Inflight Batching - in docs/source/advanced/gpt-attention.md, - - transa : bool - Is the first input transposed? Set to 'True' if you want the first - input to be transposed, 'False' otherwise. - - transb : bool - Is the second input transposed? Set to 'True' if you want the - second input to be transposed, 'False' otherwise. - - host_context_lengths: cpu Tensor = None - A host tensor that contains the lengths of the different inputs, - - max_low_rank : int - Maximum low_rank, used to determine the workspace size. - - lora_ranks : cpu Tensor with shape [batch_size] - The low_rank of each request - - lora_weights_pointers : cpu int64 Tensor with shape [batch_size, 3] - The weights pointers of each request. Consist of in_pointer, out_pointer and possibly a scales vector pointer. - - weight_index : int - The index of weight if the weight pointer pointing to multiple weights. - - Return: - The tensor produced by that layer. - - ''' - assert host_context_lengths is not None or not default_net( - ).plugin_config.remove_input_padding - - trt.get_plugin_registry().plugin_creator_list - in_hidden_size_field = trt.PluginField( - "in_hidden_size", np.array(in_hidden_size, dtype=np.int32), - trt.PluginFieldType.INT32) - out_hidden_size_field_list = [ - trt.PluginField(f"out_hidden_size_{i}", np.array(o, dtype=np.int32), - trt.PluginFieldType.INT32) - for i, o in enumerate(out_hidden_sizes) - ] - transa = 1 if transa else 0 - transa = trt.PluginField("transa", np.array(transa, dtype=np.int32), - trt.PluginFieldType.INT32) - transb = 1 if transb else 0 - transb = trt.PluginField("transb", np.array(transb, dtype=np.int32), - trt.PluginFieldType.INT32) - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Lora', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.lora_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - max_low_rank_field = trt.PluginField("max_low_rank", - np.array(max_low_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - weight_index_field = trt.PluginField("weight_index", - np.array(weight_index, dtype=np.int32), - trt.PluginFieldType.INT32) - num_lora_modules = len(out_hidden_sizes) - num_lora_modules_field = trt.PluginField( - "num_lora_modules", np.array(num_lora_modules, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - in_hidden_size_field, transa, transb, num_lora_modules_field, pf_type, - remove_input_padding, max_low_rank_field, weight_index_field - ] + out_hidden_size_field_list) - lora_plug = plg_creator.create_plugin("lora", pfc) - - plug_inputs = [input.cast(p_dtype), host_request_types - ] + lora_ranks + lora_weights_pointers - - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, lora_plug) - - if num_lora_modules == 1: - return _create_tensor(layer.get_output(0), layer).cast(input.dtype) - else: - return [ - _create_tensor(layer.get_output(i), layer).cast(input.dtype) - for i in range(num_lora_modules) - ] - - -def dora_plugin(activations: Tensor, - out_hidden_sizes: list[int], - lora_weights_pointers: list[Tensor], - host_request_types: Tensor, - host_context_lengths: Tensor | None = None) -> Tensor: - ''' - The DoRA plugin applies column-wise scaling to the output of a LoRA layer. - - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - out_hidden_sizes : list[int] - The output hidden size of each adapter in the related LoRA module. - For example, for a qkv projection out_hidden_sizes should be [q_dim, k_dim, v_dim]. - - host_request_types : Tensor = None - The tensor on the host that indicates if a request is in context or - generation phase. Its shape is [batch_size]. See Inflight Batching - in docs/source/advanced/gpt-attention.md, - - host_context_lengths: cpu Tensor = None - A host tensor that contains the lengths of the different inputs, - - Return: - The tensor produced by that layer. - - ''' - assert host_context_lengths is not None or not default_net( - ).plugin_config.remove_input_padding - - dora_plg_creator = trt.get_plugin_registry().get_creator( - 'Dora', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert dora_plg_creator is not None - - out_hidden_sizes = trt.PluginField( - f"out_hidden_sizes", np.array(out_hidden_sizes, dtype=np.int32), - trt.PluginFieldType.INT32) - - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - - lora_dtype = default_net().plugin_config.lora_plugin - type_id = trt.PluginField( - "type", np.array(int(str_dtype_to_trt(lora_dtype)), np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [type_id, remove_input_padding, out_hidden_sizes]) - - dora_plug = dora_plg_creator.create_plugin("dora", pfc, - trt.TensorRTPhase.BUILD) - - plug_inputs = [activations.cast(lora_dtype), host_request_types - ] + lora_weights_pointers - - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v3(plug_inputs, [], dora_plug) - _add_plugin_info(layer, dora_plg_creator, "dora", pfc) - output = _create_tensor(layer.get_output(0), layer).cast(activations.dtype) - return output - - -def mamba_conv1d(input: Tensor, - conv_state_or_ptr: Tensor, - conv_weight: Tensor, - conv_bias: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - dim: int, - dconv: int, - dtype: str, - pre_stride: int = 0, - post_stride: int = 0, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - apply_silu: bool = True): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - conv_state_or_ptr : Tensor (On GPU or CPU) - The conv state tensor. Its shape is [batch_size, dconv - 1, dim] - Or the CPU tensor of shape [1] for the pointer of paged states. - - conv_weight : Tensor (On GPU) - The weight tensor. Its shape is [1, dconv, dim] - - conv_bias : Tensor (On GPU) - The bias tensor. Its shape is [dim] - - host_request_types : Tensor (On CPU) - The tensor on the host that indicates if a request is in context or - generation phase. Its shape is [batch_size]. See Inflight Batching - in docs/source/advanced/gpt-attention.md, - - last_token_ids : Tensor (On GPU) - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - dim : int - The hidden dimension of conv1d - - dconv : int - The window size of conv1d - - dtype: str - data type - - pre_stride : int = 0 - The (pre) stride size of the input tensor. - The valid values of the input tensor are input[..., pre_stride: dim-post_stride] - - post_stride : int = 0 - The (post) stride size of the input tensor. - The valid values of the input tensor are input[..., pre_stride: dim-post_stride] - - host_context_lengths: Tensor (On CPU) (Optional) - A host tensor that contains the lengths of the different inputs, - - slot_mapping: Tensor (On GPU) (Optional) - Real page index in state. Its shape is [dim], used for paged state, each page shape is [dconv, dim] - - apply_silu: bool - Is there a SiLU operation after the conv1d? When True apply - SiLU activation function after the conv1d. - ''' - assert host_request_types is not None - mamba_conv1d_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'MambaConv1d', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert mamba_conv1d_plg_creator is not None - - dim = trt.PluginField("dim", np.array(dim, dtype=np.int32), - trt.PluginFieldType.INT32) - dconv = trt.PluginField("dconv", np.array(dconv, dtype=np.int32), - trt.PluginFieldType.INT32) - pre_stride = trt.PluginField("pre_stride", - np.array(pre_stride, dtype=np.int32), - trt.PluginFieldType.INT32) - post_stride = trt.PluginField("post_stride", - np.array(post_stride, dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - paged_state = trt.PluginField( - "paged_state", - np.array(np.int8(default_net().plugin_config.paged_state), - dtype=np.int8), trt.PluginFieldType.INT8) - apply_silu = trt.PluginField("apply_silu", - np.array(np.int8(apply_silu), dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - dim, dconv, pre_stride, post_stride, pf_type, remove_input_padding, - paged_state, apply_silu - ]) - mamba_conv1d_plug = mamba_conv1d_plg_creator.create_plugin( - "mamba_conv1d", pfc) - plug_inputs = [ - input, conv_state_or_ptr, conv_weight, conv_bias, host_request_types, - last_token_ids - ] - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - if default_net().plugin_config.paged_state: - plug_inputs += [slot_mapping] - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, mamba_conv1d_plug) - _add_plugin_info(layer, mamba_conv1d_plg_creator, "mamba_conv1d", pfc) - output = _create_tensor(layer.get_output(0), layer) - if default_net().plugin_config.paged_state: - return output, None - else: - present_state = _create_tensor(layer.get_output(1), layer) - return output, present_state - - -def selective_scan(input: Tensor, - state_or_ptr: Tensor, - delta: Tensor, - delta_bias: Tensor, - A: Tensor, - BC: Tensor, - D: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - dim: int, - dstate: int, - dt_rank: int, - delta_softplus: bool, - dtype: str, - z: Optional[Tensor] = None, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - nheads: int = 1, - ngroups: int = 1, - chunk_size: int = 256, - mamba_version: str = 'Mamba1'): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] - - state_or_ptr : Tensor (On GPU or CPU) - The ssm state tensor. Its shape is [batch_size, dstate, dim] - Or the CPU tensor of shape [1] for the pointer of paged states. - - delta : Tensor (On GPU) - The delta tensor. - mamba: Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - mamba2: Its shape is [batch_size, seq_len, nheads] or [num_tokens, nheads] for remove_input_padding - - delta_bias : Tensor (On GPU) - The delta bias tensor. - mamba: Its shape is [dim] - mamba2: Its shape is [nheads] - - A : Tensor (On GPU) - A matrix. - mamba: Its shape is [dstate, dim] - mamba2: Its shape is [nheads] - - BC : Tensor (On GPU) - B and C matrix. - mamba: Its shape is [batch_size, seq_len, dstate * 2] or [num_tokens, dstate * 2] for remove_input_padding - mamba2: Its shape is [batch_size, seq_len, ngroups * dstate * 2] or [num_tokens, ngroups * dstate * 2] for remove_input_padding - - D : Tensor (On GPU) - D matrix. - mamba: Its shape is [dim] - mamba2: Its shape is [nheads] - - host_request_types : Tensor (On CPU) - The tensor on the host that indicates if a request is in context or - generation phase. Its shape is [batch_size]. See Inflight Batching - in docs/source/advanced/gpt-attention.md - - last_token_ids : Tensor (On GPU) - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - dim : int - The inner dimension of SSM block - - dstate : int - The state dimension of SSM block - - dt_rank: int - The rank dimension of dt_proj - - delta_softplus : bool - Do we apply softplus to the delta. - - dtype: str - data type - - z : Tensor (On GPU) (Optional) - The z tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - host_context_lengths: Tensor (On CPU) (Optional) - A host tensor that contains the lengths of the different inputs, - - slot_mapping: Tensor (On GPU) (Optional) - Real page index in state. Its shape is [dim], used for paged state, each page shape is [dstate, dim] - - nheads: int (Optional) - The number of heads. - - ngroups: int (Optional) - The number of groups. - - chunk_size: int (Optional) - The chunk_size is used for the chunk_scan kernel. - - mamba_version: int (Optional) - Mamba version, support Mamba1 as default. - ''' - assert host_request_types is not None - selective_scan_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'SelectiveScan', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert selective_scan_plg_creator is not None - - dim = trt.PluginField("dim", np.array(dim, dtype=np.int32), - trt.PluginFieldType.INT32) - dstate = trt.PluginField("dstate", np.array(dstate, dtype=np.int32), - trt.PluginFieldType.INT32) - dt_rank = trt.PluginField("dt_rank", np.array(dt_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - nheads = trt.PluginField("nheads", np.array(nheads, dtype=np.int32), - trt.PluginFieldType.INT32) - ngroups = trt.PluginField("ngroups", np.array(ngroups, dtype=np.int32), - trt.PluginFieldType.INT32) - chunk_size = trt.PluginField("chunk_size", - np.array(chunk_size, dtype=np.int32), - trt.PluginFieldType.INT32) - delta_softplus = trt.PluginField( - "delta_softplus", np.array(np.int8(delta_softplus), dtype=np.int8), - trt.PluginFieldType.INT8) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - paged_state = trt.PluginField( - "paged_state", - np.array(np.int8(default_net().plugin_config.paged_state), - dtype=np.int8), trt.PluginFieldType.INT8) - if z is None: - z_enabled = trt.PluginField("z_enabled", np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - z_enabled = trt.PluginField("z_enabled", np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - is_mamba2 = trt.PluginField( - "is_mamba2", - np.array(1 if mamba_version == 'Mamba2' else 0, dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - dim, dstate, dt_rank, nheads, ngroups, chunk_size, delta_softplus, - pf_type, remove_input_padding, paged_state, z_enabled, is_mamba2 - ]) - selective_scan_plug = selective_scan_plg_creator.create_plugin( - "selective_scan", pfc) - - plug_inputs = [ - input, state_or_ptr, delta, delta_bias, A, BC, D, host_request_types, - last_token_ids - ] - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - if default_net().plugin_config.paged_state: - plug_inputs += [slot_mapping] - if z is not None: - plug_inputs += [z] - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, selective_scan_plug) - _add_plugin_info(layer, selective_scan_plg_creator, "selective_scan", pfc) - output = _create_tensor(layer.get_output(0), layer) - if default_net().plugin_config.paged_state: - return output, None - else: - present_state = _create_tensor(layer.get_output(1), layer) - return output, present_state - - -def rg_lru(input: Tensor, - A: Tensor, - state_or_ptr: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - dim: int, - dtype: str, - block_size: int = 0, - y: Optional[Tensor] = None, - y_bias: Optional[Tensor] = None, - gate: Optional[Tensor] = None, - gate_bias: Optional[Tensor] = None, - gate_x: Optional[Tensor] = None, - gate_x_bias: Optional[Tensor] = None, - gate_a: Optional[Tensor] = None, - gate_a_bias: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] - - A : Tensor (On GPU) - A matrix. Its shape is [dim] - - state_or_ptr : Tensor (On GPU or CPU) - The lru state tensor. Its shape is [batch_size, dstate, dim] - Or the CPU tensor of shape [1] for the pointer of paged states. - - host_request_types : Tensor (On CPU) - The tensor on the host that indicates if a request is in context or - generation phase. Its shape is [batch_size]. See Inflight Batching - in docs/source/advanced/gpt-attention.md, - - last_token_ids : Tensor (On GPU) - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - dim : int - The inner dimension of RG_LRU block - - block_size : int - The block size of the block diagonal linear layer. It is used to - support the cases that enable fused gate. - - dtype: str - data type - - y : Tensor (On GPU) (Optional) - The y tensor. Its shape is [batch_size, seq_len, dim] - - y_bias : Tensor (On GPU) (Optional) - The y_bias tensor. Its shape is [dim]. If y_bias is not None, we - will fuse GELU(y + y_bias) in this function. - - gate : Tensor (On GPU) (Optional) - The gate tensor. Its shape is [batch_size, seq_len, 2 * dim]. - If gate is not None, we will fuse the gate_x and gate_a, otherwise - use those two tensors. - - gate_bias : Tensor (On GPU) (Optional) - The gate_bias tensor. Its shape is [2 * block_num, dim // block_num]. - If gate_bias is not None, we will fuse the bias add in this function. - - gate_x : Tensor (On GPU) (Optional) - The gate_x tensor. Its shape is [batch_size, seq_len, dim] - - gate_x_bias : Tensor (On GPU) (Optional) - The gate_x_bias tensor. Its shape is [block_num, dim // block_num]. - If gate_x_bias is not None, we will fuse the bias add in this function. - - gate_a : Tensor (On GPU) (Optional) - The gate_a tensor. Its shape is [batch_size, seq_len, dim] - - gate_a_bias : Tensor (On GPU) (Optional) - The gate_a_bias tensor. Its shape is [block_num, dim // block_num]. - If gate_a_bias is not None, we will fuse the bias add in this function. - - slot_mapping: Tensor (On GPU) (Optional) - Real page index in state. Its shape is [dim], used for paged state, each page shape is [dstate, dim] - ''' - assert host_request_types is not None - lru_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LRU', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert lru_plg_creator is not None - assert (gate_x_bias is None) == (gate_a_bias is None) - enable_fuse_gate = gate is not None - has_gate_bias = (gate_bias is not None) or (gate_x_bias is not None) - if enable_fuse_gate: - assert gate is not None - assert block_size > 0 - if has_gate_bias: - assert gate_bias is not None - else: - assert gate_x is not None and gate_a is not None - if has_gate_bias: - assert gate_x_bias is not None and gate_a_bias is not None - - dim = trt.PluginField("dim", np.array(dim, dtype=np.int32), - trt.PluginFieldType.INT32) - block_size = trt.PluginField("block_size", - np.array(block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - paged_state = trt.PluginField( - "paged_state", - np.array(np.int8(default_net().plugin_config.paged_state), - dtype=np.int8), trt.PluginFieldType.INT8) - - if y is None: - y_enabled = trt.PluginField("y_enabled", np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - y_enabled = trt.PluginField("y_enabled", np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - - if y_bias is None: - y_bias_enabled = trt.PluginField("y_bias_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - y_bias_enabled = trt.PluginField("y_bias_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - - if enable_fuse_gate: - fuse_gate_enabled = trt.PluginField("fuse_gate_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - fuse_gate_enabled = trt.PluginField("fuse_gate_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - - if has_gate_bias: - gate_bias_enabled = trt.PluginField("gate_bias_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - gate_bias_enabled = trt.PluginField("gate_bias_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - dim, block_size, pf_type, remove_input_padding, paged_state, y_enabled, - y_bias_enabled, fuse_gate_enabled, gate_bias_enabled - ]) - lru_plug = lru_plg_creator.create_plugin("rg_lru", pfc) - - plug_inputs = [ - input, - A, - state_or_ptr, - host_request_types, - last_token_ids, - ] - if default_net().plugin_config.paged_state: - plug_inputs += [slot_mapping] - if y is not None: - plug_inputs += [y] - if y_bias is not None: - plug_inputs += [y_bias] - if enable_fuse_gate: - plug_inputs += [gate] - if has_gate_bias: - plug_inputs += [gate_bias] - else: - plug_inputs += [gate_x, gate_a] - if has_gate_bias: - plug_inputs += [gate_x_bias, gate_a_bias] - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, lru_plug) - _add_plugin_info(layer, lru_plg_creator, "rg_lru", pfc) - output = _create_tensor(layer.get_output(0), layer) - if default_net().plugin_config.paged_state: - return output, None - else: - present_state = _create_tensor(layer.get_output(1), layer) - return output, present_state - - -def topk(input: Tensor, - k: Union[Tensor, int], - dim: int, - largest: bool = True, - prefer_plugin: bool = True) -> Tuple[Tensor, Tensor]: - ''' - Add an topk operation. - - As explained in the ONNX documentation, - - https://github.com/onnx/onnx/blob/main/docs/Operators.md#topk - - NOTE: One distinction from the ONNX topk op, the output is always sorted - with TensorRT layer. - - Retrieve the top-K largest elements along a specified axis. - Given an input tensor of shape [a_1, a_2, ..., a_n, r] - and integer argument k, return two outputs: - Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which contains the values of the top k elements along the specified axis - Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which contains the indices of the top k elements (original indices from the input tensor). - - Parameters: - input : Tensor - The input tensor. + @staticmethod + def create_sinusoidal_positions_for_cogvlm_attention_plugin( + num_pos: int, + dim: int, + theta: float = 10000.0, + scale: float = 1.0, + scale_type: RotaryScalingType = RotaryScalingType.none, + vision_start: int = 1, + vision_length: int = 1225, + dtype=np.float32, + ): + if scale_type == RotaryScalingType.linear: + scale = 1.0 / scale + inv_freq = scale / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) + position_id = np.hstack([ + np.arange(0, vision_start + 1, dtype=dtype), + np.full(vision_length, vision_start + 1, dtype=dtype), + np.arange(vision_start + 2, + num_pos - (vision_length - 1), + dtype=dtype), + ]) + sinusoid_inp = np.expand_dims(np.einsum("i , j -> i j", + position_id, + inv_freq, + dtype=dtype), + axis=-1) + # fuse cos/sin into float2 (cos, sin). + concat = np.concatenate((np.cos(sinusoid_inp), np.sin(sinusoid_inp)), + axis=-1) - k : int - A single positive value corresponding to the number of top elements to retrieve + return inv_freq, concat.reshape(1, -1).astype(dtype) - dim: int - The dimension in which to compute the topk indices. + def create_sinusoidal_positions_long_rope_for_attention_plugin( + num_pos: int, + num_orig_pos: int, + dim: int, + theta: float = 10000.0, + scaling_short_factors: Tensor = 1.0, + scaling_long_factors: Tensor = 1.0, + short_mscale=None, + long_mscale=None, + dtype=np.float32, + ): - largest: bool - Controls whether to return largest or smallest elements + def _calc_mscale(scale): + if scale <= 1.0: + return 1.0 + return math.sqrt(1 + math.log(scale) / math.log(num_orig_pos)) - prefer_plugin : bool - Whether to use the topkLastDim plugin if dim is last dim and k is static. + if short_mscale is None: + short_mscale = _calc_mscale(num_pos / num_orig_pos) + long_mscale = short_mscale + def _compute_sinusoidal_positions(scale_factors, is_short, + for_attention_plugin): + inv_freq = 1 / (scale_factors * + (theta**(np.arange(0, dim, 2) / dim)).astype(dtype)) + sinusoid_inp = np.einsum("i , j -> i j", + np.arange(num_pos, dtype=dtype), + inv_freq, + dtype=dtype) - Returns: - The tensors (values, indices) produced by this topk operation. - ''' - dim = dim_resolve_negative(dim, input.ndim())[0] - if prefer_plugin and dim == input.ndim() - 1 and not isinstance(k, Tensor): - last_dim = input.size(-1) - if last_dim == -1: # dynamic? - last_dim = shape(input, -1) - # since we might need to flatten the input to 2d tensor, - # we need to prepare the output shape - out_shape = [] - for i in range(input.ndim() - 1): - out_shape.append(shape(input, i)) - out_shape = concat(out_shape + [k]) - if input.ndim() == 1: - input_2d = unsqueeze(input, - 0) # special handling of rank-1 dynamic tensor - elif input.ndim() != 2: - input_2d = input.view(concat([-1, last_dim]), - zero_is_placeholder=False) - else: - input_2d = input - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "TopkLastDim", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - is_largest = trt.PluginField( - "is_largest", np.array(1 if largest else 0, dtype=np.int32), - trt.PluginFieldType.INT32) - k = trt.PluginField("k", np.array(k, dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField("type_id", - np.array([int(input_2d.dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([pf_type, k, is_largest]) - topk_last_dim_plug = plg_creator.create_plugin("topk_last_dim", pfc) - plug_inputs = [input_2d] - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, topk_last_dim_plug) - _add_plugin_info(layer, plg_creator, "topk_last_dim", pfc) - values = _create_tensor(layer.get_output(0), layer) - indices = _create_tensor(layer.get_output(1), layer) - values = values.view(out_shape, zero_is_placeholder=False) - indices = indices.view(out_shape, zero_is_placeholder=False) - else: - # non-plugin path - axes = dim_to_trt_axes(dim) - layer = default_trtnet().add_topk( - input.trt_tensor, - trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN, - k=k if not isinstance(k, Tensor) else 1, - axes=axes) - if isinstance(k, Tensor): - if k.ndim() == 1: - k = squeeze(k, 0) - layer.set_input(1, k.trt_tensor) - values = _create_tensor(layer.get_output(0), layer) - indices = _create_tensor(layer.get_output(1), layer) + if for_attention_plugin: + sinusoid_inp = np.expand_dims(sinusoid_inp, axis=-1) + concat = np.concatenate( + (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), axis=-1) + else: + concat = np.concatenate( + (np.sin(sinusoid_inp), np.cos(sinusoid_inp)), axis=1) + concat = np.expand_dims(concat, axis=0) - return values, indices + mscale = short_mscale if is_short else long_mscale + concat = concat.astype(dtype) * mscale + # gpt attention plugins also need inv_freq. + if for_attention_plugin: + return inv_freq.reshape(1, -1), concat.reshape(1, -1) + else: + return concat -def scatter_nd(input: Tensor, mask: Tensor, source: Tensor) -> Tensor: - ''' - Scatter_nd is a tensor operation that writes or updates values in a tensor based on indices. + return ( + _compute_sinusoidal_positions(scaling_short_factors, True, False), + _compute_sinusoidal_positions(scaling_long_factors, False, False), + _compute_sinusoidal_positions(scaling_short_factors, True, True), + _compute_sinusoidal_positions(scaling_long_factors, False, True), + short_mscale, + ) - Parameters: - input: Tensor - The input tensor to be updated - mask: Tensor - A tensor of indices specifying the locations in data to be updated. - source: Tensor - A tensor of values to be written or scattered into data. - Returns: - New tensor with the same shape as the input tensor data, - where the values from the source tensor are scattered or written into the output tensor - at the locations specified by the mask tensor. - ''' - scatter_layer = default_trtnet().add_scatter(input.trt_tensor, - mask.trt_tensor, - source.trt_tensor, - mode=trt.ScatterMode.ND) - return _create_tensor(scatter_layer.get_output(0), scatter_layer) + @staticmethod + def create_sinusoidal_positions_long_rope( + num_pos: int, + dim: int, + theta: float, + original_max_pos: int, + short_factor: List[float], + long_factor: List[float], + dtype=np.float32, + max_seq_len: Optional[int] = None, + duplicate_data: bool = False, + ): + short_factor = np.array(short_factor, dtype=np.float32) + long_factor = np.array(long_factor, dtype=np.float32) + inv_freq = 1.0 / (theta**(np.arange(0, dim, 2, dtype=np.float32) / dim)) + t_pos = np.arange(np.max([num_pos, original_max_pos]), dtype=np.float32) -def low_latency_gemm(input: Tensor, - mat2: Tensor, - alpha: Optional[np.ndarray] = None, - strict_dtype: Optional[trt.DataType] = None) -> Tensor: - if not default_net().plugin_config.low_latency_gemm_plugin: - raise RuntimeError("Low Latency GEMM is only support with plugin") - elif default_net().plugin_config.low_latency_gemm_plugin != "fp8": - raise RuntimeError("Low Latency GEMM plugin only support fp8") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "LowLatencyGemm", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - if ((input.dtype != trt.fp8) or ((mat2.dtype) != trt.fp8)): - raise TypeError("Low Latency GEMM only support fp8 input") - if (alpha): - assert (isinstance(alpha, np.ndarray) and alpha.dtype == np.float32 - and alpha.size - == 1), "`alpha` must be passed as a float32 ndarray" - alpha = alpha if alpha else np.array(1.0, dtype=np.float32) - alpha = trt.PluginField("alpha", alpha.flatten(), - trt.PluginFieldType.FLOAT32) + # Choose proper freqs based on max_seq_len. + factor = (long_factor if max_seq_len is None + or max_seq_len > original_max_pos else short_factor) + inv_freq = inv_freq / factor + freqs = np.einsum("i,j->ij", t_pos, inv_freq) + sinusoid_inp = freqs.astype(np.float32)[..., np.newaxis] - if strict_dtype is not None: - assert isinstance(strict_dtype, trt.DataType) - p_dtype = strict_dtype - if (p_dtype not in [trt.float32, trt.float16, trt.bfloat16]): - raise ValueError( - "strict_dtype must be float32, float16 or bfloat16 in low latency gemm plugin" - ) + # Apply scaling + scale = num_pos / original_max_pos + if scale <= 1.0: + scaling_factor = 1.0 else: - raise RuntimeError( - "need to use strict dtype in low latency gemm plugin fp8") - pf_type = trt.PluginField("type_id", np.array([int(p_dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([alpha, pf_type]) - low_latency_gemm_plug = plg_creator.create_plugin( - "low_latency_gemm", pfc) - plug_inputs = [input.trt_tensor, mat2.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, - low_latency_gemm_plug) - _add_plugin_info(layer, plg_creator, "low_latency_gemm", pfc) - return _create_tensor(layer.get_output(0), layer) - - -class SideStreamIDType(IntEnum): - disable = 0 - moe = 1 - - -def low_latency_gemm_swiglu(input: Tensor, - weight: Tensor, - scale_d0: float = 1.0, - scale_d1: float = 1.0, - scale_output: float = 1.0) -> Tensor: - ''' - Add a matrix multiplication, followed by SwiGLU (`x * SiLU(gate)`) operation. - - The second SwiGLU operation takes the preceding tensor, splits it into two halves - along the last dimension, applies SiLU to the second half and multiply the results. The - behaviour is undefined if the last dimension is not even. - - Parameters: - input : Tensor - The first tensor (often called A). - - weight : Tensor - The second tensor (often called B). - - scale_d0 : float - The scale for dequantizing x, used for fp8 - - scale_d1 : float - The scale for dequantizing gate, used for fp8 - - scale_output : float - The scale for quantizing output, used for fp8 - - Returns: - The tensor produced by the inserted layer. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LowLatencyGemmSwiglu', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.low_latency_gemm_swiglu_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_scale_d0 = trt.PluginField("scale_d0", - np.array(scale_d0, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_d1 = trt.PluginField("scale_d1", - np.array(scale_d1, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_output = trt.PluginField("scale_output", - np.array(scale_output, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - pfc = trt.PluginFieldCollection( - [pf_type, pf_scale_output, pf_scale_d0, pf_scale_d1]) - low_latency_gemm_swiglu_plug = plg_creator.create_plugin( - "low_latency_gemm_swiglu", pfc) - - plug_inputs = [input.trt_tensor, weight.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, - low_latency_gemm_swiglu_plug) - - return _create_tensor(layer.get_output(0), layer) - + scaling_factor = np.sqrt(1.0 + + np.log(scale) / np.log(original_max_pos)) -def cuda_stream_sync(input_list: List[Tensor], - side_stream_id: SideStreamIDType) -> Tensor: - ''' - Wait for the side stream on the main stream. - output = input_list[0] + if duplicate_data: + sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) - Parameters: - input_list : List[Tensor] (On GPU) - The list of input tensors. - side_stream_id : int (On CPU) - The side stream ID. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "CudaStream", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None + # fuse cos/sin into float2 (cos, sin). + concat = np.concatenate( + (np.cos(sinusoid_inp) * scaling_factor, + np.sin(sinusoid_inp) * scaling_factor), + axis=-1, + ) - p_side_stream_id = trt.PluginField("side_stream_id", - np.array(side_stream_id, dtype=np.int32), - trt.PluginFieldType.INT32) - p_num_inputs = trt.PluginField("num_inputs", - np.array(len(input_list), dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField( - "type_id", np.array([int(input_list[0].dtype)], dtype=np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([p_side_stream_id, p_num_inputs, pf_type]) - plug = plg_creator.create_plugin("cuda_stream", pfc) - plug_inputs = [input.trt_tensor for input in input_list] + return None, concat.reshape(1, -1).astype(dtype) - layer = default_trtnet().add_plugin_v2(plug_inputs, plug) - _add_plugin_info(layer, plg_creator, "cuda_stream", pfc) - output = _create_tensor(layer.get_output(0), layer) - return output + @staticmethod + def create_fake_weight(dim: int, dtype=np.half): + return np.random.rand(dim).astype(dtype) + # Note: When not using deepseek_yarn, make sure to set mscale_all_dim to 0.0. + @staticmethod + def create_sinusoidal_positions_yarn( + num_pos: int, + dim: int, + base: int = 10000, + scaling_factor: float = 1.0, + original_max_position_embeddings: int = 4096, + beta_fast: int = 32, + beta_slow: int = 1, + mscale: float = 1.0, + mscale_all_dim: float = 1.0, + duplicate_data: bool = True, + dtype=None, + ): + if dtype is None: + dtype = torch.float32 -def cp_split_plugin( - input_ids: Tensor, - host_request_types: Tensor, - host_context_lengths: Tensor, # for pad-free input mode - cp_size: int = 1, - cp_rank: int = 0, -) -> Tensor: - ''' - Add an operation to perform splitting for context parallelism. + # Copy from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/main/modeling_deepseek.py + # Inverse dim formula to find dim based on number of rotations + def yarn_find_correction_dim(num_rotations, dim, base, + max_position_embeddings): + return (dim * math.log(max_position_embeddings / + (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base)) - This operation split the input_ids into cp_size chunks, and return the cp_rank-th - chunk. - When the seqlen % cp_size != 0, the chunk sizes of each rank would be - [seqlen // cp_size, seqlen // cp_size, ..., seqlen - (seqlen // cp_size) * cp_size] + # Find dim range bounds based on rotations + def yarn_find_correction_range(low_rot, high_rot, dim, base, + max_position_embeddings): + low = math.floor( + yarn_find_correction_dim(low_rot, dim, base, + max_position_embeddings)) + high = math.ceil( + yarn_find_correction_dim(high_rot, dim, base, + max_position_embeddings)) + if low < 0: + low = 0 + if high > dim - 1: + high = dim - 1 + return low, high # Clamp values just in case - It inserts a IPluginV3Layer. + def yarn_get_mscale(scale, mscale): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 - Parameters: - input : Tensor - The input tensor contains the indices to split. + def yarn_linear_ramp_mask(min, max, dim): + if min == max: + max += 0.001 # Prevent singularity - host_request_types: Tensor = None (On CPU) - The tensor on the host that indicates if a request is in context or - generation phase. Its shape is [batch_size]. See Inflight Batching - in docs/gpt_attention.md, + linear_func = (torch.arange(dim, dtype=dtype) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func - host_context_lengths: Tensor = None (On CPU) - A host tensor that contains the lengths of the different inputs + pos_freqs = base**(torch.arange(0, dim, 2, dtype=dtype) / dim) + freq_extra = 1.0 / pos_freqs + freq_inter = 1.0 / (scaling_factor * pos_freqs) - Returns: - The output split tensor. - The length of the output split tensor. - The index for rebuilding the sequence - ''' - plg_creator = trt.get_plugin_registry().get_creator( - 'CpSplit', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None + low, high = yarn_find_correction_range( + beta_fast, + beta_slow, + dim, + base, + original_max_position_embeddings, + ) + inv_freq_mask = 1 - yarn_linear_ramp_mask(low, high, dim // 2) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + t = torch.arange(num_pos, dtype=dtype) + sinusoid_inp = torch.einsum("i,j -> ij", t, inv_freq).unsqueeze(-1) - cp_size = trt.PluginField("cp_size", np.array([int(cp_size)], np.int32), - trt.PluginFieldType.INT32) - cp_rank = trt.PluginField("cp_rank", np.array([int(cp_rank)], np.int32), - trt.PluginFieldType.INT32) + _mscale = float( + yarn_get_mscale(scaling_factor, mscale) / + yarn_get_mscale(scaling_factor, mscale_all_dim)) - pfc = trt.PluginFieldCollection([cp_size, cp_rank]) - cp_split_plug = plg_creator.create_plugin("cp_split", pfc, - trt.TensorRTPhase.BUILD) - plug_inputs = [ - input_ids.trt_tensor, host_request_types.trt_tensor, - host_context_lengths.trt_tensor - ] + if duplicate_data: + emb = torch.cat((sinusoid_inp, sinusoid_inp), dim=-2) + else: + emb = sinusoid_inp - layer = default_trtnet().add_plugin_v3(plug_inputs, [], cp_split_plug) - _add_plugin_info(layer, plg_creator, "cp_split", pfc) - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(2), layer) + concat = torch.cat((torch.cos(emb) * _mscale, torch.sin(emb) * _mscale), + dim=-1) + return inv_freq.numpy(), concat.reshape((1, -1)).to(dtype).numpy() diff --git a/tensorrt_llm/graph_rewriting.py b/tensorrt_llm/graph_rewriting.py deleted file mode 100644 index 64aadd2da7e4..000000000000 --- a/tensorrt_llm/graph_rewriting.py +++ /dev/null @@ -1,645 +0,0 @@ -import inspect -import weakref -from copy import copy -from dataclasses import dataclass, field -from functools import wraps -from typing import Any, Callable, ClassVar, Dict, List, Optional, Set, Tuple, TypeVar - -import tensorrt as trt - -from ._utils import trt_gte -from .logger import logger -from .network import Network - - -class Layer: - """Layer is a wrapper for TensorRT's ILayer with several python-friendly helper functions.""" - - def __init__(self, network: Network, trt_layer: trt.ILayer): - self._network = weakref.ref(network) - self.trt_layer = trt_layer - - assert isinstance(self.network, Network) - assert isinstance(self.trt_layer, trt.ILayer) - - @property - def network(self): - return self._network() - - def get_inputs(self, *indices: int): - """Get the input tensors of the layer. - - Parameters: - idx: the indices of the input tensor, will return all inputs if left empty - - Returns: - List[Tensor] - """ - from .functional import Tensor - - indices = indices if indices else range(self.trt_layer.num_inputs) - - ret = [] - for i in indices: - assert i < self.trt_layer.num_inputs, ( - f"Invalid input index {i} for layer {self.trt_layer.name}" - ) - - tensor = self.trt_layer.get_input(i) - tensor = Tensor(trt_tensor=tensor, network=self.network, is_network_input=False) - ret.append(tensor) - return ret - - def get_outputs(self, *indices: int): - """Get the output tensor of the layer. - - Parameters: - idx: the index of the output tensor - - Returns: - List[Tensor] - """ - from .functional import Tensor - - indices = indices if indices else range(self.trt_layer.num_outputs) - - ret = [] - for i in indices: - assert i < self.trt_layer.num_outputs, ( - f"Invalid output index {i} for layer {self.trt_layer.name}" - ) - - tensor = self.trt_layer.get_output(i) - tensor = Tensor(trt_tensor=tensor, network=self.network, is_network_input=False) - ret.append(tensor) - return ret - - def is_removed(self): - return self.network.is_removed_layer(self) - - def mark_as_removed(self): - """Mark the layer as removed, this will remove the layer from the network.""" - # NOTE, since INetwork python API doesn't provide a way to remove a layer, we actually mark - # the layer as removed in the network. - self.network.mark_removed_layer(self) - - # remove the FLayerInfo if exists - FLayerInfoMemo.instance().remove(self.name) - - def __eq__(self, other: "Layer") -> bool: - if isinstance(other, Layer): - return self.trt_layer == other.trt_layer - if isinstance(other, trt.tensorrt.ILayer): - return self.trt_layer == other - return False - - def __getattr__(self, name: str) -> Any: - return getattr(self.trt_layer, name) - - # Refer to https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html?highlight=elementwise#layers - # for a complete list of TRT layers. - TRT_LAYER_TYPE_TO_LAYER = { - trt.LayerType.CONVOLUTION: trt.IConvolutionLayer, - trt.LayerType.ACTIVATION: trt.IActivationLayer, - trt.LayerType.POOLING: trt.IPoolingLayer, - trt.LayerType.LRN: trt.ILRNLayer, - trt.LayerType.SCALE: trt.IScaleLayer, - trt.LayerType.SOFTMAX: trt.ISoftMaxLayer, - trt.LayerType.DECONVOLUTION: trt.IDeconvolutionLayer, - trt.LayerType.CONCATENATION: trt.IConcatenationLayer, - trt.LayerType.ELEMENTWISE: trt.IElementWiseLayer, - trt.LayerType.UNARY: trt.IUnaryLayer, - trt.LayerType.PADDING: trt.IPaddingLayer, - trt.LayerType.SHUFFLE: trt.IShuffleLayer, - trt.LayerType.REDUCE: trt.IReduceLayer, - trt.LayerType.TOPK: trt.ITopKLayer, - trt.LayerType.GATHER: trt.IGatherLayer, - trt.LayerType.MATRIX_MULTIPLY: trt.IMatrixMultiplyLayer, - trt.LayerType.RAGGED_SOFTMAX: trt.IRaggedSoftMaxLayer, - trt.LayerType.CONSTANT: trt.IConstantLayer, - trt.LayerType.IDENTITY: trt.IIdentityLayer, - trt.LayerType.PLUGIN_V2: trt.IPluginV2Layer, - trt.LayerType.PLUGIN_V3: trt.IPluginV3Layer, - trt.LayerType.SLICE: trt.ISliceLayer, - trt.LayerType.SHAPE: trt.IShapeLayer, - trt.LayerType.PARAMETRIC_RELU: trt.IParametricReLULayer, - trt.LayerType.RESIZE: trt.IResizeLayer, - trt.LayerType.TRIP_LIMIT: trt.ITripLimitLayer, - trt.LayerType.RECURRENCE: trt.IRecurrenceLayer, - trt.LayerType.ITERATOR: trt.IIteratorLayer, - trt.LayerType.LOOP_OUTPUT: trt.ILoopOutputLayer, - trt.LayerType.SELECT: trt.ISelectLayer, - trt.LayerType.FILL: trt.IFillLayer, - trt.LayerType.QUANTIZE: trt.IQuantizeLayer, - trt.LayerType.DEQUANTIZE: trt.IDequantizeLayer, - trt.LayerType.CONDITION: trt.IConditionLayer, - trt.LayerType.CONDITIONAL_INPUT: trt.IIfConditionalInputLayer, - trt.LayerType.CONDITIONAL_OUTPUT: trt.IIfConditionalOutputLayer, - trt.LayerType.ASSERTION: trt.IAssertionLayer, - trt.LayerType.SCATTER: trt.IScatterLayer, - trt.LayerType.EINSUM: trt.IEinsumLayer, - trt.LayerType.GRID_SAMPLE: trt.IGridSampleLayer, - trt.LayerType.ONE_HOT: trt.IOneHotLayer, - trt.LayerType.NON_ZERO: trt.INonZeroLayer, - trt.LayerType.NMS: trt.INMSLayer, - trt.LayerType.REVERSE_SEQUENCE: trt.IReverseSequenceLayer, - trt.LayerType.NORMALIZATION: trt.INormalizationLayer, - trt.LayerType.CAST: trt.ICastLayer, - } - - if trt_gte(10, 8): - TRT_LAYER_TYPE_TO_LAYER[trt.LayerType.DYNAMIC_QUANTIZE] = trt.IQuantizeLayer - - def as_layer(self) -> Any: - """Convert to a actual TensorRT layer object. - - This can be IPluginV2Layer or IConvolutionLayer, so that we can access the actual layer information. - """ - if self.type in self.TRT_LAYER_TYPE_TO_LAYER: - # bypass TRT's bug of retrieving a specific ILayer type in TensorRT - self.trt_layer.__class__ = self.TRT_LAYER_TYPE_TO_LAYER[self.type] - return self.trt_layer - raise NotImplementedError(f"Unknown layer type: {self.type}") - - def __hash__(self): - return id(self.trt_layer) - - -@dataclass -class _Pattern: - name: str - # args helps to pass in/out some information - args: Dict[str, Any] = field(default_factory=dict, init=False) - - def log_info(self, msg: str): - logger.info(f"Pattern {self.name}: {msg}") - - def log_error(self, msg: str): - logger.error(f"Pattern {self.name}: {msg}") - - def log_warn(self, msg: str): - logger.warning(f"Pattern {self.name}: {msg}") - - -class PatternRewriter(_Pattern): - """A pattern rewriter is a class that can match a pattern in the graph and rewrite the matched pattern. - - There are two ways to implement a pattern rewriter, either override match() and rewrite() separately, or - override match_and_rewrite(). - """ - - def __init__( - self, - name: str, - root_layer: Optional[Set[trt.LayerType]] = None, - separate_match_rewrite=False, - ): - """Constructor. - - Parameters: - name: the name of the rewrite pattern - root_layer: the root layer types to start the pattern matching, if not provided, the pattern - will traverse all the layers in the graph. - separate_match_rewrite: if set to True, the pattern should override match() and rewrite() - separately, otherwise, the pattern should override match_and_rewrite() - """ - super().__init__(name) - self.root_layer = root_layer - self._separate_match_rewrite = separate_match_rewrite - - def match(self, layer: Layer) -> bool: - raise NotImplementedError() - - def rewrite(self, layer: Layer) -> None: - raise NotImplementedError() - - def match_and_rewrite(self, layer: Layer) -> bool: - raise NotImplementedError() - - -class PatternAnalyzer(_Pattern): - def __init__(self, name: str, root_layer: Optional[Set[trt.LayerType]]) -> None: - super().__init__(name) - self.root_layer = root_layer - - def match(self, layer: Layer) -> bool: - raise NotImplementedError() - - def analyze(self, subgraph: List[Layer]) -> None: - raise NotImplementedError() - - -class _PatternManager: - PatternType = TypeVar("PatternType") - - def __init__(self): - # records of (benefit, pattern, id) - self.patterns: Dict[str, Tuple[int, _PatternManager.PatternType]] = {} - - def add(self, label: str, pattern: "_PatternManager.PatternType", benefit: int = 0): - assert label not in self.patterns, f"Pattern {label} already exists" - self.patterns[label] = (benefit, pattern) - - def get(self, label: str) -> "_PatternManager.PatternType": - return self.patterns[label][1] - - -class RewritePatternManager(_PatternManager): - def rewrite(self, net: Network, args=None): - modified = True - # TODO: we can optimize this by asking TRT to expose a graph iterator consistent even after - # the graph is modified. - while modified: - modified = False - # Since the graph iterator is hold by the underlying INetwork, we can only rebuild the - # graph cache and match the nodes again. - for layer in net.get_layers(): - if layer.is_removed(): - continue - for profit, pattern in sorted(self.patterns.values(), key=lambda x: x[0]): - pattern.args = args - - if pattern.root_layer is not None and layer.type not in pattern.root_layer: - continue - if pattern._separate_match_rewrite: - if pattern.match(layer): - pattern.rewrite(layer) - modified = True - else: - if pattern.match_and_rewrite(layer): - modified = True - - @staticmethod - def instance(): - return _global_rewrite_pattern_manager - - -class AnalysisPatternManager(_PatternManager): - def analyze(self, graph: Network, args=None): - for layer in graph.get_layers(): - if layer.name in graph.removed_layers: - continue - for benefit, pattern in sorted(self.patterns.values(), key=lambda x: x[0]): - pattern.args = args - - if pattern.root_layer is not None and layer.type not in pattern.root_layer: - continue - if pattern.match(layer): - subgraph = pattern.match(layer) - pattern.analyze(subgraph) - - @staticmethod - def instance(): - return _global_analysis_pattern_manager - - -@dataclass -class FLayerInfo: - """The FLayerInfo is used to track the functional layers in the INetwork and help graph rewriting. - - The lifetime of a FLayer is the same as the corresponding plugin instance in the INetwork. Once the - plugin instance is removed by the graph rewriting, the FLayer will be removed as well. - - WHY this is needed? - In the current implementation, for functional methods, once it is called in Python, it will lower - to a plugin instance in the INetwork. - However, the plugin interface is black box with customized logic, we cannot retrieve necessary - information from it. This is quite different from ILayer, which provides a set of APIs to retrieve - the information. - Therefore, we need to record the high level information in the FLayerInfo, and keep - it consistent during the graph rewriting. - """ - - layer_kind: str # the method name in the functional.py - # Record the raw inputs of the functional layer to be used in the graph rewrite - # NOTE: the raw inputs contains both the constants and Tensors, the Tensors will be also updated by - # graph rewriting APIs such as `replace_all_uses_with` - raw_inputs: Dict[str, Any] - - raw_outputs: List[Any] = field(default_factory=list, init=False) - - # the corresponding ILayer name - layer_name: str = field(init=False, default="") - - # the signature of the functional layer - signature: Any = field(init=False, default=None) - - def __post_init__(self): - from .functional import Tensor - - assert self.layer_kind - - def replace_with_symbols(arg) -> Any: - if arg is None: - return None - if isinstance(arg, Tensor): - return Tensor - if isinstance(arg, (list, tuple)): - return [replace_with_symbols(x) for x in arg] - if isinstance(arg, dict): - return {k: replace_with_symbols(v) for k, v in arg.items()} - - return arg - - def amend_tensor(arg) -> Any: - if arg is None: - return None - if isinstance(arg, Tensor): - arg.network = self.network - if isinstance(arg, (list, tuple)): - [replace_with_symbols(x) for x in arg] - if isinstance(arg, dict): - {k: replace_with_symbols(v) for k, v in arg.items()} - - return arg - - self.signature = ( - self.layer_kind, - {name: replace_with_symbols(value) for name, value in self.raw_inputs.items()}, - ) - - amend_tensor(self.raw_inputs) - - def set_outputs(self, outputs: List[Any]): - self.raw_outputs = outputs - - def get_input(self, name: str) -> Any: - return self.raw_inputs[name] - - def clone_inputs(self): - """Get a shallow copy of the inputs.""" - return copy(self.raw_inputs) - - def replace_input_with(self, src, dst): - """Replace the input `src` with the input `dst` in the raw_inputs. - - src: Tensor - dst: Tensor - """ - from .functional import Tensor - - def replace(arg: Any): - if isinstance(arg, Tensor): - if arg.trt_tensor is src.trt_tensor: - return dst - return arg - elif isinstance(arg, (list, tuple)): - return [replace(x) for x in arg] - elif isinstance(arg, dict): - return {k: replace(v) for k, v in arg.items()} - return arg - - replace(self.raw_inputs) - - def replace_outputs_uses_with(self, net: Network, new_outs: List[Any]): - """Replace the output users with the new outputs. - - new_outs: List[Tensor], the new outputs to replace with - """ - from .functional import Tensor - - assert len(self.raw_outputs) == len(new_outs) - for old_out, new_out in zip(self.raw_outputs, new_outs): - assert type(old_out) is type(new_out), ( - f"rewrite error, the output type {type(old_out)} is different from the new output " - f"type {type(new_out)} not match the original output type {type(old_out)}" - ) - - def _swap_tensor_info(new, deprecated): - name = deprecated.trt_tensor.name - deprecated.trt_tensor.name = name + "_deprecated" - from .functional import cast - - new = cast(new, deprecated.dtype) - new.trt_tensor.name = name - - def _reset_network_output_tensors(network, out, new_out): - net_outputs = list() - num_outputs = network._trt_network.num_outputs - need_to_mark = False - for i in range(num_outputs): - net_outputs.append(network._trt_network.get_output(i)) - if out.trt_tensor is net_outputs[i]: - need_to_mark = True - if need_to_mark is False: - return - for output in net_outputs: - network.trt_network.unmark_output(output) - for i in range(num_outputs): - if net_outputs[i] is out.trt_tensor: - network.trt_network.mark_output(new_out.trt_tensor) - new_out.trt_tensor.dtype = out.trt_tensor.dtype - else: - network.trt_network.mark_output(net_outputs[i]) - - def replace_all_uses_with(out, new_out): - if isinstance(out, Tensor): - assert isinstance(new_out, Tensor) - out.replace_all_uses_with(new_out) - _swap_tensor_info(new_out, out) - _reset_network_output_tensors(net, out, new_out) - elif isinstance(out, list): - assert isinstance(new_out, list) - for x, y in zip(out, new_out): - replace_all_uses_with(x, y) - elif isinstance(out, dict): - assert isinstance(new_out, dict) - for k, v in out.items(): - replace_all_uses_with(v, new_out[k]) - elif isinstance(out, tuple): - assert isinstance(new_out, tuple) - for x, y in zip(out, new_out): - replace_all_uses_with(x, y) - - replace_all_uses_with(self.raw_outputs, new_outs) - - def __hash__(self) -> int: - return hash(self.signature) - - def __repr__(self) -> str: - return "".format(self.signature) - - @staticmethod - def _get_spec(arg): - """Get the spec that could impact on the Module's topology in the `forward` method.""" - from .functional import Tensor - - # For scalars, we track their value since they are constant - if arg is None: - return None - elif isinstance(arg, (bool, int, str)): - return arg - # For tensors, currently we only track their type, since they are variables - elif isinstance(arg, Tensor): - return Tensor - elif isinstance(arg, (list, tuple)): - return [FLayerInfo._get_spec(x) for x in arg] - # NOTE Free to add more types here is broken, carefully note that, from the engine building angle, - # all the constants should be captured while for the network variables, their types as placeholders - # are enough. - else: - raise TypeError(f"unsupported input type detected: {type(arg)}") - - -@dataclass -class FLayerInfoMemo: - """FLayerInfoMemo holds the FLayer of all the necessary functional layers.""" - - data: Dict[str, FLayerInfo] = field(default_factory=dict, init=False) - - cur_flayer: ClassVar[Optional[FLayerInfo]] = None - - def add(self, layer_name: str, layer: FLayerInfo) -> None: - assert layer_name not in self.data, f"FLayer {layer_name} already exists in FLayerMemo" - self.data[layer_name] = layer - - def create(self, fn: Callable, *args, **kwargs) -> FLayerInfo: - """Add a FLayer to the memo.""" - return FLayerInfo(fn.__name__, self.get_function_arg_dict(fn, *args, **kwargs)) - - def get(self, layer_name: str) -> Optional[FLayerInfo]: - return self.data.get(layer_name, None) - - def remove(self, layer_name: str) -> None: - if layer_name in self.data: - del self.data[layer_name] - - @staticmethod - def instance() -> "FLayerInfoMemo": - """A singleton instance of FLayerMemo.""" - from ._common import default_net - - return default_net().flayer_memo - - @staticmethod - def get_function_arg_dict(f: Callable, *args, **kwargs): - """Get the input argument dict of a function.""" - sig = inspect.signature(f) - - bound_args = sig.bind(*args, **kwargs) - bound_args.apply_defaults() - - return {k: v for k, v in bound_args.arguments.items() if k != "self"} - - -class FLayerScope: - """FLayerScope is used to capture the plugin within a functional method.""" - - def __init__(self, fn, *args, **kwargs): - self.layer = FLayerInfoMemo.instance().create(fn, *args, **kwargs) - - def __enter__(self): - assert FLayerInfoMemo.cur_flayer is None, "FLayerMemo is not reentrant" - # There is no FLayer hierarchy, since the functional layers are not nested - FLayerInfoMemo.cur_flayer = self.layer - - def __exit__(self, exc_type, exc_val, exc_tb): - FLayerInfoMemo.cur_flayer = None - if exc_type is None: - assert self.layer.layer_name != "", ( - f"FLayer {self.layer.layer_kind} without a plugin name detected" - ) - FLayerInfoMemo.instance().add(self.layer.layer_name, self.layer) - - -def record_signature(f): - """Helps to decorate a functional method and record its metadata with a FLayerInfo.""" - - @wraps(f) - def wrapper(*args, **kwargs): - with FLayerScope(f, *args, **kwargs): - outs = f(*args, **kwargs) - FLayerInfoMemo.cur_flayer.set_outputs(outs) - return outs - - return wrapper - - -# singletons -_global_rewrite_pattern_manager = RewritePatternManager() -_global_analysis_pattern_manager = AnalysisPatternManager() - - -class FuseAttentionWithBiasPass(PatternRewriter): - def __init__(self): - super().__init__(name="fuse_attention_with_bias", separate_match_rewrite=False) - - @staticmethod - def is_attention_plugin(layer: Layer) -> bool: - if layer.as_layer().type != trt.LayerType.PLUGIN_V2: - return False - p = layer.as_layer().plugin - conds = [ - p.plugin_namespace == "tensorrt_llm", - p.plugin_type == "GPTAttention", - p.num_outputs == 2, - ] - return all(conds) - - @staticmethod - def is_elementwise_sum(layer: Layer) -> bool: - l = layer.as_layer() # noqa: E741 - if l.type != trt.LayerType.ELEMENTWISE: - return False - return l.op == trt.ElementWiseOperation.SUM - - @staticmethod - def get_eltwise_inputs(layer: Layer): - const_inputs = [] - mutable_inputs = [] - - from .functional import Tensor - - def const_foldable(tensor: Tensor, depth=0) -> bool: - max_depth = 10 - layer = tensor.get_parent() - if layer is None or depth > max_depth: - return False - if layer.type == trt.LayerType.CONSTANT and len(layer.get_inputs()) == 0: - return True - for _ in layer.get_inputs(): - if not const_foldable(_, depth + 1): - return False - return True - - for input in layer.get_inputs(): - if const_foldable(input): - const_inputs.append(input) - else: - mutable_inputs.append(input) - return const_inputs, mutable_inputs - - def match_and_rewrite(self, layer: Layer) -> bool: - from tensorrt_llm.network import net_guard - - with net_guard(layer.network): - if not self.is_attention_plugin(layer): - return False - plugin_flayer = FLayerInfoMemo.instance().get(layer.name) - input = plugin_flayer.raw_inputs["qkv"] - if input is None or isinstance(input, list) or len(list(input.get_users())) != 1: - return False - parent_layer = input.get_parent() - if not self.is_elementwise_sum(parent_layer): - return False - eltwise_const_inputs, eltwise_mutable_inputs = self.get_eltwise_inputs(parent_layer) - if len(eltwise_const_inputs) != 1 or len(eltwise_mutable_inputs) != 1: - return False - if plugin_flayer.raw_inputs["qkv_bias"] is not None: - return False - plugin_flayer.raw_inputs["qkv"] = eltwise_mutable_inputs[0] - plugin_flayer.raw_inputs["qkv_bias"] = eltwise_const_inputs[0] - from .functional import gpt_attention - - new_outputs = gpt_attention(**plugin_flayer.raw_inputs) - plugin_flayer.replace_outputs_uses_with(layer.network, new_outputs) - return True - - -def optimize(net): - patterns = RewritePatternManager() - patterns.add( - label="fuse_attention_with_bias", - pattern=FuseAttentionWithBiasPass(), - ) - patterns.rewrite(net) diff --git a/tensorrt_llm/layers/__init__.py b/tensorrt_llm/layers/__init__.py deleted file mode 100755 index cdfb8a6f897f..000000000000 --- a/tensorrt_llm/layers/__init__.py +++ /dev/null @@ -1,79 +0,0 @@ -# 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 .activation import Mish -from .attention import (Attention, AttentionMaskParams, AttentionMaskType, - AttentionParams, BertAttention, BlockSparseAttnParams, - CogVLMAttention, DeepseekV2Attention, - KeyValueCacheParams, MropeParams, PositionEmbeddingType, - SpecDecodingParams) -from .cast import Cast -from .conv import Conv1d, Conv2d, Conv3d, ConvTranspose2d -from .embedding import Embedding, PromptTuningEmbedding -from .language_adapter import LanguageAdapter, LanguageAdapterConfig -from .linear import ColumnLinear, Linear, RowLinear -from .lora import Lora, LoraParams, LoraRuntimeParams -from .mlp import MLP, FusedGatedMLP, GatedMLP -from .moe import MOE, MoeConfig, SharedMoE -from .normalization import GroupNorm, LayerNorm, RmsNorm -from .pooling import AvgPool2d -from .recurrent import FusedRgLru, GroupedLinear, Recurrent, RgLru -from .ssm import Mamba, Mamba2 - -__all__ = [ - 'LayerNorm', - 'RmsNorm', - 'ColumnLinear', - 'Linear', - 'RowLinear', - 'AttentionMaskType', - 'PositionEmbeddingType', - 'Attention', - 'BertAttention', - 'CogVLMAttention', - 'DeepseekV2Attention', - 'GroupNorm', - 'Embedding', - 'PromptTuningEmbedding', - 'Conv2d', - 'ConvTranspose2d', - 'Conv1d', - 'Conv3d', - 'AvgPool2d', - 'Mish', - 'MLP', - 'GatedMLP', - 'FusedGatedMLP', - 'Cast', - 'AttentionParams', - 'AttentionMaskParams', - 'SpecDecodingParams', - 'MropeParams', - 'KeyValueCacheParams', - 'BlockSparseAttnParams', - 'Lora', - 'LoraParams', - 'LoraRuntimeParams', - 'MOE', - 'MoeConfig', - 'SharedMoE', - 'Mamba', - 'Mamba2', - 'Recurrent', - 'GroupedLinear', - 'RgLru', - 'FusedRgLru', - 'LanguageAdapter', - 'LanguageAdapterConfig', -] diff --git a/tensorrt_llm/layers/activation.py b/tensorrt_llm/layers/activation.py deleted file mode 100644 index 94ea283a1cf3..000000000000 --- a/tensorrt_llm/layers/activation.py +++ /dev/null @@ -1,22 +0,0 @@ -# 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 ..functional import softplus, tanh -from ..module import Module - - -class Mish(Module): - - def forward(self, input): - return input * tanh(softplus(input, beta=1.0, threshold=20.0)) diff --git a/tensorrt_llm/layers/attention.py b/tensorrt_llm/layers/attention.py deleted file mode 100755 index 29b63a4258cb..000000000000 --- a/tensorrt_llm/layers/attention.py +++ /dev/null @@ -1,2808 +0,0 @@ -# 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 math -from typing import List, Optional - -import numpy as np -import tensorrt as trt -import torch - -from .._common import default_net, precision -from .._utils import (fp32_array, int32_array, is_same_dtype, set_obj_attrs, - trt_dtype_to_np, trt_dtype_to_str) - -# isort: off -from ..functional import ( - ACT2FN, AllReduceParams, AttentionMaskType, Conditional, LayerNormType, - PositionEmbeddingType, RopeEmbeddingUtils, RotaryScalingType, Tensor, - allgather, arange, bert_attention, cast, clip, concat, constant, embedding, - expand, expand_dims, expand_mask, generate_alibi_biases, identity, - generate_alibi_slopes, generate_logn_scaling, gpt_attention, matmul, - minimum, repeat_interleave, shape, slice, softmax, split, unsqueeze, where) -# isort: on -from ..mapping import Mapping -from ..module import Module, ModuleList -from ..parameter import Parameter -from ..quantization import QuantMode -from ..quantization.functional import dequantize, quantize -from .linear import ColumnLinear, RowLinear -from .lora import LoraRuntimeParams -from .normalization import GroupNorm, LayerNorm, RmsNorm - -layernorm_map = { - LayerNormType.LayerNorm: LayerNorm, - LayerNormType.RmsNorm: RmsNorm, - LayerNormType.GroupNorm: GroupNorm, -} - - -def make_causal_mask(bsz, tgt_len, past_key_values_length, dtype): - _range = arange(start=constant(int32_array(0)), - end=tgt_len, - dtype=trt_dtype_to_str(dtype)) - mask = repeat_interleave(_range, tgt_len, 0).view(concat([tgt_len, - tgt_len])) - mask = where(mask < mask.transpose(-1, -2), 1.0, 0.0) - - zero = constant(fp32_array(0)) - zero = expand_dims(zero, [0, 1]) - zero = expand(zero, concat([tgt_len, past_key_values_length])) - mask = concat([zero, mask], dim=1) - mask *= np.finfo(trt_dtype_to_np(dtype)).min.item() - mask = mask.view(concat([1, 1, tgt_len, tgt_len + past_key_values_length])) - mask = expand(mask, - concat([bsz, 1, tgt_len, tgt_len + past_key_values_length])) - return mask - - -def compute_relative_bias(query_length, - key_length, - num_buckets, - max_distance, - bidirectional, - rel_attn_table, - tp_size=1, - tp_group=None, - tp_rank=None): - - def make_relative_position_bucket(relative_position, bidirectional, - num_buckets, max_distance): - relative_buckets = 0 - if bidirectional: - num_buckets //= 2 - relative_buckets += where(relative_position > 0, num_buckets, 0) - relative_position = relative_position.abs() - else: - relative_position = 0 - minimum(relative_position, 0) - - max_exact = num_buckets // 2 - is_small = relative_position < max_exact - - max_exact_fp = constant(fp32_array(max_exact)) - tmp = cast(relative_position, "float32") / max_exact_fp - tmp = tmp.log() - const1 = math.log(max_distance / max_exact) - const2 = constant(fp32_array(num_buckets - max_exact)) - relative_position_if_large = tmp / const1 * const2 - relative_position_if_large = cast(relative_position_if_large, "int32") - relative_position_if_large = max_exact + relative_position_if_large - relative_position_if_large = minimum(relative_position_if_large, - num_buckets - 1) - - relative_buckets += where(is_small, relative_position, - relative_position_if_large) - return relative_buckets - - context_position = arange(start=constant(int32_array(0)), - end=query_length, - dtype=trt_dtype_to_str(trt.int32)) - context_position = unsqueeze(context_position, -1) - memory_position = arange(start=constant(int32_array(0)), - end=key_length, - dtype=trt_dtype_to_str(trt.int32)) - memory_position = unsqueeze(memory_position, 0) - relative_position = memory_position - context_position - relative_position_bucket = make_relative_position_bucket( - relative_position, # shape (query_length, key_length) - bidirectional, - num_buckets, - max_distance, - ) - # shape (query_length, key_length, num_heads) - values = embedding(relative_position_bucket, - rel_attn_table, - tp_size=tp_size, - tp_group=tp_group, - tp_rank=tp_rank) - # shape (1, num_heads, query_length, key_length) - values = unsqueeze(values.permute([2, 0, 1]), 0) - return values - - -class AttentionMaskParams(object): - - def __init__(self, - self_attention_mask: Tensor = None, - self_attention_packed_mask: Tensor = None, - cross_attention_mask: Tensor = None, - cross_attention_packed_mask: Tensor = None): - self.self_attention_mask = self_attention_mask - self.self_attention_packed_mask = self_attention_packed_mask - self.cross_attention_mask = cross_attention_mask - self.cross_attention_packed_mask = cross_attention_packed_mask - - -class AttentionParams(object): - - def __init__(self, - sequence_length: Tensor = None, - context_lengths: Tensor = None, - host_context_lengths: Tensor = None, - max_context_length: int = None, - host_request_types: Tensor = None, - encoder_input_lengths: Tensor = None, - encoder_max_input_length: Tensor = None, - host_runtime_perf_knobs: Tensor = None, - host_context_progress: Tensor = None): - self.sequence_length = sequence_length - self.context_lengths = context_lengths - self.host_context_lengths = host_context_lengths - # max allowed context length. Required to - # compute scratch memory size. - self.max_context_length = max_context_length - self.host_request_types = host_request_types - - self.encoder_input_lengths = encoder_input_lengths - self.encoder_max_input_length = encoder_max_input_length - - self.host_runtime_perf_knobs = host_runtime_perf_knobs - - self.host_context_progress = host_context_progress - - # const parameters that will be reused by all layers. - self.embed_positions = None - self.rotary_inv_freq = None - self.embed_positions_for_gpt_attention = None - - # auxiliary params to support models with non-homegeneous attn layers requiring - # a different set of rope params. e.g. Gemma3. - self.embed_positions_local = None - self.rotary_inv_freq_local = None - self.embed_positions_for_gpt_attention_local = None - - # long rope const parameters - self.long_rope_embed_positions = None - self.long_rope_rotary_inv_freq = None - self.long_rope_embed_positions_for_gpt_attention = None - self.short_mscale = 1.0 - self.long_mscale = 1.0 - - def fill_attention_const_params_for_rope( - self, - embed_positions: Tensor = None, - rotary_inv_freq: Tensor = None, - embed_positions_for_gpt_attention: Tensor = None, - embed_positions_local: Tensor = None, - rotary_inv_freq_local: Tensor = None, - embed_positions_for_gpt_attention_local: Tensor = None): - self.embed_positions = embed_positions - self.rotary_inv_freq = rotary_inv_freq - self.embed_positions_for_gpt_attention = embed_positions_for_gpt_attention - self.embed_positions_local = embed_positions_local - self.rotary_inv_freq_local = rotary_inv_freq_local - self.embed_positions_for_gpt_attention_local = embed_positions_for_gpt_attention_local - return self - - def fill_attention_const_params_for_long_rope( - self, embed_positions, long_rope_embed_positions, rotary_inv_freq, - long_rope_rotary_inv_freq, embed_positions_for_gpt_attention, - long_rope_embed_positions_for_gpt_attention, short_mscale, - long_mscale): - self.embed_positions = embed_positions - self.long_rope_embed_positions = long_rope_embed_positions - self.rotary_inv_freq = rotary_inv_freq - self.long_rope_rotary_inv_freq = long_rope_rotary_inv_freq - self.embed_positions_for_gpt_attention = embed_positions_for_gpt_attention - self.long_rope_embed_positions_for_gpt_attention = long_rope_embed_positions_for_gpt_attention - self.short_mscale = short_mscale - self.long_mscale = long_mscale - return self - - def is_valid_cross_attn(self, do_cross_attention): - if do_cross_attention: - if self.encoder_input_lengths is None: - return False - if self.encoder_max_input_length is None: - return False - return True - - def is_valid(self, gpt_attention_plugin, remove_input_padding, - use_kv_cache): - if gpt_attention_plugin: - if use_kv_cache and self.sequence_length is None: - return False - if self.context_lengths is None: - return False - if self.host_request_types is None: - return False - if self.max_context_length is None: - return False - if self.host_runtime_perf_knobs is None: - return False - if self.host_context_progress is None: - return False - - if remove_input_padding: - if self.host_context_lengths is None: - return False - if not gpt_attention_plugin: - return False - - return True - - -class SpecDecodingParams: - - def __init__(self, - spec_decoding_is_generation_length_variable: bool = False, - spec_decoding_max_generation_length: int = 1, - spec_decoding_generation_lengths: Tensor = None, - spec_decoding_position_offsets: Tensor = None, - spec_decoding_packed_mask: Tensor = None, - spec_decoding_use: Tensor = None): - - self.spec_decoding_is_generation_length_variable = spec_decoding_is_generation_length_variable - self.spec_decoding_max_generation_length = spec_decoding_max_generation_length - self.spec_decoding_generation_lengths = spec_decoding_generation_lengths - self.spec_decoding_position_offsets = spec_decoding_position_offsets - self.spec_decoding_packed_mask = spec_decoding_packed_mask - self.spec_decoding_use = spec_decoding_use - - -class MropeParams: - - def __init__( - self, - mrope_rotary_cos_sin: Tensor = None, - mrope_position_deltas: Tensor = None, - ): - self.mrope_rotary_cos_sin = mrope_rotary_cos_sin - self.mrope_position_deltas = mrope_position_deltas - - -class KeyValueCacheParams: - - def __init__(self, - past_key_value: List[Tensor] = None, - host_past_key_value_lengths: Tensor = None, - host_max_attention_window_sizes: Tensor = None, - host_sink_token_length: Tensor = None, - kv_cache_block_offsets: Tensor = None, - host_kv_cache_block_offsets: Tensor = None, - host_kv_cache_pool_pointers: Tensor = None, - host_kv_cache_pool_mapping: Tensor = None, - cache_indirection: Tensor = None, - past_key_value_length: Tensor = None, - cross_kv_cache_block_offsets: Tensor = None, - host_cross_kv_cache_block_offsets: Tensor = None, - host_cross_kv_cache_pool_pointers: Tensor = None, - host_cross_kv_cache_pool_mapping: Tensor = None): - self.past_key_value = past_key_value - self.host_past_key_value_lengths = host_past_key_value_lengths - self.host_max_attention_window_sizes = host_max_attention_window_sizes - self.host_sink_token_length = host_sink_token_length - self.kv_cache_block_offsets = kv_cache_block_offsets - self.host_kv_cache_block_offsets = host_kv_cache_block_offsets - self.host_kv_cache_pool_pointers = host_kv_cache_pool_pointers - self.host_kv_cache_pool_mapping = host_kv_cache_pool_mapping - self.cross_kv_cache_block_offsets = cross_kv_cache_block_offsets - self.host_cross_kv_cache_block_offsets = host_cross_kv_cache_block_offsets - self.host_cross_kv_cache_pool_pointers = host_cross_kv_cache_pool_pointers - self.host_cross_kv_cache_pool_mapping = host_cross_kv_cache_pool_mapping - self.cache_indirection = cache_indirection - # self.past_key_value_length = past_key_value_length - - def get_first_past_key_value(self): - if self.past_key_value is None: - return None - return self.past_key_value[0] - - def fill_none_tensor_list(self, list_size): - if self.past_key_value is None: - self.past_key_value = tuple([None] * list_size) - - def is_valid(self, gpt_attention_plugin): - if gpt_attention_plugin: - if self.host_past_key_value_lengths is None: - return False - if self.host_max_attention_window_sizes is None: - return False - if self.host_sink_token_length is None: - return False - if self.cache_indirection is None: - return False - - return True - - -class BlockSparseAttnParams: - - def __init__(self, - block_size: int = 64, - homo_head_pattern: bool = False, - num_local_blocks: int = 16, - vertical_stride: int = 8): - self.block_size = block_size - self.homo_head_pattern = homo_head_pattern - self.num_local_blocks = num_local_blocks - self.vertical_stride = vertical_stride - - -class Attention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - qk_layernorm=False, - layernorm_type=LayerNormType.LayerNorm, - layernorm_share=True, - inner_layernorm=False, - eps=1e-05, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_base_local=1.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - rope_scaling_short_factors=None, - rope_scaling_long_factors=None, - rope_scaling_short_mscale=None, - rope_scaling_long_mscale=None, - original_max_position_embeddings=1024, - tp_group=None, - tp_size=1, - tp_rank=0, - quant_mode: QuantMode = QuantMode(0), - q_scaling=1.0, - cross_attention=False, - relative_attention=False, - max_distance=0, - num_buckets=0, - dense_bias=None, - clip_qkv=None, - alibi_bias_max=8, - skip_cross_kv=False, - max_attn_value=0.0, - block_sparse_params=None, - use_implicit_relative_attention=False, - reorder=False, - enable_qkv=True, - cp_group=[0], - cp_size=1, - cp_rank=0, - max_seqlen_for_logn_scaling=8192, - use_logn_scaling=False, - is_local=False): - super().__init__() - - self.local_layer_idx = local_layer_idx - self.cross_attention = cross_attention - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - assert num_attention_heads % tp_size == 0, \ - "num_attention_heads must be divisible by tp_size" - self.num_attention_heads = num_attention_heads // tp_size - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.num_kv_heads = num_kv_heads if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size - self.attention_hidden_size = self.attention_head_size * self.num_attention_heads - self.max_position_embeddings = max_position_embeddings - self.original_max_position_embeddings = original_max_position_embeddings - self.bias = bias - self.tp_group = tp_group - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dtype = dtype - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - self.cp_group = cp_group - self.cp_size = cp_size - self.cp_rank = cp_rank - self.is_local = is_local - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = q_scaling - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = position_embedding_type == PositionEmbeddingType.alibi_with_scale - self.alibi_bias_max = alibi_bias_max - self.position_embedding_type = position_embedding_type - - self.relative_attention = relative_attention - self.max_distance = max_distance - self.num_buckets = num_buckets - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_base_local = rotary_embedding_base_local - self.rotary_embedding_scaling = rotary_embedding_scaling - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.short_mscale = 1.0 - self.long_mscale = 1.0 - self.rotary_embedding_percentage = rotary_embedding_percentage - self.use_implicit_relative_attention = self.relative_attention and use_implicit_relative_attention - self.max_seqlen_for_logn_scaling = max_seqlen_for_logn_scaling - self.use_logn_scaling = use_logn_scaling - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - self.rotary_embedding_dim = 0 - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes( - self.num_attention_heads * self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale, - alibi_bias_max=self.alibi_bias_max) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - if self.use_logn_scaling: - logn_scaling = generate_logn_scaling( - self.max_seqlen_for_logn_scaling, self.max_position_embeddings) - self.register_parameter( - 'logn_scaling', - Parameter(logn_scaling, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.max_attn_value = max_attn_value - self.register_parameter('kv_cache_scaling_factor', None) - self.register_parameter('attention_output_orig_quant_scale', None) - self.register_parameter('attention_output_sf_scale', None) - - self.block_sparse_params = block_sparse_params if block_sparse_params is not None else BlockSparseAttnParams( - ) - - # The output feature size is therefore (h/tp + 2*kvh/tp) * d, where h is num_heads, - # d is head_size, kvh is the num_kv_heads and tp is tensor_parallel_size. - # In ColumnLinear op, the output dim is calculated by (h + 2*kvh) * d / tp, - # which matches the desired output size (h/tp + 2*kvh/tp) * d after splitting - - # out dim is not necessarily hidden_size + kv specific size (in MQA/GQA), but num_heads * heads_size - # example: d_model != num_heads * head_size in Flan-T5/ByT5/Gemma - if enable_qkv: - self.qkv = ColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - is_qkv=True) - self.dense = RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - # see optimize_model's add_lora for LoRA initialization - self.qkv_lora = None - self.qkv_dora = None - - # per-layer relative attention table - if self.use_implicit_relative_attention: - self.rel_attn_table = Parameter(shape=(num_attention_heads // - tp_size, num_buckets), - dtype=dtype) - - # qk layernorm - self.qk_layernorm = qk_layernorm - self.layernorm_type = layernorm_type - self.layernorm_share = layernorm_share - ln_type = layernorm_map[layernorm_type] - if self.qk_layernorm: - # layernorm_share indicates whether all the QK head in one layer shares the same norm parameters or not - if layernorm_share: - self.q_layernorm = ln_type(self.attention_head_size, - eps=eps, - dtype=dtype) - self.k_layernorm = ln_type(self.attention_head_size, - eps=eps, - dtype=dtype) - else: - assert ln_type == LayerNorm - self.q_layernorm = ln_type( - (self.num_attention_heads, self.attention_head_size), - eps=eps, - dtype=dtype, - bias=False, - tp_size=tp_size, - tp_dim=0) - self.k_layernorm = ln_type( - (self.num_attention_kv_heads, self.attention_head_size), - eps=eps, - dtype=dtype, - bias=False, - tp_size=tp_size, - tp_dim=0) - - self.inner_layernorm = ln_type(self.hidden_size, dtype=dtype, - eps=eps) if inner_layernorm else None - if clip_qkv is not None: - self.clip_qkv = fp32_array([clip_qkv]) - else: - self.clip_qkv = None - - self.skip_cross_kv = skip_cross_kv - - @staticmethod - def create_attention_const_params(model_cls, config): - # get rotary parameters. - hidden_size = config.hidden_size - num_attention_heads = config.num_attention_heads - attention_head_size = config.head_size - max_position_embeddings = config.max_position_embeddings - position_embedding_type = config.position_embedding_type - rotary_embedding_base = getattr(config, 'rotary_base', 10000.0) - rotary_embedding_scaling = getattr(config, 'rotary_scaling', None) - rotary_embedding_percentage = getattr(config, 'rotary_pct', 1.0) - # only rope need the const parameters. - if not position_embedding_type.is_rope(): - return - # attention head size - attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - # rotary embedding dim. - rotary_embedding_dim = getattr( - config, 'rotary_dim', - int(attention_head_size * rotary_embedding_percentage)) - # rotary scaling. - rotary_embedding_scale_type = RotaryScalingType.none - rotary_embedding_scale = 1.0 - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - rotary_embedding_scale = rotary_embedding_scaling.get("factor", 1.0) - - if position_embedding_type == PositionEmbeddingType.long_rope: - rope_scaling_short_factors, rope_scaling_long_factors = None, None - rope_scaling_short_mscale, rope_scaling_long_mscale = None, None - original_max_position_embeddings = 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 - - if config.architecture == "Phi3SmallForCausalLM" or config.architecture == "PhiMoEForCausalLM": - rope_scaling_short_mscale = config.longrope_short_mscale - rope_scaling_long_mscale = config.longrope_long_mscale - - embed_positions, long_rope_embed_positions, \ - (rotary_inv_freq, embed_positions_for_gpt_attention), \ - (long_rope_rotary_inv_freq, long_rope_embed_positions_for_gpt_attention), mscale \ - = RopeEmbeddingUtils.create_sinusoidal_positions_long_rope_for_attention_plugin( - max_position_embeddings, - original_max_position_embeddings, rotary_embedding_dim, - rotary_embedding_base, rope_scaling_short_factors, - rope_scaling_long_factors, rope_scaling_short_mscale, rope_scaling_long_mscale) - - if rope_scaling_short_mscale is not None: - assert rope_scaling_long_mscale is not None - short_mscale = rope_scaling_short_mscale - long_mscale = rope_scaling_long_mscale - else: - short_mscale = long_mscale = mscale - - model_cls.register_parameter( - 'embed_positions', - Parameter(embed_positions, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'long_rope_embed_positions', - Parameter(long_rope_embed_positions, - dtype='float32', - is_buffer=True)) - model_cls.register_parameter( - 'rotary_inv_freq', - Parameter(rotary_inv_freq, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'long_rope_rotary_inv_freq', - Parameter(long_rope_rotary_inv_freq, - dtype='float32', - is_buffer=True)) - model_cls.register_parameter( - 'embed_positions_for_gpt_attention', - Parameter(embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - model_cls.register_parameter( - 'long_rope_embed_positions_for_gpt_attention', - Parameter(long_rope_embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - model_cls.short_mscale = short_mscale - model_cls.long_mscale = long_mscale - elif rotary_embedding_scale_type == RotaryScalingType.yarn: - beta_fast = rotary_embedding_scaling.get("beta_fast", 32.0) - beta_slow = rotary_embedding_scaling.get("beta_slow", 1.0) - mscale = rotary_embedding_scaling.get("mscale", 1.0) - mscale_all_dim = rotary_embedding_scaling.get("mscale_all_dim", 0.0) - original_max_position_embeddings = rotary_embedding_scaling.get( - "original_max_position_embeddings", 4096) - rotary_inv_freq, embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_yarn( - max_position_embeddings, rotary_embedding_dim, - rotary_embedding_base, rotary_embedding_scale, - original_max_position_embeddings, beta_fast, beta_slow, mscale, - mscale_all_dim, False) - - embed_positions = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, - rotary_embedding_dim, - ) - model_cls.register_parameter( - 'embed_positions', - Parameter(embed_positions, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'rotary_inv_freq', - Parameter(rotary_inv_freq, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'embed_positions_for_gpt_attention', - Parameter(embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - else: - - def register_rope_params(rotary_base, rotary_embedding_scale, - rotary_embedding_scale_type, - rotary_embedding_scaling, - names_to_register): - # Rotary const weights. - embed_positions = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, - rotary_embedding_dim, - ) - - rotary_inv_freq, embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin( - max_position_embeddings, rotary_embedding_dim, rotary_base, - rotary_embedding_scale, rotary_embedding_scale_type, - rotary_embedding_scaling) - model_cls.register_parameter( - names_to_register[0], - Parameter(embed_positions, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - names_to_register[1], - Parameter(rotary_inv_freq, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - names_to_register[2], - Parameter(embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - - register_rope_params( - rotary_base=rotary_embedding_base, - rotary_embedding_scale=rotary_embedding_scale, - rotary_embedding_scale_type=rotary_embedding_scale_type, - rotary_embedding_scaling=rotary_embedding_scaling, - names_to_register=[ - 'embed_positions', 'rotary_inv_freq', - 'embed_positions_for_gpt_attention' - ]) - - # For models with non-homegeneous attention layers requiring a second set of rope params. e.g. Gemma3. - rotary_embedding_base_local = getattr(config, - 'rope_local_base_freq', None) - if rotary_embedding_base_local is not None: - register_rope_params( - rotary_base=rotary_embedding_base_local, - rotary_embedding_scale=1.0, - rotary_embedding_scale_type=RotaryScalingType.none, - rotary_embedding_scaling=None, - names_to_register=[ - 'embed_positions_local', 'rotary_inv_freq_local', - 'embed_positions_for_gpt_attention_local' - ]) - - @staticmethod - def fill_attention_params(model_cls, attention_params): - if model_cls.position_embedding_type.is_rope(): - if attention_params is None: - attention_params = AttentionParams() - if model_cls.position_embedding_type == PositionEmbeddingType.long_rope: - return attention_params.fill_attention_const_params_for_long_rope( - model_cls.embed_positions.value, - model_cls.long_rope_embed_positions.value, - model_cls.rotary_inv_freq.value, - model_cls.long_rope_rotary_inv_freq.value, - model_cls.embed_positions_for_gpt_attention.value, - model_cls.long_rope_embed_positions_for_gpt_attention.value, - model_cls.short_mscale, model_cls.long_mscale) - else: - return attention_params.fill_attention_const_params_for_rope( - model_cls.embed_positions.value, - model_cls.rotary_inv_freq.value, - model_cls.embed_positions_for_gpt_attention.value, - model_cls.embed_positions_local.value if hasattr( - model_cls, "embed_positions_local") else None, - model_cls.rotary_inv_freq_local.value if hasattr( - model_cls, "rotary_inv_freq_local") else None, - model_cls.embed_positions_for_gpt_attention_local.value - if hasattr( - model_cls, - "embed_positions_for_gpt_attention_local") else None) - # Fill nothing. - return attention_params - - def _get_output_orig_quant_scale(self): - attention_output_orig_quant_scale = self.attention_output_orig_quant_scale.value if self.attention_output_orig_quant_scale is not None else None - if attention_output_orig_quant_scale is not None and ( - default_net().plugin_config.gemm_plugin == 'nvfp4' - or self.quant_mode.has_nvfp4()): - # The scale was intended for nvfp4 quantization: max_value * scale = fp4_max * fp8_max - # So if we want to quantize the output to fp8, the scale should be divided by fp4_max - attention_output_orig_quant_scale = attention_output_orig_quant_scale / 6.0 - return attention_output_orig_quant_scale - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - attention_packed_mask=None, - use_cache=False, - spec_decoding_params=None, - mrope_params=None, - kv_cache_params=None, - attention_params=None, - encoder_output: Optional[Tensor] = None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - all_reduce_params: Optional[AllReduceParams] = None, - skip_attn=None, - ): - attention_input = hidden_states - - assert isinstance(hidden_states, (Tensor, tuple)) - - spec_decoding_params = SpecDecodingParams( - ) if spec_decoding_params is None else spec_decoding_params - - mrope_params = MropeParams() if mrope_params is None else mrope_params - logn_scaling = None - if self.use_logn_scaling: - logn_scaling = self.logn_scaling.value - - alibi_slopes = None - if self.position_embedding_type.is_alibi(): - alibi_slopes = self.alibi_slopes.value - if default_net().plugin_config.gpt_attention_plugin: - alibi_slopes = cast(alibi_slopes, hidden_states.dtype) - - qkv_lora_params = None - if lora_layer_params is not None: - if not self.cross_attention: - qkv_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_qkv") - else: - qkv_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_qkv") - - unfuse_qkv_gemm = self.qkv is None - if unfuse_qkv_gemm: - qkv_gemm = [self.q, self.k, self.v] - qkv = [gemm(hidden_states) for gemm in qkv_gemm] - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is not None: - lora = self.qkv.lora(hidden_states, qkv_lora_params) - kv_size = self.attention_head_size * self.num_attention_kv_heads - qkv_lora = split(lora, - [self.attention_hidden_size, kv_size, kv_size], - dim=1) - qkv = [tensor + lora for tensor, lora in zip(qkv, qkv_lora)] - else: - qkv = self.qkv(hidden_states, qkv_lora_params) - - if self.clip_qkv is not None: - qkv = clip(qkv, -self.clip_qkv, self.clip_qkv) - - if default_net().plugin_config.remove_input_padding: - if unfuse_qkv_gemm: - for tensor in qkv: - assert tensor.ndim() == 2 - else: - assert qkv.ndim() == 2 - - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is None and lora_layer_params is not None: - if not self.cross_attention: - q_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_q") - k_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_k") - v_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_v") - else: - q_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_q") - k_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_k") - v_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_v") - - assert (q_lora_params is not None and k_lora_params is not None and v_lora_params is not None) or \ - (q_lora_params is None and k_lora_params is None and v_lora_params is None), "q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time." - - if q_lora_params is not None and k_lora_params is not None and v_lora_params is not None: - qkv_lora_runtime_params = LoraRuntimeParams( - lora_ranks=[ - q_lora_params.lora_ranks[0], - k_lora_params.lora_ranks[0], - v_lora_params.lora_ranks[0], - ], - lora_weights_pointers=[ - q_lora_params.lora_weights_pointers[0], - k_lora_params.lora_weights_pointers[0], - v_lora_params.lora_weights_pointers[0], - ], - host_request_types=q_lora_params.host_request_types, - host_context_lengths=q_lora_params.host_context_lengths, - max_encoder_context_length=q_lora_params. - max_encoder_context_length, - host_encoder_input_lengths=q_lora_params. - host_encoder_input_lengths, - partial_lora_mask=lora_layer_params.partial_lora_mask, - ) - - q_lora, k_lora, v_lora = self.qkv_lora(hidden_states, - qkv_lora_runtime_params) - qkv_lora = concat([q_lora, k_lora, v_lora], - dim=q_lora.rank() - 1) - qkv = qkv + qkv_lora - if self.qkv_dora is not None: - qkv = self.qkv_dora(qkv, qkv_lora_runtime_params) - if self.qk_layernorm: - base_shape = shape(qkv, 0) if qkv.ndim() == 2 else concat( - [shape(qkv, 0), shape(qkv, 1)]) - qkv_sections = [ - self.num_attention_heads, self.num_attention_kv_heads, - self.num_attention_kv_heads - ] - total_heads = sum(qkv_sections) - if self.num_attention_heads != self.num_attention_kv_heads: - qkv = qkv.view( - concat([base_shape, total_heads, self.attention_head_size])) - query, key, value = split(qkv, qkv_sections, dim=qkv.ndim() - 2) - else: - qkv = qkv.view( - concat([ - base_shape, self.num_attention_heads, 3, - self.attention_head_size - ])) - query, key, value = split(qkv, 1, dim=qkv.ndim() - 2) - q_shape = concat([ - base_shape, self.num_attention_heads, - self.attention_head_size - ]) - query = query.view(q_shape) - key = key.view(q_shape) - value = value.view(q_shape) - - normalized_shape = None - if not self.layernorm_share: - normalized_shape = self.attention_head_size - query = self.q_layernorm(query, normalized_shape=normalized_shape) - key = self.k_layernorm(key, normalized_shape=normalized_shape) - qkv = concat([query, key, value], dim=query.ndim() - 2) - qkv = qkv.view( - concat([base_shape, total_heads * self.attention_head_size])) - if self.position_embedding_type == PositionEmbeddingType.chatglm: - qkv = RopeEmbeddingUtils.apply_rotary_pos_emb_chatglm( - qkv, - position_embedding, - self.num_attention_heads, - self.attention_head_size, - self.max_position_embeddings, - self.rotary_embedding_scale, - default_net().plugin_config.remove_input_padding, - ) - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - - paged_kv_cache = default_net().plugin_config.paged_kv_cache - - assert attention_params is None or attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - - if use_cache: - assert kv_cache_params is None or kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - - past_key_value = None if kv_cache_params is None else kv_cache_params.get_first_past_key_value( - ) - - # if cross attention, cross QKV only needs to be calculated once in the - # 1st decoding step --> write to cross KV cache --> remains constant - # during the entire decoding steps. - # 1st and >1st steps are distinguished by a boolean tensor `cross_kv_cache_gen` passed at runtime - # also, cross KV cache max length is set from encoder output seqlen, - # this maps to the max context length concept in decoder-only models - cross_kv = None - if self.cross_attention and encoder_output: - assert isinstance(encoder_output, Tensor) - - def compute_cross_kv(encoder_output): - if hasattr(self, 'kv'): - # We optimize the graph by adding kv in the cross attention layer, preventing computing the - # query of encoder_output. - assert qkv_lora_params is None, "Not support LoRA when we only compute key/value in cross atteniton" - # see optimization_model's optimize_cross_qkv - cross_kv = self.kv(encoder_output, qkv_lora_params) - base_shape = shape( - cross_kv, 0) if cross_kv.ndim() == 2 else concat( - [shape(cross_kv, 0), - shape(cross_kv, 1)]) - if self.qk_layernorm: - cross_kv = cross_kv.view( - concat([ - base_shape, 2 * self.num_attention_kv_heads, - self.attention_head_size - ])) - - key, value = split(cross_kv, [ - self.num_attention_kv_heads, - self.num_attention_kv_heads - ], - dim=cross_kv.ndim() - 2) - - key = self.k_layernorm(key) - cross_kv = concat([key, value], dim=key.ndim() - 2) - else: - cross_qkv = self.qkv(encoder_output, qkv_lora_params) - base_shape = shape( - cross_qkv, 0) if cross_qkv.ndim() == 2 else concat( - [shape(cross_qkv, 0), - shape(cross_qkv, 1)]) - - cross_qkv = cross_qkv.view( - concat([ - base_shape, self.num_attention_heads + - 2 * self.num_attention_kv_heads, - self.attention_head_size - ])) - - if self.qk_layernorm: - _, key, value = split(cross_qkv, [ - self.num_attention_heads, - self.num_attention_kv_heads, - self.num_attention_kv_heads - ], - dim=cross_qkv.ndim() - 2) - - key = self.k_layernorm(key) - cross_kv = concat([key, value], dim=key.ndim() - 2) - else: - _, cross_kv = split(cross_qkv, [ - self.num_attention_heads, - self.num_attention_kv_heads * 2 - ], - dim=cross_qkv.ndim() - 2) - cross_kv = cross_kv.view( - concat([ - base_shape, 2 * self.num_attention_kv_heads * - self.attention_head_size - ])) - - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is None and lora_layer_params is not None: - _, cross_k_lora, cross_v_lora = self.qkv_lora( - encoder_output, - qkv_lora_runtime_params, - is_cross_attention=True) - cross_kv_lora = concat([cross_k_lora, cross_v_lora], - dim=cross_k_lora.rank() - 1) - cross_kv = cross_kv + cross_kv_lora - if self.qkv_dora is not None: - cross_kv = self.qkv_dora(cross_kv, - qkv_lora_runtime_params, - is_cross_attention=True) - - return cross_kv - - if self.skip_cross_kv: - conditional = Conditional(cross_kv_cache_gen) - cond_in1 = conditional.add_input(encoder_output) - cond_in2 = conditional.add_input(cross_kv_reuse) - - ## True branch: context phase, compute cross qkv - cross_kv_true = compute_cross_kv(cond_in1) - - ## False branch: generation phase, no compute but need to obey shape constraints - # because TRT's IfConditional requires the output shape of two subgraphs to be identical - # our 1st attempt was to stack encoder_output [B, S, H] or [N, H] --> cross qkv [B, S, 3*H] or [N, 3*H], - # but it still introduces unnecessary concat. A better solution is to create a dummy torch tensor `cross_kv_resue` - # with the correct shape and reuse it in every generation step - cross_kv_false = cond_in2 - cross_kv = conditional.add_output(cross_kv_true, cross_kv_false) - else: - cross_kv = compute_cross_kv(encoder_output) - - if default_net().plugin_config.gpt_attention_plugin: - if self.cross_attention and (past_key_value is not None): - past_key_value = kv_cache_params.past_key_value[1] - assert self.attention_mask_type in [ - AttentionMaskType.causal, AttentionMaskType.bidirectional, - AttentionMaskType.bidirectionalglm, - AttentionMaskType.blocksparse - ], 'Plugin only support masked MHA.' - - # KV cache scales. - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - - # The output SF scale, needed when fuse_fp4_quant is enabled. - attention_output_sf_scale = self.attention_output_sf_scale.value if self.attention_output_sf_scale is not None else None - - # The rotary inv freq can be pre-computed. - rotary_inv_freq = getattr(attention_params, "rotary_inv_freq", None) - # Rotary cos/sin cache. - rotary_cos_sin = getattr(attention_params, - "embed_positions_for_gpt_attention", None) - rotary_inv_freq_local = getattr(attention_params, - "rotary_inv_freq_local", None) - rotary_cos_sin_local = getattr( - attention_params, "embed_positions_for_gpt_attention_local", - None) - - long_rope_rotary_inv_freq = getattr(attention_params, - "long_rope_rotary_inv_freq", - None) - long_rope_rotary_cos_sin = getattr( - attention_params, "long_rope_embed_positions_for_gpt_attention", - None) - - if self.position_embedding_type == PositionEmbeddingType.learned_absolute: - rotary_inv_freq = None - rotary_cos_sin = None - - # check if the cache is provided. - if self.position_embedding_type.is_rope(): - assert (rotary_inv_freq is not None) and ( - rotary_cos_sin is not None - ), "rotary_inv_freq and embed_positions_for_gpt_attention must be provided." - if self.position_embedding_type == PositionEmbeddingType.long_rope: - assert long_rope_rotary_inv_freq is not None - assert long_rope_rotary_cos_sin is not None - - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=past_key_value, - attention_mask=attention_mask, - attention_packed_mask=attention_packed_mask, - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base - if not self.is_local else self.rotary_embedding_base_local, - rotary_embedding_scale_type=self.rotary_embedding_scale_type - if not self.is_local else RotaryScalingType.none, - rotary_embedding_short_m_scale=attention_params.short_mscale, - rotary_embedding_long_m_scale=attention_params.long_mscale, - rotary_embedding_scale=self.rotary_embedding_scale - if not self.is_local else 1.0, - rotary_embedding_max_positions=self.max_position_embeddings, - rotary_embedding_original_max_positions=self. - original_max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq - if not self.is_local else rotary_inv_freq_local, - rotary_cos_sin=rotary_cos_sin - if not self.is_local else rotary_cos_sin_local, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - attention_output_orig_quant_scale=self. - _get_output_orig_quant_scale(), - attention_output_sf_scale=attention_output_sf_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - mask_type=self.attention_mask_type, - block_sparse_block_size=self.block_sparse_params.block_size, - block_sparse_homo_head_pattern=self.block_sparse_params. - homo_head_pattern, - block_sparse_num_local_blocks=self.block_sparse_params. - num_local_blocks, - block_sparse_vertical_stride=self.block_sparse_params. - vertical_stride, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets - if not self.cross_attention else - kv_cache_params.cross_kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_pool_mapping, - do_cross_attention=self.cross_attention, - cross_kv=cross_kv, - cross_kv_length=attention_params.encoder_max_input_length, - encoder_input_lengths=attention_params.encoder_input_lengths, - logn_scaling=logn_scaling, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_distance=self.max_distance, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_is_generation_length_variable=spec_decoding_params - .spec_decoding_is_generation_length_variable, - spec_decoding_max_generation_length=spec_decoding_params. - spec_decoding_max_generation_length, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - spec_decoding_use=spec_decoding_params.spec_decoding_use, - long_rope_rotary_inv_freq=long_rope_rotary_inv_freq, - long_rope_rotary_cos_sin=long_rope_rotary_cos_sin, - mrope_rotary_cos_sin=mrope_params.mrope_rotary_cos_sin, - mrope_position_deltas=mrope_params.mrope_position_deltas, - attn_logit_softcapping_scale=self.max_attn_value, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - skip_attn=skip_attn, - cp_size=self.cp_size, - cp_rank=self.cp_rank, - cp_group=self.cp_group) - - else: - # plain TensorRT mode - assert paged_kv_cache == False - - assert logn_scaling is None, "plan TensorRT mode does not support logn scaling now" - - def transpose_for_scores(x, - rotary: bool = False, - is_kv: bool = False): - _num_attention_heads = self.num_attention_kv_heads if is_kv else self.num_attention_heads - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), _num_attention_heads, self.attention_head_size - ]) - if rotary: - return x.view(new_x_shape) - else: - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - # qkv after projection is of shape - # [bs, seqlen, (num_attention_heads + 2 * num_attention_kv_heads), attention_head_size]. - # The projected and split qkv after transpose_for_scores(): - # Q[bs, num_attention_heads, seqlen, attention_head_size] - # K[bs, num_attention_kv_heads, seqlen, attention_head_size] - # V[bs, num_attention_kv_heads, seqlen, attention_head_size] - kv_size = self.attention_head_size * self.num_attention_kv_heads - if unfuse_qkv_gemm: - query, key, value = qkv[0], qkv[1], qkv[2] - else: - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - - # in cross attention mode, replace kv by encoder_output - if self.cross_attention and encoder_output is not None: - key, value = split(cross_kv, [kv_size, kv_size], dim=2) - - query = transpose_for_scores( - query, rotary=self.position_embedding_type.is_rope()) - key = transpose_for_scores( - key, is_kv=True, rotary=self.position_embedding_type.is_rope()) - value = transpose_for_scores(value, is_kv=True) - - if self.position_embedding_type.is_rope(): - if self.position_embedding_type == PositionEmbeddingType.long_rope: - sequence_length = shape(hidden_states, 1) - floor_seq_length = maximum( - sequence_length, self.original_max_position_embeddings) - - starts = concat([0, 0, 0]) - shapes = concat( - [1, floor_seq_length, self.rotary_embedding_dim]) - short = slice(attention_params.embed_positions, starts, - shapes) - long = slice(attention_params.long_rope_embed_positions, - starts, shapes) - - embed_positions = concat([short, long], dim=0) - select = where( - sequence_length - <= self.original_max_position_embeddings, 0, 1) - embed_positions = slice(embed_positions, - concat([select, 0, 0]), - sizes=shape(short)) - embed_positions = cast(embed_positions, self.dtype) - elif is_same_dtype(self.dtype, trt.bfloat16): - embed_positions = cast(attention_params.embed_positions, - trt.bfloat16) - else: - embed_positions = cast(attention_params.embed_positions, - query.dtype) - - if self.rotary_embedding_dim is not None: - # When shape(hidden_states, 1) > 1(Context phase), the embedding start from 0, - # otherwise (Generation phase) move start to position - if not use_cache: - # Only context phase is involved when kv cache is disabled. - start = 0 - else: - start = where( - shape(hidden_states, 1) > 1, 0, - shape(past_key_value, 3)) - size = where( - shape(hidden_states, 1) > 1, shape(hidden_states, 1), 1) - sincos = slice(embed_positions, concat([0, start, 0]), - concat([1, size, self.rotary_embedding_dim])) - sin, cos = split(sincos, - self.rotary_embedding_dim // 2, - dim=-1) - - key_rot_size = concat([ - shape(key, 0), - shape(key, 1), - shape(key, 2), self.rotary_embedding_dim - ]) - query_rot_size = concat([ - shape(query, 0), - shape(query, 1), - shape(query, 2), self.rotary_embedding_dim - ]) - remaining = shape(key, 3) - self.rotary_embedding_dim - key_pass_size = concat([ - shape(key, 0), - shape(key, 1), - shape(key, 2), remaining - ]) - query_pass_size = concat([ - shape(query, 0), - shape(query, 1), - shape(query, 2), remaining - ]) - k_rot = slice(key, [0, 0, 0, 0], key_rot_size) - k_pass = slice(key, [0, 0, 0, self.rotary_embedding_dim], - key_pass_size) - - q_rot = slice(query, [0, 0, 0, 0], query_rot_size) - q_pass = slice(query, [0, 0, 0, self.rotary_embedding_dim], - query_pass_size) - - k_rot = RopeEmbeddingUtils.apply_rotary_pos_emb( - k_rot, [cos, sin], self.position_embedding_type) - q_rot = RopeEmbeddingUtils.apply_rotary_pos_emb( - q_rot, [cos, sin], self.position_embedding_type) - - key = concat([k_rot, k_pass], dim=3) - query = concat([q_rot, q_pass], dim=3) - else: - key = RopeEmbeddingUtils.apply_rotary_pos_emb( - key, [cos, sin], self.position_embedding_type) - query = RopeEmbeddingUtils.apply_rotary_pos_emb( - query, [cos, sin], self.position_embedding_type) - - key = key.permute([0, 2, 1, 3]) - query = query.permute([0, 2, 1, 3]) - - if past_key_value is not None and not self.cross_attention: - if self.kv_cache_scaling_factor is not None: - past_key_value = dequantize( - past_key_value, - self.kv_cache_scaling_factor.value, - output_type=self.dtype) - - # past_key_value [bs, 2, num_heads, max_seq_len, head_dim] - past_key, past_value = split(past_key_value, 1, dim=1) - - key_shape = concat([ - shape(past_key, 0), - shape(past_key, 2), - shape(past_key, 3), - shape(past_key, 4) - ]) - past_key = past_key.view(key_shape, zero_is_placeholder=False) - past_value = past_value.view(key_shape, - zero_is_placeholder=False) - - key = concat([past_key, key], dim=2) - value = concat([past_value, value], dim=2) - - if use_cache: - key_inflated_shape = concat([ - shape(key, 0), 1, - shape(key, 1), - shape(key, 2), - shape(key, 3) - ]) - inflated_key = key.view(key_inflated_shape, - zero_is_placeholder=False) - inflated_value = value.view(key_inflated_shape, - zero_is_placeholder=False) - past_key_value = concat([inflated_key, inflated_value], dim=1) - - # TRT quantizes the tensor value by doing `cast(clip(fp_value / scale))` while - # the plugin quantizes it by doing `cast(clip(fp_value * scale))`. - if self.kv_cache_scaling_factor is not None: - past_key_value = quantize( - past_key_value, - self.kv_cache_scaling_factor.value, - dtype='fp8' - if self.quant_mode.has_fp8_kv_cache() else 'int8') - - # MQA broadcast - if self.num_attention_heads // self.num_attention_kv_heads > 1: - key = repeat_interleave( - key, - self.num_attention_heads // self.num_attention_kv_heads, 1) - value = repeat_interleave( - value, - self.num_attention_heads // self.num_attention_kv_heads, 1) - - key_length = shape(key, 2) - - # The following code creates a 2D tensor with 0s in the lower triangular (including the diagonal) and - # +INF in the upper triangular parts. This bias tensor will be added to the output of the Q*K^T matrix - # multiplication (BMM1). The +INF elements will be transformed to 0s by the Softmax operator that - # follows. The elements that corresponds to 0s in the bias are unaffected by the bias tensor. - # - # Note that when we added to another bias tensor B (for example, with AliBi), the values in the lower- - # triangular part of the B tensor are not affected and the upper-triangular ones are set to +INF. - if self.attention_mask_type == AttentionMaskType.causal and not self.cross_attention: - if self.position_embedding_type.is_alibi(): - query_length = shape(query, 2) - # bsz, tatget_length, past_key_value_length - buffer = make_causal_mask(shape(query, 0), query_length, - key_length - query_length, - trt.float32) - starts = concat([0, 0, 0, 0]) - sizes = concat([1, 1, query_length, key_length]) - generated_mask = slice(buffer, starts, sizes) - - else: - query_length = shape(query, 2) - starts = concat([0, 0, key_length - query_length, 0]) - sizes = concat([1, 1, query_length, key_length]) - if self.position_embedding_type == PositionEmbeddingType.long_rope: - buf_shape = (self.original_max_position_embeddings, - self.original_max_position_embeddings) - else: - buf_shape = (self.max_position_embeddings, - self.max_position_embeddings) - select_buf = np.expand_dims( - np.tril(np.ones(buf_shape)).astype(bool), (0, 1)) - - select_buf = np.logical_not(select_buf) - mask_buf = np.zeros_like(select_buf, np.float32) - mask_buf[select_buf] = float('-inf') - buffer = constant(mask_buf) - generated_mask = slice(buffer, starts, sizes) - - elif self.attention_mask_type == AttentionMaskType.bidirectional and not self.cross_attention: - query_length = shape(query, 2) - zero_buf = np.expand_dims( - np.zeros((self.max_position_embeddings, - self.max_position_embeddings), - dtype=np.float32), (0, 1)) - - zero_buf[:, :, :-1, -1] = 1 - zero_buf *= -10000 - - mask = constant(zero_buf) - - # context phase, query_length - mask_size = where(query_length > 1, query_length, 1) - mask_start = where(query_length > 1, - self.max_position_embeddings - mask_size, 1) - start = concat([0, 0, mask_start, mask_start]) - size = concat([1, 1, mask_size, mask_size]) - generated_mask = slice(mask, start, size) - - if attention_mask is not None: - if self.cross_attention: - batch_size = shape(attention_mask, 0) - query_len = shape(attention_mask, 1) - encoder_input_len = shape(attention_mask, 2) - attention_mask = attention_mask.view( - concat([batch_size, 1, query_len, encoder_input_len])) - attention_mask = where(attention_mask == 0, float('-inf'), - 0.0) - else: - attention_mask = expand_mask(attention_mask, - shape(query, 2)) - bias = attention_mask - if self.position_embedding_type.is_alibi(): - alibi_biases = generate_alibi_biases(alibi_slopes, key_length) - bias = alibi_biases if bias is None else bias + alibi_biases - - if self.relative_attention: - query_length = shape(query, 2) - if self.use_implicit_relative_attention: - relative_bias = compute_relative_bias( - query_length + key_length - 1, - key_length, - self.num_buckets, - self.max_distance, - False, # bidirectional - self.rel_attn_table.value.transpose(1, 0), - tp_size=self.tp_size, - tp_group=self.tp_group, - tp_rank=self.tp_rank) - else: - relative_bias = unsqueeze(self.rel_attn_table.value, 0) - start = concat([0, 0, query_length + key_length - 2, 0]) - size = concat([ - shape(relative_bias, 0), - shape(relative_bias, 1), 1, key_length - ]) - relative_bias = slice(relative_bias, start, size) - - key = key.permute([0, 1, 3, 2]) - model_type = query.dtype - with precision('float32'): - # FIXME the "with precision('float32') does not really work and lead to nan" - # in some cases - query = cast(query, 'float32') - key = cast(key, 'float32') - if norm_before_bmm1: - # Apply norm on query earlier to prevent matmul fp16 overflow. - query /= (self.q_scaling * self.norm_factor) - attention_scores = matmul(query, key) - if not norm_before_bmm1: - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - if self.max_attn_value > 0: - attention_scores = self.max_attn_value * ACT2FN['tanh']( - attention_scores / self.max_attn_value) - - if self.attention_mask_type in [ - AttentionMaskType.causal, - AttentionMaskType.bidirectional - ] and not self.cross_attention: - - bias = generated_mask if bias is None else bias + generated_mask - - if bias is not None: - bias = cast(bias, attention_scores.dtype) - attention_scores = attention_scores + bias - - if self.relative_attention: - attention_scores = attention_scores + relative_bias - - attention_probs = softmax(attention_scores, dim=-1) - attention_probs = cast(attention_probs, model_type) - - # A dummy reshape WAR for mha fusion - attention_probs = attention_probs.view( - concat([ - shape(attention_probs, 0), - shape(attention_probs, 1), - shape(attention_probs, 2), - shape(value, 2) - ])) - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([ - shape(context, 0), - shape(context, 1), self.attention_hidden_size - ])) - - dense_lora_params = None - if lora_layer_params is not None: - dense_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_dense") - - if skip_attn is not None and not default_net( - ).plugin_config.use_fp8_context_fmha: - # This case is used when we can skip this attention layer directly. - # The output would be undefined and not used if skip_attn is not None - # and set skip_attn as True during runtime - # But when use_fp8_context_fmha is enabled, the output data type of - # attention_plugin is fp8. Since TRT's conditional layer does not support - # FP8 data type yet, we cannot use it to skip the computation in such case. - - dense_conditional = Conditional(skip_attn) - skip_case = dense_conditional.add_input(attention_input) - context = dense_conditional.add_input(context) - - if self.inner_layernorm is not None: - context = self.inner_layernorm(context) - context = self.dense(context, - lora_runtime_params=dense_lora_params, - all_reduce_params=all_reduce_params) - - if skip_attn is not None and not default_net( - ).plugin_config.use_fp8_context_fmha: - context = dense_conditional.add_output(skip_case, context) - - if use_cache: - return (context, past_key_value) - else: - return context - - def set_rel_attn_table(self, max_seq_len, precomputed_relative_attention): - self.rel_attn_table = Parameter(shape=(self.num_attention_heads, - max_seq_len + 1, - max_seq_len + 1), - dtype=self.dtype) - self.rel_attn_table.value = precomputed_relative_attention - - def postprocess(self, tllm_key, weights, **kwargs): - - if tllm_key.endswith("kv_cache_scaling_factor"): - if weights is None: - return {tllm_key: torch.ones(1, ).float()} - elif isinstance(weights, torch.Tensor): - return {tllm_key: weights.float()} - elif None in weights: - return {tllm_key: torch.ones(1, ).float()} - else: - return {tllm_key: max(weights).float()} - elif tllm_key.endswith("kv_cache_rcp_scaling_factor"): - if weights is None: - return {tllm_key: torch.ones(1, ).float()} - elif isinstance(weights, torch.Tensor): - return {tllm_key: torch.reciprocal(weights.float())} - elif None in weights: - return {tllm_key: torch.ones(1, ).float()} - else: - return {tllm_key: torch.reciprocal(max(weights).float())} - else: - return {tllm_key: weights} - - -class BertAttention(Module): - - def __init__(self, - hidden_size, - num_attention_heads, - max_position_embeddings=1024, - num_layers=1, - attention_head_size=None, - num_kv_heads=None, - q_scaling=1.0, - apply_query_key_layer_scaling=False, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - cp_group=None, - cp_size=1, - cp_rank=0, - relative_attention=False, - max_distance=0, - num_buckets=0, - quant_mode=QuantMode(0)): - super().__init__() - - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size - self.attention_hidden_size = self.attention_head_size * self.num_attention_heads - self.max_position_embeddings = max_position_embeddings - self.norm_factor = math.sqrt(self.attention_head_size) - self.tp_group = tp_group - self.tp_size = tp_size - self.tp_rank = tp_rank - self.cp_group = cp_group - self.cp_size = cp_size - self.cp_rank = cp_rank - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = q_scaling - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - - self.dtype = dtype - # add quant mode to control quantization - self.quant_mode = quant_mode - - self.relative_attention = relative_attention - self.max_distance = max_distance - self.num_buckets = num_buckets - - # out dim is not necessarily hidden_size + kv specific size (in MQA/GQA), but num_heads * heads_size - # example: d_model != num_heads * head_size in Flan-T5 - self.qkv = ColumnLinear(hidden_size, - tp_size * self.attention_hidden_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - is_qkv=True) - self.dense = RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - # see optimize_model's add_lora for LoRA initialization - self.qkv_lora = None - - # per-layer relative attention table - if relative_attention: - self.rel_attn_table = Parameter(shape=(num_attention_heads // - tp_size, num_buckets), - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - input_lengths=None, - max_input_length=None, - lora_layer_params=None): - assert isinstance(hidden_states, Tensor) - - qkv_lora_params = None - if lora_layer_params is not None: - qkv_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_qkv") - - qkv = self.qkv(hidden_states, qkv_lora_params) - - if default_net().plugin_config.remove_input_padding: - assert qkv.ndim() == 2 - - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is None and lora_layer_params is not None: - q_lora_params = lora_layer_params.get_runtime_params(0, "attn_q") - k_lora_params = lora_layer_params.get_runtime_params(0, "attn_k") - v_lora_params = lora_layer_params.get_runtime_params(0, "attn_v") - - assert (q_lora_params is not None and k_lora_params is not None and v_lora_params is not None) or \ - (q_lora_params is None and k_lora_params is None and v_lora_params is None), "q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time." - - if q_lora_params is not None and k_lora_params is not None and v_lora_params is not None: - qkv_lora_params = LoraRuntimeParams( - lora_ranks=[ - q_lora_params.lora_ranks[0], - k_lora_params.lora_ranks[0], - v_lora_params.lora_ranks[0], - ], - lora_weights_pointers=[ - q_lora_params.lora_weights_pointers[0], - k_lora_params.lora_weights_pointers[0], - v_lora_params.lora_weights_pointers[0], - ], - host_request_types=q_lora_params.host_request_types, - host_context_lengths=q_lora_params.host_context_lengths) - - q_lora, k_lora, v_lora = self.qkv_lora(hidden_states, - qkv_lora_params) - qkv_lora = concat([q_lora, k_lora, v_lora], - dim=q_lora.rank() - 1) - qkv = qkv + qkv_lora - - if default_net().plugin_config.bert_attention_plugin: - # TRT plugin mode - assert input_lengths is not None - context = bert_attention( - qkv, - input_lengths, - self.num_attention_heads, - self.attention_head_size, - q_scaling=self.q_scaling, - relative_attention=self.relative_attention, - max_distance=self.max_distance, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_input_length=max_input_length, - cp_group=self.cp_group, - cp_size=self.cp_size, - cp_rank=self.cp_rank) - else: - # plain TRT mode - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - kv_size = self.attention_head_size * self.num_attention_kv_heads - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - if self.cp_size > 1 and self.cp_group is not None: - key = allgather(key, self.cp_group, gather_dim=1) - value = allgather(value, self.cp_group, gather_dim=1) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - key = key.permute([0, 1, 3, 2]) - attention_scores = matmul(query, key, use_fp32_acc=False) - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - - if self.relative_attention: - query_len = shape(attention_scores, 2) - key_len = shape(attention_scores, 3) - bias = compute_relative_bias( - query_len, - key_len, - self.num_buckets, - self.max_distance, - True, # bidirectional - self.rel_attn_table.value.transpose(1, 0), - tp_size=self.tp_size, - tp_group=self.tp_group, - tp_rank=self.tp_rank) - attention_scores = attention_scores + bias - - if attention_mask is not None: - attention_mask = expand_mask(attention_mask, shape(query, 2)) - attention_mask = cast(attention_mask, attention_scores.dtype) - attention_scores = attention_scores + attention_mask - - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([ - shape(context, 0), - shape(context, 1), self.attention_hidden_size - ])) - - dense_lora_params = None - if lora_layer_params is not None: - dense_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_dense") - context = self.dense(context, lora_runtime_params=dense_lora_params) - - return context - - -class CogVLMAttention(Attention): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - attention_mask_type=AttentionMaskType.causal, - bias=True, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - tp_group=None, - tp_size=1, - tp_rank=0, - quant_mode: QuantMode = QuantMode(0), - dense_bias=None, - ): - super().__init__(local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - dtype=dtype, - attention_mask_type=attention_mask_type, - bias=bias, - position_embedding_type=position_embedding_type, - rotary_embedding_base=rotary_embedding_base, - rotary_embedding_scaling=rotary_embedding_scaling, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=quant_mode) - - self.vis_qkv = ColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - is_qkv=True) - self.vis_dense = RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - def forward(self, - hidden_states: Tensor, - use_cache=False, - kv_cache_params=None, - attention_params=None, - vision_token_mask=None, - position_embedding=None): - - assert isinstance(hidden_states, Tensor) - assert (default_net().plugin_config.gpt_attention_plugin) - - vision_qkv = self.vis_qkv(hidden_states) - language_qkv = self.qkv(hidden_states) - qkv = where(vision_token_mask, vision_qkv, language_qkv) - - qkv = RopeEmbeddingUtils.apply_rotary_pos_emb_cogvlm( - qkv, position_embedding, self.num_attention_heads, - self.attention_head_size, self.max_position_embeddings, - self.rotary_embedding_scale, - default_net().plugin_config.remove_input_padding) - - assert attention_params is None or attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - assert kv_cache_params is None or kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - - past_key_value = None if kv_cache_params is None else kv_cache_params.get_first_past_key_value( - ) - - if default_net().plugin_config.gpt_attention_plugin: - if self.cross_attention and (past_key_value is not None): - past_key_value = kv_cache_params.past_key_value[1] - assert self.attention_mask_type in [ - AttentionMaskType.causal, AttentionMaskType.bidirectional, - AttentionMaskType.bidirectionalglm - ], 'Plugin only support masked MHA.' - - # KV cache scales. - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value if self.quant_mode.has_kv_cache_quant( - ) else None - kv_quant_orig_scale = self.kv_cache_scaling_factor.value if self.quant_mode.has_kv_cache_quant( - ) else None - - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=past_key_value, - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - position_embedding_type=self.position_embedding_type, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - attention_output_orig_quant_scale=self. - _get_output_orig_quant_scale(), - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - mask_type=self.attention_mask_type, - alibi_slopes=None, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - do_cross_attention=self.cross_attention, - cross_kv=None, - cross_kv_length=attention_params.encoder_max_input_length, - encoder_input_lengths=attention_params.encoder_input_lengths, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_distance=self.max_distance, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_position_offsets=None, - spec_decoding_packed_mask=None, - mrope_rotary_cos_sin=None, - mrope_position_deltas=None, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - ) - - vision_dense = self.vis_dense(context) - language_dense = self.dense(context) - context = where(vision_token_mask, vision_dense, language_dense) - - if use_cache: - return (context, past_key_value) - else: - return context - - -class DeepseekV2Attention(Attention): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - q_lora_rank, - kv_lora_rank, - qk_nope_head_dim=None, - qk_rope_head_dim=None, - v_head_dim=None, - eps=1e-06, - attention_mask_type=AttentionMaskType.causal, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - 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, - tp_rank=0, - quant_mode: QuantMode = QuantMode(0), - ): - super().__init__(local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_kv_heads=1, - max_position_embeddings=max_position_embeddings, - attention_head_size=kv_lora_rank + qk_rope_head_dim, - dtype=dtype, - attention_mask_type=attention_mask_type, - position_embedding_type=position_embedding_type, - rotary_embedding_base=rotary_embedding_base, - rotary_embedding_scaling=rotary_embedding_scaling, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=quant_mode, - bias=False, - dense_bias=False, - enable_qkv=False) - - self.tp_size = tp_size - - if q_lora_rank is None: - self.q_lora_rank = hidden_size - self.is_deepseek_v2_lite = True - else: - self.q_lora_rank = q_lora_rank - self.is_deepseek_v2_lite = False - - self.kv_lora_rank = kv_lora_rank - self.qk_nope_head_dim = qk_nope_head_dim - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.rotary_embedding_dim = 0 - self.rotary_scaling = rotary_scaling - self.shard_dim = 1 - - def yarn_get_mscale(scale=1, mscale=1): - if scale <= 1: - return 1.0 - return 0.1 * mscale * math.log(scale) + 1.0 - - 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"] - if mscale_all_dim: - 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_yarn( - 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) - self.register_parameter( - 'embed_positions_for_gpt_attention', - Parameter(embed_positions_for_gpt_attention, dtype='float32')) - - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - - if self.is_deepseek_v2_lite: - self.fused_a = ColumnLinear( - hidden_size, - kv_lora_rank + qk_rope_head_dim, - bias=self.dense_bias, - dtype=dtype, - ) - else: - self.fused_a = ColumnLinear( - hidden_size, - q_lora_rank + kv_lora_rank + qk_rope_head_dim, - bias=self.dense_bias, - dtype=dtype, - ) - self.q_a_layernorm = RmsNorm(q_lora_rank, dtype=dtype, eps=eps) - - self.kv_a_layernorm = RmsNorm(kv_lora_rank, dtype=dtype, eps=eps) - - self.kv_b_proj = Parameter( - shape=(self.num_attention_heads * self.qk_nope_head_dim * 2, - self.kv_lora_rank), - dtype=dtype) - self.k_b_proj_trans = Parameter( - shape=(self.num_attention_heads * self.kv_lora_rank, - self.qk_nope_head_dim), - dtype=dtype) - self.q_b_proj = Parameter( - shape=(self.num_attention_heads * - (self.qk_nope_head_dim + self.qk_rope_head_dim), - self.q_lora_rank), - dtype=dtype) - self.dense = RowLinear(tp_size * self.num_attention_heads * - self.v_head_dim, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - set_obj_attrs(self.q_b_proj, { - "weight_loader": self.weight_loader, - }) - set_obj_attrs(self.kv_b_proj, { - "weight_loader": self.weight_loader, - }) - set_obj_attrs(self.k_b_proj_trans, { - "weight_loader": self.weight_loader, - }) - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor): - # use_parallel_embedding - tp_rank = mapping.tp_rank - if self.tp_size > 1: - sharding_dim = self.sharding_dim - shard_size = param._shape[sharding_dim] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(sharding_dim, start_idx, - shard_size) - param.value = loaded_weight - - def postprocess(self, tllm_key, weights, **kwargs): - - def split_matrix_tp(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.find("q_b_proj") != -1: - q_b_proj_weight = weights.unflatten( - 0, - [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.qk_rope_head_dim, - ], - ) - - q_b_proj_weight = split_matrix_tp( - q_b_proj_weight, - self.tp_size, - self.tp_rank, - dim=0, - ) - weights = q_b_proj_weight.reshape( - self.num_attention_heads * self.tp_size * - (self.qk_nope_head_dim + self.qk_rope_head_dim) // self.tp_size, - self.q_lora_rank) - - elif tllm_key.find("kv_b_proj") != -1: - kv_b_proj_weight = weights.unflatten( - 0, - [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim, - ], - ) - kv_b_proj_weight = split_matrix_tp( - kv_b_proj_weight, - self.tp_size, - self.tp_rank, - dim=0, - ) - k_nope_weight, v_weight = kv_b_proj_weight.split( - [self.qk_nope_head_dim, self.v_head_dim], - dim=1, - ) - weights = torch.concat([ - k_nope_weight.reshape( - self.num_attention_heads * self.tp_size * - self.qk_nope_head_dim // self.tp_size, self.kv_lora_rank), - v_weight.reshape( - self.num_attention_heads * self.tp_size * self.v_head_dim // - self.tp_size, self.kv_lora_rank) - ], - dim=0) - - elif tllm_key.find("k_b_proj_trans") != -1: - kv_b_proj = weights.unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim - ]) - kv_b_proj = split(kv_b_proj, self.tp_size, self.tp_rank, dim=0) - k_nope_weight, v_weight = kv_b_proj.split( - [self.qk_nope_head_dim, self.v_head_dim], - dim=1, - ) - weights = k_nope_weight.transpose(2, 1).reshape( - self.num_attention_heads * self.kv_lora_rank, - self.qk_nope_head_dim) - - return {tllm_key: weights} - - def forward(self, - hidden_states: Tensor, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - assert default_net().plugin_config.remove_input_padding - - spec_decoding_params = SpecDecodingParams( - ) if spec_decoding_params is None else spec_decoding_params - - if default_net().plugin_config.remove_input_padding: - assert hidden_states.ndim() == 2 - - default_net().plugin_config.paged_kv_cache - - assert attention_params is None or attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - - if use_cache: - assert kv_cache_params is None or kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - - past_key_value = None if kv_cache_params is None else kv_cache_params.get_first_past_key_value( - ) - - if self.is_deepseek_v2_lite: - compressed_kv, k_pe = self.fused_a(hidden_states).split( - [self.kv_lora_rank, self.qk_rope_head_dim], -1) - compressed_kv = self.kv_a_layernorm(compressed_kv) - input_qkv = concat([hidden_states, compressed_kv, k_pe], dim=-1) - else: - compressed_q, compressed_kv, k_pe = self.fused_a( - hidden_states).split([ - self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim - ], -1) - compressed_q = self.q_a_layernorm(compressed_q) - compressed_kv = self.kv_a_layernorm(compressed_kv) - input_qkv = concat([compressed_q, compressed_kv, k_pe], dim=-1) - - if default_net().plugin_config.gpt_attention_plugin: - if self.cross_attention and (past_key_value is not None): - past_key_value = kv_cache_params.past_key_value[1] - assert self.attention_mask_type in [ - AttentionMaskType.causal, - AttentionMaskType.bidirectional, - AttentionMaskType.bidirectionalglm, - ], 'Plugin only support masked MHA.' - - # KV cache scales. - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - - rotary_cos_sin = self.embed_positions_for_gpt_attention.value - - context, past_key_value = gpt_attention( - qkv=input_qkv, - past_key_value=past_key_value, - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=1, - num_kv_heads_origin=1, - hidden_size_per_head=self.kv_lora_rank + self.qk_rope_head_dim, - q_scaling=self.q_scaling, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=None, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - attention_output_orig_quant_scale=self. - _get_output_orig_quant_scale(), - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - mask_type=self.attention_mask_type, - block_sparse_block_size=self.block_sparse_params.block_size, - block_sparse_homo_head_pattern=self.block_sparse_params. - homo_head_pattern, - block_sparse_num_local_blocks=self.block_sparse_params. - num_local_blocks, - block_sparse_vertical_stride=self.block_sparse_params. - vertical_stride, - alibi_slopes=None, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets - if not self.cross_attention else - kv_cache_params.cross_kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - do_cross_attention=self.cross_attention, - cross_kv=None, - cross_kv_length=attention_params.encoder_max_input_length, - encoder_input_lengths=attention_params.encoder_input_lengths, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_distance=self.max_distance, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_is_generation_length_variable=spec_decoding_params - .spec_decoding_is_generation_length_variable, - spec_decoding_max_generation_length=spec_decoding_params. - spec_decoding_max_generation_length, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - spec_decoding_use=spec_decoding_params.spec_decoding_use, - attn_logit_softcapping_scale=self.max_attn_value, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - is_mla_enabled_flag=True, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.kv_lora_rank, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - fused_q_proj=self.fused_q_proj.value, - q_b_proj=self.q_b_proj.value, - kv_b_proj=self.kv_b_proj.value) - - context = self.dense(context) - - if use_cache: - return (context, past_key_value) - else: - return context - - -class DiffusersAttention(Module): - - def __init__(self, - *, - query_dim: int, - cross_attention_dim: Optional[int] = None, - heads: int = 8, - kv_heads: Optional[int] = None, - dim_head: int = 64, - dropout: float = 0.0, - bias: bool = False, - upcast_attention: bool = False, - upcast_softmax: bool = False, - cross_attention_norm: Optional[str] = None, - cross_attention_norm_num_groups: int = 32, - qk_norm: Optional[str] = None, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, - norm_num_groups: Optional[int] = None, - spatial_norm_dim: Optional[int] = None, - out_bias: bool = True, - scale_qk: bool = True, - only_cross_attention: bool = False, - eps: float = 1e-5, - rescale_output_factor: float = 1.0, - residual_connection: bool = False, - out_dim: int = None, - out_context_dim: int = None, - context_pre_only=None, - pre_only=False, - elementwise_affine: bool = True, - is_causal: bool = False, - attn_forward_funcname: str = 'joint_attn_forward', - mapping=Mapping(), - dtype=None): - super().__init__() - - self.cp_size = mapping.cp_size - self.cp_group = mapping.cp_group - self.tp_group = mapping.tp_group - self.tp_size = mapping.tp_size - self.tp_rank = mapping.tp_rank - self.dtype = dtype - self.attn_forward_func = getattr(self, attn_forward_funcname) - - self.inner_dim = out_dim if out_dim is not None else dim_head * heads - self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads - self.query_dim = query_dim - self.use_bias = bias - self.is_cross_attention = cross_attention_dim is not None - self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim - - ## [TODO] Not supported yet. - # self.upcast_attention = upcast_attention - # self.upcast_softmax = upcast_softmax - # self.rescale_output_factor = rescale_output_factor - # self.residual_connection = residual_connection - # self.dropout = dropout - - self.fused_projections = False - self.out_dim = out_dim if out_dim is not None else query_dim - self.context_pre_only = context_pre_only - self.pre_only = pre_only - self.is_causal = is_causal - - self.scale_qk = scale_qk - self.scale = dim_head**-0.5 if self.scale_qk else 1.0 - - # Params for `Attention` Module - self.heads = out_dim // dim_head if out_dim is not None else heads - self.heads = self.heads // self.tp_size - self.dim_head = dim_head - # default attn settings - self.norm_factor = math.sqrt(dim_head) - self.q_scaling = 1.0 - self.max_distance = 0 - - self.added_kv_proj_dim = added_kv_proj_dim - self.only_cross_attention = only_cross_attention - if self.added_kv_proj_dim is None and self.only_cross_attention: - raise ValueError( - "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`." - ) - - if norm_num_groups is not None: - self.group_norm = GroupNorm(num_channels=query_dim, - num_groups=norm_num_groups, - eps=eps, - affine=True, - dtype=dtype) - else: - self.group_norm = None - - if spatial_norm_dim is not None: - raise NotImplementedError("SpatialNorm is not supported yet.") - else: - self.spatial_norm = None - - if qk_norm is None: - self.norm_q = None - self.norm_k = None - elif qk_norm == "layer_norm": - self.norm_q = LayerNorm(dim_head, - eps=eps, - elementwise_affine=elementwise_affine, - dtype=dtype) - self.norm_k = LayerNorm(dim_head, - eps=eps, - elementwise_affine=elementwise_affine, - dtype=dtype) - elif qk_norm == "fp32_layer_norm": - self.norm_q = LayerNorm(dim_head, - eps=eps, - elementwise_affine=False, - bias=False, - dtype=dtype) - self.norm_k = LayerNorm(dim_head, - eps=eps, - elementwise_affine=False, - bias=False, - dtype=dtype) - elif qk_norm == "rms_norm": - self.norm_q = RmsNorm(dim_head, eps=eps, dtype=dtype) - self.norm_k = RmsNorm(dim_head, eps=eps, dtype=dtype) - elif qk_norm == "rms_norm_across_heads": - # LTX applies qk norm across all heads - self.norm_q = RmsNorm(dim_head * heads, eps=eps, dtype=dtype) - self.norm_k = RmsNorm(dim_head * kv_heads, eps=eps, dtype=dtype) - elif qk_norm in ["layer_norm_across_heads", "l2"]: - raise NotImplementedError( - f"qk_norm {qk_norm} is not supported yet.") - else: - raise ValueError( - f"unknown qk_norm: {qk_norm}. Should be None,'layer_norm','fp32_layer_norm','rms_norm'" - ) - - if cross_attention_norm is None: - self.norm_cross = None - elif cross_attention_norm == "layer_norm": - self.norm_cross = LayerNorm(self.cross_attention_dim, dtype=dtype) - elif cross_attention_norm == "group_norm": - if self.added_kv_proj_dim is not None: - # The given `encoder_hidden_states` are initially of shape - # (batch_size, seq_len, added_kv_proj_dim) before being projected - # to (batch_size, seq_len, cross_attention_dim). The norm is applied - # before the projection, so we need to use `added_kv_proj_dim` as - # the number of channels for the group norm. - norm_cross_num_channels = added_kv_proj_dim - else: - norm_cross_num_channels = self.cross_attention_dim - self.norm_cross = GroupNorm( - num_channels=norm_cross_num_channels, - num_groups=cross_attention_norm_num_groups, - eps=1e-5, - affine=True, - dtype=dtype) - else: - raise ValueError( - f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'" - ) - - # [TODO] check `gather_output` - self.to_q = ColumnLinear(query_dim, - self.inner_dim, - bias=bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - if not self.only_cross_attention: - self.to_k = ColumnLinear(self.cross_attention_dim, - self.inner_kv_dim, - bias=bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - self.to_v = ColumnLinear(self.cross_attention_dim, - self.inner_kv_dim, - bias=bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - else: - self.to_k = None - self.to_v = None - - self.added_proj_bias = added_proj_bias - if self.added_kv_proj_dim is not None: - self.add_k_proj = ColumnLinear(added_kv_proj_dim, - self.inner_kv_dim, - bias=added_proj_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - self.add_v_proj = ColumnLinear(added_kv_proj_dim, - self.inner_kv_dim, - bias=added_proj_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - if self.context_pre_only is not None: - self.add_q_proj = ColumnLinear(added_kv_proj_dim, - self.inner_dim, - bias=added_proj_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - else: - self.add_q_proj = None - self.add_k_proj = None - self.add_v_proj = None - - if not self.pre_only: - self.to_out = ModuleList([ - RowLinear(self.inner_dim, - self.out_dim, - bias=out_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - dtype=dtype) - ]) - else: - self.to_out = None - - if self.context_pre_only is not None and not self.context_pre_only: - self.to_add_out = RowLinear(self.inner_dim, - self.out_dim, - bias=out_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - dtype=dtype) - else: - self.to_add_out = None - - if qk_norm is not None and added_kv_proj_dim is not None: - if qk_norm == "fp32_layer_norm": - self.norm_added_q = LayerNorm(dim_head, - elementwise_affine=False, - bias=False, - eps=eps, - dtype=dtype) - self.norm_added_k = LayerNorm(dim_head, - elementwise_affine=False, - bias=False, - eps=eps, - dtype=dtype) - elif qk_norm == "rms_norm": - self.norm_added_q = RmsNorm(dim_head, eps=eps, dtype=dtype) - self.norm_added_k = RmsNorm(dim_head, eps=eps, dtype=dtype) - else: - raise ValueError( - f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`" - ) - else: - self.norm_added_q = None - self.norm_added_k = None - - def joint_attn_forward(self, - hidden_states: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - attention_mask: Optional[Tensor] = None, - max_input_length: Optional[Tensor] = None, - *args, - **kwargs): - if attention_mask is not None: - raise NotImplementedError() - residual = identity(hidden_states) - batch_size = shape(hidden_states, 0) - - # `sample` projections. - query = self.to_q(hidden_states) - key = self.to_k(hidden_states) - value = self.to_v(hidden_states) - - head_dim = self.dim_head - inner_dim = head_dim * self.heads - - query = query.view(concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - key = key.view(concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - value = value.view(concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - - if self.norm_q is not None: - query = self.norm_q(query) - if self.norm_k is not None: - key = self.norm_k(key) - - # `context` projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = self.add_q_proj( - encoder_hidden_states) - encoder_hidden_states_key_proj = self.add_k_proj( - encoder_hidden_states) - encoder_hidden_states_value_proj = self.add_v_proj( - encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - - if self.norm_added_q is not None: - encoder_hidden_states_query_proj = self.norm_added_q( - encoder_hidden_states_query_proj) - if self.norm_added_k is not None: - encoder_hidden_states_key_proj = self.norm_added_k( - encoder_hidden_states_key_proj) - - query = concat([query, encoder_hidden_states_query_proj], dim=2) - key = concat([key, encoder_hidden_states_key_proj], dim=2) - value = concat([value, encoder_hidden_states_value_proj], dim=2) - - # Transpose from [batch_size, num_heads, seq_len, head_dim] back to - # [batch_size, seq_len, num_heads * head_dim] for attention plugin. - query = query.permute([0, 2, 1, - 3]).view(concat([batch_size, -1, inner_dim])) - key = key.permute([0, 2, 1, 3]).view(concat([batch_size, -1, - inner_dim])) - value = value.permute([0, 2, 1, - 3]).view(concat([batch_size, -1, inner_dim])) - - if default_net().plugin_config.bert_attention_plugin: - # TRT plugin mode - assert self.cp_size == 1 - shape(query, 1) - qkv = concat([query, key, value], dim=-1) - input_lengths = expand( - shape(qkv, 1).unsqueeze(0), - shape(qkv, 0).unsqueeze(0)).cast("int32") - - hidden_states = bert_attention(qkv, - input_lengths, - self.heads, - head_dim, - q_scaling=self.q_scaling, - relative_attention=False, - max_distance=self.max_distance, - max_input_length=max_input_length) - else: - # plain TRT mode - def transpose_for_scores(x): - new_x_shape = concat( - [shape(x, 0), - shape(x, 1), self.heads, head_dim]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - if self.cp_size > 1 and self.cp_group is not None: - key = allgather(key, self.cp_group, gather_dim=1) - value = allgather(value, self.cp_group, gather_dim=1) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - key = key.permute([0, 1, 3, 2]) - attention_scores = matmul(query, key, use_fp32_acc=True) - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=True).permute([0, 2, 1, 3]) - hidden_states = context.view( - concat([shape(context, 0), - shape(context, 1), inner_dim])) - - if encoder_hidden_states is not None: - # Split the attention outputs. - slice_seq_len = shape(residual, 1) - encoder_hidden_states = slice(hidden_states, - starts=concat([0, slice_seq_len, 0]), - sizes=concat([ - batch_size, - (shape(hidden_states, 1) - - slice_seq_len), inner_dim - ])) - hidden_states = slice(hidden_states, - starts=[0, 0, 0], - sizes=concat( - [batch_size, slice_seq_len, inner_dim])) - - if not self.context_pre_only: - encoder_hidden_states = self.to_add_out(encoder_hidden_states) - - # linear proj - hidden_states = self.to_out[0](hidden_states) - if encoder_hidden_states is not None: - return hidden_states, encoder_hidden_states - else: - return hidden_states - - def forward(self, - hidden_states: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - attention_mask: Optional[Tensor] = None, - max_input_length: Optional[Tensor] = None, - *args, - **kwargs): - return self.attn_forward_func( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - max_input_length=max_input_length, - *args, - **kwargs) diff --git a/tensorrt_llm/layers/cast.py b/tensorrt_llm/layers/cast.py deleted file mode 100644 index 0dd647984aff..000000000000 --- a/tensorrt_llm/layers/cast.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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 ..functional import cast -from ..module import Module - - -class Cast(Module): - - def __init__(self, output_dtype: str = 'float32') -> None: - super().__init__() - assert output_dtype in ('float32', 'float16', 'bfloat16', 'bool', - 'int32', 'int8'), TypeError( - "%s is not supported" % output_dtype) - self.output_dtype = output_dtype - - def forward(self, x): - return cast(x, self.output_dtype) diff --git a/tensorrt_llm/layers/conv.py b/tensorrt_llm/layers/conv.py deleted file mode 100644 index 49234c70962e..000000000000 --- a/tensorrt_llm/layers/conv.py +++ /dev/null @@ -1,260 +0,0 @@ -# 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 Tuple - -from ..functional import conv1d, conv2d, conv3d, conv_transpose2d -from ..module import Module -from ..parameter import Parameter - - -class Conv2d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: Tuple[int, int], - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(out_channels, in_channels // groups, - *kernel_size), - dtype=dtype) - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def forward(self, input): - return conv2d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, self.dilation, self.groups) - - -class ConvTranspose2d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: Tuple[int, int], - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - output_padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.output_padding = output_padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(in_channels, out_channels // groups, - *kernel_size), - dtype=dtype) - - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def _output_padding(self, - input, - output_size, - stride, - padding, - kernel_size, - num_spatial_dims: int, - dilation=None): - if output_size is None: - ret = self.output_padding - else: - has_batch_dim = input.dim() == num_spatial_dims + 2 - num_non_spatial_dims = 2 if has_batch_dim else 1 - if len(output_size) == num_non_spatial_dims + num_spatial_dims: - output_size = output_size[num_non_spatial_dims:] - if len(output_size) != num_spatial_dims: - raise ValueError( - "ConvTranspose{}D: for {}D input, output_size must have {} or {} elements (got {})" - .format(num_spatial_dims, input.dim(), num_spatial_dims, - num_non_spatial_dims + num_spatial_dims, - len(output_size))) - - min_sizes = [] - max_sizes = [] - for d in range(num_spatial_dims): - dim_size = ( - (input.size(d + num_non_spatial_dims) - 1) * stride[d] - - 2 * padding[d] + - (dilation[d] if dilation is not None else 1) * - (kernel_size[d] - 1) + 1) - min_sizes.append(dim_size) - max_sizes.append(min_sizes[d] + stride[d] - 1) - - for i in range(len(output_size)): - size = output_size[i] - min_size = min_sizes[i] - max_size = max_sizes[i] - if size < min_size or size > max_size: - raise ValueError(( - "requested an output size of {}, but valid sizes range " - "from {} to {} (for an input of {})").format( - output_size, min_sizes, max_sizes, - input.size()[2:])) - - res = [] - for d in range(num_spatial_dims): - res.append(output_size[d] - min_sizes[d]) - - ret = res - return ret - - def forward(self, input, output_size=None): - num_spatial_dims = 2 - output_padding = self._output_padding(input, output_size, self.stride, - self.padding, self.kernel_size, - num_spatial_dims, self.dilation) - - return conv_transpose2d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, output_padding, - self.dilation, self.groups) - - -class Conv1d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - padding: int = 0, - dilation: int = 1, - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(out_channels, in_channels // groups, - kernel_size, 1), - dtype=dtype) - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def forward(self, input): - return conv1d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, self.dilation, self.groups) - - -class Conv3d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: Tuple[int, int, int], - stride: Tuple[int, int, int] = (1, 1, 1), - padding: Tuple[int, int, int] = (0, 0, 0), - dilation: Tuple[int, int, int] = (1, 1, 1), - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(out_channels, in_channels // groups, - kernel_size[0], kernel_size[1], - kernel_size[2]), - dtype=dtype) - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def forward(self, input): - return conv3d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, self.dilation, self.groups) diff --git a/tensorrt_llm/layers/embedding.py b/tensorrt_llm/layers/embedding.py deleted file mode 100644 index fd77aa9e643a..000000000000 --- a/tensorrt_llm/layers/embedding.py +++ /dev/null @@ -1,673 +0,0 @@ -# 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 math -from typing import Optional, Sequence, Union - -import numpy as np -import torch - -from .._utils import set_obj_attrs, str_dtype_to_torch, trt_dtype_to_np -from ..functional import (ACT2FN, Tensor, arange, concat, constant, cos, div, - embedding, exp, expand, identity, meshgrid2d, outer, - pad, shape, sin, slice, unsqueeze, where) -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from .linear import ColumnLinear, Linear, RowLinear - - -class Embedding(Module): - """ - The embedding layer takes input indices (x) and the embedding lookup table (weight) as input. - And output the corresponding embeddings according to input indices. - The size of weight is [num_embeddings, embedding_dim] - - Four parameters (tp_size, tp_group, sharding_dim, tp_rank) are involved in tensor parallelism. - Only when "tp_size > 1 and tp_group is not None", tensor parallelism is enabled. - When "sharding_dim == 0", the weight is shared in the vocabulary dimension. - tp_rank must be set when sharding_dim == 0. - When "sharding_dim == 1", the weight is shard in the hidden dimension. - """ - - def __init__(self, - num_embeddings: int, - embedding_dim: int, - dtype: Optional[str] = None, - tp_size: int = 1, - tp_group: Optional[list] = None, - sharding_dim: int = 0, - tp_rank: Optional[int] = None): - super().__init__() - # num_embeddings records the total vocab size no matter using TP or not - self.num_embeddings = num_embeddings - self.embedding_dim = embedding_dim - self.tp_size = tp_size - self.tp_group = tp_group - self.sharding_dim = sharding_dim - self.tp_rank = tp_rank - self.dtype = dtype - self.tp_dim = sharding_dim - - if sharding_dim == 1: - shape = (self.num_embeddings, self.embedding_dim // self.tp_size) - elif sharding_dim == 0: - shape = (math.ceil(self.num_embeddings / self.tp_size), - self.embedding_dim) - - self.weight = Parameter(shape=shape, dtype=dtype) - - self.weight_padding_size = ((8 - shape[0] % 8) % 8, shape[1]) - - set_obj_attrs(self.weight, { - "weight_loader": self.weight_loader, - }) - - def forward(self, x): - # The embedding weight is padded to the multiple of 8. - # The reason is that when lm_head and vocab_embedding are using the same embedding weight, - # previously weights can't be depulicated in the engine because gemm will pad the weight to the multiple of 8. - # If we also pad the embedding weight to the multiple of 8, the weights can be successfully deduplicated. - # This will not affect the input and output of the gather op and perf impact is negligible. - if self.weight_padding_size[0] != 0: - padding_values = np.zeros(self.weight_padding_size, - dtype=trt_dtype_to_np( - self.weight.value.dtype)) - padding = constant(padding_values) - else: - padding = None - - return embedding(x, - self.weight.value, - tp_size=self.tp_size, - tp_group=self.tp_group, - sharding_dim=self.sharding_dim, - tp_rank=self.tp_rank, - padding=padding) - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor): - # use_parallel_embedding - tp_rank = mapping.tp_rank - if self.tp_size > 1: - sharding_dim = self.sharding_dim - shard_size = param._shape[sharding_dim] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(sharding_dim, start_idx, - shard_size) - param.value = loaded_weight - - def postprocess(self, tllm_key, weights, **kwargs): - if weights is None: - return {} - weights = weights.to(str_dtype_to_torch(self.dtype)) - return {tllm_key: weights} - - -class PromptTuningEmbedding(Embedding): - """ - PromptTuningEmbedding handles fine-tuned prompts with virtual tokens. At runtime, - a supplementary embedding dictionary is passed. Tokens whose ids are >= vocab_size are embedded - with that additional dictionary. - The prompt tuning dictionary holds multiple tasks, and each sequence is assigned a given task. - Prompt-tuned tokens from a given sequence use the adequate task dictionary, as defined by the `tasks` input. - """ - - def __init__(self, - num_embeddings, - embedding_dim, - vocab_size=None, - dtype=None, - tp_size=1, - tp_group=None, - sharding_dim=0, - tp_rank=0): - super().__init__(num_embeddings, embedding_dim, dtype, tp_size, - tp_group, sharding_dim, tp_rank) - if vocab_size is None: - vocab_size = num_embeddings - self.vocab_size = vocab_size - - def forward(self, tokens, prompt_embedding_table, tasks, task_vocab_size): - """ - Pass all tokens through both normal and prompt embedding tables. - Tokens are masked so that "normal" embedding only see "normal" tokens. Same logic for "prompt" embedding. - After those two embedding, combine results based on whether the token was "normal" or "prompt-tuned". - - Parameters: - tokens : Tensor - the ids to embed, size [batch_size, seq_len] - - prompt_embedding_table : Tensor - the additional embedding table for prompt-tuned tokens, size [num_tasks * num_tokens_per_task, hidden_size] - - tasks: Tensor - the task required by each token, size [batch_size, seq_len] - - task_vocab_size: Tensor - the number of tokens used for each task, should be equal to prompt_embedding_table's num_tokens_per_task, size [1] - - Returns: - Tokens' embedding - """ - # do not use ">=" because internally the layer works with floating points - prompt_tokens_mask = tokens > (self.vocab_size - 1) - - # clip tokens in the [0, vocab_size) range - normal_tokens = where(prompt_tokens_mask, self.vocab_size - 1, tokens) - normal_embeddings = embedding(normal_tokens, self.weight.value, - self.tp_size, self.tp_group, - self.sharding_dim, self.tp_rank) - - # put virtual tokens in the [0, max_prompt_vocab_size) range - prompt_tokens = where(prompt_tokens_mask, tokens - self.vocab_size, 0) - # add offsets to match the concatenated embedding tables - tasks = tasks * task_vocab_size - - # tasks: [batch_size, seq_len] - # prompt_tokens: [batch_size, seq_len] - # if speculative decoding is enabled the shape of prompt_tokens is [batch_size, seq_len + max_draft_len], - # so we need to expand tasks to [batch_size, seq_len + max_draft_len] - tasks = expand(tasks, shape(prompt_tokens)) - prompt_tokens = prompt_tokens + tasks - prompt_embeddings = embedding(prompt_tokens, prompt_embedding_table) - - # prompt_tokens_mask: [batch_size, seq_len] -> [batch_size, seq_len, 1] - # combine the correct sources of embedding: normal/prompt - return where(unsqueeze(prompt_tokens_mask, -1), prompt_embeddings, - normal_embeddings) - - -class LabelEmbedding(Module): - - def __init__(self, - num_classes: int, - hidden_size: int, - dropout_prob: float = 0.0, - mapping=Mapping(), - dtype=None): - super().__init__() - use_cfg_embedding = dropout_prob > 0 - self.embedding_table = Embedding(num_classes + use_cfg_embedding, - hidden_size, - tp_size=mapping.tp_size, - tp_group=mapping.tp_group, - dtype=dtype) - self.num_classes = num_classes - - def token_drop(self, labels: Tensor, force_drop_ids: Tensor): - labels = where(force_drop_ids == 1, self.num_classes, labels) - return labels - - def forward(self, labels: Tensor, force_drop_ids: Optional[Tensor] = None): - if force_drop_ids is not None: - labels = self.token_drop(labels, force_drop_ids) - embeddings = self.embedding_table(labels) - return embeddings - - -def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: Tensor): - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - omega = torch.arange(embed_dim // 2, dtype=torch.float64) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - omega = constant(omega.numpy().astype(np.float32)) - - pos = pos.view([-1]) # (M,) - out = outer(pos, omega) # (M, D/2), outer product - - emb_sin = sin(out) # (M, D/2) - emb_cos = cos(out) # (M, D/2) - - emb = concat([emb_sin, emb_cos], dim=1) # (M, D) - return emb - - -def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: Sequence[Tensor]): - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, - grid[0]) # (H*W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, - grid[1]) # (H*W, D/2) - - emb = concat([emb_h, emb_w], dim=1) # (H*W, D) - return emb - - -def get_2d_sincos_pos_embed( - embed_dim: int, - grid_size: Union[int, Sequence[int]], - cls_token: bool = False, - extra_tokens: int = 0, - interpolation_scale: float = 1.0, - base_size: int = 16, -): - if isinstance(grid_size, int): - grid_size = (grid_size, grid_size) - - grid_h = div( - div(arange(0, grid_size[0], 'float32'), - float(grid_size[0] / base_size)), interpolation_scale) - grid_w = div( - div(arange(0, grid_size[1], 'float32'), - float(grid_size[1] / base_size)), interpolation_scale) - grid_h, grid_w = meshgrid2d(grid_w, grid_h) # here w goes first - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, [ - grid_h.view([1, grid_size[1], grid_size[0]]), - grid_w.view([1, grid_size[1], grid_size[0]]) - ]) - if cls_token and extra_tokens > 0: - pos_embed = concat([ - constant( - np.zeros(shape=(extra_tokens, embed_dim), - dtype=trt_dtype_to_np(pos_embed.dtype))), pos_embed - ], - dim=0) - return pos_embed - - -class SD3PatchEmbed(Module): - """ - 2D Image to Patch Embedding with support for SD3 cropping. - """ - - def __init__( - self, - height: int = 224, - width: int = 224, - patch_size: int = 16, - in_channels: int = 3, - embed_dim: int = 768, - layer_norm: bool = False, - flatten: bool = True, - bias: bool = True, - interpolation_scale: int = 1, - pos_embed_type: str = "sincos", - pos_embed_max_size: Optional[int] = None, # For SD3 cropping - dtype=None): - from diffusers.models.embeddings import \ - get_2d_sincos_pos_embed as get_2d_sincos_pos_embed_torch - - from .conv import Conv2d - from .normalization import LayerNorm - - super().__init__() - - num_patches = (height // patch_size) * (width // patch_size) - self.flatten = flatten - self.layer_norm = layer_norm - self.pos_embed_max_size = pos_embed_max_size - - self.proj = Conv2d(in_channels, - embed_dim, - kernel_size=(patch_size, patch_size), - stride=(patch_size, patch_size), - bias=bias, - dtype=dtype) - if layer_norm: - self.norm = LayerNorm(embed_dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - else: - self.norm = None - - self.patch_size = patch_size - self.height, self.width = height // patch_size, width // patch_size - self.base_size = height // patch_size - self.interpolation_scale = interpolation_scale - - # Calculate positional embeddings based on max size or default - if pos_embed_max_size: - grid_size = pos_embed_max_size - else: - grid_size = int(num_patches**0.5) - - if pos_embed_type is None: - self.pos_embed = None - elif pos_embed_type == "sincos": - pos_embed = get_2d_sincos_pos_embed_torch( - embed_dim, - grid_size, - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - output_type="pt", - ) - self.pos_embed = Parameter( - pos_embed.detach().cpu().float().unsqueeze(0), dtype=dtype) - else: - raise ValueError( - f"Unsupported pos_embed_type: {self.pos_embed_type}") - - def cropped_pos_embed(self, height, width): - """Crops positional embeddings for SD3 compatibility.""" - if self.pos_embed_max_size is None: - raise ValueError("`pos_embed_max_size` must be set for cropping.") - - height = height // self.patch_size - width = width // self.patch_size - if height > self.pos_embed_max_size: - raise ValueError( - f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - if width > self.pos_embed_max_size: - raise ValueError( - f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - - top = (self.pos_embed_max_size - height) // 2 - left = (self.pos_embed_max_size - width) // 2 - spatial_pos_embed = identity(self.pos_embed.value).view( - [1, self.pos_embed_max_size, self.pos_embed_max_size, -1]) - spatial_pos_embed = slice(spatial_pos_embed, - starts=[0, top, left, 0], - sizes=concat([ - shape(spatial_pos_embed, 0), height, - width, - shape(spatial_pos_embed, 3) - ])) - spatial_pos_embed = spatial_pos_embed.view( - concat( - [1, -1, - shape(spatial_pos_embed, - spatial_pos_embed.ndim() - 1)])) - return spatial_pos_embed - - def forward(self, latent): - # [TODO] to support height and width for runtime - if self.pos_embed_max_size is not None: - height, width = latent.shape[-2:] - else: - height, width = latent.shape[-2] // self.patch_size, latent.shape[ - -1] // self.patch_size - latent = self.proj(latent) - if self.flatten: - latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC - if self.layer_norm: - latent = self.norm(latent) - if self.pos_embed is None: - return latent.cast(latent.dtype) - # Interpolate or crop positional embeddings as needed - if self.pos_embed_max_size: - pos_embed = self.cropped_pos_embed(height, width) - else: - if self.height != height or self.width != width: - pos_embed = get_2d_sincos_pos_embed( - embed_dim=self.pos_embed.value.shape[-1], - grid_size=(height, width), - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - ) - pos_embed = unsqueeze(pos_embed.cast('float32'), axis=0) - else: - pos_embed = self.pos_embed.value - - pos_embed = pos_embed.cast(latent.dtype) - output = (latent + pos_embed).cast(latent.dtype) - return output - - -def get_timestep_embedding( - timesteps: Tensor, - embedding_dim: int, - flip_sin_to_cos: bool = False, - downscale_freq_shift: float = 1, - scale: float = 1, - max_period: int = 10000, -) -> Tensor: - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - - Args - timesteps (Tensor): - a 1-D Tensor of N indices, one per batch element. These may be fractional. - embedding_dim (int): - the dimension of the output. - flip_sin_to_cos (bool): - Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) - downscale_freq_shift (float): - Controls the delta between frequencies between dimensions - scale (float): - Scaling factor applied to the embeddings. - max_period (int): - Controls the maximum frequency of the embeddings - Returns - Tensor: an [N x dim] Tensor of positional embeddings. - """ - assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - exponent = -math.log(max_period) * np.arange( - start=0, stop=half_dim, dtype=np.float32) - exponent = exponent / (half_dim - downscale_freq_shift) - exponent = constant(exponent) - - emb = exp(exponent) - emb = unsqueeze(timesteps, -1).cast('float32') * unsqueeze(emb, 0) - - # scale embeddings - emb = scale * emb - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = concat([cos(emb), sin(emb)], dim=-1) - else: - emb = concat([sin(emb), cos(emb)], dim=-1) - - # zero pad - if embedding_dim % 2 == 1: - emb = pad(emb, (0, 1, 0, 0)) - return emb - - -class TimestepEmbedding(Module): - - def __init__(self, - in_channels: int, - time_embed_dim: int, - act_fn: str = "silu", - out_dim: int = None, - post_act_fn: Optional[str] = None, - cond_proj_dim=None, - sample_proj_bias=True, - mapping=None, - dtype=None): - super().__init__() - tp_group = mapping.tp_group - tp_size = mapping.tp_size - self.linear_1 = ColumnLinear(in_channels, - time_embed_dim, - sample_proj_bias, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype, - gather_output=False) - - if cond_proj_dim is not None: - self.cond_proj = Linear(cond_proj_dim, - in_channels, - bias=False, - dtype=dtype) - else: - self.cond_proj = None - - self.act = ACT2FN[act_fn] - - if out_dim is not None: - time_embed_dim_out = out_dim - else: - time_embed_dim_out = time_embed_dim - self.linear_2 = RowLinear(time_embed_dim, - time_embed_dim_out, - sample_proj_bias, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - - if post_act_fn is None: - self.post_act = None - else: - self.post_act = ACT2FN[post_act_fn] - - def forward(self, sample, condition=None): - if condition is not None: - sample = sample + self.cond_proj(condition) - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - - if self.post_act is not None: - sample = self.post_act(sample) - return sample - - -class Timesteps(Module): - - def __init__(self, - num_channels: int, - flip_sin_to_cos: bool, - downscale_freq_shift: float, - scale: int = 1): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - self.scale = scale - - def forward(self, timesteps) -> Tensor: - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - scale=self.scale, - ) - return t_emb - - -class PixArtAlphaTextProjection(Module): - """ - Projects caption embeddings. Also handles dropout for classifier-free guidance. - - Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py - """ - - def __init__(self, - in_features, - hidden_size, - out_features=None, - act_fn="gelu_tanh", - mapping=None, - dtype=None): - super().__init__() - if out_features is None: - out_features = hidden_size - tp_group = mapping.tp_group - tp_size = mapping.tp_size - self.linear_1 = ColumnLinear(in_features=in_features, - out_features=hidden_size, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype, - gather_output=False) - self.act_1 = ACT2FN[act_fn] - self.linear_2 = RowLinear(in_features=hidden_size, - out_features=out_features, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - - def forward(self, caption): - hidden_states = self.linear_1(caption) - hidden_states = self.act_1(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states - - -class CombinedTimestepTextProjEmbeddings(Module): - - def __init__(self, - embedding_dim, - pooled_projection_dim, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.time_proj = Timesteps(num_channels=256, - flip_sin_to_cos=True, - downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, - time_embed_dim=embedding_dim, - mapping=mapping, - dtype=dtype) - self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, - embedding_dim, - act_fn="silu", - mapping=mapping, - dtype=dtype) - - def forward(self, timestep: Tensor, pooled_projection: Tensor): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder( - timesteps_proj.cast(dtype=pooled_projection.dtype)) # (N, D) - - pooled_projections = self.text_embedder(pooled_projection) - - conditioning = timesteps_emb + pooled_projections - self.register_network_output('output', conditioning) - return conditioning - - -class CombinedTimestepLabelEmbeddings(Module): - - def __init__(self, - num_classes, - embedding_dim, - class_dropout_prob=0.0, - mapping=Mapping(), - dtype=None): - super().__init__() - self.time_proj = Timesteps(num_channels=256, - flip_sin_to_cos=True, - downscale_freq_shift=1) - self.timestep_embedder = TimestepEmbedding(in_channels=256, - time_embed_dim=embedding_dim, - mapping=mapping, - dtype=dtype) - self.class_embedder = LabelEmbedding(num_classes, - embedding_dim, - class_dropout_prob, - mapping=mapping, - dtype=dtype) - - def forward(self, - timestep: Tensor, - class_labels: Tensor, - hidden_dtype: Optional[str] = 'float32'): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder( - timesteps_proj.cast(dtype=hidden_dtype)) # (N, D) - class_labels = self.class_embedder(class_labels) # (N, D) - conditioning = timesteps_emb + class_labels # (N, D) - return conditioning diff --git a/tensorrt_llm/layers/language_adapter.py b/tensorrt_llm/layers/language_adapter.py deleted file mode 100755 index b9fa790b7d58..000000000000 --- a/tensorrt_llm/layers/language_adapter.py +++ /dev/null @@ -1,94 +0,0 @@ -# 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 dataclasses import asdict, dataclass, field - -from ..mapping import Mapping -from ..module import Module -from ..quantization import QuantMode -from .moe import MOE, MoeConfig - - -@dataclass -class LanguageAdapterConfig: - num_languages: int | None = None - top_k: int = 1 - normalization_mode: MoeConfig.ExpertScaleNormalizationMode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED - ffn_hidden_size: int = 1024 - language_list: list = field(default_factory=list) - - def validate(self) -> "LanguageAdapterConfig": - if (self.num_languages is None or self.num_languages == 0 - or self.top_k == 0): - raise ValueError( - "LanguageAdapterConfig's num_languages and top_k must not be set to 0" - ) - if (self.normalization_mode - != MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED): - raise ValueError( - "LanguageAdapterConfig's normalization_mode must be set to DEVICE_LIMITED since it skips both softmax and renormalization" - ) - if (len(self.language_list) == 0): - raise ValueError( - "LanguageAdapterConfig's language_list must not be empty") - return self - - @classmethod - def from_dict(cls, config: dict): - return cls(**config) - - def to_dict(self): - return asdict(self) - - def to_MOE_config(self): - return MoeConfig( - num_experts=self.num_languages, - top_k=self.top_k, - normalization_mode=self.normalization_mode, - ) - - -class LanguageAdapter(Module): - """ - Language Adapter module that uses MOE plugin with static expert selection passed in as a parameter in request. - A language MLP is selected by user for each request. - see https://arxiv.org/pdf/2005.00052 for more details. - """ - - def __init__( - self, - language_adapter_config: LanguageAdapterConfig, - hidden_size: int, - hidden_act: str, - mapping: Mapping = Mapping(), - has_mlp_bias: bool = True, - dtype=None, - quant_mode=QuantMode(0), - ): - super().__init__() - self.config = language_adapter_config - self.config.validate() - - self.layers = MOE( - hidden_size=hidden_size, - ffn_hidden_size=language_adapter_config.ffn_hidden_size, - hidden_act=hidden_act, - dtype=dtype, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - quant_mode=quant_mode, - static_routing=True, - moe_config=self.config.to_MOE_config(), - mapping=mapping) diff --git a/tensorrt_llm/layers/linear.py b/tensorrt_llm/layers/linear.py deleted file mode 100644 index b5539ded43cf..000000000000 --- a/tensorrt_llm/layers/linear.py +++ /dev/null @@ -1,551 +0,0 @@ -# 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 abc import ABCMeta, abstractmethod -from typing import Optional - -import numpy as np -import tensorrt as trt -import torch - -from .._common import default_net, default_trtnet -from .._utils import set_obj_attrs, str_dtype_to_torch, str_dtype_to_trt -from ..functional import (AllReduceFusionOp, AllReduceParams, Tensor, - _add_plugin_info, _create_tensor, allgather, - allreduce, cast, gemm_allreduce, low_latency_gemm, - matmul) -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from ..plugin import TRT_LLM_PLUGIN_NAMESPACE -from .lora import LoraRuntimeParams - - -def _gemm_plugin(input: Tensor, - mat2: Tensor, - transa: bool = False, - transb: bool = False, - pad_lda: int = 0, - pad_ldb: int = 0, - pad_ldc: int = 0, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - strict_dtype: Optional[trt.DataType] = None) -> Tensor: - ''' - output = op(mat2)op(input) - - Parameters: - input : Tensor (On GPU) - The input tensor. - - mat2 : Tensor (On GPU) - The mat2 tensor. - - transa : bool - Is the input tensor transposed? Set to 'True' if you want the - input tensor to be transposed, 'False' otherwise. - - transb : bool - Is the mat2 tensor transposed? Set to 'True' if you want the - mat2 tensor to be transposed, 'False' otherwise. - - pad_lda: int - Padding to the lead dimension of input tensor. It is used to - support the strided GEMM that only uses the sub-tensor for - computation. The GEMM plugin computation is - [N, K] x [K, M+pad_lda] -> [N, M] if transa, - [N, K] x [K+pad_lda, M] -> [N, M] if not transa. - - pad_ldb: int - Padding to the lead dimension of mat2 tensor. It is used to - support the strided GEMM that only uses the sub-tensor for - computation. The GEMM plugin computation is - [N, K+pad_ldb] x [K, M] -> [N, M] if transb, - [N+pad_ldb, K] x [K, M] -> [N, M] if not transb. - - pad_ldc: int - Padding to the lead dimension of output tensor. It is used to - support the strided GEMM that only uses the sub-tensor for - computation. The GEMM plugin computation is - [N, K] x [K, M] -> [N+pad_ldc, M]. - - use_fp8: bool - Do we use fp8 GEMM. - - alpha: float - Alpha for fp8 GEMM. - - strict_dtype: trt.DataType - Set the data type for the GEMM plugin. If it is None, the data - type is the gemm_plugin type set in the plugin_config. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "Gemm", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - if use_fp8: - assert ( - isinstance(alpha, np.ndarray) and alpha.dtype == np.float32 - and alpha.size == 1 - ), "`alpha` must be passed as a float32 ndarray if `use_fp8` is enabled for _gemm_plugin" - assert input.dtype == trt.fp8 - assert mat2.dtype == trt.fp8 - - transa = 1 if transa else 0 - transa = trt.PluginField("transa", np.array(transa, dtype=np.int32), - trt.PluginFieldType.INT32) - transb = 1 if transb else 0 - transb = trt.PluginField("transb", np.array(transb, dtype=np.int32), - trt.PluginFieldType.INT32) - pad_lda = trt.PluginField("pad_lda", np.array(pad_lda, dtype=np.int32), - trt.PluginFieldType.INT32) - pad_ldb = trt.PluginField("pad_ldb", np.array(pad_ldb, dtype=np.int32), - trt.PluginFieldType.INT32) - pad_ldc = trt.PluginField("pad_ldc", np.array(pad_ldc, dtype=np.int32), - trt.PluginFieldType.INT32) - use_fp8 = 1 if use_fp8 else 0 - use_fp8 = trt.PluginField("use_fp8", np.array(use_fp8, dtype=np.int32), - trt.PluginFieldType.INT32) - alpha = alpha if alpha else np.array(1.0, dtype=np.float32) - alpha = trt.PluginField("alpha", alpha.flatten(), - trt.PluginFieldType.FLOAT32) - - if strict_dtype is not None: - assert isinstance(strict_dtype, trt.DataType) - p_dtype = strict_dtype - else: - p_dtype = str_dtype_to_trt(default_net().plugin_config.gemm_plugin) - assert p_dtype != trt.fp8, "need to use strict dtype in gemm plugin fp8" - pf_type = trt.PluginField("type_id", np.array([int(p_dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection( - [transa, transb, pad_lda, pad_ldb, pad_ldc, pf_type, use_fp8, alpha]) - gemm_plug = plg_creator.create_plugin("gemm", pfc) - plug_inputs = [input.trt_tensor, mat2.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "gemm", pfc) - return _create_tensor(layer.get_output(0), layer) - - -class LinearBase(Module, metaclass=ABCMeta): - - def __init__( - self, - local_in_features, - local_out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - share_weight=None, - strict_dtype=False, - pad_lda=0, - pad_ldc=0, - prefer_managed_weight=True, - ): - super().__init__() - self.in_features = local_in_features - self.out_features = local_out_features - self.dtype = dtype - self.pad_lda = pad_lda - self.pad_ldc = pad_ldc - self.prefer_managed_weight = prefer_managed_weight - - self.share_weight = share_weight - if not share_weight: - self.weight = Parameter( - shape=(self.out_features, self.in_features), - dtype=dtype, - prefer_managed=self.prefer_managed_weight, - ) - set_obj_attrs( - self.weight, - { - "weight_loader": self.weight_loader, - }, - ) - else: - self.weight = share_weight - - self.tp_size = tp_size - self.tp_group = tp_group - self.strict_dtype = str_dtype_to_trt( - self.dtype) if strict_dtype else None - - if bias: - self.bias = Parameter(shape=(self.out_features, ), dtype=dtype) - assert pad_ldc == 0, "not support pad_ldc with bias" - else: - self.register_parameter("bias", None) - - # see optimize_model's add_lora for LoRA initialization - self.lora = None - self.dora = None - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor) -> None: - tp_rank = mapping.tp_rank - shard_size = param._shape[self.tp_split_dim()] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(self.tp_split_dim(), start_idx, - shard_size) - - param.value = loaded_weight - - @classmethod - @abstractmethod - def tp_split_dim(cls) -> int: - pass - - def get_weight(self) -> Tensor: - if default_net( - ).plugin_config.manage_weights and self.prefer_managed_weight: - gemm_plugin = default_net().plugin_config.gemm_plugin - # nvfp4 plugin does not use this code path - use_gemm_plugin = gemm_plugin is not None and gemm_plugin != 'nvfp4' - use_low_latency_gemm_plugin = default_net( - ).plugin_config.low_latency_gemm_plugin == 'fp8' - use_gemm_allreduce_plugin = default_net( - ).plugin_config.gemm_allreduce_plugin is not None - return self.weight.get_managed_tensor(network=default_net()) - else: - return self.weight.get_constant_tensor(network=default_net()) - - def multiply_and_lora( - self, - x, - weight, - gemm_plugin: Optional[str] = None, - low_latency_gemm_plugin: Optional[str] = None, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - ): - hidden_state = x - if low_latency_gemm_plugin: - strict_dtype = str_dtype_to_trt(self.dtype) if isinstance( - self.dtype, str) else self.dtype - x = low_latency_gemm(x, weight, alpha, strict_dtype) - elif gemm_plugin and gemm_plugin != 'nvfp4': # nvfp4 gemm plugin has its own implementation - if gemm_plugin == 'fp8': - strict_dtype = str_dtype_to_trt(self.dtype) if isinstance( - self.dtype, str) else self.dtype - else: - strict_dtype = self.strict_dtype - x = _gemm_plugin(x, - weight, - transb=True, - pad_lda=self.pad_lda, - pad_ldc=self.pad_ldc, - use_fp8=use_fp8, - alpha=alpha, - strict_dtype=strict_dtype) - else: - x = matmul(x, weight, transb=True) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora( - hidden_state - if lora_hidden_state is None else lora_hidden_state, - lora_runtime_params=lora_runtime_params, - ) - if self.dora is not None: - x = self.dora(x, lora_runtime_params=lora_runtime_params) - - return x - - @abstractmethod - def collect_and_bias(self, x: Tensor) -> Tensor: - pass - - def multiply_collect( - self, - x, - weight, - gemm_plugin: Optional[str] = None, - low_latency_gemm_plugin: Optional[str] = None, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - **kwargs): - x = self.multiply_and_lora( - x, - weight, - gemm_plugin=gemm_plugin, - low_latency_gemm_plugin=low_latency_gemm_plugin, - use_fp8=use_fp8, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - ) - return self.collect_and_bias(x, **kwargs) - - def forward(self, - x, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - **kwargs) -> Tensor: - return self.multiply_collect( - x, - self.get_weight(), - gemm_plugin=default_net().plugin_config.gemm_plugin, - use_fp8=False, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - **kwargs) - - -class Linear(LinearBase): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - share_weight=None, - strict_dtype=False, - pad_lda=0, - pad_ldc=0, - prefer_managed_weight=True, - is_qkv=False, - ): - super().__init__( - local_in_features=in_features, - local_out_features=out_features // tp_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - share_weight=share_weight, - strict_dtype=strict_dtype, - pad_lda=pad_lda, - pad_ldc=pad_ldc, - prefer_managed_weight=prefer_managed_weight, - ) - self.gather_output = gather_output - self.is_qkv = is_qkv - self.tp_dim = 0 - if bias: - set_obj_attrs( - self.bias, - { - "weight_loader": self.weight_loader, - }, - ) - - @classmethod - def tp_split_dim(cls) -> int: - return 0 - - def collect_and_bias(self, x, **kwargs): - - if self.bias is not None: - bias = cast(self.bias.value, x.dtype) - x = x + bias - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=-1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - using_head_as_leading_dim = kwargs.get("using_head_as_leading_dim", - False) - config = kwargs.get("config", None) - if self.is_qkv: - if isinstance(weights, list): - head_size = config.hidden_size // config.num_attention_heads if config.head_size is None else config.head_size - if getattr(config, "remove_duplicated_kv_heads", False): - if config.remove_duplicated_kv_heads: - k, v = weights[1:] - k = k.reshape([ - k.shape[0] // head_size // 2, 2, head_size, - self.in_features - ]) - v = v.reshape([ - v.shape[0] // head_size // 2, 2, head_size, - self.in_features - ]) - assert (k[:, 0] == k[:, 1]).all() - assert (v[:, 0] == v[:, 1]).all() - k = k[:, 0].reshape([-1, self.in_features]) - v = v[:, 0].reshape([-1, self.in_features]) - weights[1] = k - weights[2] = v - # Duplicate kv heads in case of invalid TP size - tp_size = config.mapping.tp_size - num_kv_heads = config.num_key_value_heads - if num_kv_heads < tp_size: - for qkv_idx in range(3): - v = weights[qkv_idx] - if qkv_idx > 0: - assert tp_size % num_kv_heads == 0 - reps = tp_size // num_kv_heads - if tllm_key.endswith("bias"): - v = v.reshape(num_kv_heads, - head_size)[:, None, :].expand( - num_kv_heads, reps, head_size) - v = v.reshape(num_kv_heads * reps * head_size) - else: - v = v.reshape(num_kv_heads, head_size, - -1)[:, None, :, :].expand( - num_kv_heads, reps, head_size, - v.shape[1]) - v = v.reshape(num_kv_heads * reps * head_size, - -1) - weights[qkv_idx] = v.chunk( - tp_size, self.tp_dim)[config.mapping.tp_rank] - - weights = torch.cat(weights) - if using_head_as_leading_dim: - # Reorder [n_head, 3, head_dim, ...] into [3, n_head, head_dim, ...] - assert config.num_attention_heads == config.num_key_value_heads, "using_head_as_leading_dim require head_size to be multiple of 3." - num_heads = config.num_attention_heads - head_dim = self.out_features // (3 * num_heads) - w = weights.reshape(num_heads, 3, head_dim, -1) - w = w.transpose(0, 1) - if w.shape[-1] > 1: - weights = w.reshape(-1, self.in_features) # Weight - else: - weights = w.reshape(-1) # Bias - - if isinstance(weights, list): - weights = torch.cat(weights) - - weights = weights.to(str_dtype_to_torch(self.dtype)) - return {tllm_key: weights} - - -ColumnLinear = Linear - - -class RowLinear(LinearBase): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - strict_dtype: bool = False, - pad_lda=0, - prefer_managed_weight=True, - is_expert=False, - ): - super().__init__( - local_in_features=in_features // tp_size, - local_out_features=out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - strict_dtype=strict_dtype, - pad_lda=pad_lda, - prefer_managed_weight=prefer_managed_weight, - ) - - self.tp_dim = 1 - self.tp_size = tp_size - self.is_expert = is_expert - - @classmethod - def tp_split_dim(cls) -> int: - return 1 - - def multiply_collect( - self, - x, - weight, - gemm_plugin: Optional[str] = None, - low_latency_gemm_plugin: Optional[str] = None, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - **kwargs): - - gemm_allreduce_plugin = default_net( - ).plugin_config.gemm_allreduce_plugin - if gemm_allreduce_plugin: - if lora_runtime_params != None or lora_hidden_state != None: - raise RuntimeError( - "gemm_allreduce_plugin not supported with lora.") - - output_dtype = self.dtype - if isinstance(output_dtype, str): - output_dtype = str_dtype_to_trt(output_dtype) - - x = gemm_allreduce( - a=x, - b=weight, - transa=False, # row-major - transb=True, # col-major - alpha=alpha, - group=self.tp_group, # ranks participating - output_dtype=output_dtype, - fp8_inputs_override=use_fp8) - - if self.bias is not None: - bias = cast(self.bias.value, x.dtype) - if self.is_expert: - x = x + bias / self.tp_size - else: - x = x + bias - return x - else: - return super().multiply_collect(x, weight, gemm_plugin, - low_latency_gemm_plugin, use_fp8, - alpha, lora_runtime_params, - lora_hidden_state, **kwargs) - - def collect_and_bias(self, x, **kwargs): - all_reduce_params: Optional[AllReduceParams] = kwargs.get( - "all_reduce_params", None) - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = ( - need_bias and (all_reduce_params is not None) - and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM)) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - if not self.is_expert: - x = allreduce(x, - self.tp_group, - all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias - else: - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias / self.tp_size - return x - - if self.bias is not None: - bias = cast(self.bias.value, x.dtype) - x = x + bias - - return x diff --git a/tensorrt_llm/layers/lora.py b/tensorrt_llm/layers/lora.py deleted file mode 100644 index 26e93040dd5c..000000000000 --- a/tensorrt_llm/layers/lora.py +++ /dev/null @@ -1,181 +0,0 @@ -# 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 List - -import numpy as np - -from .._common import default_net -from ..functional import Tensor, constant, dora_plugin, lora_plugin, where -from ..module import Module - - -class LoraRuntimeParams(object): - - def __init__( - self, - lora_ranks: List[Tensor] = None, - lora_weights_pointers: List[Tensor] = None, - host_request_types: Tensor = None, - host_context_lengths: Tensor = None, - max_encoder_context_length: Tensor = None, - host_encoder_input_lengths: Tensor = None, - weight_index: int = 0, - partial_lora_mask: Tensor = None, - ): - - self.lora_ranks = lora_ranks - self.lora_weights_pointers = lora_weights_pointers - self.host_request_types = host_request_types - self.host_context_lengths = host_context_lengths - self.max_encoder_context_length = max_encoder_context_length - self.host_encoder_input_lengths = host_encoder_input_lengths - self.weight_index = weight_index - self.partial_lora_mask = partial_lora_mask # Partial LoRA for https://arxiv.org/abs/2401.16420 - - -class Lora(Module): - - def __init__(self, - in_hidden_size: int = 0, - out_hidden_sizes: List[int] = [0], - max_low_rank: int = 0) -> None: - super().__init__() - - self.in_hidden_size = in_hidden_size - self.out_hidden_sizes = out_hidden_sizes - self.max_low_rank = max_low_rank - - def forward(self, - x, - lora_runtime_params: LoraRuntimeParams = None, - is_cross_attention: bool = False): - if default_net().plugin_config.lora_plugin: - result = lora_plugin( - x, - in_hidden_size=self.in_hidden_size, - out_hidden_sizes=self.out_hidden_sizes, - host_request_types=lora_runtime_params.host_request_types, - transb=True, - # For cross attention, host_encoder_input_lengths should be used instead of host_context_lengths - host_context_lengths=lora_runtime_params.host_context_lengths - if not is_cross_attention else - lora_runtime_params.host_encoder_input_lengths, - max_low_rank=self.max_low_rank, - lora_ranks=lora_runtime_params.lora_ranks, - lora_weights_pointers=lora_runtime_params.lora_weights_pointers, - weight_index=lora_runtime_params.weight_index, - ) - if lora_runtime_params.partial_lora_mask is not None: - zero_tensor = constant(np.array([0.0], dtype=np.float16)) - if isinstance(result, List): - result = [ - where(lora_runtime_params.partial_lora_mask, r, - zero_tensor) for r in result - ] - elif isinstance(result, Tensor): - result = where(lora_runtime_params.partial_lora_mask, - result, zero_tensor) - else: - assert False - else: - assert False, "Not support lora without plugin" - - return result - - -class Dora(Module): - - def __init__(self, out_hidden_sizes: List[int] = [0]) -> None: - super().__init__() - self.out_hidden_sizes = out_hidden_sizes - - def forward(self, - x, - lora_runtime_params: LoraRuntimeParams = None, - is_cross_attention: bool = False): - assert lora_runtime_params.weight_index == 0, "DoRA does not support weight_index != 0" - if default_net().plugin_config.lora_plugin and default_net( - ).plugin_config.dora_plugin: - result = dora_plugin( - x, - out_hidden_sizes=self.out_hidden_sizes, - host_request_types=lora_runtime_params.host_request_types, - host_context_lengths=lora_runtime_params.host_context_lengths - if not is_cross_attention else - lora_runtime_params.host_encoder_input_lengths, - lora_weights_pointers=lora_runtime_params.lora_weights_pointers, - ) - else: - assert False, "Not support dora without plugin" - - return result - - -class LoraParams(object): - - def __init__( - self, - lora_ranks=None, # : List[dict[Tensor]] - lora_weights_pointers=None, # : List[dict[Tensor]] - host_context_lengths: Tensor = None, - max_encoder_context_length: Tensor = None, # For cross attention - host_request_types: Tensor = None, - host_encoder_input_lengths: Tensor = None, # For cross attention - weight_index: int = 0, - partial_lora_mask: Tensor = None, - ): - - self.lora_ranks = lora_ranks - self.lora_weights_pointers = lora_weights_pointers - - self.host_context_lengths = host_context_lengths - self.max_encoder_context_length = max_encoder_context_length - self.host_request_types = host_request_types - self.host_encoder_input_lengths = host_encoder_input_lengths - self.weight_index = weight_index - - self.partial_lora_mask = partial_lora_mask # Partial LoRA for https://arxiv.org/abs/2401.16420 - - def get_layer_params(self, layer_idx: int): - return LoraParams( - lora_ranks=[self.lora_ranks[layer_idx]], - lora_weights_pointers=[self.lora_weights_pointers[layer_idx]], - host_context_lengths=self.host_context_lengths, - max_encoder_context_length=self.max_encoder_context_length, - host_request_types=self.host_request_types, - host_encoder_input_lengths=self.host_encoder_input_lengths, - weight_index=self.weight_index, - partial_lora_mask=self.partial_lora_mask) - - def get_runtime_params(self, layer_idx: int, lora_module: str): - if f"{lora_module}_lora_ranks" in self.lora_ranks[layer_idx]: - return LoraRuntimeParams( - lora_ranks=[ - self.lora_ranks[layer_idx][f"{lora_module}_lora_ranks"] - ], - lora_weights_pointers=[ - self.lora_weights_pointers[layer_idx] - [f"{lora_module}_lora_weights_pointers"] - ], - host_context_lengths=self.host_context_lengths, - max_encoder_context_length=self.max_encoder_context_length, - host_request_types=self.host_request_types, - host_encoder_input_lengths=self.host_encoder_input_lengths, - weight_index=self.weight_index, - partial_lora_mask=self.partial_lora_mask, - ) - else: - return None diff --git a/tensorrt_llm/layers/mlp.py b/tensorrt_llm/layers/mlp.py deleted file mode 100644 index cbd2a764059a..000000000000 --- a/tensorrt_llm/layers/mlp.py +++ /dev/null @@ -1,565 +0,0 @@ -# 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 - -import tensorrt as trt - -from .._common import default_net -from ..functional import (ACT2FN, AllReduceParams, cast, chunk, concat, - gemm_swiglu, is_gated_activation, - low_latency_gemm_swiglu) -from ..mapping import Mapping -from ..module import Module -from ..quantization import QuantMode -from ..quantization.functional import quantize -from ..quantization.layers import FP8Linear, FP8RowLinear -from .linear import ColumnLinear, RowLinear -from .lora import LoraRuntimeParams -from .normalization import LayerNorm - - -def fc_gate_lora(hidden_states, lora, fused_gate_up_lora, lora_layer_params): - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - mlp_gate_up_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate_up") - - if mlp_gate_up_lora_params is not None: - assert fused_gate_up_lora is not None - mlp_gate_up_lora = fused_gate_up_lora(hidden_states, - mlp_gate_up_lora_params) - return mlp_gate_up_lora - - elif mlp_fc_lora_params is not None and mlp_gate_lora_params is not None: - mlp_in_lora_params = LoraRuntimeParams( - lora_ranks=[ - mlp_fc_lora_params.lora_ranks[0], - mlp_gate_lora_params.lora_ranks[0] - ], - lora_weights_pointers=[ - mlp_fc_lora_params.lora_weights_pointers[0], - mlp_gate_lora_params.lora_weights_pointers[0] - ], - host_request_types=mlp_fc_lora_params.host_request_types, - host_context_lengths=mlp_fc_lora_params.host_context_lengths) - - mlp_fc_lora, mlp_gate_lora = lora(hidden_states, mlp_in_lora_params) - mlp_in_result = concat([mlp_gate_lora, mlp_fc_lora], - dim=mlp_fc_lora.rank() - 1) - return mlp_in_result - return None - - -def fc_gate_dora(hidden_states, dora, fused_gate_up_dora, lora_layer_params): - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - mlp_gate_up_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate_up") - - if mlp_gate_up_lora_params is not None: - assert fused_gate_up_dora is not None - return fused_gate_up_dora(hidden_states, mlp_gate_up_lora_params) - - if mlp_fc_lora_params is not None and mlp_gate_lora_params is not None: - mlp_in_lora_params = LoraRuntimeParams( - lora_ranks=[ - mlp_fc_lora_params.lora_ranks[0], - mlp_gate_lora_params.lora_ranks[0] - ], - lora_weights_pointers=[ - mlp_fc_lora_params.lora_weights_pointers[0], - mlp_gate_lora_params.lora_weights_pointers[0] - ], - host_request_types=mlp_fc_lora_params.host_request_types, - host_context_lengths=mlp_fc_lora_params.host_context_lengths) - - return dora(hidden_states, mlp_in_lora_params) - return None - - -class MLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - is_expert=False, - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act in [ - 'swiglu', 'gegelu' - ] else ffn_hidden_size - self.inner_layernorm = LayerNorm(ffn_hidden_size, dtype=dtype, - eps=eps) if inner_layernorm else None - - self.fc = ColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.proj = RowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - is_expert=is_expert) - - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.dtype = dtype - self.bias = bias - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - self.eps = eps - self.is_expert = is_expert - # see optimize_model's add_lora for LoRA initialization - self.lora = None - self.dora = None - - def forward(self, hidden_states, lora_layer_params=None, gegelu_limit=None): - if lora_layer_params is not None: - assert lora_layer_params.get_runtime_params( - 0, "mlp_gate_up" - ) is None, f"LoRA module 'mlp_gate_up' is not supported in {self}" - if is_gated_activation(self.hidden_act): - inter = self.fc(hidden_states) - lora_result = fc_gate_lora(hidden_states, self.lora, None, - lora_layer_params) - if lora_result is not None: - inter = inter + lora_result - if self.dora is not None: - inter = fc_gate_dora(inter, self.dora, - self.fused_gate_up_dora, - lora_layer_params) - else: - mlp_fc_lora_params = None - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - inter = self.fc(hidden_states, mlp_fc_lora_params) - - mlp_proj_lora_params = None - if lora_layer_params is not None: - mlp_proj_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_4h_to_h") - - if self.hidden_act == 'gegelu': - inter = ACT2FN[self.hidden_act](inter, gegelu_limit) - else: - inter = ACT2FN[self.hidden_act](inter) - if self.inner_layernorm is not None: - inter = self.inner_layernorm(inter) - output = self.proj(inter, lora_runtime_params=mlp_proj_lora_params) - return output - - -class GatedMLP(MLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - is_expert=False, - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode, - inner_layernorm=inner_layernorm, - eps=eps, - is_expert=is_expert) - - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.tp_group = tp_group - self.tp_size = tp_size - - self.gate = ColumnLinear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None): - if lora_layer_params is not None: - assert lora_layer_params.get_runtime_params( - 0, "mlp_gate_up" - ) is None, f"LoRA module 'mlp_gate_up' is not supported in {self}" - - mlp_fc_lora_params = None - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - - mlp_gate_lora_params = None - if lora_layer_params is not None: - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - - mlp_proj_lora_params = None - if lora_layer_params is not None: - mlp_proj_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_4h_to_h") - - inter = self.fc(hidden_states, mlp_fc_lora_params) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states, mlp_gate_lora_params) - intermediate = inter * gate - if self.inner_layernorm is not None: - intermediate = self.inner_layernorm(intermediate) - output = self.proj(intermediate, - lora_runtime_params=mlp_proj_lora_params, - all_reduce_params=all_reduce_params) - return output - - -class FusedGatedMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - is_expert=False, - ): - super().__init__() - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.bias = bias - self.dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - - self.fused_fc = ColumnLinear( - self.hidden_size, - self.ffn_hidden_size * 2, - bias=self.bias, - dtype=self.dtype, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - ) - self.inner_layernorm = LayerNorm(ffn_hidden_size, dtype=dtype, - eps=eps) if inner_layernorm else None - self.proj = RowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - is_expert=is_expert) - - # see optimize_model's add_lora for LoRA initialization - self.lora = None # used for split up and gate proj - self.fused_gate_up_lora = None # used for merged up_gate proj - self.dora = None - self.fused_gate_up_dora = None - - def fc_gate_plugin(self, hidden_states, lora_layer_params=None): - # Combine the following pattern - # - # SiLU(FC(x)) * Gate(x) - # - # into: - # - # SwiGLU(FusedFC(x)) - if default_net( - ).plugin_config.low_latency_gemm_swiglu_plugin is not None: - p_dtype = default_net().plugin_config.low_latency_gemm_swiglu_plugin - else: - p_dtype = default_net().plugin_config.gemm_swiglu_plugin - use_fp8 = p_dtype == 'fp8' - assert use_fp8, "gemm_swiglu_plugin and low_latency_gemm_swiglu_plugin only supports fp8 now" - - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - - if mlp_fc_lora_params is not None or mlp_gate_lora_params is not None: - raise NotImplementedError( - f"LoRA of splitting fc and gate is not yet implemented for gemm_swiglu_plugin" - ) - - if self.hidden_act != 'silu': - raise NotImplementedError( - f"Activation {self.hidden_act} not yet implemented for gemm_swiglu_plugin" - ) - - if self.bias: - raise NotImplementedError( - f"bias not yet implemented for gemm_swiglu_plugin fp8") - - assert isinstance( - self.fused_fc, - FP8Linear), "fp8 gemm_swiglu only supports fp8 weights" - assert isinstance( - self.proj, - FP8RowLinear), "fp8 gemm_swiglu only supports fp8 weights" - assert self.fused_fc.weight.shape == ( - self.hidden_size, self.ffn_hidden_size * 2 // - self.tp_size), "fp8 gemm_swiglu only supports (k, n) weights" - - scale_d0 = (self.fused_fc.weights_scaling_factor.raw_value.item() * - self.fused_fc.activation_scaling_factor.raw_value.item()) - scale_d1 = scale_d0 - scale_output = 1.0 / self.proj.activation_scaling_factor.raw_value.item( - ) - activation_scaling_factor = cast( - self.fused_fc.activation_scaling_factor.value, self.dtype) - if hidden_states.dtype != trt.fp8: - hidden_states = quantize(hidden_states, activation_scaling_factor, - 'fp8') - - if default_net( - ).plugin_config.low_latency_gemm_swiglu_plugin is not None: - inter = low_latency_gemm_swiglu(hidden_states, - self.fused_fc.weight.value, - scale_d0, scale_d1, scale_output) - else: - inter = gemm_swiglu(hidden_states, self.fused_fc.weight.value, None, - scale_d0, scale_d1, scale_output) - - lora_result = fc_gate_lora(hidden_states, self.lora, - self.fused_gate_up_lora, lora_layer_params) - if lora_result is not None: - inter = inter + lora_result - - return inter - - def fc_gate(self, hidden_states, lora_layer_params=None): - # Combine the following pattern - # - # SiLU(FC(x)) * Gate(x) - # - # into: - # - # SwiGLU(FusedFC(x)) - # - # Upside is we don't need to modify 4 different weight loading paths just to concat weights - - inter = self.fused_fc(hidden_states) - - lora_result = fc_gate_lora(hidden_states, self.lora, - self.fused_gate_up_lora, lora_layer_params) - if lora_result is not None: - inter = inter + lora_result - if self.dora is not None: - inter = fc_gate_dora(inter, self.dora, self.fused_gate_up_lora, - lora_layer_params) - - if self.hidden_act == 'silu': - inter = ACT2FN['swiglu'](inter) - elif self.hidden_act == 'gelu': - inter = ACT2FN['geglu'](inter) - else: - raise NotImplementedError( - f"Activation {self.hidden_act} not yet implemented for {self.__class__.__name__}." - ) - return inter - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None): - if default_net().plugin_config.gemm_swiglu_plugin or default_net( - ).plugin_config.low_latency_gemm_swiglu_plugin: - inter = self.fc_gate_plugin(hidden_states, lora_layer_params) - else: - inter = self.fc_gate(hidden_states, lora_layer_params) - - if self.inner_layernorm is not None: - inter = self.inner_layernorm(inter) - - mlp_proj_lora_params = None - if lora_layer_params is not None: - mlp_proj_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_4h_to_h") - output = self.proj(inter, - lora_runtime_params=mlp_proj_lora_params, - all_reduce_params=all_reduce_params) - return output - - -class LinearGELU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - approximate: str = 'tanh', - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.proj = ColumnLinear(dim_in, - dim_out, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - if approximate != 'tanh': - raise NotImplementedError('GELU only support tanh now.') - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states = ACT2FN['gelu_pytorch_tanh'](hidden_states) - return hidden_states - - -class LinearGEGLU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - approximate: str = 'tanh', - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.proj = ColumnLinear(dim_in, - dim_out * 2, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - if approximate != 'tanh': - raise NotImplementedError('GELU only support tanh now.') - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states, gate = chunk(hidden_states, - 2, - dim=(hidden_states.ndim() - 1)) - return hidden_states * ACT2FN['gelu_pytorch_tanh'](gate) - - -class LinearApproximateGELU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.proj = ColumnLinear(dim_in, - dim_out, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - - def forward(self, x): - x = self.proj(x) - return x * ACT2FN['sigmoid'](1.702 * x) - - -class LinearSwiGLU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.proj = ColumnLinear(dim_in, - dim_out * 2, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - self.hidden_act = 'silu' - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states, gate = chunk(hidden_states, - 2, - dim=(hidden_states.ndim() - 1)) - return hidden_states * ACT2FN[self.hidden_act](gate) - - -class LinearActivation(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - bias: bool = True, - activation: str = "silu", - mapping=Mapping(), - dtype=None): - super().__init__() - - self.proj = ColumnLinear(dim_in, - dim_out, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - self.hidden_act = activation - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - return ACT2FN[self.activation](hidden_states) diff --git a/tensorrt_llm/layers/moe.py b/tensorrt_llm/layers/moe.py deleted file mode 100755 index c6a0427b3a97..000000000000 --- a/tensorrt_llm/layers/moe.py +++ /dev/null @@ -1,1553 +0,0 @@ -# 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 dataclasses import asdict, dataclass -from enum import IntEnum -from typing import List, Optional, Type, Union - -import numpy as np -import tensorrt as trt -import torch - -from tensorrt_llm._torch.utils import ActivationType -from tensorrt_llm._utils import (get_init_params, str_dtype_to_torch, - str_dtype_to_trt) -from tensorrt_llm.layers.lora import LoraParams - -from .._common import default_net, default_trtnet -from .._utils import QuantModeWrapper, get_sm_version, int32_array -from ..functional import (AllReduceParams, SideStreamIDType, Tensor, - _add_plugin_info, _create_tensor, abs, allreduce, - cast, concat, constant, cuda_stream_sync, div, expand, - gather_nd, gt, is_gated_activation) -from ..functional import max as trt_max -from ..functional import (maximum, non_gated_version, nonzero, reduce_scatter, - repeat_interleave, scatter, scatter_nd, shape, - sigmoid, softmax, split, sub, sum, topk, unsqueeze, - where) -from ..mapping import Mapping -from ..module import Module, ModuleList -from ..parameter import Parameter -from ..plugin import TRT_LLM_PLUGIN_NAMESPACE -from ..quantization import GroupwiseQuantAlgo, QuantMode -from ..quantization.functional import (get_weight_scale_interleave_factor, - postprocess_weight_only, - preprocess_weights_for_mixed_gemm, - quantize) -from .linear import RowLinear -from .mlp import MLP, GatedMLP - -activation_str_to_int_map = { - # [WARNING] Keep the below in sync with cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h - "gelu": int(ActivationType.Gelu), - "gelu_new": int(ActivationType.Gelu), - "relu": int(ActivationType.Relu), - "silu": int(ActivationType.Silu), - "swiglu": int(ActivationType.Swiglu), - "geglu": int(ActivationType.Geglu), - "swiglu_bias": int(ActivationType.SwigluBias), - "identity": int(ActivationType.Identity), - "relu2": int(ActivationType.Relu2), -} - - -class MoeGroupwiseQuantParams(): - - def __init__(self, - group_size=-1, - zero=False, - pre_quant_scale=False, - use_w4a8_awq=False, - act_scale_1=None, - weight_scale_1=None, - weight_zero_1=None, - alpha_1=None, - act_scale_2=None, - weight_scale_2=None, - weight_zero_2=None, - alpha_2=None) -> None: - self.group_size = group_size - self.quant_algo = zero * GroupwiseQuantAlgo.ZERO + pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA - - self.quant_params = [] - - if group_size == -1: - return - - assert weight_scale_1 - assert weight_scale_2 - self.quant_params += [weight_scale_1, weight_scale_2] - if pre_quant_scale: - assert act_scale_1 - assert act_scale_2 - self.quant_params += [act_scale_1, act_scale_2] - if zero: - assert weight_zero_1 - assert weight_zero_2 - self.quant_params += [weight_zero_1, weight_zero_2] - if use_w4a8_awq: - assert alpha_1 - assert alpha_2 - self.quant_params += [alpha_1, alpha_2] - - -@dataclass -class MoeConfig: - - class ExpertScaleNormalizationMode(IntEnum): - NONE = 0 - RENORMALIZE = 1 - SPARSE_MIXER = 2 - DEVICE_LIMITED = 3 - DEVICE_LIMITED_RENORM = 4 - - num_experts: int = 0 - shared_expert_intermediate_size: int = 0 - - top_k: int = 0 - normalization_mode: ExpertScaleNormalizationMode = ExpertScaleNormalizationMode.RENORMALIZE - sparse_mixer_epsilon: float = 0.01 - tp_mode: int = 0 - - device_limited_n_group: int = 0 - device_limited_topk_group: int = 0 - device_limited_routed_scaling_factor: float = 1.0 - - def validate(self) -> "MoeConfig": - if (self.num_experts == 0) != (self.top_k == 0): - raise ValueError( - "Both or neither MoeConfig's num_experts and top_k must be set to 0" - ) - return self - - def has_moe(self) -> bool: - return self.num_experts > 1 - - @classmethod - def from_dict(cls, config: dict): - return cls(**config) - - def to_dict(self): - return asdict(self) - - -def _moe_plugin(moe_config, - hidden_states, - hidden_states_raw, - token_selected_experts, - token_final_scales, - expert_weights_1, - expert_weights_2, - expert_bias_1, - expert_bias_2, - expert_scale_1, - expert_scale_2, - expert_scale_3, - expert_scale_4, - expert_scale_5, - expert_scale_6, - groupwise_quant_params, - hidden_size, - ffn_hidden_size, - act_fn, - dtype, - weight_dtype, - output_dtype, - lora_params: LoraParams, - lora_max_low_rank, - quant_mode=QuantMode(0), - tp_size=1, - ep_size=1, - tp_rank=0, - ep_rank=0, - side_stream_id=SideStreamIDType.disable): - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - - if isinstance(weight_dtype, str): - weight_dtype = str_dtype_to_trt(weight_dtype) - - if isinstance(output_dtype, str): - output_dtype = str_dtype_to_trt(output_dtype) - - def from_parameter(x): - if isinstance(x, Parameter): - return x.value - return x - - expert_weights_1 = from_parameter(expert_weights_1) - expert_weights_2 = from_parameter(expert_weights_2) - expert_bias_1 = from_parameter(expert_bias_1) - expert_bias_2 = from_parameter(expert_bias_2) - expert_scale_1 = from_parameter(expert_scale_1) - expert_scale_2 = from_parameter(expert_scale_2) - expert_scale_3 = from_parameter(expert_scale_3) - expert_scale_4 = from_parameter(expert_scale_4) - expert_scale_5 = from_parameter(expert_scale_5) - expert_scale_6 = from_parameter(expert_scale_6) - - # Create the plugin with our required state - num_experts = moe_config.num_experts - p_remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int32(default_net().plugin_config.remove_input_padding), - dtype=np.int32), trt.PluginFieldType.INT32) - # We pass the full number of experts (not divided by ep_size) even for EP mode - p_num_experts = trt.PluginField("number_of_experts", - np.array(num_experts, dtype=np.int32), - trt.PluginFieldType.INT32) - p_experts_per_token = trt.PluginField( - "experts_per_token", np.array(moe_config.top_k, dtype=np.int32), - trt.PluginFieldType.INT32) - p_expert_hidden_size = trt.PluginField( - "expert_hidden_size", np.array(hidden_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_expert_inter_size = trt.PluginField( - "expert_inter_size", np.array(ffn_hidden_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_groupwise_quant_algo = trt.PluginField( - "groupwise_quant_algo", - np.array(groupwise_quant_params.quant_algo, dtype=np.int32), - trt.PluginFieldType.INT32) - p_group_size = trt.PluginField( - "group_size", np.array(groupwise_quant_params.group_size, - dtype=np.int32), trt.PluginFieldType.INT32) - p_activation_type = trt.PluginField( - "activation_type", - np.array(activation_str_to_int_map[act_fn], dtype=np.int32), - trt.PluginFieldType.INT32) - p_type_id = trt.PluginField("type_id", np.array([int(dtype)], - dtype=np.int32), - trt.PluginFieldType.INT32) - - p_weight_type_id = trt.PluginField( - "weight_type_id", np.array([int(weight_dtype)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_output_type_id = trt.PluginField( - "output_type_id", np.array([int(output_dtype)], dtype=np.int32), - trt.PluginFieldType.INT32) - - if isinstance(quant_mode, QuantModeWrapper): - # We only need to get one quant mode here for specific moe layer - quant_mode = quant_mode[0] - p_quant_mode = trt.PluginField("quant_mode", - np.array([int(quant_mode)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_use_final_scales = trt.PluginField( - "use_final_scales", - np.array([int(token_final_scales is not None)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_use_bias = trt.PluginField( - "use_bias", np.array([int(expert_bias_1 is not None)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_tp_size = trt.PluginField("tp_size", np.array(tp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_tp_rank = trt.PluginField("tp_rank", np.array(tp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - p_ep_size = trt.PluginField("ep_size", np.array(ep_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_ep_rank = trt.PluginField("ep_rank", np.array(ep_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_force_determinism = trt.PluginField( - "force_determinism", np.array([int(False)], dtype=np.int32), - trt.PluginFieldType.INT32) - - p_side_stream_id = trt.PluginField("side_stream_id", - np.array(side_stream_id, dtype=np.int32), - trt.PluginFieldType.INT32) - - use_lora = default_net().plugin_config.lora_plugin is not None - p_use_lora = trt.PluginField("use_lora", np.array([int(use_lora)], - np.int32), - trt.PluginFieldType.INT32) - if use_lora: - p_lora_type_id = trt.PluginField( - "lora_type_id", - np.array([ - int(str_dtype_to_trt(default_net().plugin_config.lora_plugin)) - ], np.int32), trt.PluginFieldType.INT32) - p_max_low_rank = trt.PluginField( - "max_low_rank", np.array(lora_max_low_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc_inputs = [ - p_remove_input_padding, p_num_experts, p_experts_per_token, - p_expert_hidden_size, p_expert_inter_size, p_groupwise_quant_algo, - p_group_size, p_activation_type, p_type_id, p_weight_type_id, - p_output_type_id, p_quant_mode, p_use_bias, p_use_final_scales, - p_tp_size, p_tp_rank, p_ep_size, p_ep_rank, p_force_determinism, - p_side_stream_id, p_use_lora - ] - - if use_lora: - pfc_inputs += [p_lora_type_id, p_max_low_rank] - - pfc = trt.PluginFieldCollection(pfc_inputs) - - # Create the plugin with our constant inputs to the constructor - plugin_creator = trt.get_plugin_registry().get_plugin_creator( - 'MixtureOfExperts', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plugin_creator is not None - moe_plugin = plugin_creator.create_plugin("mixture_of_experts", pfc) - - # Instantiate the plugin with our specific inputs - plugin_inputs = [hidden_states, expert_weights_1, expert_weights_2] - plugin_inputs += [token_selected_experts] - - # Add conditional inputs - - # Final scales do a final rescale of the output of the experts - if token_final_scales is not None: - plugin_inputs += [token_final_scales] - - # Expert biases - if expert_bias_1: - assert expert_bias_2 - plugin_inputs += [expert_bias_1, expert_bias_2] - - # Add conditional inputs - if (quant_mode.is_weight_only() and not quant_mode.has_per_group_scaling()): - assert expert_scale_1 - assert expert_scale_2 - plugin_inputs += [expert_scale_1, expert_scale_2] - elif quant_mode.has_fp8_qdq(): - # FP8 always has scales 1-3 - assert expert_scale_1 - assert expert_scale_2 - assert expert_scale_3 - plugin_inputs += [expert_scale_1, expert_scale_2, expert_scale_3] - - if expert_scale_4 is not None: - assert output_dtype == trt.fp8 - plugin_inputs += [expert_scale_4] - - # Lora needs an extra parameter to be able to dequant the input back to backbone type - if use_lora: - assert expert_scale_5 - plugin_inputs += [expert_scale_5] - elif quant_mode.has_per_group_scaling(): - plugin_inputs += groupwise_quant_params.quant_params - # Lora needs an extra parameter to be able to dequant the input back to backbone type - if use_lora: - assert expert_scale_5 - plugin_inputs += [expert_scale_5] - elif quant_mode.has_nvfp4(): - assert expert_scale_1 - assert expert_scale_2 - assert expert_scale_3 - assert expert_scale_4 - assert expert_scale_5 - assert expert_scale_6 - plugin_inputs += [ - expert_scale_1, expert_scale_2, expert_scale_3, expert_scale_4, - expert_scale_5, expert_scale_6 - ] - - # Lora parameters - if use_lora: - # Check if lora_params is not None - moe_h_4h_params = lora_params.get_runtime_params(0, "moe_h_to_4h") - if moe_h_4h_params is not None: - moe_h_4h_weight_ptrs = moe_h_4h_params.lora_weights_pointers - moe_h_4h_lora_ranks = moe_h_4h_params.lora_ranks - plugin_inputs += (moe_h_4h_weight_ptrs + moe_h_4h_lora_ranks) - - moe_4h_h_params = lora_params.get_runtime_params(0, "moe_4h_to_h") - if moe_4h_h_params is not None: - moe_4h_h_weight_ptrs = moe_4h_h_params.lora_weights_pointers - moe_4h_h_lora_ranks = moe_4h_h_params.lora_ranks - plugin_inputs += (moe_4h_h_weight_ptrs + moe_4h_h_lora_ranks) - - if is_gated_activation(act_fn): - moe_gate_params = lora_params.get_runtime_params(0, "moe_gate") - if moe_gate_params is not None: - moe_gate_weight_ptrs = moe_gate_params.lora_weights_pointers - moe_gate_lora_ranks = moe_gate_params.lora_ranks - plugin_inputs += (moe_gate_weight_ptrs + moe_gate_lora_ranks) - - host_request_types = lora_params.host_request_types - plugin_inputs += [host_request_types] - - if default_net().plugin_config.remove_input_padding: - plugin_inputs += [lora_params.host_context_lengths] - - # A control flow tensor required to synchronize the side stream - if side_stream_id != SideStreamIDType.disable: - plugin_inputs += [hidden_states_raw] - - # Pass the inputs to the plugin - plugin_inputs = [i.trt_tensor for i in plugin_inputs] - layer = default_trtnet().add_plugin_v2(plugin_inputs, moe_plugin) - _add_plugin_info(layer, plugin_creator, "mixture_of_experts", pfc) - if not default_net().strongly_typed: - for ii in range(layer.num_inputs): - if layer.get_input(ii).dtype == str_dtype_to_trt("int8"): - layer.get_input(ii).set_dynamic_range(-127, 127) - - # Fetch the output tensor - output = _create_tensor(layer.get_output(0), layer) - - # If the side stream is enabled, also return the synchronization tensor for the side stream - if side_stream_id != SideStreamIDType.disable: - output = (output, _create_tensor(layer.get_output(1), layer)) - return output - - -def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1], - w_packed_int4x2.shape[2] * 2, - dtype=torch.int8) - w_unpacked[:, :, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, :, 1::2] = w_packed_int4x2 // 16 - w_unpacked = w_unpacked.view(-1, 8)[:, [0, 4, 1, 5, 2, 6, 3, 7]].view( - w_unpacked.shape) - return w_unpacked.contiguous() - - -# This exists so that MOE can have the same name format as a regular MLP, just with different shaped weight tensors -class MOEWeightWrapper(Module): - - def __init__(self, in_features: int, out_features: int, - experts_per_node: int, quant_mode: QuantMode, - groupwise_quant_algo: int, group_size: int, - dtype: Union[str, - trt.DataType], weight_dtype: Union[str, - trt.DataType], - has_bias: bool, wrapper_tllm_to_externel_key_dict: dict, - tp_size: int, tp_dim: int): - super().__init__() - self.quant_mode = quant_mode - self.groupwise_quant_algo = groupwise_quant_algo - self.group_size = group_size - self.expert_shape = (experts_per_node, out_features, in_features) - self.dtype = dtype - self.weight_dtype = weight_dtype - self.has_bias = has_bias - self.tllm_to_externel_key_dict = wrapper_tllm_to_externel_key_dict - self.tp_size = tp_size - self.tp_dim = 1 - tp_dim if quant_mode.has_per_group_scaling( - ) else tp_dim - self.is_padded = False - - if quant_mode.is_weight_only( - ) and not quant_mode.has_per_group_scaling(): - bytes_per_col_scale = 2 if quant_mode.is_int4_weight_only() else 1 - # We use a different shape here because the quantized weights have their own layout - self.expert_shape = (experts_per_node, in_features, - out_features // bytes_per_col_scale) - self.per_channel_scale = Parameter(shape=(experts_per_node, - out_features), - dtype=dtype) - else: - self.register_parameter('per_channel_scale', None) - - if quant_mode.has_nvfp4(): - self.expert_shape = (experts_per_node, out_features, in_features) - weight_dtype = trt.fp4 - - if not quant_mode.has_per_group_scaling(): - self.weight = Parameter(shape=self.expert_shape, - dtype=weight_dtype, - prefer_managed=True) - - if has_bias: - self.bias = Parameter(shape=(experts_per_node, out_features), - dtype=dtype) - else: - self.register_parameter('bias', None) - - self.scaling_vector_size = 16 - if quant_mode.has_fp8_qdq(): - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(experts_per_node, 1), - dtype=trt.float32) - elif quant_mode.has_nvfp4(): - self.weights_block_scaling_factor_interleaved = Parameter( - shape=(experts_per_node, out_features, - in_features // self.scaling_vector_size), - dtype=trt.fp8) - self.weights_block_scaling_factor = Parameter( - shape=(experts_per_node, out_features, - in_features // self.scaling_vector_size), - dtype=trt.fp8) - self.activation_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - # alpha = 1.0 / (weight_global_scale * act_global_scale) - self.alpha = Parameter(shape=(experts_per_node, ), - dtype=trt.float32) - elif quant_mode.has_per_group_scaling(): - self.weight = Parameter( - shape=(experts_per_node, in_features, - out_features // 4), # int4 <--> fp16/bf16 - dtype=dtype) - if groupwise_quant_algo & GroupwiseQuantAlgo.W4A8_ALPHA: - scale_interleave_factor = get_weight_scale_interleave_factor( - in_features, group_size) - else: - scale_interleave_factor = 1 - scale_shape = (experts_per_node, - in_features // group_size // scale_interleave_factor, - out_features * scale_interleave_factor) - self.weights_scaling_factor = Parameter(shape=scale_shape, - dtype=dtype) - if groupwise_quant_algo & GroupwiseQuantAlgo.ZERO: - self.zero = Parameter(shape=scale_shape, dtype=dtype) - else: - self.register_parameter('zero', None) - if groupwise_quant_algo & GroupwiseQuantAlgo.PRE_QUANT_SCALE: - self.prequant_scaling_factor = Parameter(shape=(1, in_features), - dtype=dtype) - else: - self.register_parameter('prequant_scaling_factor', None) - if groupwise_quant_algo & GroupwiseQuantAlgo.W4A8_ALPHA: - self.alpha = Parameter(shape=(experts_per_node, 1), - dtype=trt.float32) - else: - self.register_parameter('alpha', None) - self.tllm_to_externel_key_dict.update( - {"weight": ["qweight", "qzeros", "scales"]}) - else: - self.register_parameter('weights_scaling_factor', None) - self.register_parameter('weights_block_scaling_factor', None) - self.register_parameter('weights_block_scaling_factor_interleaved', - None) - self.register_parameter('activation_scaling_factor', None) - self.register_parameter('activation_global_scaling_factor', None) - self.register_parameter('alpha', None) - self.register_parameter('zero', None) - self.register_parameter('prequant_scaling_factor', None) - - def postprocess(self, tllm_key, weights, **kwargs): - - def stack_weights(tllm_key, weights): - if "fc" in tllm_key: - weights = torch.cat([ - torch.stack(weights[:len(weights) // 2]), - torch.stack(weights[len(weights) // 2:]) - ], - dim=-2) - elif "proj" in tllm_key: - weights = torch.stack(weights) - return weights - - def postprocess_awq(tllm_key, weights): - if not tllm_key.endswith("weight"): - return {} - weights = [weights[i::3] for i in range(3)] - for idx, w in enumerate(weights): - if "fc" in tllm_key: - weights[idx] = torch.cat([ - torch.stack(w[:len(w) // 2]), - torch.stack(w[len(w) // 2:]) - ], - dim=-1) - elif "proj" in tllm_key: - weights[idx] = torch.stack(w) - qweight_int32, qzeros_int32, scales_fp16 = weights - qweight = unpack_int32_into_int8(qweight_int32) - 8 - qweight -= (qweight >> 4) << 4 - qweight = qweight.view(torch.uint8) - qweight = (qweight[:, :, 1::2] * 16 + qweight[:, :, ::2]).view( - torch.int8) - qweight = preprocess_weights_for_mixed_gemm( - qweight, torch.quint4x2, - torch.float16).view(str_dtype_to_torch(self.dtype)) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.to( - str_dtype_to_torch(self.dtype)) - - results = { - tllm_key: qweight, - tllm_key.replace("weight", "weights_scaling_factor"): - scales_fp16, - tllm_key.replace("weight", "zero"): zeros_x_scales_fp16, - } - return results - - if self.quant_mode.has_per_group_scaling(): - return postprocess_awq(tllm_key, weights) - elif tllm_key.endswith("weight"): - if isinstance(weights, torch.Tensor): - weights = [weights] - else: - if self.quant_mode.has_fp8_qdq(): - experts_per_node = self.weights_scaling_factor.shape[0] - if 'fc' in tllm_key: - # Example weights: - # ['model.layers.0.block_sparse_moe.experts.0.w3.weight', - # 'model.layers.0.block_sparse_moe.experts.0.w3.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w3.weight', - # 'model.layers.0.block_sparse_moe.experts.7.w3.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.0.w1.weight', - # 'model.layers.0.block_sparse_moe.experts.0.w1.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w1.weight', - # 'model.layers.0.block_sparse_moe.experts.7.w1.weight_scale'] - assert experts_per_node * 4 == len(weights) - - def w3_weight_idx(expert_id): - return 2 * expert_id - - def w3_weight_scale_idx(expert_id): - return 2 * expert_id + 1 - - def w1_weight_idx(expert_id): - return 2 * (expert_id + experts_per_node) - - def w1_weight_scale_idx(expert_id): - return 2 * (expert_id + experts_per_node) + 1 - - weights_requantized = [None] * (2 * experts_per_node) - # Since w1/w3 share weight_scale by picking the max, we need to requantize weight - for i in range(experts_per_node): - w3_weight = weights[w3_weight_idx(i)].float() - w3_weight_scale = weights[w3_weight_scale_idx( - i)].float() - w1_weight = weights[w1_weight_idx(i)].float() - w1_weight_scale = weights[w1_weight_scale_idx( - i)].float() - - max_weight_scale = max(w3_weight_scale, - w1_weight_scale) - - weights_requantized[i] = (w3_weight * - w3_weight_scale / - max_weight_scale).to( - torch.float8_e4m3fn) - weights_requantized[i + experts_per_node] = ( - w1_weight * w1_weight_scale / - max_weight_scale).to(torch.float8_e4m3fn) - - weights = weights_requantized - else: - assert 'proj' in tllm_key, f"tllm_key is {tllm_key}, which does not contain fc or proj" - # Example weights: - # ['model.layers.0.block_sparse_moe.experts.0.w2.weight', - # 'model.layers.0.block_sparse_moe.experts.0.w2.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w2.weight', - # 'model.layers.0.block_sparse_moe.experts.7.w2.weight_scale', - assert 2 * experts_per_node == len(weights) - - def w2_weight_idx(expert_id): - return 2 * expert_id - - # No need to requantize, simply skip weight_scale - weights = [ - weights[w2_weight_idx(i)] - for i in range(experts_per_node) - ] - - weights = stack_weights(tllm_key, weights) - - if not self.quant_mode.has_any_quant(): - # When each rank holds single expert, weights will be a list - if isinstance(weights, list): - weights = stack_weights(tllm_key, weights) - weights = weights.to(str_dtype_to_torch(self.dtype)) - - # FP8 scaling factors - if tllm_key.endswith("activation_scaling_factor"): - # Use max input range. - weights = max(weights).float().reshape((1, )) - - if tllm_key.endswith("weights_scaling_factor"): - if tllm_key.split('.')[-2] == 'fc': - # Example weights: - # ['model.layers.0.block_sparse_moe.experts.0.w3.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w3.weight_scale', - # 'model.layers.0.block_sparse_moe.experts.0.w1.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w1.weight_scale'] - experts_per_node = self.weights_scaling_factor.shape[0] - assert experts_per_node * 2 == len(weights) - - def w3_weight_scale_idx(expert_id): - return expert_id - - def w1_weight_scale_idx(expert_id): - return expert_id + experts_per_node - - # w1 and w3 share the weight scale by picking the max - weights = [ - max(weights[w3_weight_scale_idx(i)], - weights[w1_weight_scale_idx(i)]) - for i in range(experts_per_node) - ] - - weights = stack_weights(tllm_key, weights) - - # FP4 scaling factors - if tllm_key.endswith("weights_block_scaling_factor"): - weights = stack_weights(tllm_key, weights) - if tllm_key.endswith("weights_block_scaling_factor_interleaved"): - weights = stack_weights(tllm_key, weights) - weights = torch.ops.trtllm.block_scale_interleave( - weights.to(torch.float8_e4m3fn).view( - torch.uint8).cpu().contiguous()).reshape( - weights.shape).view(torch.float8_e4m3fn) - if tllm_key.endswith("activation_global_scaling_factor"): - # Use max input range. - weights = max(weights).float().reshape((1, )) - if tllm_key.endswith("alpha"): - # weights are: [e0_w3_weight_scale, e0_w3_input_scale, e1_w3_weight_scale, e1_w3_input_scale - # ..., e7_w3_weight_scale, e7_w3_input_scale, e0_w1_weight_scale, e0_w1_input_scale, ...] - weights_global_scale = weights[::2] - activation_global_scale = weights[1::2] - if 'fc' in tllm_key: - weights_global_scale = torch.stack( - weights_global_scale[:len(weights_global_scale) // 2]) - else: - weights_global_scale = torch.stack(weights_global_scale) - weights = (weights_global_scale * - max(activation_global_scale).float()).reshape((-1, )) - - # Weight only - if self.quant_mode.is_weight_only(): - if "per_channel_scale" in tllm_key: - return {} - weights = weights.to(str_dtype_to_torch(self.dtype)) - return postprocess_weight_only( - tllm_key, weights, torch.int8 if - self.quant_mode.is_int8_weight_only() else torch.quint4x2, self) - - return weights - - -class MixtureOfExperts(Module): - - def __init__(self, - moe_config: MoeConfig, - hidden_size: int, - ffn_hidden_size: int, - hidden_act: str, - mapping: Mapping = Mapping(), - bias: bool = True, - dtype=None, - tp_group: List[int] = None, - tp_size: int = 1, - quant_mode=QuantMode(0), - use_all_reduce=True, - pre_quant_scale=False, - zero=False, - use_w4a8_awq=False, - use_int8_weight=False, - group_size: int = -1, - static_routing=False): - super().__init__() - - self.moe_config = moe_config - self.num_experts = moe_config.num_experts - self.top_k = moe_config.top_k - - self.hidden_act = hidden_act - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.expert_inter_size = ffn_hidden_size - self.dtype = dtype - self.weight_dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.mapping = mapping - self.quant_mode = quant_mode - self.bias = bias - self.use_all_reduce = use_all_reduce - self.zero = zero - self.pre_quant_scale = pre_quant_scale - self.use_w4a8_awq = use_w4a8_awq - self.use_int8_weight = use_int8_weight - self.group_size = group_size - - if self.use_int8_weight and self.group_size > 0: - raise NotImplementedError("INT8-GPTQ is not implemented for MoE.") - - self.static_routing = static_routing - - self.experts_per_node = self.num_experts - if self.mapping.has_moe_ep(): - if self.num_experts % self.mapping.moe_ep_size != 0: - raise ValueError( - f"MixtureOfExperts - Number of experts {self.num_experts} is not a multiple of EP size {self.mapping.moe_ep_size}" - ) - self.experts_per_node = self.experts_per_node // self.mapping.moe_ep_size - - if self.mapping.has_moe_tp(): - if self.ffn_hidden_size % self.mapping.moe_tp_size != 0: - raise ValueError( - f"MixtureOfExperts - FFN Hidden Size {self.ffn_hidden_size} is not a multiple of TP size {self.mapping.moe_tp_size}" - ) - self.expert_inter_size = self.ffn_hidden_size // self.mapping.moe_tp_size - - if quant_mode.has_fp8_rowwise(): - raise ValueError( - "MixtureOfExperts - MOE Does not support FP8 rowwise quantize") - - if quant_mode.has_fp8_qdq() and self.bias: - # TODO We will need to revisit this if we have a use case for it - raise ValueError( - f"MixtureOfExperts - Bias is not supported with FP8") - - if quant_mode.is_weight_only(): - self.weight_dtype = trt.int8 - elif quant_mode.has_fp8_qdq(): - self.weight_dtype = trt.fp8 - - rank_experts = self.mapping.ep_experts(self.num_experts) - self.wrapper_tllm_to_externel_key_dict = { - "mlp": - "block_sparse_moe", - "proj": [f"experts.{expert}.w2" for expert in rank_experts], - "fc": [f"experts.{expert}.w3" for expert in rank_experts] + - [f"experts.{expert}.w1" for expert in rank_experts] - } - - if quant_mode.has_fp8_qdq(): - self.wrapper_tllm_to_externel_key_dict.update({ - "weight": [ - "weight", "weight_scale" - ], # We need weight_scale to do requantization for w1/w3 fusion - "weights_scaling_factor": - "weight_scale", - "activation_scaling_factor": - "input_scale" - }) - - if quant_mode.has_nvfp4(): - self.wrapper_tllm_to_externel_key_dict.update({ - "weights_block_scaling_factor_interleaved": - "weight_scale", - "weights_block_scaling_factor": - "weight_scale", - "activation_global_scaling_factor": - "input_scale", - "alpha": ["weight_scale_2", "input_scale"], - }) - - # Since output dimension is usually low (in the order of 10s), no TP at - # all is more efficient as no allreduce required in the end. - # Note that if we see models that have large number of experts, we may - # need to consider add TP back here. - # TODO: Arctic has large # experts, we may need to add TP back here. - if not self.static_routing: - self.router = RowLinear( - hidden_size, - self.num_experts, - bias=False, - dtype= - "float32", # Routing is sensitive since it conditions what experts are used - tp_group=None, - tp_size=1, - strict_dtype=True) - self.router.tllm_to_externel_key_dict = { - "mlp": "block_sparse_moe", - "router": "gate" - } - - self.init_experts() - - self.max_low_rank = None - - def init_experts(self): - # Note we use horizontal fusion for gated activation to do the operation in one GEMM invocation - # The left matrix is a linear projection (no activation applied) - # The right matrix is the gating value (activation applied) - # The naming convention is the inverse of GatedMLP, but the same as `tensorrt_llm/functional.py` - fc_out_size = self.expert_inter_size * 2 if is_gated_activation( - self.hidden_act) else self.expert_inter_size - groupwise_quant_algo = self.zero * GroupwiseQuantAlgo.ZERO + self.pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + self.use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA - - self.fc = MOEWeightWrapper(self.hidden_size, fc_out_size, - self.experts_per_node, self.quant_mode, - groupwise_quant_algo, self.group_size, - self.dtype, self.weight_dtype, self.bias, - self.wrapper_tllm_to_externel_key_dict, - self.mapping.moe_tp_size, 0) - self.proj = MOEWeightWrapper(self.expert_inter_size, self.hidden_size, - self.experts_per_node, self.quant_mode, - groupwise_quant_algo, self.group_size, - self.dtype, self.weight_dtype, self.bias, - self.wrapper_tllm_to_externel_key_dict, - self.mapping.moe_tp_size, 1) - - def default_routing(self, logits): - topk_values, topk_indices = topk(softmax(cast(logits, trt.float32), - dim=-1), - k=self.moe_config.top_k, - dim=-1) - return topk_indices, topk_values - - def renormalize(self, logits): - # Get top-k experts and renormalize their scores - token_scores, token_selected_experts = topk(cast(logits, trt.float32), - k=self.moe_config.top_k, - dim=-1) - token_final_scales = softmax(token_scores, dim=-1) - return token_selected_experts, token_final_scales - - def group_limited_greedy(self, logits): - n_group = self.moe_config.device_limited_n_group - scores = softmax(cast(logits, trt.float32), -1) - scores_shape = [shape(scores, i) for i in range(scores.ndim())] - group_scores = scores.view( - concat(scores_shape[:-1] + - [n_group, scores_shape[-1] // n_group])).max(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, - cast(group_idx, 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 = scores * score_mask * \ - self.moe_config.device_limited_routed_scaling_factor - return scores - - def sparse_mixer(self, logits): - router_logits = cast(logits, trt.float32) - - topk_values = [] - topk_indices = [] - - assert self.top_k == 2, "Sparse mixer only supports top_k = 2" - - def mask_and_softmax(router_logits): - # Get max of remaining values - max_values = trt_max(router_logits, dim=-1, keepdim=True) - - # Calculate mask for epsilon condition - abs_values = abs(router_logits) - max_abs = maximum(abs_values, max_values) - diff = sub(max_values, router_logits) - ratio = div(diff, max_abs) - - # Apply epsilon mask - eps_mask = gt(ratio, 2 * self.moe_config.sparse_mixer_epsilon) - router_logits = where(eps_mask, -float('inf'), router_logits) - curr_values, curr_indices = topk(softmax(router_logits), - k=1, - dim=-1) - return curr_indices, curr_values - - curr_indices, curr_values = mask_and_softmax(router_logits) - topk_values.append(curr_values) - topk_indices.append(curr_indices) - - # Mask the last selected expert to -inf - router_logits = scatter(router_logits, -1, curr_indices, - curr_values * 0 - float('inf')) - - curr_indices, curr_values = mask_and_softmax(router_logits) - topk_values.append(curr_values) - topk_indices.append(curr_indices) - - # Concatenate results - values = concat(topk_values, dim=1) - indices = concat(topk_indices, dim=1) - - return indices, values - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - last_local_layer_residual=None, - side_stream_id: Optional[SideStreamIDType] = SideStreamIDType. - disable, - static_routing_input: Optional[Tensor] = None): - moe_router_lora_params = None - if lora_layer_params is not None: - moe_router_lora_params = lora_layer_params.get_runtime_params( - 0, "moe_router") - - if not self.static_routing: - routing_input = cast(hidden_states, trt.float32) - routing = self.router(routing_input, moe_router_lora_params) - else: - routing = None - - # token_selected_experts is shape (num_tokens, experts_per_token). - # It is a list of selected expert indices for each token - # token_final_scales is shape (num_tokens, experts_per_token). May be None - # It contains a final scaling/weighting factor applied to the output of each selected expert before summing the results - if self.static_routing: - token_selected_experts = static_routing_input - token_final_scales = None - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED: - token_final_scales, token_selected_experts = topk( - self.group_limited_greedy(routing), - k=self.moe_config.top_k, - dim=-1) - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM: - token_final_scales, token_selected_experts = topk( - self.group_limited_greedy(routing), - k=self.moe_config.top_k, - dim=-1) - token_final_scales /= sum(token_final_scales, dim=-1, keepdim=True) - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE: - token_selected_experts, token_final_scales = self.renormalize( - routing) - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.SPARSE_MIXER: - token_selected_experts, token_final_scales = self.sparse_mixer( - routing) - else: - token_selected_experts, token_final_scales = self.default_routing( - routing) - - output = self.forward_experts(hidden_states, token_selected_experts, - token_final_scales, lora_layer_params, - side_stream_id) - if side_stream_id != SideStreamIDType.disable: - output, side_stream_sync_tensor = output - if self.use_all_reduce: - output = self.forward_allreduce(output, all_reduce_params, - last_local_layer_residual) - if side_stream_id != SideStreamIDType.disable: - # All tensors that the side channel receives as input must be synced - # on the main stream, to prevent their memory from being released or - # reused by the main stream before the side stream has finished. - tensors_to_sync = (side_stream_sync_tensor, hidden_states, - token_selected_experts, token_final_scales, - lora_layer_params) - tensors_to_sync = tuple(t for t in tensors_to_sync if t is not None) - output = (output, tensors_to_sync) - return output - - def forward_experts(self, hidden_states, token_selected_experts, - token_final_scales, lora_layer_params, side_stream_id): - - groupwise_quant_params = MoeGroupwiseQuantParams() - if self.quant_mode.has_fp8_qdq(): - assert self.fc.weight.value.dtype == trt.fp8, ( - "mlp fc weight dtype should be fp8 in the fp8 quantization mode." - ) - assert self.proj.weight.value.dtype == trt.fp8, ( - "mlp proj weight dtype should be fp8 in the fp8 quantization mode." - ) - hidden_states_quant = hidden_states - if hidden_states_quant.dtype != trt.fp8: - hidden_states_quant = quantize( - hidden_states, self.fc.activation_scaling_factor.value, - 'fp8') - - dtype_quant = trt.fp8 - weight_dtype_quant = trt.fp8 - - fc1_dequant = self.fc.weights_scaling_factor.value * self.fc.activation_scaling_factor.value - fc2_quant = div(1.0, self.proj.activation_scaling_factor.value) - fc2_dequant = self.proj.weights_scaling_factor.value * self.proj.activation_scaling_factor.value - fc1_act_dequant = self.fc.activation_scaling_factor.value - - scale_1 = fc1_dequant - scale_2 = fc2_quant - scale_3 = fc2_dequant - scale_4 = None - scale_5 = fc1_act_dequant - scale_6 = None - - output_dtype_quant = self.dtype - - if output_dtype_quant == trt.fp8 and scale_4 is None: - raise RuntimeError( - "Cannot output FP8 value without knowing quantization parameter" - ) - elif self.quant_mode.has_nvfp4(): - # We pass through the weights unchanged, the quantization is done in the plugin - hidden_states_quant = hidden_states - dtype_quant = trt.fp4 - weight_dtype_quant = trt.fp4 - output_dtype_quant = self.dtype - - scale_1 = div(1.0, self.fc.activation_global_scaling_factor.value) - scale_2 = self.fc.weights_block_scaling_factor_interleaved - scale_3 = self.fc.alpha - scale_4 = div(1.0, self.proj.activation_global_scaling_factor.value) - scale_5 = self.proj.weights_block_scaling_factor_interleaved - scale_6 = self.proj.alpha - elif self.quant_mode.has_per_group_scaling(): - hidden_states_quant = hidden_states - dtype_quant = trt.fp8 if self.use_w4a8_awq else self.dtype - weight_dtype_quant = self.weight_dtype - output_dtype_quant = self.dtype - - scale_1 = None - scale_2 = None - scale_3 = None - scale_4 = None - scale_5 = None - scale_6 = None - pre_quant_scale_1 = self.fc.prequant_scaling_factor.value if self.fc.prequant_scaling_factor else None - zero_1 = self.fc.zero.value if self.fc.zero else None - alpha_1 = self.fc.alpha.value if self.fc.alpha else None - pre_quant_scale_2 = self.proj.prequant_scaling_factor.value if self.proj.prequant_scaling_factor else None - zero_2 = self.proj.zero.value if self.proj.zero else None - alpha_2 = self.proj.alpha.value if self.proj.alpha else None - groupwise_quant_params = MoeGroupwiseQuantParams( - self.group_size, - self.zero, - self.pre_quant_scale, - self.use_w4a8_awq, - pre_quant_scale_1, - self.fc.weights_scaling_factor.value, - zero_1, - alpha_1, - pre_quant_scale_2, - self.proj.weights_scaling_factor.value, - zero_2, - alpha_2, - ) - else: - hidden_states_quant = hidden_states - dtype_quant = self.dtype - weight_dtype_quant = self.weight_dtype - output_dtype_quant = self.dtype - - scale_1 = self.fc.per_channel_scale - scale_2 = self.proj.per_channel_scale - scale_3 = None - scale_4 = None - scale_5 = None - scale_6 = None - output = _moe_plugin(self.moe_config, - hidden_states_quant, - hidden_states, - token_selected_experts, - token_final_scales, - expert_weights_1=self.fc.weight.value, - expert_weights_2=self.proj.weight.value, - expert_bias_1=self.fc.bias, - expert_bias_2=self.proj.bias, - expert_scale_1=scale_1, - expert_scale_2=scale_2, - expert_scale_3=scale_3, - expert_scale_4=scale_4, - expert_scale_5=scale_5, - expert_scale_6=scale_6, - groupwise_quant_params=groupwise_quant_params, - hidden_size=self.hidden_size, - ffn_hidden_size=self.expert_inter_size, - act_fn=self.hidden_act, - dtype=dtype_quant, - weight_dtype=weight_dtype_quant, - output_dtype=output_dtype_quant, - lora_params=lora_layer_params, - lora_max_low_rank=self.max_low_rank, - quant_mode=self.quant_mode, - tp_size=self.mapping.moe_tp_size, - tp_rank=self.mapping.moe_tp_rank, - ep_size=self.mapping.moe_ep_size, - ep_rank=self.mapping.moe_ep_rank, - side_stream_id=side_stream_id) - - return output - - def forward_allreduce(self, - output, - all_reduce_params: Optional[AllReduceParams], - last_local_layer_residual=None): - - if last_local_layer_residual is not None: - if self.mapping.tp_rank == 0: - output = output + last_local_layer_residual - else: - # we need to add this line here to minimize the numerical difference - output = output + 0 - # reshape to (-1) - output = output.view(concat([-1])) - if self.tp_size > 1 and self.tp_group is not None: - output = reduce_scatter(output, self.tp_group) - # reshape to (-1, hidden_size // tp_size) - output = output.view(concat([-1, self.hidden_size // self.tp_size])) - return output - if self.tp_size > 1 and self.tp_group is not None: - output = allreduce(output, - self.tp_group, - all_reduce_params=all_reduce_params) - return output - - def load_weights(self, moe: "MixtureOfExperts"): - ''' - Load weights from base MOE layer - ''' - raise NotImplementedError("Subclass shall override this") - - def to(self, - moe_cls: Type["MixtureOfExperts"], - quant_config=None) -> "MixtureOfExperts": - - if isinstance(moe_cls, MoeOOTB): - if self.moe_config.normalization_mode in [ - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED, - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - ]: - raise ValueError( - 'MoeOOTB doesn\'t support group_limited_greedy yet.') - from ..quantization.quantize import quantize - if isinstance(self, moe_cls): - return self - - new_moe = moe_cls(**get_init_params(self)) - # If config is not None, set quantization from config - if quant_config is not None: - quantize(new_moe, quant_config) - - new_moe.load_weights(self) - if not self.static_routing: - new_moe.router = self.router - return new_moe - - -MOE = MixtureOfExperts - - -# TODO: Support `group_limited_greedy` in MoeOOTB. -class MoeOOTB(MOE): - - def init_experts(self): - if self.quant_mode.is_weight_only(): - raise ValueError( - f"OOTB MOE does not support weight only quantization now, current quant mode: {self.quant_mode}" - ) - - if get_sm_version() >= 100: - raise RuntimeError( - "MoeOOTB does not support SM version >= 100, please use SM version < 100" - ) - - ClsMLP = GatedMLP if is_gated_activation(self.hidden_act) else MLP - - tp_size = 1 - tp_group = None - self.experts = ModuleList([ - ClsMLP(self.hidden_size, self.expert_inter_size, - non_gated_version(self.hidden_act), self.bias, self.dtype, - tp_group, tp_size, self.quant_mode) - for _ in range(self.experts_per_node) - ]) - - def moe_to_expert_lora_params(self, lora_layer_params, expert_idx): - - def get_params(module): - ranks = lora_layer_params.get_runtime_params(0, - module).lora_ranks[0] - weights_pointers = lora_layer_params.get_runtime_params( - 0, module).lora_weights_pointers[0] - return ranks, weights_pointers - - if lora_layer_params is None: - return None - fc_lora_ranks, fc_lora_weights_pointers = get_params("moe_h_to_4h") - proj_lora_ranks, proj_lora_weights_pointers = get_params("moe_4h_to_h") - gate_lora_ranks = None - gate_lora_weights_pointers = None - if is_gated_activation(self.hidden_act): - gate_lora_ranks, gate_lora_weights_pointers = get_params("moe_gate") - return LoraParams( - lora_ranks=[{ - "mlp_h_to_4h_lora_ranks": fc_lora_ranks, - "mlp_4h_to_h_lora_ranks": proj_lora_ranks, - "mlp_gate_lora_ranks": gate_lora_ranks, - }], - lora_weights_pointers=[{ - "mlp_h_to_4h_lora_weights_pointers": - fc_lora_weights_pointers, - "mlp_4h_to_h_lora_weights_pointers": - proj_lora_weights_pointers, - "mlp_gate_lora_weights_pointers": - gate_lora_weights_pointers, - }], - host_context_lengths=lora_layer_params.host_context_lengths, - max_encoder_context_length=lora_layer_params. - max_encoder_context_length, - host_request_types=lora_layer_params.host_request_types, - host_encoder_input_lengths=lora_layer_params. - host_encoder_input_lengths, - weight_index=expert_idx, - ) - - def forward_experts(self, hidden_states, token_selected_experts, - token_final_scales, lora_layer_params, side_stream_id): - assert side_stream_id == SideStreamIDType.disable, "MoeOOTB does not support using side stream" - # TODO: https://nvbugspro.nvidia.com/bug/4781396 after this nvbug is fixed, we will remove this check. - if lora_layer_params is not None: - for module in ["mlp_h_to_4h", "mlp_4h_to_h", "mlp_gate"]: - if lora_layer_params.get_runtime_params(0, module) is not None: - raise RuntimeError( - f"MoE OOTB does not support {module} LoRA module, please enable MoE plugin" - ) - - topk_indices = token_selected_experts - topk_values = token_final_scales - - hidden_size = shape(hidden_states, -1) - # [B*sq, hidden] - inputs_merged = hidden_states.view(concat([-1, hidden_size])) - flat_topk_indices = topk_indices.view( - concat([-1, shape(topk_indices, -1)])) - flat_topk_values = topk_values.view(concat([-1, - shape(topk_values, -1)])) - - # Create output space - zero_buffer = inputs_merged * 0.0 - output = zero_buffer - - expert_indices_stack = [] - indices_stack = [] - # When topk indices are equal to expert index, the expert will inference the tokens. - # Bundle all indices and experts index, then do mask once. - for i, expert in enumerate(self.experts): - if self.mapping.has_moe_ep(): - index = i + self.experts_per_node * self.mapping.moe_ep_rank - else: - index = i - expert_indices_stack.append( - flat_topk_indices.view(concat([1, shape(flat_topk_indices)]))) - - indices_stack.append(constant(int32_array(index))) - - all_expert_indices = concat(expert_indices_stack, dim=0) - indices = expand( - concat(indices_stack).view(concat([len(self.experts), 1, 1])), - shape(all_expert_indices)) - - # Create all experts mask - all_expert_mask = all_expert_indices == indices - - experts_weights = cast( - sum(flat_topk_values * - cast(all_expert_mask, flat_topk_values.dtype), - dim=-1, - keepdim=True), self.dtype) - - all_expert_mask = cast( - sum(cast(all_expert_mask, flat_topk_values.dtype), - dim=-1, - keepdim=True), 'bool') - all_expert_mask = repeat_interleave(all_expert_mask, shape(output, -1), - 2) - - # split the mask and weights for each expert - experts_mask = split(all_expert_mask, 1, dim=0) - expert_weights = split(experts_weights, 1, dim=0) - - for i, expert in enumerate(self.experts): - if self.mapping.has_moe_ep(): - index = i + self.experts_per_node * self.mapping.moe_ep_rank - else: - index = i - # get mask token index - non_zero_index = nonzero(experts_mask[i].view( - concat([-1, hidden_size]))) - non_zero_index = non_zero_index.transpose(1, 0) - input_for_expert = gather_nd(inputs_merged, non_zero_index, 0) - input_for_expert = input_for_expert.view(concat([-1, hidden_size]), - zero_is_placeholder=False) - - # Expert inference - expert_output = expert( - input_for_expert, - lora_layer_params=self.moe_to_expert_lora_params( - lora_layer_params, index)) - - # scatter expert output to real position - expert_finialized_output = zero_buffer - expert_finialized_output = scatter_nd( - expert_finialized_output, non_zero_index, - expert_output.view([-1])) * expert_weights[i] - - output += expert_finialized_output - - output = output.view(shape(hidden_states)) - - return output - - def load_weights(self, moe: MOE): - for i, expert in enumerate(self.experts): - is_gated_act = is_gated_activation(self.hidden_act) - # Gated weight pack in expert1 weights - # expert_weights_1 - experts_weight_1_raw = moe.fc.weight.raw_value - fc1_weight_scale = None - fc1_activation_scale = None - fc2_weight_scale = None - fc2_activation_scale = None - - if self.quant_mode.has_fp8_qdq(): - fc1_weight_scale = moe.fc.weights_scaling_factor.raw_value - fc1_activation_scale = moe.fc.activation_scaling_factor.raw_value - fc2_weight_scale = moe.proj.weights_scaling_factor.raw_value - fc2_activation_scale = moe.proj.activation_scaling_factor.raw_value - - if self.quant_mode.is_weight_only(): - expert.fc.weight.value = experts_weight_1_raw[ - i, :, -self.expert_inter_size:] - if is_gated_act: - expert.gate.weight.value = experts_weight_1_raw[ - i, :, :self.expert_inter_size] - else: - expert.fc.weight.value = experts_weight_1_raw[ - i, -self.expert_inter_size:, :] - if is_gated_act: - expert.gate.weight.value = experts_weight_1_raw[ - i, :self.expert_inter_size, :] - - if self.quant_mode.has_fp8_qdq(): - expert.fc.activation_scaling_factor.value = fc1_activation_scale - expert.fc.weights_scaling_factor.value = fc1_weight_scale[i] - expert.proj.activation_scaling_factor.value = fc2_activation_scale - expert.proj.weights_scaling_factor.value = fc2_weight_scale[i] - if is_gated_act: - expert.gate.activation_scaling_factor.value = fc1_activation_scale - expert.gate.weights_scaling_factor.value = fc1_weight_scale[ - i] - - if self.quant_mode.has_nvfp4(): - expert.fc.activation_global_scaling_factor.value = moe.fc.activation_global_scaling_factor.raw_value - expert.fc.weights_block_scaling_factor.value = moe.fc.weights_block_scaling_factor.raw_value[ - i, -self.expert_inter_size:, :] - expert.fc.weights_block_scaling_factor_interleaved.value = moe.fc.weights_block_scaling_factor_interleaved.raw_value[ - i, -self.expert_inter_size:, :] - expert.fc.alpha.value = np.array(moe.fc.alpha.raw_value[i]) - if is_gated_act: - expert.gate.activation_global_scaling_factor.value = moe.fc.activation_global_scaling_factor.raw_value - expert.gate.weights_block_scaling_factor.value = moe.fc.weights_block_scaling_factor.raw_value[ - i, :-self.expert_inter_size, :] - expert.gate.weights_block_scaling_factor_interleaved.value = moe.fc.weights_block_scaling_factor_interleaved.raw_value[ - i, :-self.expert_inter_size, :] - expert.gate.alpha.value = np.array( - moe.fc.alpha.raw_value[i]) - - expert.proj.activation_global_scaling_factor.value = moe.proj.activation_global_scaling_factor.raw_value - expert.proj.weights_block_scaling_factor.value = moe.proj.weights_block_scaling_factor.raw_value[ - i] - expert.proj.weights_block_scaling_factor_interleaved.value = moe.proj.weights_block_scaling_factor_interleaved.raw_value[ - i] - expert.proj.alpha.value = np.array(moe.proj.alpha.raw_value[i]) - - # expert_weights_2 - experts_weight_2_raw = moe.proj.weight.raw_value - expert.proj.weight.value = experts_weight_2_raw[i, :, :] - - has_bias = self.bias - if has_bias: - experts_bias_1_raw = moe.fc.bias.raw_value - expert.fc.bias.value = experts_bias_1_raw[ - i, -self.expert_inter_size:] - experts_bias_2_raw = moe.proj.bias.raw_value - expert.proj.bias.value = experts_bias_2_raw[i, :] - if is_gated_act: - expert.gate.bias.value = experts_bias_1_raw[ - i, :self.expert_inter_size] - - -# Add SharedMoE class -class SharedMoE(MOE): - - def __init__(self, - moe_config: MoeConfig, - hidden_size: int, - ffn_hidden_size: int, - hidden_act: str, - mapping: Mapping = Mapping(), - bias: bool = True, - dtype=None, - tp_group: List[int] = None, - tp_size: int = 1, - quant_mode=QuantMode(0), - use_shared_gate: bool = False, - use_side_stream: bool = False): - super().__init__( - moe_config=moe_config, - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - mapping=mapping, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode, - use_all_reduce=False, - ) - self.shared_expert = MLP( - hidden_size=hidden_size, - ffn_hidden_size=moe_config.shared_expert_intermediate_size, - hidden_act=hidden_act, - bias=False, - dtype=self.dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=self.quant_mode, - is_expert=True, - ) - self.use_shared_gate = use_shared_gate - if use_shared_gate: - self.shared_expert_gate = RowLinear( - hidden_size, - 1, - bias=False, - dtype=dtype, - tp_group=None, - tp_size=1, - ) - else: - self.shared_expert_gate = None - self.use_side_stream = use_side_stream - - def forward(self, hidden_states, lora_layer_params=None): - side_stream_id = SideStreamIDType.moe if self.use_side_stream else SideStreamIDType.disable - if self.use_side_stream: - routed_output, tensors_to_sync = super().forward( - hidden_states, - lora_layer_params=lora_layer_params, - side_stream_id=side_stream_id, - ) - else: - routed_output = super().forward( - hidden_states, - lora_layer_params=lora_layer_params, - ) - shared_output = self.shared_expert( - hidden_states, - lora_layer_params=lora_layer_params, - ) - if self.shared_expert_gate is not None: - gate_lora_params = None - if lora_layer_params is not None: - gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_router") - shared_output = sigmoid( - self.shared_expert_gate(hidden_states, - gate_lora_params)) * shared_output - if self.use_side_stream: - # tensors_to_sync are included in the inputs to ensure that their - # memory space is not reused for other tensors on the main stream - # until the side stream has finished - shared_output = cuda_stream_sync([shared_output, *tensors_to_sync], - side_stream_id) - hidden_states = routed_output + shared_output - if self.tp_size > 1 and self.tp_group is not None: - hidden_states = allreduce(hidden_states, self.tp_group) - return hidden_states diff --git a/tensorrt_llm/layers/normalization.py b/tensorrt_llm/layers/normalization.py deleted file mode 100644 index a09699c5f08f..000000000000 --- a/tensorrt_llm/layers/normalization.py +++ /dev/null @@ -1,341 +0,0 @@ -# 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 - -from ..functional import (ACT2FN, Tensor, chunk, group_norm, layer_norm, - rms_norm, unsqueeze) -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from .embedding import CombinedTimestepLabelEmbeddings, Embedding -from .linear import Linear - - -class LayerNorm(Module): - - def __init__(self, - normalized_shape, - eps=1e-05, - elementwise_affine=True, - bias=True, - dtype=None, - tp_size=1, - tp_dim=-1): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) - - self.eps = eps - self.dtype = dtype - self.tp_size = tp_size - self.tp_dim = tp_dim - - def forward(self, x, normalized_shape=None): - weight = 1. if self.weight is None else self.weight.value - bias = 0. if self.bias is None else self.bias.value - if normalized_shape is None: - normalized_shape = self.normalized_shape - return layer_norm(x, normalized_shape, weight, bias, self.eps) - - -class RmsNorm(Module): - - def __init__(self, - normalized_shape, - num_groups=1, - eps=1e-06, - elementwise_affine=True, - dtype=None): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - self.num_groups = num_groups - num_channels = normalized_shape[-1] - if num_channels % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - self.eps = eps - self.dtype = dtype - - def forward(self, x, normalized_shape=None): - weight = None if self.weight is None else self.weight.value - if normalized_shape is None: - normalized_shape = self.normalized_shape - return rms_norm(x, normalized_shape, self.num_groups, weight, self.eps) - - -class GroupNorm(Module): - - def __init__(self, - num_groups, - num_channels, - eps=1e-05, - affine=True, - dtype=None): - super().__init__() - - if num_channels % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - - self.num_groups = num_groups - self.num_channels = num_channels - self.affine = affine - - if self.affine: - self.weight = Parameter(shape=(self.num_channels, ), dtype=dtype) - self.bias = Parameter(shape=(self.num_channels, ), dtype=dtype) - else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) - - self.eps = eps - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - return group_norm(x, self.num_groups, weight, bias, self.eps) - - -class AdaLayerNorm(Module): - - def __init__(self, - embedding_dim: int, - num_embeddings: Optional[int] = None, - output_dim: Optional[int] = None, - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-5, - chunk_dim: int = 0, - mapping=Mapping(), - dtype=None): - super().__init__() - self.chunk_dim = chunk_dim - output_dim = output_dim or embedding_dim * 2 - if num_embeddings is not None: - self.emb = Embedding(num_embeddings, embedding_dim, dtype=dtype) - else: - self.emb = None - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - output_dim, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - self.norm = LayerNorm(output_dim // 2, - eps=norm_eps, - elementwise_affine=norm_elementwise_affine, - dtype=dtype) - - def forward(self, - x: Tensor, - timestep: Optional[Tensor] = None, - temb: Optional[Tensor] = None): - assert timestep is not None or temb is not None - if self.emb is not None and timestep is not None: - temb = self.emb(timestep) - temb = self.linear(self.silu(temb)) - if self.chunk_dim == 1: - shift, scale = chunk(temb, 2, dim=1) - shift = unsqueeze(shift, 1) - scale = unsqueeze(scale, 1) - else: - scale, shift = chunk(temb, 2, dim=0) - x = self.norm(x) * (1 + scale) + shift - return x - - -class AdaLayerNormZero(Module): - - def __init__(self, - embedding_dim: int, - num_embeddings: Optional[int] = None, - norm_type: str = "layer_norm", - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - if num_embeddings is not None: - self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, - embedding_dim, - dtype=dtype) - else: - self.emb = None - - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - 6 * embedding_dim, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - elif norm_type == "fp32_layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - bias=False, - dtype=dtype) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward(self, - x: Tensor, - timestep: Optional[Tensor] = None, - class_labels: Optional[Tensor] = None, - hidden_dtype: str = None, - emb: Optional[Tensor] = None): - assert emb is not None or self.emb is not None - if self.emb is not None: - emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype) - emb = self.linear(self.silu(emb)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = chunk( - emb, 6, dim=1) - x = self.norm(x) * (1 + unsqueeze(scale_msa, 1)) + unsqueeze( - shift_msa, 1) - return x, gate_msa, shift_mlp, scale_mlp, gate_mlp - - -class AdaLayerNormZeroSingle(Module): - - def __init__(self, - embedding_dim: int, - norm_type: str = "layer_norm", - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - 3 * embedding_dim, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward(self, x: Tensor, emb: Optional[Tensor] = None): - emb = self.linear(self.silu(emb)) - shift_msa, scale_msa, gate_msa = chunk(emb, 3, dim=1) - x = self.norm(x) * (1 + unsqueeze(scale_msa, 1)) + unsqueeze( - shift_msa, 1) - return x, gate_msa - - -class AdaLayerNormContinuous(Module): - - def __init__(self, - embedding_dim: int, - conditioning_embedding_dim: int, - elementwise_affine: bool = True, - eps: float = 1e-5, - bias: bool = True, - norm_type: str = "layer_norm", - mapping=Mapping(), - dtype=None): - super().__init__() - self.silu = ACT2FN['silu'] - self.linear = Linear(conditioning_embedding_dim, - embedding_dim * 2, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - eps=eps, - elementwise_affine=elementwise_affine, - bias=bias, - dtype=dtype) - elif norm_type == "rms_norm": - self.norm = RmsNorm(embedding_dim, - eps=eps, - elementwise_affine=elementwise_affine, - dtype=dtype) - else: - raise ValueError(f"unknown norm_type {norm_type}") - - def forward(self, x: Tensor, conditioning_embedding: Tensor): - # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) - emb = self.linear(self.silu(conditioning_embedding).cast(x.dtype)) - scale, shift = chunk(emb, 2, dim=1) - x = self.norm(x) * unsqueeze((1 + scale), 1) + unsqueeze(shift, 1) - return x - - -class SD35AdaLayerNormZeroX(Module): - - def __init__(self, - embedding_dim: int, - norm_type: str = "layer_norm", - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - 9 * embedding_dim, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm'." - ) - - def forward(self, hidden_states: Tensor, emb: Tensor): - emb = self.linear(self.silu(emb).cast(hidden_states.dtype)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 = chunk( - emb, 9, dim=1) - norm_hidden_states = self.norm(hidden_states) - hidden_states = norm_hidden_states * ( - 1 + unsqueeze(scale_msa, 1)) + unsqueeze(shift_msa, 1) - norm_hidden_states2 = norm_hidden_states * ( - 1 + unsqueeze(scale_msa2, 1)) + unsqueeze(shift_msa2, 1) - return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 diff --git a/tensorrt_llm/layers/pooling.py b/tensorrt_llm/layers/pooling.py deleted file mode 100644 index 78153e6b1ef9..000000000000 --- a/tensorrt_llm/layers/pooling.py +++ /dev/null @@ -1,38 +0,0 @@ -# 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, Tuple - -from ..functional import avg_pool2d -from ..module import Module - - -class AvgPool2d(Module): - - def __init__(self, - kernel_size: Tuple[int], - stride: Optional[Tuple[int]] = None, - padding: Optional[Tuple[int]] = (0, 0), - ceil_mode: bool = False, - count_include_pad: bool = True) -> None: - super().__init__() - self.kernel_szie = kernel_size - self.stride = stride - self.padding = padding - self.ceil_mode = ceil_mode - self.count_include_pad = count_include_pad - - def forward(self, input): - return avg_pool2d(input, self.kernel_szie, self.stride, self.padding, - self.ceil_mode, self.count_include_pad) diff --git a/tensorrt_llm/layers/recurrent.py b/tensorrt_llm/layers/recurrent.py deleted file mode 100644 index fde8dc24eb24..000000000000 --- a/tensorrt_llm/layers/recurrent.py +++ /dev/null @@ -1,316 +0,0 @@ -# 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 - -import torch - -from .._utils import set_obj_attrs -from ..functional import Tensor, allgather, cast, concat, matmul, rg_lru, shape -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from .linear import ColumnLinear, RowLinear -from .ssm import MambaConv1d - - -class GroupedLinear(Module): - - def __init__(self, - in_features, - out_features, - num_blocks, - bias=True, - dtype=None, - use_fp8=False, - tp_group=None, - tp_size=1, - gather_output=True, - strict_dtype=False, - fuse_bias=False): - super().__init__() - assert in_features % num_blocks == 0 and out_features % num_blocks == 0 - assert num_blocks % tp_size == 0 - assert not (gather_output and fuse_bias) - self.in_features = in_features // tp_size - self.out_features = out_features // tp_size - self.num_blocks = num_blocks // tp_size - self.dtype = dtype - self.use_fp8 = use_fp8 - self.fuse_bias = fuse_bias - - self.weight = Parameter(shape=(self.num_blocks, - self.in_features // self.num_blocks, - self.out_features // self.num_blocks), - dtype=('fp8' if use_fp8 else dtype)) - set_obj_attrs(self.weight, { - "weight_loader": self.weight_loader, - }) - - self.tp_size = tp_size - self.tp_group = tp_group - self.gather_output = gather_output - self.strict_dtype = self.dtype if strict_dtype else None - - if bias: - self.bias = Parameter(shape=(self.num_blocks, - self.out_features // self.num_blocks), - dtype=dtype) - set_obj_attrs(self.bias, { - "weight_loader": self.weight_loader, - }) - else: - self.register_parameter('bias', None) - - def multiply_gather(self, x, weight): - grouped_shape = [] - out_shape = [] - ndim = x.ndim() - for i in range(x.ndim() - 1): - grouped_shape.append(shape(x, i)) - out_shape.append(shape(x, i)) - grouped_shape.extend( - [self.num_blocks, self.in_features // self.num_blocks]) - out_shape.append(self.out_features) - x = x.view(concat(grouped_shape)).permute([i for i in range(ndim - 2)] + - [-2, -3, -1]) - x = matmul(x, weight) - x = x.permute([i for i in range(ndim - 2)] + [-2, -3, -1]) - - if self.bias is not None and not self.fuse_bias: - bias = cast(self.bias.value, x.dtype) - x = x + bias - x = x.view(concat(out_shape)) - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=-1) - - return x - - def forward(self, x): - return self.multiply_gather(x, self.weight.value) - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor): - tp_rank = mapping.tp_rank - output_dim = 0 - shard_size = param._shape[output_dim] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) - param.value = loaded_weight - - -class RgLru(Module): - - def __init__(self, - lru_width, - num_heads=1, - dtype=None, - tp_group=None, - tp_size=1): - super().__init__() - self.lru_width = lru_width - self.dtype = dtype - self.num_heads = num_heads - self.tp_group = tp_group - self.tp_size = tp_size - - self.recurrent_param = Parameter(shape=(self.lru_width // - self.tp_size, ), - dtype=self.dtype) - self.input_gate = GroupedLinear(self.lru_width, - self.lru_width, - self.num_heads, - dtype=self.dtype, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - fuse_bias=True) - self.recurrent_gate = GroupedLinear(self.lru_width, - self.lru_width, - self.num_heads, - dtype=self.dtype, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - fuse_bias=True) - - def forward(self, - x: Tensor, - y: Tensor, - y_bias: Tensor, - lru_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - slot_mapping: Optional[Tensor] = None): - gate_x = self.input_gate(x) - gate_a = self.recurrent_gate(x) - out, lru_state = rg_lru(input=x, - gate_x=gate_x, - gate_x_bias=self.input_gate.bias.value, - gate_a=gate_a, - gate_a_bias=self.recurrent_gate.bias.value, - y=y, - y_bias=y_bias, - state_or_ptr=lru_state, - A=self.recurrent_param.value, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - dim=self.lru_width // self.tp_size, - dtype=self.dtype, - slot_mapping=slot_mapping) - return out, lru_state - - -class FusedRgLru(Module): - - def __init__(self, - lru_width, - num_heads=1, - dtype=None, - tp_group=None, - tp_size=1): - super().__init__() - self.lru_width = lru_width - self.tp_size = tp_size - self.dtype = dtype - self.dim = self.lru_width // self.tp_size - self.block_size = self.lru_width // num_heads - - self.recurrent_param = Parameter(shape=(self.lru_width // tp_size, ), - dtype=dtype) - self.gate = GroupedLinear(self.lru_width, - self.lru_width * 2, - num_heads, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - fuse_bias=True) - - def forward(self, - x: Tensor, - y: Tensor, - y_bias: Tensor, - lru_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - slot_mapping: Optional[Tensor] = None): - gate = self.gate(x) - out, lru_state = rg_lru(input=x, - gate=gate, - gate_bias=self.gate.bias.value, - block_size=self.block_size, - y=y, - y_bias=y_bias, - state_or_ptr=lru_state, - A=self.recurrent_param.value, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - dim=self.dim, - dtype=self.dtype, - slot_mapping=slot_mapping) - return out, lru_state - - -class Recurrent(Module): - - def __init__( - self, - width, - lru_width, - d_conv=4, - num_heads=1, - dtype=None, - tp_group=None, - tp_size=1, - ): - super().__init__() - self.width = width - self.lru_width = lru_width - self.d_conv = d_conv - self.dtype = dtype - - self.linear_x = ColumnLinear(self.width, - self.lru_width, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.linear_y = ColumnLinear(self.width, - self.lru_width, - bias=False, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.y_bias = Parameter(shape=(self.lru_width // tp_size, ), - dtype=dtype) - - self.conv1d = MambaConv1d(self.lru_width // tp_size, - self.d_conv, - dtype=self.dtype, - apply_silu=False) - - self.rg_lru = RgLru(self.lru_width, - num_heads=num_heads, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - self.linear_out = RowLinear(self.lru_width, - self.width, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - def forward(self, - hidden_states: Tensor, - conv_state: Tensor, - lru_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - hidden_states: [B, L, D] or [T, D] - conv_state: [B, W, D] or [1] of type int64 for paged state - lru_state: [B, N] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - # y branch - y = self.linear_y(hidden_states) - - # x branch - x = self.linear_x(hidden_states) - x_conv, conv_state = self.conv1d(x, conv_state, host_request_types, - last_token_ids, host_context_lengths, - slot_mapping, conv_indices) - - # rg-lru - out, lru_state = self.rg_lru(x_conv, y, self.y_bias.value, lru_state, - host_request_types, last_token_ids, - slot_mapping) - - # linear out - out = self.linear_out(out) - return out, conv_state, lru_state diff --git a/tensorrt_llm/layers/ssm.py b/tensorrt_llm/layers/ssm.py deleted file mode 100644 index 41c7a5b8f2f3..000000000000 --- a/tensorrt_llm/layers/ssm.py +++ /dev/null @@ -1,358 +0,0 @@ -# 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 math -from typing import Optional - -from .._common import default_net -from ..functional import (ACT2FN, Tensor, concat, conv2d, gather, mamba_conv1d, - permute, selective_scan, shape, split, view) -from ..module import Module -from ..parameter import Parameter -from .linear import ColumnLinear, Linear, RowLinear -from .normalization import RmsNorm - - -class MambaConv1d(Module): - - def __init__(self, - d_inner, - d_conv=4, - pre_stride=0, - post_stride=0, - dtype=None, - apply_silu=True): - super().__init__() - self.d_inner = d_inner - self.d_conv = d_conv - self.pre_stride = pre_stride - self.post_stride = post_stride - self.dtype = dtype - self.weight = Parameter(shape=(self.d_inner, 1, self.d_conv, 1), - dtype=dtype) - self.bias = Parameter(shape=(self.d_inner, ), dtype=dtype) - self.apply_silu = apply_silu - - def forward(self, - x: Tensor, - conv_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - x: [B, L, D] or [T, D] - conv_state: [B, W, D] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - if default_net().plugin_config.mamba_conv1d_plugin: - transposed_weight = permute( - view(self.weight.value, shape=[self.d_inner, 1, self.d_conv]), - (1, 2, 0)) - x_conv, conv_state = mamba_conv1d( - x, conv_state, transposed_weight, self.bias.value, - host_request_types, last_token_ids, self.d_inner, self.d_conv, - self.dtype, self.pre_stride, self.post_stride, - host_context_lengths, slot_mapping, self.apply_silu) - else: - assert not default_net().plugin_config.paged_state - assert len( - x.shape - ) == 3, "remove_input_padding is not supported by OOTB for Mamba." - if self.pre_stride > 0: - _, x = split(x, - [self.pre_stride, self.d_inner + self.post_stride], - dim=-1) - if self.post_stride > 0: - x, _ = split(x, [self.d_inner, self.post_stride], dim=-1) - x = x.permute([0, 2, 1]) - - # In context phase, conv_state is a zero tensor, and it is used for padding - # In generation phase, conv_state is a tensor of the past x - x_pad = concat([conv_state, x], dim=2) - - # Update conv_state - conv_state = gather(x_pad, 2, conv_indices) - - # Convolution - x_pad = x_pad.view( - concat([shape(x_pad, 0), - shape(x_pad, 1), - shape(x_pad, 2), 1])) - x_conv = conv2d(x_pad, - self.weight.value, - self.bias.value, - groups=self.d_inner) - if self.apply_silu: - x_conv = ACT2FN['silu'](x_conv) - x_conv = x_conv.view( - concat([shape(x_conv, 0), - shape(x_conv, 1), - shape(x_conv, 2)])) - - # Get dt, B and C - x_conv = x_conv.permute([0, 2, 1]) - return x_conv, conv_state - - -class Mamba(Module): - - def __init__(self, - d_model, - d_inner, - d_state=16, - d_conv=4, - dt_rank="auto", - bias=False, - dtype=None): - super().__init__() - self.d_model = d_model - self.d_state = d_state - self.d_conv = d_conv - self.d_inner = d_inner - self.dt_rank = math.ceil(self.d_model / - 16) if dt_rank == "auto" else dt_rank - self.dtype = dtype - - self.A = Parameter(shape=(self.d_state, self.d_inner), dtype="float32") - - self.D = Parameter(shape=(self.d_inner, ), dtype="float32") - self.dt_bias = Parameter(shape=(self.d_inner, ), dtype="float32") - - self.in_proj_x = Linear(self.d_model, - self.d_inner, - bias=bias, - dtype=dtype, - gather_output=False) - self.in_proj_z = Linear(self.d_model, - self.d_inner, - bias=bias, - dtype=dtype, - gather_output=False) - - self.conv1d = MambaConv1d(self.d_inner, self.d_conv, dtype=self.dtype) - - self.x_proj = Linear(self.d_inner, - self.dt_rank + self.d_state * 2, - bias=False, - dtype=dtype, - gather_output=False) - - self.dt_proj = Linear(self.dt_rank, - self.d_inner, - bias=False, - dtype=dtype, - gather_output=False, - pad_lda=self.d_state * 2) - - self.out_proj = Linear(self.d_inner, - self.d_model, - bias=bias, - dtype=dtype, - gather_output=False) - - def forward(self, - hidden_states: Tensor, - conv_state: Tensor, - ssm_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - hidden_states: [B, L, D] or [T, D] - conv_state: [B, W, D] or [1] of type int64 for paged state - ssm_state: [B, N, D] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - # in_proj - x = self.in_proj_x(hidden_states) - z = self.in_proj_z(hidden_states) - - x_conv, conv_state = self.conv1d(x, conv_state, host_request_types, - last_token_ids, host_context_lengths, - slot_mapping, conv_indices) - - # Get dt, B and C - x_dbl = self.x_proj(x_conv) - if default_net().plugin_config.gemm_plugin: - dt = self.dt_proj(x_dbl) - else: - dt, _ = split(x_dbl, [self.dt_rank, self.d_state * 2], dim=-1) - dt = self.dt_proj(dt) - - # selective scan - y, ssm_state = selective_scan(x_conv, - ssm_state, - dt, - self.dt_bias.value, - self.A.value, - x_dbl, - self.D.value, - host_request_types, - last_token_ids, - self.d_inner, - self.d_state, - self.dt_rank, - delta_softplus=True, - dtype=self.dtype, - z=z, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping) - # out_proj - out = self.out_proj(y) - return out, conv_state, ssm_state - - -class Mamba2(Module): - - def __init__(self, - d_model, - d_inner, - d_state=16, - d_conv=4, - headdim=64, - ngroups=1, - chunk_size=256, - bias=False, - rmsnorm=True, - dtype=None, - tp_group=None, - tp_size=1): - super().__init__() - self.d_model = d_model - self.d_state = d_state - self.d_conv = d_conv - assert d_inner % tp_size == 0 - self.d_inner = d_inner // tp_size - self.headdim = headdim - assert ngroups % tp_size == 0 - self.ngroups = ngroups // tp_size - self.chunk_size = chunk_size - self.rmsnorm = rmsnorm - self.dtype = dtype - assert d_inner % headdim == 0 - nheads = d_inner // headdim - assert nheads % tp_size == 0 - self.nheads = nheads // tp_size - # conv1d needs alignment to 8 fp16s - self.pad_ldc = (self.nheads + 7) // 8 * 8 - self.nheads - pad_ldc = self.pad_ldc * tp_size - - self.A = Parameter(shape=(self.nheads, ), dtype="float32") - self.D = Parameter(shape=(self.nheads, ), dtype="float32") - self.dt_bias = Parameter(shape=(self.nheads, ), dtype="float32") - - d_in_proj = 2 * d_inner + 2 * ngroups * d_state + nheads - self.in_proj = ColumnLinear(d_model, - d_in_proj, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - pad_ldc=pad_ldc) - - self.conv_dim = (d_inner + 2 * ngroups * d_state) // tp_size - self.conv1d = MambaConv1d(self.conv_dim, - self.d_conv, - pre_stride=self.d_inner, - post_stride=self.nheads + self.pad_ldc, - dtype=self.dtype) - - if rmsnorm: - self.norm = RmsNorm(normalized_shape=self.d_inner, - num_groups=self.ngroups, - eps=1e-5, - dtype=dtype) - - self.out_proj = RowLinear(d_inner, - d_model, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - def forward(self, - hidden_states: Tensor, - conv_state: Tensor, - ssm_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - hidden_states: [B, L, D] or [T, D] - conv_state: [B, W, D_conv] or [1] of type int64 for paged state - ssm_state: [B, H, N, D] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - # in_proj - zxbcdt = self.in_proj(hidden_states) - - # conv1d - xbc_conv, conv_state = self.conv1d(zxbcdt, conv_state, - host_request_types, last_token_ids, - host_context_lengths, slot_mapping, - conv_indices) - - # mamba scan - y, ssm_state = selective_scan(xbc_conv, - ssm_state, - zxbcdt, - self.dt_bias.value, - self.A.value, - xbc_conv, - self.D.value, - host_request_types, - last_token_ids, - self.d_inner, - self.d_state, - dt_rank=0, - delta_softplus=True, - dtype=self.dtype, - z=zxbcdt, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - nheads=self.nheads, - ngroups=self.ngroups, - chunk_size=self.chunk_size, - mamba_version='Mamba2') - - # norm - if self.rmsnorm: - y = self.norm(y) - - # out_proj - out = self.out_proj(y) - return out, conv_state, ssm_state diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 1682eb51e71c..019b97913d38 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -4,7 +4,6 @@ from ..executor import CompletionOutput, LoRARequest, RequestError from ..sampling_params import GuidedDecodingParams, SamplingParams from ..scheduling_params import SchedulingParams -from .build_cache import BuildCacheConfig from .llm import LLM, RequestOutput # yapf: disable from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType, @@ -24,9 +23,8 @@ SADecodingConfig, SAEnhancerConfig, SaveHiddenStatesDecodingConfig, SchedulerConfig, SkipSoftmaxAttentionConfig, TorchCompileConfig, - TorchLlmArgs, TrtLlmArgs, UserProvidedDecodingConfig) -from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo, - QuantConfig) + TorchLlmArgs, UserProvidedDecodingConfig) +from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder from .mpi_session import MpiCommSession from .thinking_budget import (ThinkingBudgetLogitsProcessor, @@ -56,11 +54,9 @@ 'MTPDecodingConfig', 'SchedulerConfig', 'CapacitySchedulerPolicy', - 'BuildConfig', 'QuantConfig', 'QuantAlgo', 'CalibConfig', - 'BuildCacheConfig', 'RequestError', 'MpiCommSession', 'ExtendedRuntimePerfKnobConfig', @@ -78,7 +74,6 @@ 'DraftTargetDecodingConfig', 'LlmArgs', 'TorchLlmArgs', - 'TrtLlmArgs', 'AutoDecodingConfig', 'AttentionDpConfig', 'LoRARequest', diff --git a/tensorrt_llm/llmapi/build_cache.py b/tensorrt_llm/llmapi/build_cache.py deleted file mode 100644 index e689145863c9..000000000000 --- a/tensorrt_llm/llmapi/build_cache.py +++ /dev/null @@ -1,312 +0,0 @@ -import contextlib -import datetime -import enum -import hashlib -import json -import os -import shutil -from dataclasses import dataclass -from pathlib import Path -from typing import Any, List, Optional - -import filelock -from pydantic import Field, model_validator - -import tensorrt_llm -from tensorrt_llm.builder import BuildConfig -from tensorrt_llm.llmapi.utils import (StrictBaseModel, enable_llm_debug, - print_colored) -from tensorrt_llm.logger import logger - - -def get_build_cache_config_from_env() -> tuple[bool, str]: - """ - Get the build cache configuration from the environment variables - """ - build_cache_enabled = os.environ.get('TLLM_LLMAPI_BUILD_CACHE') == '1' - build_cache_root = os.environ.get( - 'TLLM_LLMAPI_BUILD_CACHE_ROOT', - '/tmp/.cache/tensorrt_llm/llmapi/') # nosec B108 - return build_cache_enabled, build_cache_root - - -class BuildCacheConfig(StrictBaseModel): - """ - Configuration for the build cache. - - Note: - The build-cache assumes the weights of the model are not changed during the execution. If the weights are - changed, you should remove the caches manually. - """ - - cache_root: Optional[Path] = Field( - default=None, - description= - "The root directory for the build cache. Falls back to env var if not provided." - ) - max_records: int = Field( - default=10, - gt=0, - description="The maximum number of records to store in the cache.") - max_cache_storage_gb: float = Field( - default=256.0, - description="The maximum amount of storage (in GB) to use for the cache." - ) - - @model_validator(mode="after") - def set_default_cache_root(self) -> "BuildCacheConfig": - """Set cache_root from environment variable if not provided.""" - if self.cache_root is None: - _, default_cache_root = get_build_cache_config_from_env() - self.cache_root = Path(default_cache_root) - return self - - -class BuildCache: - """ - The BuildCache class is a class that manages the intermediate products from the build steps. - - NOTE: currently, only engine-building is supported - TODO[chunweiy]: add support for other build steps, such as quantization, convert_checkpoint, etc. - """ - # The version of the cache, will be used to determine if the cache is compatible - CACHE_VERSION = 0 - - def __init__(self, config: Optional[BuildCacheConfig] = None): - config = config or BuildCacheConfig() - self.cache_root = config.cache_root - self.max_records = config.max_records - self.max_cache_storage_gb = config.max_cache_storage_gb - - def free_storage_in_gb(self) -> float: - ''' Get the free storage capacity of the cache. ''' - # measure the root directory - if self.cache_root.parent.exists(): - usage = shutil.disk_usage(self.cache_root.parent) - return usage.free / 1024**3 - return 0 - - def get_engine_building_cache_stage(self, - build_config: BuildConfig, - model_path: Optional[Path] = None, - force_rebuild: bool = False, - **kwargs) -> 'CachedStage': - ''' - Get the build step for engine building. - ''' - build_config_str = json.dumps(self.prune_build_config_for_cache_key( - build_config.model_dump(mode="json")), - sort_keys=True) - - kwargs_str = json.dumps(kwargs, sort_keys=True) - - return CachedStage(parent=self, - kind=CacheRecord.Kind.Engine, - cache_root=self.cache_root, - force_rebuild=force_rebuild, - inputs=[build_config_str, model_path, kwargs_str]) - - def prune_caches(self, has_incoming_record: bool = False): - ''' - Clean up the cache records to make sure the cache size is within the limit - - Args: - has_incoming_record (bool): If the cache has incoming record, the existing records will be further pruned to - reserve space for the incoming record - ''' - if not self.cache_root.exists(): - return - self._clean_up_cache_dir() - records = [] - for dir in self.cache_root.iterdir(): - records.append(self._load_cache_record(dir)) - records.sort(key=lambda x: x.time, reverse=True) - max_records = self.max_records - 1 if has_incoming_record else self.max_records - # prune the cache to meet max_records and max_cache_storage_gb limitation - while len(records) > max_records or sum( - r.storage_gb for r in records) > self.max_cache_storage_gb: - record = records.pop() - # remove the directory and its content - shutil.rmtree(record.path) - - @staticmethod - def prune_build_config_for_cache_key(build_config: dict) -> dict: - black_list = ['dry_run'] - dic = build_config.copy() - for key in black_list: - if key in dic: - dic.pop(key) - return dic - - def load_cache_records(self) -> List["CacheRecord"]: - ''' - Load all the cache records from the cache directory - ''' - records = [] - if not self.cache_root.exists(): - return records - - for dir in self.cache_root.iterdir(): - records.append(self._load_cache_record(dir)) - return records - - def _load_cache_record(self, cache_dir) -> "CacheRecord": - ''' - Get the cache record from the cache directory - ''' - metadata = json.loads((cache_dir / 'metadata.json').read_text()) - storage_gb = sum(f.stat().st_size for f in cache_dir.glob('**/*') - if f.is_file()) / 1024**3 - return CacheRecord(kind=CacheRecord.Kind.__members__[metadata['kind']], - storage_gb=storage_gb, - path=cache_dir, - time=datetime.datetime.fromisoformat( - metadata['datetime'])) - - def _clean_up_cache_dir(self): - ''' - Clean up the files in the cache directory, remove anything that is not in the cache - ''' - # get all the files and directies in the cache_root - if not self.cache_root.exists(): - return - for file_or_dir in self.cache_root.iterdir(): - if not self.is_cache_valid(file_or_dir): - logger.info(f"Removing invalid cache directory {dir}") - if file_or_dir.is_file(): - file_or_dir.unlink() - else: - shutil.rmtree(file_or_dir) - - def is_cache_valid(self, cache_dir: Path) -> bool: - ''' - Check if the cache directory is valid - ''' - if not cache_dir.exists(): - return False - - metadata_path = cache_dir / 'metadata.json' - if not metadata_path.exists(): - return False - - metadata = json.loads(metadata_path.read_text()) - if metadata.get('version') != BuildCache.CACHE_VERSION: - return False - - content = cache_dir / 'content' - if not content.exists(): - return False - - return True - - -@dataclass -class CachedStage: - ''' - CachedStage is a class that represents a stage in the build process, it helps to manage the intermediate product. - - The cache is organized as follows: - - this_cache_dir/ # name is like "engine-" - metadata.json # the metadata of the cache - content/ # the actual product of the build step, such trt-llm engine directory - ''' - # The parent should be kept alive by CachedStep instance - parent: BuildCache - cache_root: Path - # The inputs will be used to determine if the step needs to be re-run, so all the variables should be put here - inputs: List[Any] - kind: "CacheRecord.Kind" - # If force_rebuild is set to True, the cache will be ignored - force_rebuild: bool = False - - def get_hash_key(self): - lib_version = tensorrt_llm.__version__ - input_strs = [str(i) for i in self.inputs] - return hashlib.md5( - f"{lib_version}-{input_strs}".encode()).hexdigest() # nosec B324 - - def get_cache_path(self) -> Path: - ''' - The path to the product of the build step, will be overwritten if the step is re-run - ''' - return self.cache_root / f"{self.kind.value}-{self.get_hash_key()}" - - def get_engine_path(self) -> Path: - return self.get_cache_path() / 'content' - - def get_cache_metadata(self) -> dict: - res = { - "version": BuildCache.CACHE_VERSION, - "datetime": datetime.datetime.now().isoformat(), - "kind": self.kind.name, - } - return res - - def is_cached(self) -> bool: - ''' - Check if the product of the build step is in the cache - ''' - if self.force_rebuild: - return False - try: - if self.get_cache_path().exists(): - metadata = json.loads( - (self.get_cache_path() / 'metadata.json').read_text()) - if metadata["version"] == BuildCache.CACHE_VERSION: - return True - except: - pass - - return False - - @contextlib.contextmanager - def write_guard(self): - ''' Guard the cache writing process. - - The cache writing process should be atomic, so the filelock is used to protect the cache writing process. And - the cache metadata will be written to the cache directory. - - Args: - final_engien_dir: the final engine directory - ''' - self.parent.prune_caches(has_incoming_record=True) - - target_dir = self.get_cache_path() - if enable_llm_debug(): - print_colored(f"Writing cache to {target_dir}\n", "yellow") - - # To avoid the cache modification conflict, a dummy directory is used to write the cache, and then rename it to - # the target directory - dummy_target_dir = Path(f"{target_dir.parent}/{target_dir.name}.dummy") - - dummy_target_dir.mkdir(parents=True, exist_ok=True) - # TODO[chunweiy]: deal with the cache modification conflict - lock = filelock.FileLock(dummy_target_dir / '.filelock', timeout=10) - - with open(dummy_target_dir / 'metadata.json', 'w') as f: - f.write(json.dumps(self.get_cache_metadata())) - - with lock: - yield dummy_target_dir / 'content' - - # If engine building is successful, rename the dummy directory to the target directory - if target_dir.exists(): - shutil.rmtree(target_dir) - shutil.move(dummy_target_dir, target_dir) - - -@dataclass(unsafe_hash=True) -class CacheRecord: - ''' - CacheRecord is a class that represents a record in the cache directory. - ''' - - class Kind(enum.Enum): - Engine = 'engine' - Checkpoint = 'checkpoint' - - kind: Kind - storage_gb: float - path: Path - time: datetime.datetime diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index cd19e015308a..2bf39e1335bf 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1,9 +1,7 @@ import atexit import json import os -import shutil import socket -import tempfile import time import weakref from collections.abc import Mapping @@ -24,9 +22,7 @@ from tensorrt_llm.metrics.enums import MetricNames from .._utils import nvtx_range_debug -from ..bindings import executor as tllm from ..bindings import steady_clock_now -from ..builder import EngineConfig from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..executor import (DetokenizedGenerationResultBase, GenerationExecutor, @@ -44,14 +40,12 @@ from ..logger import logger from ..sampling_params import LogitsProcessor, SamplingParams from ..scheduling_params import SchedulingParams -from .llm_args import (TORCH_LLMARGS_EXPLICIT_DOCSTRING, - TRT_LLMARGS_EXPLICIT_DOCSTRING, PeftCacheConfig, - PybindMirror, TorchLlmArgs, TrtLlmArgs) +from .llm_args import TORCH_LLMARGS_EXPLICIT_DOCSTRING, TorchLlmArgs from .llm_utils import (CachedModelLoader, KvCacheRetentionConfig, - LlmBuildStats, ModelLoader, _ModelRuntimeContext) + LlmBuildStats, ModelLoader) from .mpi_session import MpiPoolSession, external_mpi_comm_available from .thinking_budget import add_thinking_budget_logits_processor -from .tokenizer import TokenizerBase, _xgrammar_tokenizer_info +from .tokenizer import TokenizerBase # TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import from .utils import (append_docstring, exception_handler, get_device_count, logger_debug, set_api_status) @@ -203,15 +197,6 @@ def _contains_bart_forced_tokens_logits_processor(processor: Any) -> bool: return False -TRT_LLM_DOCSTRING = TRT_LLMARGS_EXPLICIT_DOCSTRING + """ - - Attributes: - tokenizer (tensorrt_llm.llmapi.tokenizer.TokenizerBase, optional): The tokenizer loaded by LLM instance, if any. - workspace (pathlib.Path): The directory to store intermediate files. - llm_id (str): The unique ID of the LLM instance. - disaggregated_params (dict): The disaggregated parameters of the LLM instance. -""" - TORCH_LLM_DOCSTRING = TORCH_LLMARGS_EXPLICIT_DOCSTRING + """ Attributes: @@ -279,8 +264,9 @@ def __init__(self, LlmArgs as AutoDeployLlmArgs llm_args_cls = AutoDeployLlmArgs else: - logger.info("Using LLM with TensorRT backend") - llm_args_cls = TrtLlmArgs + raise ValueError( + f"Unknown backend: {backend!r}. Supported backends are " + "'pytorch' and '_autodeploy'.") # check the kwargs and raise ValueError directly valid_keys = set( @@ -350,17 +336,12 @@ def __init__(self, self._executor: Optional[GenerationExecutor] = None self._encode_only: bool = False self._encoder_executor = None - if self._on_trt_backend: - self._workspace = tempfile.TemporaryDirectory( - suffix="-llm-workspace", dir=self.args.workspace) - else: - self._workspace = None + self._workspace = None self._hf_model_dir: Optional[Path] = None self._hf_model_config = None self._generation_config = None - self.runtime_context: Optional[_ModelRuntimeContext] = None self.llm_build_stats = LlmBuildStats() self._build_model() @@ -644,11 +625,10 @@ def generate_async( # With pytorch backend, py_executor has logic to handle max_tokens of 1, # so set to 1 to avoid allocating unnecessary KV cache blocks for single request - # TODO: Also support for trt backend is_ctx_only = disaggregated_params is not None and disaggregated_params.request_type == "context_only" is_gen_only = disaggregated_params is not None and disaggregated_params.request_type == "generation_only" - if is_ctx_only and not self._on_trt_backend: + if is_ctx_only: sampling_params.max_tokens = 1 if isinstance(inputs, PreprocessedInputs): @@ -1426,10 +1406,6 @@ def _build_model(self): self.llm_build_stats)) self._engine_dir, self._hf_model_dir = model_loader() - @property - def _on_trt_backend(self) -> bool: - return isinstance(self.args, TrtLlmArgs) - def _try_load_tokenizer(self) -> Optional[TokenizerBase]: if self.args.skip_tokenizer_init: return None @@ -1438,9 +1414,6 @@ def _try_load_tokenizer(self) -> Optional[TokenizerBase]: assert isinstance(self.args.tokenizer, TokenizerBase) return self.args.tokenizer - if self.runtime_context is not None: - return self.runtime_context.tokenizer - # TODO smor- need to refine what is the desired behavior if lora is enabled # in terms of the tokenizer initialization process if hasattr(self.args, "backend") and self.args.backend in [ @@ -1543,199 +1516,6 @@ def __del__(self): self.shutdown() -@append_docstring(TRT_LLM_DOCSTRING) -class _TrtLLM(BaseLLM): - """LLM class is the main class for running a LLM model using TensorRT LLM backend. - - Parameters: - """ - - def __init__(self, - model: Union[str, Path], - tokenizer: Optional[Union[str, Path, TokenizerBase, - PreTrainedTokenizerBase]] = None, - tokenizer_mode: Literal['auto', 'slow'] = 'auto', - skip_tokenizer_init: bool = False, - trust_remote_code: bool = False, - tensor_parallel_size: int = 1, - dtype: str = "auto", - revision: Optional[str] = None, - tokenizer_revision: Optional[str] = None, - **kwargs: Any) -> None: - super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init, - trust_remote_code, tensor_parallel_size, dtype, - revision, tokenizer_revision, **kwargs) - - @property - def workspace(self) -> Path: - return Path(self._workspace.name) if self._on_trt_backend else None - - def save(self, engine_dir: str) -> None: - """Save the built engine to the given path. - - Args: - engine_dir (str): The path to save the engine. - """ - logger.info(f"Save model to {engine_dir}") - if self._engine_dir is None: - raise RuntimeError("The engine is not built yet.") - - if self._engine_dir.absolute() == os.path.abspath(engine_dir): - return - - if not self.mpi_session or not self.mpi_session.is_comm_session(): - shutil.copytree(self._engine_dir, engine_dir, dirs_exist_ok=True) - else: - # NFS is fragile, so we copy files one by one - target_engine_dir = Path(engine_dir) - target_engine_dir.mkdir(parents=True, exist_ok=True) - # copy files one by one - for file in self._engine_dir.iterdir(): - logger_debug( - f"Copying {file} to {target_engine_dir / file.name}\n") - shutil.copy(file, target_engine_dir / file.name) - - def _build_model(self): - super()._build_model() - # update the model_dir to a local dir for the runtime, such as tokenizer loading. - if self._engine_dir is not None: - self.args.model = self._engine_dir - - # Tokenizer and config loading should be after calling model_loader(), since model_loader() may download the model from HF hub. - # It should also be before bindings ExecutorConfig, which may depend on tokenizer info. - self._tokenizer = self._try_load_tokenizer() - # Load HF config from the original HF model dir when available, - # since self.args.model now points to the engine dir (whose - # config.json uses TRT-LLM schema, not HF schema). - if self._hf_model_dir is not None: - self._hf_model_config = ModelLoader.load_hf_model_config( - self._hf_model_dir, - trust_remote_code=self.args.trust_remote_code) - else: - self._hf_model_config = self._try_load_hf_model_config() - self._generation_config = self._try_load_generation_config() - - # Multimodal special handling: - # 1. Default load_tokenizer may fail because MM has different tokenizer configuration. Hence we initialize it inside input processor - # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__ - self.input_processor = create_input_processor( - self._hf_model_dir, - self.tokenizer, - trust_remote_code=self.args.trust_remote_code) - self._tokenizer = self.input_processor.tokenizer - - max_batch_size = self.args.max_batch_size - max_num_tokens = self.args.max_num_tokens - max_seq_len = self.args.max_seq_len - - build_config = self.args.build_config - - max_batch_size = max_batch_size or build_config.max_batch_size - max_num_tokens = max_num_tokens or build_config.max_num_tokens - max_seq_len = max_seq_len or build_config.max_seq_len - - self._executor_config = tllm.ExecutorConfig( - max_beam_width=self.args.max_beam_width, - scheduler_config=PybindMirror.maybe_to_pybind( - self.args.scheduler_config), - batching_type=PybindMirror.maybe_to_pybind(self.args.batching_type) - or tllm.BatchingType.INFLIGHT, - max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - gather_generation_logits=self.args.gather_generation_logits, - fail_fast_on_attention_window_too_large=getattr( - self.args, 'fail_fast_on_attention_window_too_large', False)) - - # also set executor_config.max_seq_len in TRT workflow, to deduce default max_tokens - if max_seq_len is not None: - self._executor_config.max_seq_len = max_seq_len - else: - engine_config = EngineConfig.from_json_file(self._engine_dir / - "config.json") - self._executor_config.max_seq_len = engine_config.build_config.max_seq_len - - if self.args.kv_cache_config is not None: - self._executor_config.kv_cache_config = PybindMirror.maybe_to_pybind( - self.args.kv_cache_config) - if os.getenv("FORCE_DETERMINISTIC", "0") == "1": - # Disable KV cache reuse for deterministic mode - self._executor_config.kv_cache_config.enable_block_reuse = False - self._executor_config.kv_cache_config.enable_partial_reuse = False - if self.args.peft_cache_config is not None: - self._executor_config.peft_cache_config = PybindMirror.maybe_to_pybind( - self.args.peft_cache_config) - - lora_config = None - if self.args.build_config.plugin_config.lora_plugin: - engine_config = EngineConfig.from_json_file(self._engine_dir / - "config.json") - lora_config = engine_config.build_config.lora_config - if self.args.lora_config is not None: - logger.info( - "Overriding lora_config from engine with lora_config from LLM args" - ) - lora_config = self.args.lora_config - - max_lora_rank = lora_config.max_lora_rank - num_lora_modules = engine_config.pretrained_config.num_hidden_layers * \ - len(lora_config.lora_target_modules + lora_config.missing_qkv_modules) - - peft_cache_config_model = PeftCacheConfig.from_pybind( - self._executor_config.peft_cache_config - ) if self._executor_config.peft_cache_config is not None else PeftCacheConfig( - ) - if lora_config.max_loras is not None: - peft_cache_config_model.num_device_module_layer = \ - max_lora_rank * num_lora_modules * lora_config.max_loras - if lora_config.max_cpu_loras is not None: - peft_cache_config_model.num_host_module_layer = \ - max_lora_rank * num_lora_modules * lora_config.max_cpu_loras - self._executor_config.peft_cache_config = peft_cache_config_model._to_pybind( - ) - - if self.args.decoding_config is not None: - self._executor_config.decoding_config = self.args.decoding_config - if self.args.guided_decoding_backend == 'xgrammar': - self._executor_config.guided_decoding_config = tllm.GuidedDecodingConfig( - backend=tllm.GuidedDecodingConfig.GuidedDecodingBackend. - XGRAMMAR, - **_xgrammar_tokenizer_info(self.tokenizer)) - elif self.args.guided_decoding_backend is not None: - raise ValueError( - f"Unsupported guided decoding backend {self.args.guided_decoding_backend}" - ) - - self._executor_config.normalize_log_probs = self.args.normalize_log_probs - self._executor_config.enable_chunked_context = self.args.enable_chunked_prefill - self._executor_config.max_beam_width = self.args.max_beam_width or self.args.build_config.max_beam_width - if self.args.extended_runtime_perf_knob_config is not None: - self._executor_config.extended_runtime_perf_knob_config = PybindMirror.maybe_to_pybind( - self.args.extended_runtime_perf_knob_config) - if self.args.cache_transceiver_config is not None: - self._executor_config.cache_transceiver_config = PybindMirror.maybe_to_pybind( - self.args.cache_transceiver_config) - self._executor_config.llm_parallel_config = self.args.parallel_config - return_logits = (self.args.gather_generation_logits - or (self.args.build_config - and self.args.build_config.gather_context_logits)) - - self._executor = self._executor_cls.create( - self._engine_dir, - executor_config=self._executor_config, - batched_logits_processor=self.args.batched_logits_processor, - model_world_size=self.args.parallel_config.world_size, - mpi_session=self.mpi_session, - reuse_mpi_comm=external_mpi_comm_available( - self.args.parallel_config.world_size), - return_logits=return_logits, - postproc_worker_config=PostprocWorkerConfig( - num_postprocess_workers=self.args.num_postprocess_workers, - postprocess_tokenizer_dir=self.args.postprocess_tokenizer_dir, - post_processor_hook=self.args.post_processor_hook, - ), - is_llm_executor=True) - - @append_docstring(TORCH_LLM_DOCSTRING) class _TorchLLM(BaseLLM): """LLM class is the main class for running a LLM model using PyTorch backend. @@ -1758,7 +1538,7 @@ def __init__(self, backend = kwargs.pop("backend", "pytorch") - # Validate that users don't pass TrtLlmArgs-specific arguments + # Validate that only arguments supported by the PyTorch backend are passed. self._validate_args_for_torch_backend(kwargs) super().__init__(model, @@ -1888,23 +1668,20 @@ def _build_model(self): llm_args=self.args) def _validate_args_for_torch_backend(self, kwargs: dict) -> None: - """Validate that users don't pass TrtLlmArgs-specific arguments when using PyTorch backend. + """Validate that only arguments supported by the PyTorch backend are passed. """ - trtllm_fields = set(TrtLlmArgs.model_fields.keys()) torchllm_fields = set(TorchLlmArgs.model_fields.keys()) - trtllm_specific_fields = trtllm_fields - torchllm_fields - - # Check if any TrtLlmArgs-specific arguments are passed - trtllm_specific_args = [] - for key in kwargs: - if key in trtllm_specific_fields: - trtllm_specific_args.append(key) + # Check if any arguments not supported by the PyTorch backend are passed. + unsupported_args = [ + key for key in kwargs + if key not in torchllm_fields and key not in ('_mpi_session', + 'backend') + ] - if trtllm_specific_args: + if unsupported_args: raise ValueError( - f"The following arguments are specific to TensorRT backend and cannot be used with PyTorch backend: {trtllm_specific_args}.\n" - f"Please use 'from tensorrt_llm._tensorrt_engine import LLM' instead to use the TensorRT backend." + f"The following arguments are not supported by the PyTorch backend: {unsupported_args}." ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fd3276163831..935cad900b29 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -70,16 +70,12 @@ # isort: on # yapf: enable -from ..builder import BuildConfig, EngineConfig from ..logger import logger from ..mapping import CpType, Mapping -from ..models.automodel import AutoConfig -from ..models.modeling_utils import (PretrainedConfig, QuantAlgo, QuantConfig, - SpeculativeDecodingMode) +from ..models.modeling_utils import QuantAlgo, QuantConfig from ..sampling_params import BatchedLogitsProcessor from ..usage.config import UsageContext # noqa: F401 from ..usage.config import TelemetryConfig, TelemetryField -from .build_cache import BuildCacheConfig from .tokenizer import TokenizerBase, tokenizer_factory from .utils import (StrictBaseModel, generate_api_docs_as_docstring, get_type_repr) @@ -1557,12 +1553,6 @@ class CalibConfig(StrictBaseModel): "The maximum sequence length to initialize tokenizer for calibration.") -class _ModelFormatKind(Enum): - HF = 0 - TLLM_CKPT = 1 - TLLM_ENGINE = 2 - - class DecodingBaseConfig(StrictBaseModel): max_draft_len: Optional[NonNegativeInt] = Field( default=None, description="The maximum number of draft tokens.") @@ -3797,7 +3787,7 @@ class DwdpConfig(StrictBaseModel): class BaseLlmArgs(StrictBaseModel): - """Base class for both TorchLlmArgs and TrtLlmArgs. It contains all the arguments that are common to both.""" + """Base class for the LLM arguments. It contains all the common arguments.""" model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") # Explicit arguments @@ -4069,8 +4059,7 @@ class BaseLlmArgs(StrictBaseModel): exclude_json_schema=True, # hide from API references validate_default=True, status="deprecated", - telemetry=TelemetryField.categorical('pytorch', 'tensorrt', - '_autodeploy')) + telemetry=TelemetryField.categorical('pytorch', '_autodeploy')) return_perf_metrics: bool = Field(default=False, description="Return perf metrics.", @@ -4121,16 +4110,11 @@ def coerce_env_overrides_to_str(cls, v): return {str(k): str(val) for k, val in v.items()} _parallel_config: Optional[_ParallelConfig] = PrivateAttr(default=None) - _model_format: Optional[_ModelFormatKind] = PrivateAttr(default=None) @property def parallel_config(self) -> _ParallelConfig: return self._parallel_config - @property - def model_format(self) -> _ModelFormatKind: - return self._model_format - @property def speculative_model(self) -> Optional[Union[str, Path]]: return self.speculative_config.speculative_model if self.speculative_config is not None else None @@ -4317,330 +4301,6 @@ def get_runtime_sizes(self, ) -> Tuple[int, int, int, int]: ) -class TrtLlmArgs(BaseLlmArgs): - enable_tqdm: bool = Field(default=False, - description="Enable tqdm for progress bar.") - - workspace: Optional[str] = Field(default=None, - description="The workspace for the model.") - - fail_fast_on_attention_window_too_large: bool = Field( - default=True, - description= - "Fail fast when attention window is too large to fit even a single sequence in the KV cache.", - status="deprecated") - - # Once set, the model will reuse the build_cache - enable_build_cache: Union[BuildCacheConfig, - bool] = Field(default=False, - description="Enable build cache.") - - extended_runtime_perf_knob_config: Optional[ - ExtendedRuntimePerfKnobConfig] = Field( - default=None, description="Extended runtime perf knob config.") - - # Quantization and calibration configurations - calib_config: CalibConfig = Field(default_factory=CalibConfig, - description="Calibration config.") - - quant_config: QuantConfig = Field(default_factory=QuantConfig, - description="Quantization config.") - - embedding_parallel_mode: Literal[ - 'NONE', 'SHARDING_ALONG_VOCAB', 'SHARDING_ALONG_HIDDEN'] = Field( - default='SHARDING_ALONG_VOCAB', - description="The embedding parallel mode.") - - fast_build: bool = Field(default=False, description="Enable fast build.") - - # BuildConfig is introduced to give users a familiar interface to configure the model building. - build_config: Optional[BuildConfig] = Field(default=None, - description="Build config.") - - # Prompt adapter arguments - enable_prompt_adapter: bool = Field(default=False, - description="Enable prompt adapter.") - - max_prompt_adapter_token: int = Field( - default=0, description="The maximum number of prompt adapter tokens.") - - batching_type: Optional[BatchingType] = Field(default=None, - description="Batching type.") - - normalize_log_probs: bool = Field( - default=False, description="Normalize log probabilities.") - - # Private attributes - # This is used to hold the options for convert_checkpoint - _convert_checkpoint_options: Dict[str, - Any] = PrivateAttr(default_factory=dict) - - @model_validator(mode="after") - def init_build_config(self): - """Creating a default BuildConfig if none is provided""" - build_config = getattr(self, "build_config", None) - if build_config is None: - kwargs = {} - if self.max_batch_size: - kwargs["max_batch_size"] = self.max_batch_size - if self.max_num_tokens: - kwargs["max_num_tokens"] = self.max_num_tokens - if self.max_seq_len: - kwargs["max_seq_len"] = self.max_seq_len - if self.max_beam_width: - kwargs["max_beam_width"] = self.max_beam_width - if self.max_input_len: - kwargs["max_input_len"] = self.max_input_len - self.build_config = BuildConfig(**kwargs) - return self - - @model_validator(mode="after") - def validate_build_config_remaining(self): - is_trt_llm_args = isinstance(self, TrtLlmArgs) - - # TODO: remove the checker when manage weights support all data types - if is_trt_llm_args and self.fast_build and (self.quant_config.quant_algo - is QuantAlgo.FP8): - self.build_config.plugin_config.manage_weights = True - - if self.parallel_config.world_size == 1 and self.build_config: - self.build_config.plugin_config.nccl_plugin = None - - if self.enable_lora and self.backend != 'pytorch': - self.build_config.plugin_config.lora_plugin = 'auto' - if self.lora_config is not None: - self.build_config.lora_config.max_lora_rank = self.lora_config.max_lora_rank - - if hasattr(self, - 'enable_prompt_adapter') and self.enable_prompt_adapter: - self.build_config.max_prompt_embedding_table_size = self.max_prompt_adapter_token * self.build_config.max_batch_size - - return self - - @model_validator(mode="after") - def validate_speculative_config(self): - if self.speculative_config: - if not self.speculative_config.supports_backend(self.backend): - raise ValueError( - f"Speculation type {self.speculative_config.decoding_type} does not " - f"support backend {self.backend}") - - # Below, we only need to set speculative_decoding_mode/decoding_config for speculation - # on the TRT backend. - if isinstance(self.speculative_config, LookaheadDecodingConfig): - max_draft_len = self.speculative_config.calculate_speculative_resource( - )[2] - assert max_draft_len > 0 - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.LOOKAHEAD_DECODING - self.build_config.max_draft_len = max( - self.build_config.max_draft_len, max_draft_len) - self.decoding_config = DecodingConfig( - decoding_mode=DecodingMode.Lookahead(), - lookahead_decoding_config=PybindMirror.maybe_to_pybind( - self.speculative_config)) - - elif isinstance(self.speculative_config, MedusaDecodingConfig): - assert self.speculative_config.max_draft_len > 0 - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.MEDUSA - self.build_config.max_draft_len = self.speculative_config.max_draft_len - self.decoding_config = DecodingConfig( - decoding_mode=DecodingMode.Medusa(), - medusa_choices=self.speculative_config.medusa_choices) - - elif isinstance(self.speculative_config, Eagle3DecodingConfig): - raise ValueError( - "speculative_config.decoding_type 'Eagle3' is only supported on the PyTorch backend. " - "Use decoding_type 'Eagle' for the TensorRT backend.") - - elif isinstance(self.speculative_config, EagleDecodingConfig): - assert self.speculative_config.max_draft_len > 0 - assert self.speculative_config.speculative_model is not None, "EAGLE draft model must be specified." - self.build_config.max_draft_len = self.speculative_config.max_draft_len - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.EAGLE - eagle_config = _EagleConfig( - self.speculative_config.eagle_choices, - self.speculative_config.greedy_sampling, - self.speculative_config.posterior_threshold, - self.speculative_config.use_dynamic_tree, - self.speculative_config.dynamic_tree_max_topK) - self.decoding_config = DecodingConfig( - decoding_mode=DecodingMode.Eagle(), - eagle_config=eagle_config) - elif isinstance(self.speculative_config, PARDDecodingConfig): - raise ValueError( - "speculative_config.decoding_type 'PARD' is only supported on the PyTorch backend." - ) - elif isinstance(self.speculative_config, DFlashDecodingConfig): - raise ValueError( - "speculative_config.decoding_type 'DFlash' is only supported on the PyTorch backend." - ) - else: - raise ValueError( - f"Unrecognized speculative config type {type(self.speculative_config)}" - ) - - else: - self.decoding_config = None - - return self - - def _load_config_from_engine(self, engine_dir: Path): - engine_config = EngineConfig.from_json_file(engine_dir / "config.json") - self._pretrained_config = engine_config.pretrained_config - self.build_config = engine_config.build_config - - # load and check parallel_config - mapping = self._pretrained_config.mapping - if self.parallel_config.tp_size not in (1, mapping.tp_size): - raise ValueError( - f"tp_size {self.parallel_config.tp_size} is not consistent with the engine's tp_size {mapping.tp_size}" - ) - if self.parallel_config.pp_size not in (1, mapping.pp_size): - raise ValueError( - f"pp_size {self.parallel_config.pp_size} is not consistent with the engine's pp_size {mapping.pp_size}" - ) - if self.parallel_config.cp_size not in (1, mapping.cp_size): - raise ValueError( - f"cp_size {self.parallel_config.cp_size} is not consistent with the engine's cp_size {mapping.cp_size}" - ) - self._parallel_config = _ParallelConfig( - tp_size=mapping.tp_size, - pp_size=mapping.pp_size, - cp_size=mapping.cp_size, - gpus_per_node=mapping.gpus_per_node, - moe_cluster_size=mapping.moe_cluster_size, - moe_tp_size=mapping.moe_tp_size, - moe_ep_size=mapping.moe_ep_size) - - def _load_config_from_ckpt(self, ckpt_dir: Path): - pretrained_config = PretrainedConfig.from_json_file(ckpt_dir / - "config.json") - tp_size = pretrained_config.mapping.tp_size - pp_size = pretrained_config.mapping.pp_size - cp_size = pretrained_config.mapping.cp_size - moe_cluster_size = pretrained_config.mapping.moe_cluster_size - moe_tp_size = pretrained_config.mapping.moe_tp_size - moe_ep_size = pretrained_config.mapping.moe_ep_size - gpus_per_node = pretrained_config.mapping.gpus_per_node - # load parallel_config - if self.parallel_config.tp_size != 1 and self.parallel_config.tp_size != tp_size: - raise ValueError( - f"tp_size {self.parallel_config.tp_size} is not consistent with the checkpoint's tp_size {tp_size}" - ) - if self.parallel_config.pp_size != 1 and self.parallel_config.pp_size != pp_size: - raise ValueError( - f"pp_size {self.parallel_config.pp_size} is not consistent with the checkpoint's pp_size {pp_size}" - ) - if self.parallel_config.cp_size != 1 and self.parallel_config.cp_size != cp_size: - raise ValueError( - f"cp_size {self.parallel_config.cp_size} is not consistent with the checkpoint's cp_size {cp_size}" - ) - self._parallel_config = _ParallelConfig( - tp_size=tp_size, - pp_size=pp_size, - cp_size=cp_size, - gpus_per_node=gpus_per_node, - moe_cluster_size=moe_cluster_size, - moe_tp_size=moe_tp_size, - moe_ep_size=moe_ep_size) - - @model_validator(mode="after") - def validate_model_format_misc(self): - """Load the model format, and do the following: - - 1. Load the build_config if got an engine. - 2. Load the parallel_config if got a checkpoint. - """ - model_obj = _ModelWrapper(self.model) - - if model_obj.is_local_model and self.backend not in [ - 'pytorch', '_autodeploy' - ]: - # Load parallel_config from the engine. - model_format = get_model_format( - self.model, trust_remote_code=self.trust_remote_code) - - if model_format is _ModelFormatKind.TLLM_ENGINE: - if self.build_config is not None: - logger.warning( - "The build_config is ignored for model format of TLLM_ENGINE." - ) - self._load_config_from_engine(model_obj.model_dir) - runtime_defaults = self._pretrained_config.runtime_defaults - if runtime_defaults: - self.kv_cache_config.fill_empty_fields_from_runtime_defaults( - runtime_defaults) - - # Load parallel_config from the checkpoint. - elif model_format is _ModelFormatKind.TLLM_CKPT: - # We need to create a temporary instance to call _load_config_from_ckpt - self._load_config_from_ckpt(model_obj.model_dir) - else: - model_format = _ModelFormatKind.HF - - # Store the model format in the values - self._model_format = model_format - return self - - @model_validator(mode="after") - def validate_build_config_with_runtime_params(self): - """Sync runtime parameters with build_config limits. - - This validator runs AFTER validate_model_format_misc so that when - loading from an engine, we have the real build_config loaded. - """ - if self.build_config is None: - raise ValueError("build_config is not initialized") - - # These can be lower than build_config limits - for field in ("max_batch_size", "max_num_tokens"): - runtime_val = getattr(self, field) - build_val = getattr(self.build_config, field) - if runtime_val is not None and runtime_val > build_val: - logger.warning( - f"{field} [{runtime_val}] clamped to build_config.{field} [{build_val}]" - ) - setattr(self, field, build_val) - - # These must match build_config exactly - for field in ("max_seq_len", "max_beam_width", "max_input_len"): - runtime_val = getattr(self, field) - build_val = getattr(self.build_config, field) - if runtime_val is not None and runtime_val != build_val: - logger.warning( - f"{field} [{runtime_val}] overridden by build_config.{field} [{build_val}]" - ) - setattr(self, field, build_val) - - return self - - @model_validator(mode="after") - def setup_embedding_parallel_mode(self): - if self.embedding_parallel_mode == 'NONE': - self._convert_checkpoint_options['use_parallel_embedding'] = False - elif self.embedding_parallel_mode == 'SHARDING_ALONG_VOCAB': - self._convert_checkpoint_options['use_parallel_embedding'] = True - self._convert_checkpoint_options['embedding_sharding_dim'] = 0 - elif self.embedding_parallel_mode == 'SHARDING_ALONG_HIDDEN': - self._convert_checkpoint_options['use_parallel_embedding'] = True - self._convert_checkpoint_options['embedding_sharding_dim'] = 1 - # No else clause needed since validation already happened - return self - - @model_validator(mode="after") - def validate_enable_build_cache(self): - if not self.enable_build_cache: - return self - self.enable_build_cache = BuildCacheConfig() if isinstance( - self.enable_build_cache, bool) else self.enable_build_cache - return self - - @model_validator(mode="after") - def validate_kv_cache_dtype(self): - assert self.kv_cache_config.dtype == "auto", "KvCacheConfig.dtype is not supported by the TensorRT backend." - return self - - class LoadFormat(Enum): AUTO = 0 # Initialize all weights randomly. @@ -5260,11 +4920,6 @@ def extra_resource_managers(self) -> Dict[str, object]: def extra_resource_managers(self, value: Dict[str, object]) -> None: self._extra_resource_managers = value - @model_validator(mode="after") - def set_model_format(self): - self._model_format = _ModelFormatKind.HF - return self - @model_validator(mode="after") def validate_encoder_modes(self) -> 'TorchLlmArgs': if self.encode_only and self.mm_encoder_only: @@ -5682,16 +5337,6 @@ def update_llm_args_with_extra_dict( "kv_cache_dtype": "dtype", "enable_block_reuse": "enable_block_reuse", } - # Scalars that live both at the top level of LlmArgs and inside - # `build_config`. The build_config patch propagates the winning source - # to the nested location. - build_config_dual_loc_keys = ( - "max_batch_size", - "max_num_tokens", - "max_beam_width", - "max_seq_len", - ) - explicit_cli_keys = explicit_cli_keys or set() if 'hf_revision' in llm_args_dict: @@ -5765,9 +5410,7 @@ def update_llm_args_with_extra_dict( field_mapping = { "quant_config": QuantConfig, "calib_config": CalibConfig, - "build_config": BuildConfig, "decoding_config": DecodingConfig, - "enable_build_cache": BuildCacheConfig, "lora_config": LoraConfig, "moe_config": MoeConfig, "nvfp4_gemm_config": Nvfp4GemmConfig, @@ -5787,36 +5430,6 @@ def update_llm_args_with_extra_dict( llm_args = llm_args | llm_args_dict - # build_config only works for TensorRT backend, it will be ignored in PyTorch backend - if "build_config" in llm_args: - # Ensure build_config is a BuildConfig object, not a dict - if isinstance(llm_args["build_config"], dict): - llm_args["build_config"] = BuildConfig(**llm_args["build_config"]) - - # Propagate dual-location scalars into build_config: explicit CLI flag - # wins; otherwise YAML's top-level scalar; otherwise leave alone. Warn - # only when the explicit CLI value actually differs from the YAML - # build_config value being replaced (a genuine override). - for key in build_config_dual_loc_keys: - if key in explicit_cli_keys and key in llm_args: - # Warn only on a genuine override of a YAML build_config value; - # otherwise just record where the value came from. - if getattr(llm_args["build_config"], key) != llm_args[key]: - logger.warning( - f"Explicit CLI flag --{key}={llm_args[key]} overrides " - f"the value set in the YAML build_config; CLI takes " - f"precedence.") - else: - logger.info( - f"build_config.{key} set to {llm_args[key]} from explicit CLI flag" - ) - setattr(llm_args["build_config"], key, llm_args[key]) - elif key in llm_args_dict: - setattr(llm_args["build_config"], key, llm_args_dict[key]) - logger.info( - f"build_config.{key} set to {llm_args_dict[key]} from YAML top-level scalar" - ) - return llm_args @@ -5835,40 +5448,8 @@ def update_llm_args_with_extra_options( return llm_args -def get_model_format(model_dir: str, - trust_remote_code: bool = False) -> _ModelFormatKind: - """Get the format of the model.""" - if not (Path(model_dir) / 'config.json').exists(): - raise ValueError( - f"Failed to infer model format because no config.json exists in {model_dir}" - ) - - with open(Path(model_dir) / 'config.json') as f: - config = json.load(f) - - try: - if 'pretrained_config' in config and 'build_config' in config: - model_format = _ModelFormatKind.TLLM_ENGINE - EngineConfig.from_json_file(Path(model_dir) / 'config.json') - elif 'architecture' in config and 'dtype' in config: - model_format = _ModelFormatKind.TLLM_CKPT - PretrainedConfig.from_checkpoint(model_dir) - else: - model_format = _ModelFormatKind.HF - AutoConfig.from_hugging_face(model_dir, - trust_remote_code=trust_remote_code) - except Exception as e: - raise ValueError( - f"Inferred model format {model_format}, but failed to load config.json: {e}" - ) - else: - return model_format - - LlmArgs = TorchLlmArgs -TRT_LLMARGS_EXPLICIT_DOCSTRING = generate_api_docs_as_docstring(TrtLlmArgs, - indent=' ' * 4) TORCH_LLMARGS_EXPLICIT_DOCSTRING = generate_api_docs_as_docstring(TorchLlmArgs, indent=' ' * 4) diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 154f60a4b8ef..57b485ccbe5c 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -1,39 +1,26 @@ import json import os -import shutil import tempfile -import time import weakref -from dataclasses import asdict, dataclass, field, is_dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple, Union -import torch import transformers -from pydantic import BaseModel -from tqdm import tqdm -from .._utils import (global_mpi_rank, local_mpi_rank, mpi_barrier, - mpi_broadcast, mpi_rank, release_gc) +from .._utils import global_mpi_rank, local_mpi_rank, mpi_rank # yapf: disable from ..bindings.executor import (BatchingType, CapacitySchedulerPolicy, ContextChunkingPolicy, ExecutorConfig, KvCacheRetentionConfig) # yapf: enable -from ..builder import BuildConfig, Engine, build -from ..llmapi.llm_args import TrtLlmArgs from ..logger import logger -from ..mapping import Mapping -from ..models.automodel import MODEL_MAP, AutoConfig, AutoModelForCausalLM -from ..models.modeling_utils import PretrainedConfig, QuantAlgo, QuantConfig +from ..models.modeling_utils import QuantAlgo, QuantConfig from ..models.quant_config_utils import \ update_quant_config_from_compressed_tensors -from ..module import Module from ..quantization.modelopt_config import (is_modelopt_quant_config, read_modelopt_quant_config, warn_if_inline_diverges) -from .build_cache import (BuildCache, BuildCacheConfig, CachedStage, - get_build_cache_config_from_env) # yapf: disable from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, @@ -41,67 +28,14 @@ KvCacheConfig, LlmArgs, LookaheadDecodingConfig, MedusaDecodingConfig, MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig, TorchLlmArgs, - UserProvidedDecodingConfig, _ModelFormatKind, - _ModelWrapper, _ParallelConfig, - update_llm_args_with_extra_dict, + UserProvidedDecodingConfig, _ModelWrapper, + _ParallelConfig, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) # yapf: enable -from .mpi_session import MPINodeState, MpiSession +from .mpi_session import MpiSession from .tokenizer import TransformersTokenizer, load_hf_tokenizer # TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import -from .utils import (download_hf_model, download_hf_pretrained_config, - enable_llm_debug, get_directory_size_in_gb, logger_debug, - print_colored, print_traceback_on_error) - - -@dataclass -class _ModelInfo: - dtype: Optional[str] = None - architecture: Optional[str] = None - - @property - def model_name(self) -> str: - if self.architecture is None: - raise RuntimeError("The architecture is not set yet.") - - return self.architecture - - @classmethod - def from_pretrained_config(cls, config: PretrainedConfig): - return cls(dtype=config.dtype, architecture=config.architecture) - - @classmethod - def from_builder_config_json(cls, config: dict): - if 'version' in config: - # The Dict format is { 'builder_config':..., 'plugin_config':...} - dtype = config['plugin_config']['gpt_attention_plugin'] - else: - dtype = config['pretrained_config']['dtype'] - - return cls(dtype=dtype, architecture=config['builder_config']['name']) - - @classmethod - def from_module(cls, module: Module): - raise NotImplementedError() - - -@dataclass -class _ModelRuntimeContext: - """_ModelRuntimeContext holds the minimum runtime resources for running a model. - - It could be a runtime cache in MPI nodes. - """ - engine: Optional[Engine] = None - mapping: Optional[Mapping] = None - model_info: Optional[_ModelInfo] = None - - # This is only used when build-cache is enabled - engine_path: Optional[str] = None - - @property - def model_arch(self) -> str: - # "LlamaForCausalLM" or "OPTForCausalLM" and so on - return self.engine.config.pretrained_config['architecture'] +from .utils import download_hf_model, print_traceback_on_error class ModelLoader: @@ -123,208 +57,16 @@ def __init__(self, self.llm_args.speculative_model ) if self.llm_args.speculative_model is not None else None - if isinstance(self.llm_args, TrtLlmArgs): - self.convert_checkpoint_options = self.llm_args._convert_checkpoint_options self.rank = mpi_rank() self.global_rank = global_mpi_rank() self.mapping = llm_args.parallel_config.to_mapping() - self._build_pipeline = [] - # For model from hub, the _model_dir is None, and will updated once downloaded self._model_dir: Optional[ Path] = self.model_obj.model_dir if self.model_obj.is_local_model else None self._speculative_model_dir: Optional[ Path] = self.speculative_model_obj.model_dir if self.speculative_model_obj is not None and self.speculative_model_obj.is_local_model else None - self._model_info: Optional[_ModelInfo] = None - self._model_format = self.llm_args.model_format - - if isinstance(self.llm_args, TrtLlmArgs): - assert self.llm_args.build_config - self.build_config = self.llm_args.build_config - - self._gather_build_steps() - - def _gather_build_steps(self): - # Prepare the model processing pipeline - if isinstance(self.llm_args.model, Module): - # Build engine from user provided model - self._build_pipeline.append( - ("Build TensorRT LLM engine", - self._build_engine_from_inmemory_model)) - return - - if (self.model_obj.is_hub_model - and self._model_format is not _ModelFormatKind.TLLM_ENGINE): - # Download HF model if necessary - if self.model_obj.model_name is None: - raise ValueError( - "Either model_dir or model should be provided to ModelConfig." - ) - self._build_pipeline.append( - ("Downloading HF model", self._download_hf_model)) - - if self._model_format is _ModelFormatKind.HF: - # HF -> TRT checkpoints -> engine - self._build_pipeline.append( - ("Loading HF model to memory", self._load_model_from_hf)) - self._build_pipeline.append( - ("Building TRT-LLM engine", self._build_engine)) - elif self._model_format is _ModelFormatKind.TLLM_CKPT: - # TRT checkpoints -> engine - self._build_pipeline.append(("Loading TRT checkpoints to memory", - self._load_model_from_ckpt)) - self._build_pipeline.append( - ("Build TRT-LLM engine", self._build_engine)) - elif self._model_format is _ModelFormatKind.TLLM_ENGINE: - # Nothing need to do - pass - else: - raise ValueError(f"Unknown model format {self._model_format}") - - class BuildPipeline: - - def __init__(self, enable_tqdm: bool, labels: List[str], - step_handlers: List[Callable], - llm_build_stats: "LlmBuildStats"): - assert len(labels) == len(step_handlers) - self.labels = labels - self.step_handlers = step_handlers - self.llm_build_stats = llm_build_stats - - self.to_log = mpi_rank() == 0 - self.counter = 0 - - self.progress_bar = tqdm( - total=len(labels)) if enable_tqdm and self.to_log else None - - def __call__(self): - start_time = time.time() - - for i in range(len(self.labels)): - self.step_forward() - - if self.to_log: - if self.progress_bar: - self.progress_bar.close() - else: - overall_latency = time.time() - start_time - print_colored("Loading model done.\n", 'bold_green') - print_colored( - 'Total latency: {:.3f}s\n'.format(overall_latency), - 'grey') - - def step_forward(self): - n_steps = len(self.labels) - - label = self.labels[self.counter] - - # display step information - if self.to_log: - if self.progress_bar: - self.progress_bar.set_description(self.labels[self.counter]) - else: - print_colored("Loading Model: ") - print_colored(f"[{self.counter+1}/{n_steps}]\t", - 'bold_green') - print_colored(f"{label}\n") - - # execute the step - start_time = time.time() - self.step_handlers[self.counter]() - # release resource after each step - release_gc() - - if self.progress_bar: - self.progress_bar.update(1) - - latency = time.time() - start_time - if self.to_log and not self.progress_bar: - print_colored("Time: {:.3f}s\n".format(latency), 'grey') - - self.llm_build_stats.build_steps_info.append((label, latency)) - - self.counter += 1 - - def __call__(self, engine_dir: Optional[Path] = None) -> Path: - """The engine_dir is the path to save the built engine. - """ - if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: - return self.model_obj.model_dir - - if self.llm_args.parallel_config.is_multi_gpu: - torch.cuda.set_device(self.global_rank % self.mapping.gpus_per_node) - - pipeline = ModelLoader.BuildPipeline( - self.llm_args.enable_tqdm, - [label for label, _ in self._build_pipeline], - [handler for _, handler in self._build_pipeline], - llm_build_stats=self.llm_build_stats, - ) - pipeline() - - assert engine_dir - - runtime_context = _ModelRuntimeContext( - engine=self._engine, - mapping=self.mapping, - model_info=self._model_info, - ) - self.save(runtime_context, self.model_obj.model_dir, engine_dir) - return engine_dir - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - for attr_name in dir(self): - if not callable(getattr( - self, attr_name)) and not attr_name.startswith("__"): - if attr_name not in ('model_format', 'workspace'): - setattr(self, attr_name, None) - - release_gc() - - @property - def workspace(self) -> str: - return self._workspace - - @property - def model_format(self) -> _ModelFormatKind: - return self._model_format - - def save( - self, - model: _ModelRuntimeContext, - model_dir: str, - engine_dir: str, - ): - """Save the built engine on a single GPU to the given path.""" - model.engine.save(engine_dir) - if model.mapping.rank == 0: - tokenizer = ModelLoader.load_hf_tokenizer( - model_dir, - trust_remote_code=self.llm_args.trust_remote_code, - use_fast=self.llm_args.tokenizer_mode != 'slow') - if tokenizer is not None: - tokenizer.save_pretrained(engine_dir) - - def _download_hf_model(self): - """Download HF model from third-party model hub like www.modelscope.cn or huggingface.""" - model_dir = None - # Only the rank0 are allowed to download model - if mpi_rank() == 0: - assert self._workspace is not None - assert isinstance(self.model_obj.model_name, str) - # this will download only once when multiple MPI processes are running - model_dir = download_hf_model(self.model_obj.model_name, - revision=self.llm_args.revision) - print_colored(f"Downloaded model to {model_dir}\n", 'grey') - # Make all the processes got the same model_dir - self._model_dir = mpi_broadcast(model_dir, root=0) - self.model_obj.model_dir = self._model_dir # mark as a local model - assert self.model_obj.is_local_model def _apply_modelopt_quant_config(self, hf_quant_config: Dict[str, Any], explicit_kv_cache_quant_algo) -> None: @@ -541,136 +283,6 @@ def _update_from_hf_quant_config(self) -> bool: return False - def _load_model_from_hf(self): - """Load a TRT-LLM model from a HF model.""" - assert self._model_dir is not None - - model_cls = AutoModelForCausalLM.get_trtllm_model_class( - self._model_dir, self.llm_args.trust_remote_code, - self.llm_args.decoding_config.decoding_mode - if hasattr(self.llm_args, "speculative_model") - and self.llm_args.speculative_model else None) - - prequantized = self._update_from_hf_quant_config() - - # FP4 Gemm force to use plugin. - if self.llm_args.quant_config.quant_mode.has_nvfp4(): - self.llm_args.build_config.plugin_config.gemm_plugin = "nvfp4" - - if self.llm_args.load_format == 'dummy': - config = model_cls.config_class.from_hugging_face( - str(self._model_dir), - dtype=self.llm_args.dtype, - mapping=self.mapping, - quant_config=self.llm_args.quant_config, - **self.convert_checkpoint_options, - ) - self.model = model_cls(config) - elif self.llm_args.quant_config._requires_calibration and not prequantized: - assert self.workspace is not None - checkpoint_dir = f"{self.workspace}/quantized-checkpoint" - if self.rank == 0: - model_cls.quantize( - self._model_dir, - checkpoint_dir, - dtype=self.llm_args.dtype, - mapping=self.mapping, - quant_config=self.llm_args.quant_config, - **self.llm_args.calib_config.model_dump(), - trust_remote_code=self.llm_args.trust_remote_code, - ) - if self.llm_args.parallel_config.is_multi_gpu: - mpi_barrier() - self.model = model_cls.from_checkpoint(checkpoint_dir, - rank=self.mapping.rank) - else: - self.model = model_cls.from_hugging_face( - str(self._model_dir), - dtype=self.llm_args.dtype, - mapping=self.mapping, - quant_config=self.llm_args.quant_config, - load_model_on_cpu= - True, # TODO:TRTLLM-195 to enhance the weights loading memory usage and choose best location - trust_remote_code=self.llm_args.trust_remote_code, - speculative_model_dir=self._speculative_model_dir, - speculative_config=self.llm_args.speculative_config - if not isinstance(self.llm_args.speculative_config, - LookaheadDecodingConfig) else None, - **self.convert_checkpoint_options, - ) - - self.pretrained_config = self.model.config - self._model_info = _ModelInfo.from_pretrained_config( - self.pretrained_config) - - @print_traceback_on_error - def _load_model_from_ckpt(self): - """Load a TRT-LLM model from checkpoint.""" - self.pretrained_config = PretrainedConfig.from_json_file( - os.path.join(self._model_dir, 'config.json')) - self.pretrained_config.mapping = self.mapping - - #TODO: TRTLLM-1091, change the architecture in the checkpoint to TRT-LLM one, not HF one. - architecture = self.pretrained_config.architecture - assert architecture in MODEL_MAP, \ - f"Unsupported model architecture: {architecture}" - model_cls = MODEL_MAP[architecture] - - if self.llm_args.load_format == 'dummy': - self.model = model_cls(self.pretrained_config) - else: - self.model = model_cls.from_checkpoint( - self._model_dir, config=self.pretrained_config) - self._model_info = _ModelInfo.from_pretrained_config( - self.pretrained_config) - - # load parallel embedding related options - self.convert_checkpoint_options[ - 'use_parallel_embedding'] = self.pretrained_config.use_parallel_embedding - - def _build_engine_from_inmemory_model(self): - assert isinstance(self.llm_args.model, Module) - self._model_info = _ModelInfo.from_module(self.model) - - @print_traceback_on_error - def _build_engine(self): - assert isinstance( - self.build_config, - BuildConfig), f"build_config is not set yet: {self.build_config}" - - logger_debug(f"rank{mpi_rank()} begin to build engine...\n", "green") - - # avoid side effects by copying the original build_config - copied_build_config = self.build_config.model_copy(deep=True) - - copied_build_config.update_kv_cache_type(self._model_info.architecture) - assert self.model is not None, "model has not been loaded yet." - - self._engine = build(self.model, copied_build_config) - self.mapping = self.model.config.mapping - - # delete the model explicitly to free all the build-time resources - self.model = None - logger_debug(f"rank{mpi_rank()} build engine done\n", "green") - - def _save_engine_for_runtime(self): - """Persist the engine to disk for the cpp runtime. - - Currently, the cpp runtime can accept an engine path, - that requires the engine should always be saved to disk. - - This explicit saving will be removed in the future when the cpp runtime can accept the engine buffer directly. - But this is necessary for a build cache, but it can be optimized to async IO. - """ - if self.build_cache_enabled: - self._model_dir = self.engine_cache_stage.cache_dir - self._model_format = _ModelFormatKind.TLLM_ENGINE - return - - def _load_engine_buffer(self): - # Load engine buffer from disk - self._engine = Engine.from_dir(self._model_dir) - @staticmethod def load_hf_tokenizer(model_dir, trust_remote_code: bool = True, @@ -777,9 +389,6 @@ def _download_hf_model_if_needed(self, def __call__(self) -> Tuple[Path, Union[Path, None]]: - if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: - return Path(self.llm_args.model), None - # Download speculative model from HuggingFace if needed (all backends) if (self.llm_args.speculative_config is not None and self.llm_args.speculative_config.speculative_model is not None): @@ -792,196 +401,25 @@ def __call__(self) -> Tuple[Path, Union[Path, None]]: if self.llm_args.backend == "_autodeploy": return None, "" - self.engine_cache_stage: Optional[CachedStage] = None self._hf_model_dir = None self.model_loader = ModelLoader(self.llm_args) - if self.llm_args.backend is not None: - if self.llm_args.backend not in ["pytorch", "_autodeploy"]: - raise ValueError( - f'backend {self.llm_args.backend} is not supported.') - - self._hf_model_dir = self._download_hf_model_if_needed( - self.model_loader.model_obj, revision=self.llm_args.revision) - - if self.llm_args.quant_config.quant_algo is not None: - logger.warning( - "QuantConfig for pytorch backend is ignored. You can load " - "quantized model with hf_quant_config.json directly.") - # Currently, this is to make updated quant_config visible by llm.args.quant_config - # TODO: Unify the logics with those in tensorrt_llm/_torch/model_config.py - self.model_loader._update_from_hf_quant_config() - - return None, self._hf_model_dir - - if self.model_loader.model_obj.is_hub_model: - # This will download the config.json from HF model hub, this helps to create a PretrainedConfig for - # cache key. - self._hf_model_dir = download_hf_pretrained_config( - self.model_loader.model_obj.model_name, - revision=self.llm_args.revision) - - elif self.model_loader.model_obj.is_local_model: - self._hf_model_dir = self.model_loader.model_obj.model_dir if self.llm_args.model_format is _ModelFormatKind.HF else None - - if self.build_cache_enabled: - print_colored("Build cache is enabled.\n", 'yellow') - - self.engine_cache_stage = self._get_engine_cache_stage() - if self.engine_cache_stage.is_cached(): - self.llm_build_stats.cache_hitted = True - print_colored( - f"Reusing cached engine in {self.engine_cache_stage.get_engine_path()}\n\n", - 'grey') - self.model_loader.model_obj.model_dir = self.engine_cache_stage.get_engine_path( - ) - self.llm_build_stats.engine_dir = self.model_loader.model_obj.model_dir - return self.llm_build_stats.engine_dir, self._hf_model_dir - - return self._build_model(), self._hf_model_dir - - def get_engine_dir(self) -> Path: - if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: - return self.model_obj.model_dir - - # generate a new path for writing the engine - if self.build_cache_enabled: - cache_stage = self._get_engine_cache_stage() - return cache_stage.get_engine_path() - - return self.workspace / "tmp.engine" - - @property - def build_cache_enabled(self) -> bool: - _enable_build_cache, _ = get_build_cache_config_from_env() - - return (self.llm_args.enable_build_cache - or _enable_build_cache) and (self.llm_args.model_format - is _ModelFormatKind.HF) - - def _get_engine_cache_stage(self) -> CachedStage: - """Get the cache stage for engine building.""" - build_cache = BuildCache(self.llm_args.enable_build_cache) - - assert self._hf_model_dir is not None, "HF model dir is required for cache key." - - def serialize(d) -> str: - if hasattr(d, "to_dict"): - dic = d.to_dict() - elif is_dataclass(d): - dic = asdict(d) - elif isinstance(d, BaseModel): - dic = d.model_dump(mode="json") - else: - raise ValueError(f"Could not serialize type: {type(d)}") - return json.dumps(dic, sort_keys=True) - - parallel_config = self.llm_args.parallel_config - - force_rebuild = False - if self.llm_args.model_format is not _ModelFormatKind.HF: - force_rebuild = True - - return build_cache.get_engine_building_cache_stage( - build_config=self.llm_args.build_config, - model_path=self._hf_model_dir, - force_rebuild=force_rebuild, - # Other configs affecting the engine building - parallel_config=serialize(parallel_config), - pretrained_config=serialize(self.get_pretrained_config()), - quant_config=serialize(self.llm_args.quant_config), - ) - - def get_pretrained_config(self) -> PretrainedConfig: - """Get the PretrainedConfig for cache key. - - NOTE, this is not the HF model's config, but the TRT-LLM's config. We use this as a generic information for - HF and other models. - """ - assert self._hf_model_dir is not None - return AutoConfig.from_hugging_face( - self._hf_model_dir, - mapping=self.llm_args.parallel_config.to_mapping(), - quant_config=self.llm_args.quant_config, - dtype=self.llm_args.dtype, - trust_remote_code=self.llm_args.trust_remote_code) - - def _build_model(self) -> Path: - model_format = self.llm_args.model_format - - def build_task(engine_dir: Path): - if model_format is not _ModelFormatKind.TLLM_ENGINE: - model_loader_kwargs = { - 'llm_args': self.llm_args, - 'workspace': str(self.workspace), - 'llm_build_stats': self.llm_build_stats, - } - - if self.llm_args.parallel_config.is_multi_gpu: - assert self.mpi_session - - #mpi_session cannot be pickled so remove from self.llm_args - if self.llm_args.mpi_session: - del self.llm_args.mpi_session - - # The engine_dir:Path will be stored to MPINodeState.state - build_infos = self.mpi_session.submit_sync( - CachedModelLoader._node_build_task, - engine_dir=engine_dir, - **model_loader_kwargs) - self.llm_build_stats.build_steps_info = build_infos[0] - - else: # single-gpu - with ModelLoader(**model_loader_kwargs) as model_loader: - model_loader(engine_dir=engine_dir) - - release_gc() - - has_storage = True - if self.build_cache_enabled: - try: - # TODO[chunweiy]: Cover the case when the model is from HF model hub. - if self.model_loader.model_obj.is_local_model: - # This is not perfect, but will make build-cache much more robust. - free_storage = self.engine_cache_stage.parent.free_storage_in_gb( - ) - model_size = get_directory_size_in_gb( - self.model_loader.model_obj.model_dir) - require_size = model_size * 1.3 - has_storage = free_storage >= require_size - - if not has_storage: - print_colored( - "Build cache is disabled since the cache storage is too small.\n ", - 'yellow') - print_colored( - f"Free storage: {free_storage}GB, Required storage: {require_size}GB\n", - 'grey') - except ValueError: - has_storage = False - except Exception as e: - logger.error(e) - has_storage = False - - if enable_llm_debug(): - print_colored(f"Has cache storage: {has_storage}\n", 'yellow') - - if has_storage: - with self.engine_cache_stage.write_guard() as engine_dir: - build_task(engine_dir) - self.llm_build_stats.cache_hitted = True + if self.llm_args.backend not in ["pytorch", "_autodeploy"]: + raise ValueError( + f'backend {self.llm_args.backend} is not supported.') - else: - print_colored( - "The cache directory is too small, build-cache is disabled.\n", - 'grey') - self.llm_build_stats.cache_hitted = False - self.llm_build_stats.cache_info = "The cache root directory is too small." + self._hf_model_dir = self._download_hf_model_if_needed( + self.model_loader.model_obj, revision=self.llm_args.revision) - if not (has_storage and self.build_cache_enabled): - build_task(self.get_engine_dir()) + if self.llm_args.quant_config.quant_algo is not None: + logger.warning( + "QuantConfig for pytorch backend is ignored. You can load " + "quantized model with hf_quant_config.json directly.") + # Currently, this is to make updated quant_config visible by llm.args.quant_config + # TODO: Unify the logics with those in tensorrt_llm/_torch/model_config.py + self.model_loader._update_from_hf_quant_config() - return self.get_engine_dir() + return None, self._hf_model_dir @print_traceback_on_error @staticmethod @@ -994,27 +432,6 @@ def _node_download_hf_model( else: return None - @print_traceback_on_error - @staticmethod - def _node_build_task( - llm_args: LlmArgs, - workspace: Optional[str | tempfile.TemporaryDirectory] = None, - llm_build_stats: Optional['LlmBuildStats'] = None, - engine_dir: Optional[Path] = None, - ): - if MPINodeState.is_initialized(): - raise RuntimeError("The MPI node is already initialized.") - - with ModelLoader(llm_args, - workspace=workspace, - llm_build_stats=llm_build_stats) as model_loader: - model_loader(engine_dir=engine_dir) - return model_loader.llm_build_stats.build_steps_info - - def save(self, engine_dir: Path): - # copy the engine directory to the target directory - shutil.copytree(self.get_engine_dir(), engine_dir) - @dataclass class LlmBuildStats: @@ -1038,10 +455,7 @@ class LlmBuildStats: 'LlmArgs', 'LlmBuildStats', 'ModelLoader', - '_ModelRuntimeContext', - '_ModelInfo', '_ParallelConfig', - '_ModelFormatKind', '_ModelWrapper', 'BatchingType', 'ExecutorConfig', @@ -1055,8 +469,6 @@ class LlmBuildStats: 'UserProvidedDecodingConfig', 'ContextChunkingPolicy', 'CapacitySchedulerPolicy', - 'BuildConfig', - 'BuildCacheConfig', 'QuantConfig', 'CalibConfig', 'CudaGraphConfig', diff --git a/tensorrt_llm/logger.py b/tensorrt_llm/logger.py index 49bd2a415342..7ed205e8f24f 100644 --- a/tensorrt_llm/logger.py +++ b/tensorrt_llm/logger.py @@ -17,8 +17,6 @@ import sys from typing import Dict, Optional -import tensorrt as trt - try: from polygraphy.logger import G_LOGGER except ImportError: @@ -75,7 +73,6 @@ def _extract_module(qualname: str) -> str: "flash_mla": "flashmla", "quantization": "quantize", "scaffolding": "scaffold", - "_tensorrt_engine": "trt_engn", "visual_gen": "vis_gen", "__pycache__": "pycache", "tokenizer": "tokenizr", @@ -186,7 +183,6 @@ def __init__(self): min_severity = self.DEFAULT_LEVEL self._min_severity = min_severity - self._trt_logger = trt.Logger(severity_map[min_severity][0]) self._logger = logging.getLogger(self.PREFIX) self._logger.propagate = False handler = logging.StreamHandler(stream=sys.stdout) @@ -204,7 +200,7 @@ def __init__(self): # Set the underlying Python logger to the minimum of all configured # levels so that per-module overrides more verbose than the global # level are not silently dropped by Python's logging framework. - global_py_level = severity_map[min_severity][1] + global_py_level = severity_map[min_severity][0] if self._module_levels: # Map our numeric levels to Python logging levels. _numeric_to_py = { @@ -224,7 +220,7 @@ def __init__(self): self._polygraphy_logger = G_LOGGER if self._polygraphy_logger is not None: - self._polygraphy_logger.module_severity = severity_map[min_severity][2] + self._polygraphy_logger.module_severity = severity_map[min_severity][1] # For log_once self._appeared_keys = set() @@ -268,10 +264,6 @@ def _func_wrapper(self, severity): else: raise AttributeError(f"No such severity: {severity}") - @property - def trt_logger(self) -> trt.ILogger: - return self._trt_logger - def log(self, severity, *msg): module = _get_caller_module() if not self.is_severity_enabled(severity, module): @@ -338,20 +330,21 @@ def set_level(self, min_severity): ) return self._min_severity = min_severity - self._trt_logger.min_severity = severity_map[min_severity][0] - self._logger.setLevel(severity_map[min_severity][1]) + self._logger.setLevel(severity_map[min_severity][0]) if self._polygraphy_logger is not None: - self._polygraphy_logger.module_severity = severity_map[min_severity][2] + self._polygraphy_logger.module_severity = severity_map[min_severity][1] severity_map = { - "internal_error": [trt.Logger.INTERNAL_ERROR, logging.CRITICAL], - "error": [trt.Logger.ERROR, logging.ERROR], - "warning": [trt.Logger.WARNING, logging.WARNING], - "info": [trt.Logger.INFO, logging.INFO], - "verbose": [trt.Logger.VERBOSE, logging.DEBUG], - "debug": [trt.Logger.VERBOSE, logging.DEBUG], - "trace": [trt.Logger.VERBOSE, logging.DEBUG], + # [0] Python logging level; [1] polygraphy level (appended below when + # polygraphy is installed). + "internal_error": [logging.CRITICAL], + "error": [logging.ERROR], + "warning": [logging.WARNING], + "info": [logging.INFO], + "verbose": [logging.DEBUG], + "debug": [logging.DEBUG], + "trace": [logging.DEBUG], } if G_LOGGER is not None: diff --git a/tensorrt_llm/lora_helper.py b/tensorrt_llm/lora_helper.py index f2222505cfe3..a2f85efb9359 100644 --- a/tensorrt_llm/lora_helper.py +++ b/tensorrt_llm/lora_helper.py @@ -67,27 +67,6 @@ def get_default_trtllm_modules_to_hf_modules(): } -def use_lora( - model, - lora_config: "LoraConfig", - trtllm_modules_to_hf_modules: Optional[Dict[str, str]] = None, -): - """Use LoRA with the given model and configuration. - - This function is a wrapper that delegates to the appropriate loading function - based on the LoRA checkpoint source. - """ - if lora_config.lora_ckpt_source == "nemo": - from .lora_manager import load_nemo_lora - load_nemo_lora(model, lora_config) - elif lora_config.lora_ckpt_source == "hf": - from .lora_manager import load_hf_lora - load_hf_lora(model, lora_config, trtllm_modules_to_hf_modules) - else: - raise ValueError( - f"Unsupported lora_ckpt_source: {lora_config.lora_ckpt_source}") - - class LoraConfig(StrictBaseModel): lora_dir: List[str] = Field( default_factory=list, diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index 933c63d1f615..f703a85769d5 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -17,15 +17,14 @@ from tensorrt_llm.bindings import internal as tb_internal -from ._utils import pad_vocab_size, release_gc, str_dtype_to_torch, torch_to_numpy -from .layers.linear import ColumnLinear +from ._utils import release_gc, str_dtype_to_torch, torch_to_numpy from .lora_helper import ( LoraConfig, get_default_trtllm_modules_to_hf_modules, get_missing_qkv_modules_from_lora_modules, ) from .mapping import Mapping -from .models.convert_utils import get_model_path, load_state_dict, split_matrix_tp +from .models.convert_utils import get_model_path, load_state_dict if TYPE_CHECKING: from .runtime import ModelConfig @@ -426,23 +425,13 @@ def get_target_modules(self): return self.lora_target_modules -def load_nemo_lora(model, lora_config: LoraConfig): - lora_loader = NemoLoraLoader(lora_config.lora_dir) - - if not lora_loader.is_valid: - raise ValueError(f"Failed to load NeMo LoRA from {lora_config.lora_dir}") - - if len(lora_config.lora_target_modules) == 0: - lora_config.lora_target_modules = lora_loader.lora_target_modules - - def load_torch_hf_lora(lora_config: LoraConfig): - """This is a shortned version of load_hf_lora that is used for torch models. + """Load an HF LoRA checkpoint for the PyTorch workflow. - Main problem is model.config in legacy code is custom (defined in the legacy code) whereas - pivot model config is the transformer's one. + Populates lora_config (trtllm_modules_to_hf_modules and inferred + lora_target_modules) from the HF adapter directory. The actual weights are + loaded later by LoraManager when requests arrive with LoRA UIDs. """ - # TODO smor- need to comibe with load_hf_lora if not lora_config.trtllm_modules_to_hf_modules: lora_config.trtllm_modules_to_hf_modules = get_default_trtllm_modules_to_hf_modules() @@ -530,94 +519,6 @@ def load_torch_lora(lora_config: LoraConfig): ) -def load_hf_lora( - model, - lora_config: LoraConfig, - trtllm_modules_to_hf_modules: Optional[Dict[str, str]] = None, -): - trtllm_modules_to_hf_modules = ( - trtllm_modules_to_hf_modules or get_default_trtllm_modules_to_hf_modules() - ) - lora_config.trtllm_modules_to_hf_modules = trtllm_modules_to_hf_modules - - lora_loader = HfLoraLoader(lora_config.lora_dir) - - if len(lora_config.lora_target_modules) == 0: - lora_config.lora_target_modules = lora_loader.get_target_modules( - trtllm_modules_to_hf_modules - ) - if len(lora_config.lora_target_modules) == 0: - raise ValueError( - "lora_target_modules is empty. " - "Please specify lora_target_modules or provide lora_dir to infer lora_target_modules." - ) - - missing_qkv_modules = LoraManager.get_missing_qkv_modules(lora_config.lora_target_modules) - lora_config.lora_target_modules.extend(missing_qkv_modules) - - if lora_loader.is_valid: - config = model.config - torch_dtype = str_dtype_to_torch(config.dtype) - # the lora checkpoint might finetune the embedding - if lora_loader.vocab_size != 0: - config.vocab_size = lora_loader.vocab_size - mapping = config.mapping - if mapping.is_first_pp_rank() and lora_loader.embed_tokens is not None: - weight = lora_loader.embed_tokens - if config.use_parallel_embedding: - weight = split_matrix_tp( - weight, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim, - ) - if model.transformer.vocab_embedding.weight.raw_value.shape != weight.shape: - model.transformer.vocab_embedding = model.transformer.vocab_embedding.__class__( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - dtype=config.dtype, - tp_size=mapping.tp_size if config.use_parallel_embedding else 1, - tp_group=mapping.tp_group if config.use_parallel_embedding else None, - sharding_dim=config.embedding_sharding_dim, - tp_rank=mapping.tp_rank, - ) - model.transformer.vocab_embedding.weight.value = weight.to(torch_dtype) - if mapping.is_last_pp_rank() and lora_loader.lm_head is not None: - weight = lora_loader.lm_head - vocab_size = lora_loader.vocab_size - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - weight = torch.from_numpy( - np.pad( - torch_to_numpy(weight), - ((0, pad_width), (0, 0)), - "constant", - constant_values=0, - ) - ) - else: - vocab_size_padded = vocab_size - if model.lm_head.weight.raw_value.shape != weight.shape: - model.lm_head = ColumnLinear( - config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - model.lm_head.weight.value = split_matrix_tp( - weight, - mapping.tp_size, - mapping.tp_rank, - dim=0, - ).to(torch_dtype) - - def unpack_nemo_weights(nemo_archive_path: str) -> Tuple[Dict, Dict[str, torch.Tensor]]: """Unpack model config and weights from a NeMo .nemo archive file. diff --git a/tensorrt_llm/models/__init__.py b/tensorrt_llm/models/__init__.py index 96bd4eff96d2..3af378d4678f 100755 --- a/tensorrt_llm/models/__init__.py +++ b/tensorrt_llm/models/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -12,212 +12,19 @@ # 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 .baichuan.model import BaichuanForCausalLM -from .bert.model import (BertForQuestionAnswering, - BertForSequenceClassification, BertModel, - RobertaForQuestionAnswering, - RobertaForSequenceClassification, RobertaModel) -from .bloom.model import BloomForCausalLM, BloomModel -from .chatglm.config import ChatGLMConfig -from .chatglm.model import ChatGLMForCausalLM, ChatGLMModel -from .clip.model import CLIPVisionTransformer -from .cogvlm.config import CogVLMConfig -from .cogvlm.model import CogVLMForCausalLM -from .commandr.model import CohereForCausalLM -from .dbrx.config import DbrxConfig -from .dbrx.model import DbrxForCausalLM -from .deepseek_v1.model import DeepseekForCausalLM -from .deepseek_v2.model import DeepseekV2ForCausalLM -from .dit.model import DiT -from .eagle.model import EagleForCausalLM -from .enc_dec.model import DecoderModel, EncoderModel, WhisperEncoder -from .falcon.config import FalconConfig -from .falcon.model import FalconForCausalLM, FalconModel -from .gemma.config import (GEMMA2_ARCHITECTURE, GEMMA3_ARCHITECTURE, - GEMMA_ARCHITECTURE, GemmaConfig) -from .gemma.model import GemmaForCausalLM -from .gpt.config import GPTConfig -from .gpt.model import GPTForCausalLM, GPTModel -from .gptj.config import GPTJConfig -from .gptj.model import GPTJForCausalLM, GPTJModel -from .gptneox.model import GPTNeoXForCausalLM, GPTNeoXModel -from .grok.model import GrokForCausalLM -from .llama.config import LLaMAConfig -from .llama.model import LLaMAForCausalLM, LLaMAModel -from .mamba.model import MambaForCausalLM -from .medusa.config import MedusaConfig -from .medusa.model import MedusaForCausalLm -from .mllama.model import MLLaMAForCausalLM -from .mmdit_sd3.model import SD3Transformer2DModel -from .modeling_utils import (PretrainedConfig, PretrainedModel, - SpeculativeDecodingMode) -from .mpt.model import MPTForCausalLM, MPTModel -from .multimodal_encoders.config import LlavaNextVisionConfig -from .multimodal_encoders.model import LlavaNextVisionWrapper -from .nemotron_nas.model import DeciLMForCausalLM -from .opt.model import OPTForCausalLM, OPTModel -from .phi3.model import Phi3ForCausalLM, Phi3Model -from .phi.model import PhiForCausalLM, PhiModel -from .qwen.model import QWenForCausalLM -from .recurrentgemma.model import RecurrentGemmaForCausalLM -from .redrafter.model import ReDrafterForLLaMALM, ReDrafterForQWenLM -from .stdit.model import STDiT3Model + +from .modeling_utils import (LayerQuantConfig, PretrainedConfig, QuantAlgo, + QuantConfig, SpeculativeDecodingMode) + +# Architecture registry is intentionally empty: the PyTorch backend resolves +# model classes via tensorrt_llm._torch.models. +MODEL_MAP = {} __all__ = [ - 'BertModel', - 'BertForQuestionAnswering', - 'BertForSequenceClassification', - 'RobertaModel', - 'RobertaForQuestionAnswering', - 'RobertaForSequenceClassification', - 'BloomModel', - 'BloomForCausalLM', - 'CLIPVisionTransformer', - 'DiT', - 'SD3Transformer2DModel', - 'STDiT3', - 'DeepseekForCausalLM', - 'FalconConfig', - 'DeepseekV2ForCausalLM', - 'FalconForCausalLM', - 'FalconModel', - 'GPTConfig', - 'GPTModel', - 'GPTForCausalLM', - 'OPTForCausalLM', - 'OPTModel', - 'LLaMAConfig', - 'LLaMAForCausalLM', - 'LLaMAModel', - 'LlavaNextVisionWrapper', - 'LlavaNextVisionConfig', - 'MedusaConfig', - 'MedusaForCausalLm', - 'ReDrafterForLLaMALM', - 'ReDrafterForQWenLM', - 'GPTJConfig', - 'GPTJModel', - 'GPTJForCausalLM', - 'GPTNeoXModel', - 'GPTNeoXForCausalLM', - 'PhiModel', - 'PhiConfig', - 'Phi3Model', - 'Phi3Config', - 'PhiForCausalLM', - 'Phi3ForCausalLM', - 'ChatGLMConfig', - 'ChatGLMForCausalLM', - 'ChatGLMModel', - 'BaichuanForCausalLM', - 'QWenConfig' - 'QWenForCausalLM', - 'QWenModel', - 'EncoderModel', - 'DecoderModel', 'PretrainedConfig', - 'PretrainedModel', - 'WhisperEncoder', - 'MambaForCausalLM', - 'MambaConfig', - 'MPTForCausalLM', - 'MPTModel', - 'SkyworkForCausalLM', - 'GemmaConfig', - 'GemmaForCausalLM', - 'DbrxConfig', - 'DbrxForCausalLM', - 'RecurrentGemmaForCausalLM', - 'CogVLMConfig', - 'CogVLMForCausalLM', - 'EagleForCausalLM', 'SpeculativeDecodingMode', - 'CohereForCausalLM', - 'MLLaMAForCausalLM', + 'QuantConfig', + 'LayerQuantConfig', + 'QuantAlgo', + 'MODEL_MAP', ] - -MODEL_MAP = { - 'GPT2LMHeadModel': GPTForCausalLM, - 'GPT2LMHeadCustomModel': GPTForCausalLM, - 'GPTBigCodeForCausalLM': GPTForCausalLM, - 'Starcoder2ForCausalLM': GPTForCausalLM, - 'FuyuForCausalLM': GPTForCausalLM, - 'Kosmos2ForConditionalGeneration': GPTForCausalLM, - 'JAISLMHeadModel': GPTForCausalLM, - 'GPTForCausalLM': GPTForCausalLM, - 'NemotronForCausalLM': GPTForCausalLM, - 'OPTForCausalLM': OPTForCausalLM, - 'BloomForCausalLM': BloomForCausalLM, - 'RWForCausalLM': FalconForCausalLM, - 'FalconForCausalLM': FalconForCausalLM, - 'PhiForCausalLM': PhiForCausalLM, - 'Phi3ForCausalLM': Phi3ForCausalLM, - 'Phi3VForCausalLM': Phi3ForCausalLM, - 'Phi3SmallForCausalLM': Phi3ForCausalLM, - 'PhiMoEForCausalLM': Phi3ForCausalLM, - 'Phi4MMForCausalLM': Phi3ForCausalLM, - 'MambaForCausalLM': MambaForCausalLM, - 'GPTNeoXForCausalLM': GPTNeoXForCausalLM, - 'GPTJForCausalLM': GPTJForCausalLM, - 'MptForCausalLM': MPTForCausalLM, - 'MPTForCausalLM': MPTForCausalLM, - 'GLMModel': ChatGLMForCausalLM, - 'ChatGLMModel': ChatGLMForCausalLM, - 'ChatGLMForCausalLM': ChatGLMForCausalLM, - 'ChatGLMForConditionalGeneration': ChatGLMForCausalLM, - 'LlamaForCausalLM': LLaMAForCausalLM, - 'LlavaLlamaModel': LLaMAForCausalLM, - 'LlavaNextForConditionalGeneration': LlavaNextVisionWrapper, - 'ExaoneForCausalLM': LLaMAForCausalLM, - 'MistralForCausalLM': LLaMAForCausalLM, - 'MixtralForCausalLM': LLaMAForCausalLM, - 'ArcticForCausalLM': LLaMAForCausalLM, - 'Grok1ModelForCausalLM': GrokForCausalLM, - 'InternLMForCausalLM': LLaMAForCausalLM, - 'InternLM2ForCausalLM': LLaMAForCausalLM, - 'InternLMXComposer2ForCausalLM': LLaMAForCausalLM, - 'GraniteForCausalLM': LLaMAForCausalLM, - 'GraniteMoeForCausalLM': LLaMAForCausalLM, - 'MedusaForCausalLM': MedusaForCausalLm, - 'MedusaLlamaForCausalLM': MedusaForCausalLm, - 'ReDrafterForLLaMALM': ReDrafterForLLaMALM, - 'ReDrafterForQWenLM': ReDrafterForQWenLM, - 'BaichuanForCausalLM': BaichuanForCausalLM, - 'BaiChuanForCausalLM': BaichuanForCausalLM, - 'SkyworkForCausalLM': LLaMAForCausalLM, - GEMMA_ARCHITECTURE: GemmaForCausalLM, - GEMMA2_ARCHITECTURE: GemmaForCausalLM, - GEMMA3_ARCHITECTURE: GemmaForCausalLM, - 'QWenLMHeadModel': QWenForCausalLM, - 'QWenForCausalLM': QWenForCausalLM, - 'Qwen2ForCausalLM': QWenForCausalLM, - 'Qwen2MoeForCausalLM': QWenForCausalLM, - 'Qwen2ForSequenceClassification': QWenForCausalLM, - 'Qwen2VLForConditionalGeneration': QWenForCausalLM, - 'Qwen2VLModel': QWenForCausalLM, - 'Qwen3ForCausalLM': QWenForCausalLM, - 'Qwen3MoeForCausalLM': QWenForCausalLM, - 'WhisperEncoder': WhisperEncoder, - 'EncoderModel': EncoderModel, - 'DecoderModel': DecoderModel, - 'DbrxForCausalLM': DbrxForCausalLM, - 'RecurrentGemmaForCausalLM': RecurrentGemmaForCausalLM, - 'CogVLMForCausalLM': CogVLMForCausalLM, - 'DiT': DiT, - 'SD3Transformer2DModel': SD3Transformer2DModel, - 'STDiT3': STDiT3Model, - 'DeepseekForCausalLM': DeepseekForCausalLM, - 'DeciLMForCausalLM': DeciLMForCausalLM, - 'DeepseekV2ForCausalLM': DeepseekV2ForCausalLM, - 'EagleForCausalLM': EagleForCausalLM, - 'CohereForCausalLM': CohereForCausalLM, - 'MLLaMAModel': MLLaMAForCausalLM, # For modelopt - 'MllamaForConditionalGeneration': - MLLaMAForCausalLM, # For mllama load by Auto - 'BertForQuestionAnswering': BertForQuestionAnswering, - 'BertForSequenceClassification': BertForSequenceClassification, - 'BertModel': BertModel, - 'RobertaModel': RobertaModel, - 'RobertaForQuestionAnswering': RobertaForQuestionAnswering, - 'RobertaForSequenceClassification': RobertaForSequenceClassification, -} diff --git a/tensorrt_llm/models/baichuan/__init__.py b/tensorrt_llm/models/baichuan/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/baichuan/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/baichuan/config.py b/tensorrt_llm/models/baichuan/config.py deleted file mode 100644 index 599b98089f75..000000000000 --- a/tensorrt_llm/models/baichuan/config.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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 ...logger import logger -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class BaichuanConfig(PretrainedConfig): - - def __init__(self, model_version: Optional[str] = None, **kwargs): - super().__init__(**kwargs) - if model_version is None: - model_version = BaichuanConfig.guess_model_version(self) - self.model_version = model_version - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in BaichuanConfig - output['model_version'] = self.model_version - return output - - @staticmethod - def guess_model_version( - config: Union['transformers.PretrainedConfig', - 'BaichuanConfig']) -> str: - logger.warning( - "Model version is not set, trying to guess from loaded config") - size = '7' if config.num_attention_heads == 32 else '13' - version = '1' if config.vocab_size == 64000 else '2' - return f'v{version}_{size}b' - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - model_version = kwargs.pop('model_version', None) - if model_version is None: - model_version = BaichuanConfig.guess_model_version(hf_config) - - if model_version == 'v1_7b' or model_version == 'v2_7b': - position_embedding_type = 'rope_gpt_neox' - max_position_embeddings = hf_config.max_position_embeddings - else: - position_embedding_type = 'alibi' - max_position_embeddings = hf_config.model_max_length - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture='BaichuanForCausalLM', - dtype=dtype, - vocab_size=hf_config.vocab_size, - max_position_embeddings=max_position_embeddings, - hidden_size=hf_config.hidden_size, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - num_key_value_heads=hf_config.num_attention_heads, - hidden_act=hf_config.hidden_act, - intermediate_size=hf_config.intermediate_size, - norm_epsilon=hf_config.rms_norm_eps, - position_embedding_type=position_embedding_type, - model_version=model_version, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/baichuan/convert.py b/tensorrt_llm/models/baichuan/convert.py deleted file mode 100644 index 576d5e1a34c8..000000000000 --- a/tensorrt_llm/models/baichuan/convert.py +++ /dev/null @@ -1,931 +0,0 @@ -# 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 copy -import functools -import os -import time -from collections import defaultdict -from typing import Any, Dict, List, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from datasets import Dataset -from tqdm import tqdm -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.pytorch_utils import Conv1D - -from ...logger import logger -from ...quantization import QuantAlgo, QuantMode -from ..convert_utils import (load_calib_dataset, smooth_gemm, - smooth_gemm_fc1_gate, weight_only_quantize_dict) -from .config import BaichuanConfig - - -def generate_int8( - weights: torch.Tensor, - act_range: Dict[str, torch.Tensor], - is_qkv: bool = False, - multi_query_mode: bool = False, -): - """ - This function has two purposes: - - compute quantized weights, scaled either per-tensor or per-column - - compute scaling factors - - Depending on the GEMM API (CUTLASS/CUBLAS) the required scaling factors differ. - CUTLASS uses two sets of scaling factors. One for the activation X, one for the weight W. - CUBLAS only has one (we can't do per-row scaling). So we must provide pre-multiplied scaling factor. - - Here is the list of what we need (T means per-tensor, C per-column): - - scale_x_orig_quant puts fp activation into the quantized range (i.e. [-128, 127], for int8). Used before the GEMM. (T) - - scale_y_quant_orig puts quantized activation into the fp range. Used if the GEMM outputs int8. (T) - - scale_w_quant_orig puts weights from quant range to fp range (used with CUTLASS) (T, C) - - scale_y_accum_quant puts the GEMM result (XW) from accumulation range (int32) - to quant range (int8) (used for CUBLAS) (T, C) - - Note that we don't do anything special about row-parallel GEMM. Theoretically, we could have per-GPU scaling factors too, - but then the model would change depending on the number of GPUs used. - - For QKV projection, the behavior is special. Even if we have a single matrix to perform QKV projection, we consider it - as three different matrices: Q, K, and V. So per-tensor actually means one scaling factor for each Q, K and V. - For our GEMM implementation to respect this behavior, we use per-column mode and replicate values along columns. - """ - - # compute weight scaling factors for fp->int8 and int8->fp - if is_qkv and not multi_query_mode: - scale_w_orig_quant_t = 127. / act_range["w"].reshape(3, -1).max( - dim=-1, keepdims=True)[0].cpu().numpy() - scale_w_orig_quant_c = 127. / act_range["w"].reshape(3, - -1).cpu().numpy() - elif is_qkv and multi_query_mode: - hidden_dim = weights.shape[0] - local_dim = act_range["w"].shape[0] - kv_dim = (local_dim - hidden_dim) // 2 - scale_w_q = act_range["w"][0:hidden_dim] - scale_w_k = act_range["w"][hidden_dim:hidden_dim + kv_dim] - scale_w_v = act_range["w"][-kv_dim:] - - scale_w_qkv_t = torch.concat([ - scale_w_q.max(dim=0, keepdim=True)[0], - scale_w_k.max(dim=0, keepdim=True)[0], - scale_w_v.max(dim=0, keepdim=True)[0] - ]) - - scale_w_orig_quant_t = 127. / scale_w_qkv_t.cpu().numpy() - scale_w_orig_quant_c = 127. / act_range["w"].cpu().numpy() - else: - scale_w_orig_quant_t = 127. / act_range["w"].max().cpu().numpy() - scale_w_orig_quant_c = 127. / act_range["w"].cpu().numpy() - scale_w_quant_orig_t = 1.0 / scale_w_orig_quant_t - scale_w_quant_orig_c = 1.0 / scale_w_orig_quant_c - - # compute the rest of needed scaling factors - scale_x_orig_quant_t = np.array(127. / act_range["x"].max().item()) - scale_y_orig_quant_t = np.array(127. / act_range["y"].max().item()) - scale_y_quant_orig_t = np.array(act_range["y"].max().item() / 127.) - scale_y_accum_quant_t = scale_y_orig_quant_t / (scale_x_orig_quant_t * - scale_w_orig_quant_t) - scale_y_accum_quant_c = scale_y_orig_quant_t / (scale_x_orig_quant_t * - scale_w_orig_quant_c) - if is_qkv and not multi_query_mode: - scale_y_accum_quant_t = np.broadcast_to(scale_y_accum_quant_t, - scale_w_orig_quant_c.shape) - scale_w_quant_orig_t = np.broadcast_to(scale_w_quant_orig_t, - scale_w_orig_quant_c.shape) - if is_qkv and multi_query_mode: - scale_q_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[0], - scale_w_q.shape) - scale_k_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[1], - scale_w_k.shape) - scale_v_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[2], - scale_w_v.shape) - scale_y_accum_quant_t = np.concatenate( - [scale_q_y_accum_t, scale_k_y_accum_t, scale_v_y_accum_t]) - scale_w_quant_orig_t = np.concatenate([ - np.broadcast_to(scale_w_quant_orig_t[0], scale_w_q.shape), - np.broadcast_to(scale_w_quant_orig_t[1], scale_w_k.shape), - np.broadcast_to(scale_w_quant_orig_t[2], scale_w_v.shape) - ]) - - to_i8 = lambda x: x.round().clip(-127, 127).astype(np.int8) - - if is_qkv and multi_query_mode: - scale_w_quant_orig_t_expand = np.ones([weights.shape[-1]]) - scale_w_quant_orig_t_expand[:hidden_dim] = scale_w_quant_orig_t[0] - scale_w_quant_orig_t_expand[hidden_dim:hidden_dim + - kv_dim] = scale_w_quant_orig_t[1] - scale_w_quant_orig_t_expand[-kv_dim:] = scale_w_quant_orig_t[2] - weight_int8 = to_i8(weights * scale_w_quant_orig_t_expand) - else: - weight_int8 = to_i8(weights * scale_w_orig_quant_t) - return { - "weight.int8": weight_int8, - "weight.int8.col": to_i8(weights * scale_w_orig_quant_c), - "scale_x_orig_quant": scale_x_orig_quant_t.astype(np.float32), - "scale_w_quant_orig": scale_w_quant_orig_t.astype(np.float32), - "scale_w_quant_orig.col": scale_w_quant_orig_c.astype(np.float32), - "scale_y_accum_quant": scale_y_accum_quant_t.astype(np.float32), - "scale_y_accum_quant.col": scale_y_accum_quant_c.astype(np.float32), - "scale_y_quant_orig": scale_y_quant_orig_t.astype(np.float32), - } - - -@torch.no_grad() -def capture_activation_range( - model: AutoModelForCausalLM, - tokenizer: AutoTokenizer, - dataset: Dataset, - num_samples: int = 512, -): - model.eval() - device = next(model.parameters()).device - act_scales: defaultdict[Any, Dict[str, - torch.Tensor]] = defaultdict(lambda: { - "x": None, - "y": None, - "w": None - }) - - test_token_num = 923 - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor( - name: str, - tensor: torch.Tensor, - act_scales: defaultdict[Any, Dict[str, torch.Tensor]], - key: str, - ): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook( - m: nn.Module, - x: torch.Tensor, - y: torch.Tensor, - name: str, - ): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - line_encoded = tokenizer(line, - return_tensors="pt", - max_length=test_token_num, - padding=True, - truncation=True).input_ids.to(device) - model(line_encoded) - - for h in hooks: - h.remove() - - return act_scales - - -@torch.no_grad() -def smooth_baichuan_model( - model: AutoModelForCausalLM, - scales: Dict[Any, Dict[str, torch.Tensor]], - alpha: float, - baichuan_smoother: Dict[str, torch.Tensor], -): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - class_name = module.__class__.__name__ - if 'Layer' not in class_name: - continue - print(f'smoothing module: {name}, class_name: {class_name}') - # qkv_proj - layer_name_qkv = name + ".self_attn.W_pack" - - smoother = smooth_gemm(module.self_attn.W_pack.weight, - scales[layer_name_qkv]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_qkv]["x"] / smoother - scales[layer_name_qkv]["w"] = module.self_attn.W_pack.weight.abs().max( - dim=1)[0].float() - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - baichuan_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0].float() - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0].float() - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0].float() - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - baichuan_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0].float() - - -def get_tllm_linear_sq_weight( - vals: Dict[str, np.ndarray], - prefix: str, - shape: List[int], - tensor_parallel: int, - quant_mode: QuantMode, - is_qkv: bool = False, - last_prefix: Optional[str] = None, - bias: Optional[torch.Tensor] = None, - smoother_value=None, - smoother_shape=None, - rank: int = 0, - cat_dim: int = 0, - multi_query_mode: bool = False, -): - per_token = quant_mode.has_per_token_dynamic_scaling() - per_channel = quant_mode.has_per_channel_scaling() - results: Dict[str, torch.Tensor] = {} - - def multi_query_split( - data: np.array, - start: int, - size: int, - tp_size: int, - cur_rank: int, - ): - q, k, v = np.split(data, [start, start + size], axis=-1) - q_split = np.split(q, tp_size, axis=-1) - k_split = np.split(k, tp_size, axis=-1) - v_split = np.split(v, tp_size, axis=-1) - return [ - np.concatenate((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array(cur_per_channel_value, - dtype=np.float32).reshape(col_shape)).contiguous() - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def split(weight: torch.Tensor, - tp_size: int, - rank: int = 0, - dim: int = 0) -> torch.Tensor: - if tp_size == 1: - return weight - elif weight.ndim == 1: - return torch.chunk(weight, tp_size)[rank].contiguous() - else: - return torch.chunk(weight, tp_size, dim=dim)[rank].contiguous() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_weight(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> Optional[torch.Tensor]: - if f'{prefix}.weight' not in params: - return None - return params[f'{prefix}.weight'].to(dtype).detach().cpu() - - -def quantize(hf_model_dir: str, - output_dir: str, - config: BaichuanConfig, - device: str = 'cuda', - calib_dataset: str = 'ccdv/cnn_dailymail', - trust_remote_code: bool = True): - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map='auto' if device != 'cpu' else 'cpu', - dtype='auto' - if not config.quantization._use_plugin_sq else torch.float16, - trust_remote_code=trust_remote_code) - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, use_fast=False, trust_remote_code=trust_remote_code) - dataset = load_calib_dataset(calib_dataset) - - act_range = capture_activation_range(hf_model, tokenizer, dataset) - smoother = {} - if config.quantization._use_plugin_sq: - smooth_baichuan_model(hf_model, act_range, - config.quantization.smoothquant_val, smoother) - - quant_mode = config.quantization.quant_mode - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - inter_size = config.intermediate_size - num_key_value_heads = config.num_key_value_heads - multi_query_mode = (num_key_value_heads != num_attention_heads) - - for rank in range(config.mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = {} - tensor_parallel = config.mapping.tp_size - - for l in config.mapping.pp_layers(config.num_hidden_layers): - prefix = f'model.layers.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - # self_attn.W_pack -> attention.qkv - qkv_weight = get_weight(model_params, prefix + 'self_attn.W_pack', - dtype) - qkv_weight = qkv_weight.t().numpy() - qkv_out_dim = qkv_weight.shape[1] - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.W_pack'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - quant_mode, - is_qkv=True, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1, - multi_query_mode=multi_query_mode)) - - if config.quantization.kv_cache_quant_algo == QuantAlgo.INT8: - qkv_weight = get_weight(model_params, - prefix + 'self_attn.W_pack', dtype) - qkv_weight = qkv_weight.t().numpy() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.W_pack'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.from_numpy( - np.array([int8_weights['scale_y_quant_orig']], - dtype=np.float32)).contiguous() - - # attn.out_proj -> attention.dense - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - attn_dense_weight = attn_dense_weight.t().numpy() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'self_attn.o_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'self_attn.o_proj')], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - - # mlp.gate_proj -> mlp.fc - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - mlp_fc_weight = mlp_fc_weight.t().numpy() - int8_weights = generate_int8( - mlp_fc_weight, act_range.get(prefix + 'mlp.gate_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', [1, inter_size // tensor_parallel], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - - # mlp.down_proj -> mlp.proj - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - mlp_proj_weight = mlp_proj_weight.t().numpy() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'mlp.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'mlp.down_proj'], - smoother_shape=[1, inter_size // tensor_parallel], - rank=rank, - cat_dim=0)) - - # mlp.up_proj -> mlp.gate - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - mlp_gate_weight = mlp_gate_weight.t().numpy() - int8_weights = generate_int8(mlp_gate_weight, - act_range.get(prefix + 'mlp.up_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', [1, inter_size // tensor_parallel], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - - # input layer_norm - input_ln_weight = get_weight(model_params, - prefix + 'input_layernorm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', - dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - embed_w = get_weight(model_params, 'model.embed_tokens', dtype) - if config.mapping.is_first_pp_rank(): - # Embedding - if config.use_parallel_embedding: - embed_w = split_matrix(embed_w, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=config.embedding_sharding_dim) - weights['transformer.vocab_embedding.weight'] = embed_w - - lm_head_w = get_weight(model_params, 'lm_head', dtype) - if config.mapping.is_last_pp_rank(): - # lm_head weight and bias - weights['lm_head.weight'] = split_matrix(lm_head_w.clone(), - config.mapping.tp_size, - config.mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -def load_weights_from_hf_model(hf_model: AutoModelForCausalLM, - config: BaichuanConfig): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_hidden_layers = config.num_hidden_layers - hf_key = [ - "model.embed_tokens.weight", # vocab_embedding - "lm_head.weight", # lm_head - "model.norm.weight", # ln_f - "self_attn.W_pack.weight", # attention.qkv - "self_attn.o_proj.weight", # attention.dense - "mlp.up_proj.weight", # mlp.gate - "mlp.down_proj.weight", # mlp.proj - "mlp.gate_proj.weight", # mlp.fc - "input_layernorm.weight", # input_layernorm - "post_attention_layernorm.weight", # post_layernorm - ] - - def load(key_id: int, layer_idx: int = -1, tp_dim: int = -1): - layer_prefix = "" if layer_idx == -1 else f"model.layers.{layer_idx}." - v: torch.Tensor = model_params[layer_prefix + hf_key[key_id]] - if key_id == 3: - q_emb = v.shape[0] // 3 - model_emb = v.shape[1] - v = v.reshape(3, q_emb, model_emb) - if v.shape[1] % config.mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(config.mapping.tp_size)) - v = v.split(v.shape[1] // config.mapping.tp_size, - dim=1)[config.mapping.tp_rank] - v = v.reshape(3 * (q_emb // config.mapping.tp_size), model_emb) - if tp_dim >= 0: - if v.shape[tp_dim] % config.mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(config.mapping.tp_size)) - v = v.split(v.shape[tp_dim] // config.mapping.tp_size, - dim=tp_dim)[config.mapping.tp_rank] - v = v.to(dtype).contiguous().detach().cpu() - return v - - # Convert vocab_embedding - if config.mapping.is_first_pp_rank(): - embed_w = load(0) - if config.use_parallel_embedding: - embed_w = split_matrix(embed_w, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=config.embedding_sharding_dim) - weights['transformer.vocab_embedding.weight'] = embed_w - - # Convert lm_head - v = load(1, -1, 0) - if config.model_version.startswith('v2'): - v = torch.nn.functional.normalize(v) - if config.mapping.is_last_pp_rank(): - weights['lm_head.weight'] = v - - # Convert ln_f - if config.mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = load(2) - - # Convert layers - layers_range = config.mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f"transformer.layers.{l}." - weights[prefix + 'attention.qkv.weight'] = load(3, l) - weights[prefix + 'attention.dense.weight'] = load(4, l, 1) - weights[prefix + 'mlp.gate.weight'] = load(5, l, 0) - weights[prefix + 'mlp.proj.weight'] = load(6, l, 1) - weights[prefix + 'mlp.fc.weight'] = load(7, l, 0) - weights[prefix + 'input_layernorm.weight'] = load(8, l) - weights[prefix + 'post_layernorm.weight'] = load(9, l) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - - return weight_only_quantize_dict(weights, - quant_algo=config.quantization.quant_algo, - plugin=True) - - -def load_weights_from_gptq(config: BaichuanConfig, quant_ckpt_path: str): - logger.info('Loading weights from groupwise GPTQ Baichuan safetensors...') - weights = {} - tik = time.time() - - gptq_baichuan = safetensors.safe_open(quant_ckpt_path, - framework="pt", - device=0) - gptq_prefix = "model." - gptq_suffix_list = [".qweight", ".qzeros", ".scales"] - gptq_key_list = [ - "embed_tokens.weight", # vocab_embedding - "lm_head.weight", # lm_head - "norm.weight", # ln_f - "self_attn.W_pack", # attention.qkv - "_proj", # - "self_attn.o_proj", # attention.dense - "mlp.up_proj", # mlp.gate - "mlp.down_proj", # mlp.proj - "mlp.gate_proj", # mlp.fc - "input_layernorm.weight", # input_layernorm - "post_attention_layernorm.weight", # post_layernorm - ] - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = getattr(torch, config.dtype) - - def load(key: str, no_prefix: bool = False): - if no_prefix: - return gptq_baichuan.get_tensor(key) - else: - return gptq_baichuan.get_tensor(gptq_prefix + key) - - def torch_split(tensor: torch.Tensor, dim: int): - if tensor.shape[dim] % config.mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(config.mapping.tp_size)) - assert False, "Invalid TP size" - return tensor.split(tensor.shape[dim] // config.mapping.tp_size, - dim=dim)[config.mapping.tp_rank] - - def unpack_int32_into_int8(w_packed: torch.Tensor): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(prefix: str, - tensors: List[torch.Tensor], - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in tensors - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in tensors - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_GPTQ_FOR_LLAMA = 1 # GPTQ-for-LLaMA added 1 to zeros - - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - # return processed interleaved weight, original scales and zeros * scales - weights[prefix + ".weight"] = qweight_interleaved - weights[prefix + ".weights_scaling_factor"] = scales_fp16 - weights[prefix + ".zero"] = zeros_x_scales_fp16 - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = load(gptq_key_list[0]) - if config.mapping.is_first_pp_rank(): - if config.use_parallel_embedding: - v = split_matrix(v, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=config.embedding_sharding_dim) - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - - # 2. lm_head - original_v = load(gptq_key_list[1], True) - if config.model_version.startswith('v2'): - # baichuan v2 models use NormHead - logger.info(f'Normalizing lm_head.weight for {config.model_version}') - v = torch_split(torch.nn.functional.normalize(original_v), 0) - else: - v = torch_split(original_v, 0) - if config.mapping.is_last_pp_rank(): - weights['lm_head.weight'] = v.to(torch_dtype) - - # 3. ln_f - v = load(gptq_key_list[2]) - if config.mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - - # 4. Weights inside each layer - num_hidden_layers = config.num_hidden_layers - layers_range = config.mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - hf_prefix = f"layers.{l}." - tllm_prefix = f"transformer.layers.{l}." - logger.info(f'Process weights in layer: {layer_idx}') - - # 4.1 attention.qkv - qkv_weight_list = [] - for suf in gptq_suffix_list: - qkv_list = [] - comp_part = load(hf_prefix + gptq_key_list[3] + suf) - qkv = torch.chunk(comp_part, 3, 1) - for i in range(3): - comp_part = qkv[i] - comp_part = torch_split(comp_part, 1) - qkv_list.append(comp_part) - qkv_weight_list.append(torch.cat(qkv_list, dim=1)) - - process_and_assign_weight(tllm_prefix + "attention.qkv", - qkv_weight_list) - - # 4.2 attention.dense - v = [ - load(hf_prefix + gptq_key_list[5] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "attention.dense", v, 0) - - # 4.3 mlp.gate - v = [ - load(hf_prefix + gptq_key_list[6] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "mlp.gate", v, 1) - - # 4.4 mlp.proj - v = [ - load(hf_prefix + gptq_key_list[7] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "mlp.proj", v, 0) - - # 4.5 mlp.fc - v = [ - load(hf_prefix + gptq_key_list[8] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "mlp.fc", v, 1) - - # 4.6 input_layernorm - v = load(hf_prefix + gptq_key_list[9]) - weights[tllm_prefix + 'input_layernorm.weight'] = v.to(torch_dtype) - - # 4.7 pst_layernorm - v = load(hf_prefix + gptq_key_list[10]) - weights[tllm_prefix + 'post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/baichuan/model.py b/tensorrt_llm/models/baichuan/model.py deleted file mode 100644 index 00dab61b69ec..000000000000 --- a/tensorrt_llm/models/baichuan/model.py +++ /dev/null @@ -1,253 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (Attention, AttentionMaskType, ColumnLinear, Embedding, - GatedMLP, RmsNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) -from .config import BaichuanConfig -from .convert import load_weights_from_hf_model - - -class BaichuanDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx): - super().__init__() - self.layer_idx = layer_idx - self.config = config - hidden_size = config.hidden_size - dtype = config.dtype - position_embedding_type = config.position_embedding_type - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - quant_mode = config.quant_mode - - self.input_layernorm = RmsNorm(normalized_shape=hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - bias=False, - position_embedding_type=position_embedding_type, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=quant_mode) - - self.mlp = GatedMLP(hidden_size=hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=dtype, - bias=False, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - self.post_layernorm = RmsNorm(normalized_shape=hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class BaichuanModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - hidden_size = config.hidden_size - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(BaichuanDecoderLayer, config) - self.ln_f = RmsNorm(normalized_shape=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): - 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) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_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 BaichuanForCausalLM(DecoderModelForCausalLM): - config_class = BaichuanConfig - - def __init__(self, config: PretrainedConfig): - transformer = BaichuanModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - lm_head = ColumnLinear(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) - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a BaichuanForCausalLM object from give parameters - ''' - import transformers - - assert hf_model_or_dir is not None - if isinstance(hf_model_or_dir, transformers.PreTrainedModel): - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_or_dir, - trust_remote_code=trust_remote_code, - dtype='auto') - hf_config_or_dir = hf_model_or_dir - - config = BaichuanConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from .convert import quantize - - config = BaichuanConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - quantize(hf_model_dir, - output_dir, - config=config, - device=device, - calib_dataset=calib_dataset) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) diff --git a/tensorrt_llm/models/bert/__init__.py b/tensorrt_llm/models/bert/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/bert/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/bert/config.py b/tensorrt_llm/models/bert/config.py deleted file mode 100644 index 5c2627aa7e23..000000000000 --- a/tensorrt_llm/models/bert/config.py +++ /dev/null @@ -1,117 +0,0 @@ -# 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 - -import torch -import transformers - -from ..._utils import torch_dtype_to_str -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class BERTConfig(PretrainedConfig): - - def __init__(self, - *, - is_roberta: bool = False, - type_vocab_size, - pad_token_id=None, - num_labels=None, - **kwargs): - self.is_roberta = is_roberta - self.type_vocab_size = type_vocab_size - self.pad_token_id = pad_token_id - self.num_labels = num_labels - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - output['is_roberta'] = self.is_roberta - output['type_vocab_size'] = self.type_vocab_size - output['pad_token_id'] = self.pad_token_id - output['num_labels'] = self.num_labels - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained(hf_config_dir) - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - head_dim = getattr( - hf_config, "head_dim", - hf_config.hidden_size // hf_config.num_attention_heads) - head_size = getattr(hf_config, "kv_channels", head_dim) - num_labels = getattr(hf_config, "num_labels", None) - - if (hf_config.position_embedding_type == 'absolute'): - position_embedding_type = 'learned_absolute' - else: - raise NotImplementedError( - f"{hf_config.position_embedding_type} hasn't been supported") - - if hf_config.model_type == "bert": - is_roberta = False - else: - is_roberta = True - - if dtype == 'auto': - dtype = getattr(hf_config, 'torch_dtype', None) - if dtype is None: - dtype = 'float16' - if isinstance(dtype, torch.dtype): - dtype = torch_dtype_to_str(dtype) - if dtype == 'float32': - dtype = 'float16' - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - hidden_size=hf_config.hidden_size, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - vocab_size=hf_config.vocab_size, - hidden_act=hf_config.hidden_act, - logits_dtype='float32', - norm_epsilon=hf_config.layer_norm_eps, - position_embedding_type=position_embedding_type, - max_position_embeddings=hf_config.max_position_embeddings, - num_key_value_heads=num_key_value_heads, - intermediate_size=hf_config.intermediate_size, - head_size=head_size, - quantization=quant_config, - mapping=mapping, - #BERT model args - is_roberta=is_roberta, - type_vocab_size=hf_config.type_vocab_size, - pad_token_id=hf_config.pad_token_id, - num_labels=num_labels, - **kwargs) diff --git a/tensorrt_llm/models/bert/convert.py b/tensorrt_llm/models/bert/convert.py deleted file mode 100644 index 90fe2a62bfe6..000000000000 --- a/tensorrt_llm/models/bert/convert.py +++ /dev/null @@ -1,309 +0,0 @@ -# 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 time -from typing import Union - -import torch - -# isort: off -from transformers import (AutoModel, AutoModelForQuestionAnswering, - AutoModelForSequenceClassification) -from transformers import (BertPreTrainedModel, RobertaPreTrainedModel) -# isort: on -from ...logger import logger -from ..convert_utils import split, split_qkv_bias_tp, split_qkv_tp -from .config import BERTConfig - - -def extract_layer_idx(name): - ss = name.split('.') - for s in ss: - if s.isdigit(): - return s - return None - - -def _load_weights_from_hf_bert_model(hf_model: Union[BertPreTrainedModel, - RobertaPreTrainedModel], - model_config: BERTConfig, - torch_dtype: torch.dtype = torch.float16): - weights = {} - no_match = {} - mapping = model_config.mapping - # use different prefix because BertModel is used both individually and as part of model - trtllm_prefix = "" if (model_config.architecture - in ["BertModel", "RobertaModel"]) else "bert." - for k, v in hf_model.state_dict().items(): - key = None - v = v.to(torch_dtype).cpu() - if 'embeddings.word_embeddings.weight' in k: - key = f'{trtllm_prefix}embedding.vocab_embedding.weight' - elif 'embeddings.position_embeddings.weight' in k: - key = f'{trtllm_prefix}embedding.position_embedding.weight' - elif 'embeddings.token_type_embeddings.weight' in k: - key = f'{trtllm_prefix}embedding.token_embedding.weight' - elif 'embeddings.LayerNorm.weight' in k: - key = f'{trtllm_prefix}embedding.embedding_ln.weight' - elif 'embeddings.LayerNorm.bias' in k: - key = f'{trtllm_prefix}embedding.embedding_ln.bias' - else: - layer_idx = extract_layer_idx(k) - if layer_idx is None: - no_match[k] = v - continue - idx = int(layer_idx) - if 'attention.output.dense.weight' in k: - #TODO: add TP support - key = f'{trtllm_prefix}layers.{idx}.attention.dense.weight' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=1) - elif 'attention.output.dense.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.dense.bias' - elif 'attention.output.LayerNorm.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.input_layernorm.weight' - elif 'attention.output.LayerNorm.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.input_layernorm.bias' - elif 'intermediate.dense.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.fc.weight' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=0) - elif 'intermediate.dense.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.fc.bias' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=0) - elif 'output.dense.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.proj.weight' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=1) - elif 'output.dense.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.proj.bias' - elif 'output.LayerNorm.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.post_layernorm.weight' - elif 'output.LayerNorm.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.post_layernorm.bias' - elif 'attention.self.query.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.q.weight' - elif 'attention.self.query.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.q.bias' - elif 'attention.self.key.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.k.weight' - elif 'attention.self.key.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.k.bias' - elif 'attention.self.value.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.v.weight' - elif 'attention.self.value.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.v.bias' - else: - no_match[k] = v - continue - weights[key] = v - - for idx in range(model_config.num_hidden_layers): - qkv_key = f'{trtllm_prefix}layers.{idx}.attention.qkv' - q_key = f'{trtllm_prefix}layers.{idx}.attention.q' - k_key = f'{trtllm_prefix}layers.{idx}.attention.k' - v_key = f'{trtllm_prefix}layers.{idx}.attention.v' - for postfix in ['weight', 'bias']: - v = torch.cat( - (weights[f'{q_key}.{postfix}'], weights[f'{k_key}.{postfix}'], - weights[f'{v_key}.{postfix}']), - dim=0) - v_clone = v.clone() - split_v = v_clone - if postfix == 'weight': - split_v = split_qkv_tp(v_clone, - model_config.num_attention_heads, - model_config.hidden_size, - mapping.tp_size, mapping.tp_rank) - - elif postfix == 'bias': - split_v = split_qkv_bias_tp(v_clone, - model_config.num_attention_heads, - model_config.hidden_size, - mapping.tp_size, mapping.tp_rank) - else: - assert True, f"Unknown postfix={postfix}!" - #add qkv weight/bias - weights[f'{qkv_key}.{postfix}'] = split_v - #remove separate q, k , v - del weights[f'{q_key}.{postfix}'] - del weights[f'{k_key}.{postfix}'] - del weights[f'{v_key}.{postfix}'] - return (weights, no_match) - - -def _load_weights_from_hf_bert_qa_model( - hf_model: Union[BertPreTrainedModel, RobertaPreTrainedModel], - model_config: BERTConfig, - torch_dtype: torch.dtype = torch.float16): - weights, no_match = _load_weights_from_hf_bert_model( - hf_model, model_config, torch_dtype) - - weights['qa_outputs.weight'] = no_match['qa_outputs.weight'] - - weights['qa_outputs.bias'] = no_match['qa_outputs.bias'] - del no_match['qa_outputs.weight'] - del no_match['qa_outputs.bias'] - - return (weights, no_match) - - -def _load_weights_from_hf_bert_cls_model( - hf_model: Union[BertPreTrainedModel, RobertaPreTrainedModel], - model_config: BERTConfig, - torch_dtype: torch.dtype = torch.float16): - - weights, no_match = _load_weights_from_hf_bert_model( - hf_model, model_config, torch_dtype) - - if model_config.is_roberta: - # roberta Version - weights['classifier.dense.weight'] = no_match['classifier.dense.weight'] - weights['classifier.dense.bias'] = no_match['classifier.dense.bias'] - weights['classifier.out_proj.weight'] = no_match[ - 'classifier.out_proj.weight'] - weights['classifier.out_proj.bias'] = no_match[ - 'classifier.out_proj.bias'] - del no_match['classifier.dense.weight'] - del no_match['classifier.dense.bias'] - del no_match['classifier.out_proj.weight'] - del no_match['classifier.out_proj.bias'] - else: - weights['pooler.dense.weight'] = no_match['bert.pooler.dense.weight'] - weights['pooler.dense.bias'] = no_match['bert.pooler.dense.bias'] - weights['classifier.weight'] = no_match['classifier.weight'] - weights['classifier.bias'] = no_match['classifier.bias'] - del no_match['bert.pooler.dense.weight'] - del no_match['bert.pooler.dense.bias'] - del no_match['classifier.weight'] - del no_match['classifier.bias'] - - return (weights, no_match) - - -def load_hf_bert_base(model_dir: str, - load_model_on_cpu: bool = False, - dtype: torch.dtype = torch.float16): - """ - load huggingface BertModel and RobertaModel model - """ - model = AutoModel.from_pretrained( - model_dir, - trust_remote_code=True, - ) - if not load_model_on_cpu: - model.cuda().to(dtype) - model.eval() - return model - - -def load_hf_bert_qa(model_dir: str, - load_model_on_cpu: bool = False, - dtype: torch.dtype = torch.float16): - """ - load huggingface BertForQuestionAnswering and RobertaForQuestionAnswering - """ - model = AutoModelForQuestionAnswering.from_pretrained( - model_dir, - trust_remote_code=True, - ) - if not load_model_on_cpu: - model.cuda().to(dtype) - model.eval() - return model - - -def load_hf_bert_cls(model_dir: str, - load_model_on_cpu: bool = False, - dtype: torch.dtype = torch.float16): - """ - load huggingface BertForSequenceClassification and RobertaForSequenceClassification - """ - model = AutoModelForSequenceClassification.from_pretrained( - model_dir, - trust_remote_code=True, - ) - if not load_model_on_cpu: - model.cuda().to(dtype) - model.eval() - return model - - -def load_weights_from_hf_model( - hf_model, - config: BERTConfig, -): - """ - load trtllm weights from hf model - - return a dict of weights, with trtllm weights naming - - """ - #TODO: add quantization support - weights = {} - tik = time.time() - - torch_dtype = getattr(torch, config.dtype) - - #NOTE: Bert - no_match = None - if config.architecture in [ - "BertForQuestionAnswering", "RobertaForQuestionAnswering" - ]: - weights, no_match = _load_weights_from_hf_bert_qa_model( - hf_model=hf_model, model_config=config, torch_dtype=torch_dtype) - elif config.architecture in ["BertModel", "RobertaModel"]: - weights, no_match = _load_weights_from_hf_bert_model( - hf_model=hf_model, model_config=config, torch_dtype=torch_dtype) - elif config.architecture in [ - "BertForSequenceClassification", "RobertaForSequenceClassification" - ]: - weights, no_match = _load_weights_from_hf_bert_cls_model( - hf_model=hf_model, model_config=config, torch_dtype=torch_dtype) - else: - assert False, f"Unknown BERT model {config.architecture}" - - if no_match is not None: - logger.warning( - f"These weights from huggingface model are not used:\n {[key for key in no_match.keys()]}" - ) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: BERTConfig, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail'): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - logger.warning(f"FP8 Support for Bert will come soon!") diff --git a/tensorrt_llm/models/bert/model.py b/tensorrt_llm/models/bert/model.py deleted file mode 100644 index a276ac8edf62..000000000000 --- a/tensorrt_llm/models/bert/model.py +++ /dev/null @@ -1,548 +0,0 @@ -# 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, OrderedDict, Union - -import numpy as np -import tensorrt as trt -import torch -import transformers - -from tensorrt_llm.models.modeling_utils import PretrainedModel - -from ..._common import default_net -from ...functional import (ACT2FN, Tensor, concat, constant, cumsum, expand, - index_select, select, shape, slice, unsqueeze) -from ...layers import MLP, BertAttention, Embedding, LayerNorm, Linear -from ...mapping import Mapping -from ...module import Module, ModuleList -from ..modeling_utils import QuantConfig -from .config import BERTConfig -from .convert import (load_hf_bert_base, load_hf_bert_cls, load_hf_bert_qa, - load_weights_from_hf_model) - - -class BertEmbedding(Module): - - def __init__(self, - vocab_size, - hidden_size, - max_position_embeddings, - type_vocab_size, - dtype=None): - super().__init__() - self.vocab_embedding = Embedding(vocab_size, hidden_size, dtype=dtype) - self.position_embedding = Embedding(max_position_embeddings, - hidden_size, - dtype=dtype) - self.token_embedding = Embedding(type_vocab_size, - hidden_size, - dtype=dtype) - self.max_position_embeddings = max_position_embeddings - - self.embedding_ln = LayerNorm(normalized_shape=hidden_size, dtype=dtype) - - def forward(self, input_ids, position_ids, token_type_ids): - x = self.vocab_embedding(input_ids) - x = x + self.position_embedding(position_ids) - x = x + self.token_embedding(token_type_ids) - x = self.embedding_ln(x) - return x - - -class BertEncoderLayer(Module): - - def __init__(self, - hidden_size, - num_attention_heads, - max_position_embeddings, - hidden_act='relu', - tp_group=None, - tp_size=1, - dtype=None): - super().__init__() - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - self.attention = BertAttention( - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=hidden_act, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - def forward(self, - hidden_states, - attention_mask=None, - input_lengths=None, - max_input_length=None): - residual = hidden_states - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - input_lengths=input_lengths, - max_input_length=max_input_length) - - hidden_states = residual + attention_output - - hidden_states = self.input_layernorm(hidden_states) - - residual = hidden_states - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - hidden_states = self.post_layernorm(hidden_states) - - return hidden_states - - -class BertBase(PretrainedModel): - ''' - Base class that provides from_huggingface() and prepare_inputs() methods - ''' - config_class = BERTConfig - - def __init__(self, config: BERTConfig): - super().__init__(config) - - @classmethod - def load_hf_bert(cls, model_dir: str, load_model_on_cpu: bool, - dtype: torch.dtype): - """ - Use as the abstractmethod, load corresponding HF model. - Subclass must implement this method! - """ - - assert cls.__name__ != "BertBase", f"Never call from BertBase class!" - - if cls.__name__ == "BertModel": - return load_hf_bert_base(model_dir, load_model_on_cpu, dtype) - elif cls.__name__ == "BertForQuestionAnswering": - return load_hf_bert_qa(model_dir, load_model_on_cpu, dtype) - elif cls.__name__ == "BertForSequenceClassification": - return load_hf_bert_cls(model_dir, load_model_on_cpu, dtype) - else: - assert False, f"Unknown class {cls.__name__}!" - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'float16', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - """ - Create a BertModel object from give parameters - """ - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - tllm_config = BERTConfig.from_hugging_face( - hf_config_or_dir=hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - #NOTE: override architecture info - RobertaCls_mapping = { - "BertModel": "RobertaModel", - "BertForQuestionAnswering": "RobertaForQuestionAnswering", - "BertForSequenceClassification": "RobertaForSequenceClassification", - } - if tllm_config.is_roberta: - setattr(tllm_config, 'architecture', - RobertaCls_mapping[cls.__name__]) - else: - setattr(tllm_config, 'architecture', cls.__name__) - - torch_dtype = torch.float16 if dtype == 'float16' else torch.float32 - if not use_preloading: - hf_model = cls.load_hf_bert(model_dir=hf_model_dir, - load_model_on_cpu=load_model_on_cpu, - dtype=torch_dtype) - weights = load_weights_from_hf_model(hf_model=hf_model, - config=tllm_config) - model = cls(tllm_config) - model.load(weights) - - return model - - # Override the PretrainedModel's meothd, can unify in the future. - def prepare_inputs(self, max_batch_size, max_input_len, **kwargs): - remove_input_padding = default_net().plugin_config.remove_input_padding - # opt_shape is set to half of max batch_size and seq_len by default - # tune this according to real data distribution - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - inlen_range = [1, (max_input_len + 1) // 2, max_input_len] - num_tokens_range = [ - 1, - (max_input_len * max_batch_size + 1) // 2, - max_input_len * max_batch_size, - ] - if not remove_input_padding: - input_ids = Tensor( - name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - # also called segment_ids - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - else: - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("num_tokens", [num_tokens_range])]), - ) - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', [num_tokens_range])]), - ) - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', [num_tokens_range])]), - ) - max_input_length = Tensor( - name="max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("max_input_length", [inlen_range])]), - ) - input_lengths = Tensor(name='input_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', [bs_range]) - ])) - - inputs = { - 'input_ids': input_ids, - 'input_lengths': input_lengths, - 'token_type_ids': token_type_ids, - } - - if remove_input_padding: - inputs['position_ids'] = position_ids - inputs['max_input_length'] = max_input_length - - return inputs - - -class BertModel(BertBase): - - def __init__(self, config: BERTConfig): - super().__init__(config) - - self.config = config - self.max_position_embeddings = config.max_position_embeddings - self.padding_idx = config.pad_token_id - self.is_roberta = config.is_roberta - self.embedding = BertEmbedding( - vocab_size=config.vocab_size, - hidden_size=config.hidden_size, - max_position_embeddings=config.max_position_embeddings, - type_vocab_size=config.type_vocab_size, - dtype=config.dtype) - - self.layers = ModuleList([ - BertEncoderLayer( - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - max_position_embeddings=config.max_position_embeddings, - hidden_act=config.hidden_act, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - dtype=config.dtype) for _ in range(config.num_hidden_layers) - ]) - - def forward(self, - input_ids=None, - input_lengths=None, - position_ids=None, - token_type_ids=None, - hidden_states=None, - max_input_length=None): - # remove_input_padding requires these fields as explicit input - mask = None - if not default_net().plugin_config.remove_input_padding: - seq_len_2d = concat([1, shape(input_ids, 1)]) - - # create position ids - position_ids_buffer = constant( - np.expand_dims( - np.arange(self.max_position_embeddings).astype(np.int32), - 0)) - tmp_position_ids = slice(position_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - tmp_position_ids = expand(tmp_position_ids, shape(input_ids)) #BxL - tmp_input_lengths = unsqueeze(input_lengths, 1) #Bx1 - tmp_input_lengths = expand(tmp_input_lengths, - shape(input_ids)) #BxL - mask = tmp_position_ids < tmp_input_lengths # BxL - mask = mask.cast('int32') - - if position_ids is None: - if self.is_roberta: - # see create_position_ids_from_input_ids() in https://github.com/huggingface/transformers/blob/main/src/transformers/models/roberta/modeling_roberta.py - position_ids = (tmp_position_ids + 1) * mask - position_ids = position_ids + self.padding_idx - else: - position_ids = slice(position_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - position_ids = expand(position_ids, shape(input_ids)) - - # create token_type_ids - if token_type_ids is None: - token_type_ids_buffer = constant( - np.expand_dims( - np.zeros(self.max_position_embeddings).astype(np.int32), - 0)) - token_type_ids = slice(token_type_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - token_type_ids = expand(token_type_ids, shape(input_ids)) - - hidden_states = self.embedding(input_ids, position_ids, token_type_ids) - self.register_network_output('embedding_output', hidden_states) - - for idx, layer in enumerate(self.layers): - hidden_states = layer(hidden_states=hidden_states, - input_lengths=input_lengths, - attention_mask=mask, - max_input_length=max_input_length) - # keep the last layer output name as hidden_states - if ((idx == (self.config.num_hidden_layers - 1)) and - (self.config.architecture in ["BertModel", "RobertaModel"])): - hidden_states.mark_output('hidden_states', self.config.dtype) - else: - self.register_network_output(f"layer_{idx}_output", - hidden_states) - - return hidden_states - - -RobertaModel = BertModel - - -class BertForQuestionAnswering(BertBase): - - def __init__(self, config: BERTConfig): - super().__init__(config) - self.bert = BertModel(config) - self.num_labels = config.num_labels - self.qa_outputs = Linear(config.hidden_size, - config.num_labels, - dtype=config.dtype) - - def forward(self, - input_ids=None, - input_lengths=None, - token_type_ids=None, - position_ids=None, - hidden_states=None, - max_input_length=None): - - remove_input_padding = default_net().plugin_config.remove_input_padding - if remove_input_padding: - assert token_type_ids is not None and \ - position_ids is not None and \ - max_input_length is not None, \ - "token_type_ids, position_ids, max_input_length is required " \ - "in remove_input_padding mode" - hidden_states = self.bert.forward(input_ids=input_ids, - input_lengths=input_lengths, - token_type_ids=token_type_ids, - position_ids=position_ids, - hidden_states=hidden_states, - max_input_length=max_input_length) - - logits = self.qa_outputs(hidden_states) - logits.mark_output('logits', self.config.logits_dtype) - - return logits - - -RobertaForQuestionAnswering = BertForQuestionAnswering - - -class BertPooler(Module): - - def __init__(self, hidden_size, dtype): - super().__init__() - self.dense = Linear(hidden_size, hidden_size, dtype=dtype) - self.activation = ACT2FN['tanh'] - - def forward(self, hidden_states, input_lengths, remove_input_padding): - if not remove_input_padding: - # We "pool" the model by simply taking the hidden state corresponding - # to the first token. - first_token_tensor = select(hidden_states, 1, 0) - else: - # when remove_input_padding is enabled, the shape of hidden_states is [num_tokens, hidden_size] - # We can take the first token of each sequence according to input_lengths, - # and then do pooling similar to padding mode. - # For example, if input_lengths is [8, 5, 6], then the indices of first tokens - # should be [0, 8, 13] - first_token_indices = cumsum( - concat([ - 0, - slice(input_lengths, - starts=[0], - sizes=(shape(input_lengths) - - constant(np.array([1], dtype=np.int32)))) - ]), 0) - first_token_tensor = index_select(hidden_states, 0, - first_token_indices) - - pooled_output = self.dense(first_token_tensor) - pooled_output = self.activation(pooled_output) - return pooled_output - - -class RobertaClassificationHead(Module): - """Head for sentence-level classification tasks.""" - - def __init__(self, hidden_size, dtype, num_labels): - super().__init__() - self.dense = Linear(hidden_size, hidden_size, dtype=dtype) - self.out_proj = Linear(hidden_size, num_labels) - - def forward(self, hidden_states, input_lengths, remove_input_padding): - - if not remove_input_padding: - # We "pool" the model by simply taking the hidden state corresponding - # to the first token. - first_token_tensor = select(hidden_states, 1, 0) - else: - # when remove_input_padding is enabled, the shape of hidden_states is [num_tokens, hidden_size] - # We can take the first token of each sequence according to input_lengths, - # and then do pooling similar to padding mode. - # For example, if input_lengths is [8, 5, 6], then the indices of first tokens - # should be [0, 8, 13] - first_token_indices = cumsum( - concat([ - 0, - slice(input_lengths, - starts=[0], - sizes=(shape(input_lengths) - - constant(np.array([1], dtype=np.int32)))) - ]), 0) - first_token_tensor = index_select(hidden_states, 0, - first_token_indices) - - x = self.dense(first_token_tensor) - x = ACT2FN['tanh'](x) - x = self.out_proj(x) - return x - - -class BertForSequenceClassification(BertBase): - - def __init__(self, config: BERTConfig): - super().__init__(config) - - self.config = config - self.is_roberta = config.is_roberta - self.bert = BertModel(config) - self.num_labels = config.num_labels - - if not config.is_roberta: - self.pooler = BertPooler(hidden_size=config.hidden_size, - dtype=config.dtype) - self.classifier = Linear(config.hidden_size, - config.num_labels, - dtype=config.dtype) - else: - self.classifier = RobertaClassificationHead( - hidden_size=config.hidden_size, - num_labels=config.num_labels, - dtype=config.dtype) - - def forward(self, - input_ids, - input_lengths, - token_type_ids=None, - position_ids=None, - hidden_states=None, - max_input_length=None): - - remove_input_padding = default_net().plugin_config.remove_input_padding - - # required as explicit input in remove_input_padding mode - # see examples/models/core/bert/run_remove_input_padding.py for how to create them from input_ids and input_lengths - if remove_input_padding: - assert token_type_ids is not None and \ - position_ids is not None and \ - max_input_length is not None, \ - "token_type_ids, position_ids, max_input_length is required " \ - "in remove_input_padding mode" - - hidden_states = self.bert.forward(input_ids=input_ids, - input_lengths=input_lengths, - token_type_ids=token_type_ids, - position_ids=position_ids, - hidden_states=hidden_states, - max_input_length=max_input_length) - - if not self.is_roberta: - pooled_output = self.pooler( - hidden_states=hidden_states, - input_lengths=input_lengths, - remove_input_padding=remove_input_padding) - logits = self.classifier(pooled_output) - else: - logits = self.classifier(hidden_states=hidden_states, - input_lengths=input_lengths, - remove_input_padding=remove_input_padding) - - logits.mark_output('logits', self.config.logits_dtype) - return logits - - -RobertaForSequenceClassification = BertForSequenceClassification diff --git a/tensorrt_llm/models/bloom/__init__.py b/tensorrt_llm/models/bloom/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/bloom/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/bloom/model.py b/tensorrt_llm/models/bloom/model.py deleted file mode 100644 index a189f7c3b740..000000000000 --- a/tensorrt_llm/models/bloom/model.py +++ /dev/null @@ -1,165 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm, PositionEmbeddingType) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class BloomDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - num_layers=config.num_hidden_layers, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - position_embedding_type=PositionEmbeddingType.alibi, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - reorder=True, - quant_mode=config.quant_mode) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act='gelu', - dtype=dtype, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class BloomModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.ln_embed = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(BloomDecoderLayer, config) - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_params=None): - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.ln_embed(hidden_states) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_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 BloomForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - transformer = BloomModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(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) - - super().__init__(config, transformer, lm_head) diff --git a/tensorrt_llm/models/chatglm/__init__.py b/tensorrt_llm/models/chatglm/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/chatglm/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/chatglm/config.py b/tensorrt_llm/models/chatglm/config.py deleted file mode 100644 index 9ae0be3152dc..000000000000 --- a/tensorrt_llm/models/chatglm/config.py +++ /dev/null @@ -1,182 +0,0 @@ -# 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 ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - -GLM_VERSIONS = ['glm4', 'chatglm3', 'chatglm2', 'chatglm', 'glm'] -GLM_ARCH1_VERSIONS = ['chatglm', 'glm'] -GLM_ARCH2_VERSIONS = ['glm4', 'chatglm3', 'chatglm2'] - - -class ChatGLMConfig(PretrainedConfig): - - def __init__(self, - *, - chatglm_version: str = 'chatglm3', - add_bias_linear: bool = False, - add_qkv_bias: bool = True, - apply_query_key_layer_scaling: bool = False, - apply_residual_connection_post_layernorm: bool = False, - rmsnorm: bool = True, - rotary_pct: float = 0.5, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - self.chatglm_version = chatglm_version - self.add_bias_linear = add_bias_linear - self.add_qkv_bias = add_qkv_bias - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm - self.rmsnorm = rmsnorm - self.rotary_pct = rotary_pct - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in ChatGLMConfig - output['chatglm_version'] = self.chatglm_version - output['add_bias_linear'] = self.add_bias_linear - output['add_qkv_bias'] = self.add_qkv_bias - output[ - 'apply_query_key_layer_scaling'] = self.apply_query_key_layer_scaling - output[ - 'apply_residual_connection_post_layernorm'] = self.apply_residual_connection_post_layernorm - output['rmsnorm'] = self.rmsnorm - output['rotary_pct'] = self.rotary_pct - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - # load hugging face config - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - logits_dtype = kwargs.pop('logits_dtype', 'float32') - use_parallel_embedding = kwargs.pop('use_parallel_embedding', False) - embedding_sharding_dim = kwargs.pop('embedding_sharding_dim', 0) - chatglm_version = kwargs.pop('chatglm_version', None) - - # get chatglm version - if chatglm_version is None: - print("Inferring chatglm version from path...") - for v in GLM_VERSIONS: - if v in hf_config._name_or_path: - chatglm_version = v - break - if 'glm_4' in hf_config._name_or_path.replace("-", "_"): - chatglm_version = 'glm4' - assert chatglm_version in GLM_VERSIONS - print(f"Chatglm version: {chatglm_version}") - - if chatglm_version == 'glm': - hf_config.num_kv_heads = hf_config.num_attention_heads - hf_config.ffn_hidden_size = hf_config.hidden_size * 4 - hf_config.hidden_act = 'gelu' - hf_config.layernorm_epsilon = 1e-5 - hf_config.max_position_embeddings = hf_config.max_sequence_length - hf_config.add_bias_linear = True - hf_config.add_qkv_bias = True - hf_config.apply_query_key_layer_scaling = False - hf_config.apply_residual_connection_post_layernorm = False - hf_config.rmsnorm = False - hf_config.rope_ratio = 1.0 - elif chatglm_version == 'chatglm': - hf_config.num_kv_heads = hf_config.num_attention_heads - hf_config.ffn_hidden_size = hf_config.inner_hidden_size - hf_config.hidden_act = 'gelu' - hf_config.max_position_embeddings = hf_config.max_sequence_length - hf_config.add_bias_linear = True - hf_config.add_qkv_bias = True - hf_config.apply_query_key_layer_scaling = False - hf_config.apply_residual_connection_post_layernorm = False - hf_config.rmsnorm = False - hf_config.rope_ratio = 1.0 - else: - hf_config.vocab_size = hf_config.padded_vocab_size - hf_config.num_kv_heads = hf_config.multi_query_group_num - hf_config.hidden_act = 'swiglu' - hf_config.max_position_embeddings = hf_config.seq_length - hf_config.rmsnorm = getattr(hf_config, 'rmsnorm', 1.0) - hf_config.rope_ratio = getattr(hf_config, 'rope_ratio', 1.0) - - if chatglm_version == 'glm': - position_embedding_type = 'learned_absolute' - elif chatglm_version == 'chatglm': - position_embedding_type = 'chatglm' - elif chatglm_version in GLM_ARCH2_VERSIONS: - position_embedding_type = 'rope_gptj' - - rotary_base = 10000.0 - rotary_embedding_scaling = None - if chatglm_version == 'chatglm2': - if hf_config.rope_ratio > 1: - rotary_embedding_scaling = { - 'type': 'linear', - 'factor': hf_config.rope_ratio - } - elif chatglm_version == 'chatglm3' or chatglm_version == 'glm4': - rotary_base *= hf_config.rope_ratio - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - logits_dtype=logits_dtype, - num_hidden_layers=hf_config.num_layers, - num_attention_heads=hf_config.num_attention_heads, - num_key_value_heads=hf_config.num_kv_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.ffn_hidden_size, - norm_epsilon=hf_config.layernorm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type=position_embedding_type, - max_position_embeddings=hf_config.max_position_embeddings, - rotary_pct=0.5, - rotary_base=rotary_base, - rotary_scaling=rotary_embedding_scaling, - hidden_act=hf_config.hidden_act, - use_parallel_embedding=use_parallel_embedding, - embedding_sharding_dim=embedding_sharding_dim, - quantization=quant_config, - mapping=mapping, - chatglm_version=chatglm_version, - add_bias_linear=hf_config.add_bias_linear, - add_qkv_bias=hf_config.add_qkv_bias, - apply_query_key_layer_scaling=False, - apply_residual_connection_post_layernorm=hf_config. - apply_residual_connection_post_layernorm, - rmsnorm=hf_config.rmsnorm, - ) diff --git a/tensorrt_llm/models/chatglm/convert.py b/tensorrt_llm/models/chatglm/convert.py deleted file mode 100644 index 6113de88300d..000000000000 --- a/tensorrt_llm/models/chatglm/convert.py +++ /dev/null @@ -1,727 +0,0 @@ -import copy -import functools -import os -import time -from collections import defaultdict -from typing import Dict, Optional - -import numpy as np -import safetensors -import torch -from tqdm import tqdm -from transformers import AutoModel, AutoTokenizer - -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.models import ChatGLMConfig -from tensorrt_llm.models.convert_utils import (generate_int8, get_weight, - get_weight_and_bias, - load_calib_dataset, smooth_gemm) -from tensorrt_llm.quantization import QuantAlgo - -from .config import GLM_ARCH1_VERSIONS, GLM_ARCH2_VERSIONS - - -def split(weight: torch.Tensor, - tp_size: int, - rank: int = 0, - dim: int = 0) -> torch.Tensor: - if tp_size == 1: - return weight - elif weight.ndim == 1: - return torch.chunk(weight, tp_size)[rank].contiguous() - else: - return torch.chunk(weight, tp_size, dim=dim)[rank].contiguous() - - -def tile_kv_weight_bias(v: torch.Tensor, kv_num_head: int, tp_size: int): - head_size = v.shape[0] // kv_num_head - reps = tp_size // kv_num_head - if v.ndim == 1: - v = v.reshape(kv_num_head, head_size)[:, None, :] - v = v.expand(kv_num_head, reps, head_size).reshape(-1).clone() - else: - hidden_size = v.shape[1] - v = v.reshape(kv_num_head, head_size, hidden_size)[:, None, :, :] - v = v.expand(kv_num_head, reps, head_size, - hidden_size).reshape(-1, hidden_size).clone() - return v - - -def split_qkv(v: torch.Tensor, tp_size: int, rank: int, hidden_size: int, - num_heads: int, num_kv_heads: int): - head_size = hidden_size // num_heads - if tp_size == 1: - return v - - assert v.shape[0] == hidden_size + head_size * num_kv_heads * 2 - query = v[:hidden_size] - key = v[hidden_size:hidden_size + head_size * num_kv_heads] - value = v[hidden_size + head_size * num_kv_heads:hidden_size + - head_size * num_kv_heads * 2] - - if num_kv_heads < tp_size: - key = tile_kv_weight_bias(key, num_kv_heads, tp_size) - value = tile_kv_weight_bias(value, num_kv_heads, tp_size) - assert (key.shape[0] % (tp_size * head_size)) == 0 - assert (value.shape[0] % (tp_size * head_size)) == 0 - - q_tmp = torch.chunk(query, tp_size, dim=0)[rank] - k_tmp = torch.chunk(key, tp_size, dim=0)[rank] - v_tmp = torch.chunk(value, tp_size, dim=0)[rank] - return torch.concatenate([q_tmp, k_tmp, v_tmp], dim=0).contiguous() - - -def split_embedding( - param: torch.Tensor, - tp_size: int, - tp_rank: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_size, tp_rank, dim=sharding_dim) - - -def swap_and_split_mlp(weight: torch.Tensor, - tp_size: int, - tp_rank: int, - dim: int = 0) -> torch.Tensor: - """Swap the positions of gate and fc weights, and split weights for tensor parallel. - """ - gate_weight, fc_weight = torch.chunk(weight, 2, dim=dim) - fc_w = split(fc_weight, tp_size, tp_rank, dim=dim) - gate_w = split(gate_weight, tp_size, tp_rank, dim=dim) - return torch.cat([fc_w, gate_w], dim=dim).contiguous() - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -@torch.no_grad() -def capture_activation_range( - model, - tokenizer, - dataset, - num_samples=64, - seq_len=512, -): - - model.eval() - device = next(model.parameters()).device - scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - def stat_tensor(name, tensor, key): - tensor = tensor.view(-1, tensor.shape[-1]).detach() - comming_max = tensor.abs().max(dim=0)[0].float() - if scales[name][key] is None: - scales[name][key] = comming_max - else: - scales[name][key] = torch.max(scales[name][key], comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, "x") - stat_tensor(name, y, "y") - # TODO: we don't need to do it every forward because inference does not change weight - if scales[name]["w"] is None: - scales[name]["w"] = m.weight.abs().clip(1e-8, None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, torch.nn.Linear): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="Calibration"): - input_ids = tokenizer( - dataset[i], - return_tensors="pt", - max_length=seq_len, - truncation=True, - ) - model(input_ids.input_ids.to(device)) - - for h in hooks: - h.remove() - - return scales - - -@torch.no_grad() -def smooth_chatglm_model( - model, - act_range, - alpha, - model_smoother, -): - for name, module in model.named_modules(): - if not module._get_name() == "GLMBlock": - continue - - # QKV multiplication weight - layer_name = name + '.self_attention.query_key_value' - print(f'Smoothing module: {layer_name}') - weight = module.self_attention.query_key_value.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - module.input_layernorm.weight, - None, - alpha, - ) - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - # Dense multiplication weight - layer_name = name + ".self_attention.dense" - print(f'Smoothing module: {layer_name}') - weight = module.self_attention.dense.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - None, - None, - alpha, - ) - model_smoother[layer_name] = smoother.float() - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - # Multilayer perceptron h -> 4h weight - layer_name = name + ".mlp.dense_h_to_4h" - print(f'Smoothing module: {layer_name}') - weight = module.mlp.dense_h_to_4h.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - module.post_attention_layernorm.weight, - None, - alpha, - ) - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - # Multilayer perceptron 4h -> h weight - layer_name = name + ".mlp.dense_4h_to_h" - print(f'Smoothing module: {layer_name}') - weight = module.mlp.dense_4h_to_h.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - None, - None, - alpha, - ) - model_smoother[layer_name] = smoother.float() - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - smoother_value=None, - smoother_shape=None): - results = {} - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - - cur_weights = original_weights - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float16)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - - results[prefix + - 'per_channel_scale'] = cur_per_channel_value.reshape(col_shape) - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - cur_weights = original_weights - - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - results[prefix + 'smoother'] = smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - return results - - -def load_weights_from_hf_model(hf_model: AutoModel, - config: ChatGLMConfig, - act_range: Optional[dict] = None, - smoother: Optional[dict] = None): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - num_kv_heads = config.num_key_value_heads - num_hidden_layers = config.num_hidden_layers - - chatglm_version = config.chatglm_version - mapping = config.mapping - use_parallel_embedding = config.use_parallel_embedding - sharding_dim = config.embedding_sharding_dim - - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - use_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - - if use_weight_only: - 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 - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - if chatglm_version in GLM_ARCH1_VERSIONS: - prefix = f'transformer.layers.{l}' - elif chatglm_version in GLM_ARCH2_VERSIONS: - prefix = f'transformer.encoder.layers.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - - # Attention QKV - attention_attr_name = '' - if chatglm_version in GLM_ARCH1_VERSIONS: - attention_attr_name = 'attention' - elif chatglm_version in GLM_ARCH2_VERSIONS: - attention_attr_name = 'self_attention' - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, f'{prefix}.{attention_attr_name}.query_key_value', - dtype) - - if use_smooth_quant: - qkv_act_range = act_range.get( - f'{prefix}.{attention_attr_name}.query_key_value') - qkv_vals_int8 = generate_int8(qkv_weight.t(), - qkv_act_range, - is_qkv=True, - multi_query_mode=True) - weights.update( - get_tllm_linear_sq_weight( - vals=qkv_vals_int8, - prefix=f'{tllm_prex}.attention.qkv.', - shape=[1, qkv_weight.size(0)], - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None)) - if qkv_bias is not None: - qkv_b = split_qkv(qkv_bias, - mapping.tp_size, - mapping.tp_rank, - hidden_size, - num_attention_heads, - num_kv_heads=num_kv_heads) - weights[f'{tllm_prex}.attention.qkv.bias'] = qkv_b - else: - qkv_w = split_qkv(qkv_weight, - mapping.tp_size, - mapping.tp_rank, - hidden_size, - num_attention_heads, - num_kv_heads=num_kv_heads) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = split_qkv(qkv_bias, - mapping.tp_size, - mapping.tp_rank, - hidden_size, - num_attention_heads, - num_kv_heads=num_kv_heads) - - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', - qkv_b, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_act_range = act_range.get( - f'{prefix}.{attention_attr_name}.query_key_value') - qkv_vals_int8 = generate_int8(qkv_weight.t(), - qkv_act_range, - is_qkv=True, - multi_query_mode=True) - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = qkv_vals_int8[ - 'scale_y_quant_orig'].contiguous() - - # Attention dense - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, f'{prefix}.{attention_attr_name}.dense', dtype) - - if use_smooth_quant: - dense_act_range = act_range.get( - f'{prefix}.{attention_attr_name}.dense') - dense_smoother = smoother.get( - f'{prefix}.{attention_attr_name}.dense') - dense_vals_int8 = generate_int8(attn_dense_weight.t(), - dense_act_range, - is_qkv=False, - multi_query_mode=True) - weights.update( - get_tllm_linear_sq_weight( - vals=dense_vals_int8, - prefix=f'{tllm_prex}.attention.dense.', - shape=[1, hidden_size], - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix= - f'{tllm_prex}.attention.quantization_scaling_factor', - smoother_value=dense_smoother, - smoother_shape=[1, hidden_size])) - if attn_dense_bias is not None: - weights[f'{tllm_prex}.attention.dense.bias'] = attn_dense_bias - else: - attn_dense_w = split(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, - f'{tllm_prex}.attention.dense', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP FC - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_h_to_4h', dtype) - - if use_smooth_quant: - fc_act_range = act_range.get(f'{prefix}.mlp.dense_h_to_4h') - fc_vals_int8 = generate_int8(mlp_fc_weight.t(), - fc_act_range, - is_qkv=False, - multi_query_mode=True) - cur_weights = get_tllm_linear_sq_weight( - vals=fc_vals_int8, - prefix=f'{tllm_prex}.mlp.fc.', - shape=[1, mlp_fc_weight.size(0)], - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - ) - cur_weights[f'{tllm_prex}.mlp.fc.weight'] = swap_and_split_mlp( - cur_weights[f'{tllm_prex}.mlp.fc.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0, - ) - if per_channel: - cur_weights[ - f'{tllm_prex}.mlp.fc.per_channel_scale'] = swap_and_split_mlp( - cur_weights[f'{tllm_prex}.mlp.fc.per_channel_scale'], - mapping.tp_size, - mapping.tp_rank, - dim=1, - ) - weights.update(cur_weights) - if chatglm_version in GLM_ARCH1_VERSIONS: - if mlp_fc_bias is not None: - mlp_fc_b = split(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b - elif chatglm_version in GLM_ARCH2_VERSIONS: - if mlp_fc_bias is not None: - mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size, - mapping.tp_rank) - weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b - else: - if chatglm_version in GLM_ARCH1_VERSIONS: - mlp_fc_w = split(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - if mlp_fc_bias is None: - mlp_fc_b = None - else: - mlp_fc_b = split(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - elif chatglm_version in GLM_ARCH2_VERSIONS: - mlp_fc_w = swap_and_split_mlp(mlp_fc_weight, mapping.tp_size, - mapping.tp_rank) - if mlp_fc_bias is None: - mlp_fc_b = None - else: - mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size, - mapping.tp_rank) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', - mlp_fc_b, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP Proj - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_4h_to_h', dtype) - - if use_smooth_quant: - proj_act_range = act_range.get(f'{prefix}.mlp.dense_4h_to_h') - proj_smoother = smoother.get(f'{prefix}.mlp.dense_4h_to_h') - proj_vals_int8 = generate_int8(mlp_proj_weight.t(), - proj_act_range, - is_qkv=False, - multi_query_mode=True) - weights.update( - get_tllm_linear_sq_weight( - vals=proj_vals_int8, - prefix=f'{tllm_prex}.mlp.proj.', - shape=[1, hidden_size], - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.mlp.quantization_scaling_factor', - smoother_value=proj_smoother, - smoother_shape=[1, config.intermediate_size])) - if mlp_proj_bias is not None: - weights[f'{tllm_prex}.mlp.proj.bias'] = mlp_proj_bias - else: - mlp_proj_w = split(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - if input_ln_bias is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.post_attention_layernorm', dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - if post_ln_bias is not None: - weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_bias - - if mapping.is_first_pp_rank(): - if chatglm_version == 'glm': - embed_w = get_weight(model_params, 'word_embeddings', dtype) - pos_embed_w = get_weight(model_params, - 'transformer.position_embeddings', dtype) - weights['transformer.position_embedding.weight'] = split_embedding( - pos_embed_w, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - block_embed_w = get_weight(model_params, - 'transformer.block_position_embeddings', - dtype) - weights['transformer.block_embedding.weight'] = split_embedding( - block_embed_w, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - elif chatglm_version == 'chatglm': - embed_w = get_weight(model_params, 'transformer.word_embeddings', - dtype) - elif chatglm_version in GLM_ARCH2_VERSIONS: - embed_w = get_weight(model_params, - 'transformer.embedding.word_embeddings', dtype) - - weights['transformer.vocab_embedding.weight'] = split_embedding( - embed_w, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - - if mapping.is_last_pp_rank(): - if chatglm_version == 'glm': - lm_head_weight = get_weight(model_params, 'word_embeddings', - dtype).clone() - elif chatglm_version == 'chatglm': - lm_head_weight = get_weight(model_params, - 'transformer.word_embeddings', - dtype).clone() - elif chatglm_version in GLM_ARCH2_VERSIONS: - lm_head_weight = get_weight(model_params, - 'transformer.output_layer', dtype) - - weights['lm_head.weight'] = split(lm_head_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - if chatglm_version in GLM_ARCH1_VERSIONS: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'transformer.final_layernorm', - dtype) - elif chatglm_version in GLM_ARCH2_VERSIONS: - ln_f_w, ln_f_b = get_weight_and_bias( - model_params, 'transformer.encoder.final_layernorm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: ChatGLMConfig, - calib_dataset: str = 'cnn_dailymail', - device: str = 'auto', - trust_remote_code: bool = True): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - if config.chatglm_version == 'glm': - device_map = 'cuda' if device != "cpu" else 'cpu' - else: - device_map = 'auto' if device != "cpu" else 'cpu' - hf_model = AutoModel.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - dtype='auto' if config.chatglm_version != 'glm' else getattr( - torch, config.dtype), - device_map=device_map) - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - ) - dataset = load_calib_dataset(calib_dataset) - - act_range = capture_activation_range(hf_model, - tokenizer, - dataset, - num_samples=64) - smoother = {} - if use_smooth_quant: - smooth_chatglm_model(hf_model, act_range, quant_config.smoothquant_val, - smoother) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model( - hf_model, - config=config, - act_range=act_range, - smoother=smoother, - ) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights diff --git a/tensorrt_llm/models/chatglm/model.py b/tensorrt_llm/models/chatglm/model.py deleted file mode 100644 index 8382419e8eb1..000000000000 --- a/tensorrt_llm/models/chatglm/model.py +++ /dev/null @@ -1,372 +0,0 @@ -# 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 - -import torch -from transformers import AutoModel - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import Tensor, concat, shape -from ...layers import (MLP, Attention, AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, KeyValueCacheParams, LayerNorm, - RmsNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import GLM_ARCH1_VERSIONS, GLM_ARCH2_VERSIONS, ChatGLMConfig -from .convert import load_weights_from_hf_model - - -class ChatGLMDecoderLayer(Module): - - def __init__(self, config: ChatGLMConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - self.chatglm_version = config.chatglm_version - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - layernorm_epsilon = config.norm_epsilon - - self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm - self.alpha = (2 * config.num_hidden_layers)**0.5 - norm_cls = RmsNorm if config.rmsnorm else LayerNorm - - if config.chatglm_version == 'glm': - attention_mask_type = AttentionMaskType.bidirectionalglm - elif config.chatglm_version == 'chatglm': - attention_mask_type = AttentionMaskType.bidirectional - elif config.chatglm_version in GLM_ARCH2_VERSIONS: - attention_mask_type = AttentionMaskType.causal - - self.input_layernorm = norm_cls( - normalized_shape=hidden_size, - eps=layernorm_epsilon, - elementwise_affine=True, - dtype=dtype, - ) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - num_layers=config.num_hidden_layers, - apply_query_key_layer_scaling=config.apply_query_key_layer_scaling, - attention_mask_type=attention_mask_type, - bias=config.add_qkv_bias, - dense_bias=config.add_bias_linear, - dtype=config.dtype, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - rotary_embedding_percentage=config.rotary_pct, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=config.quant_mode, - q_scaling=1.0, - cross_attention=False, - relative_attention=False, - max_distance=0, - num_buckets=0, - cp_rank=config.mapping.cp_rank, - cp_size=config.mapping.cp_size, - cp_group=config.mapping.cp_group, - ) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = MLP( - hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - bias=config.add_bias_linear, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode, - ) - - self.post_layernorm = norm_cls( - normalized_shape=hidden_size, - eps=layernorm_epsilon, - elementwise_affine=True, - dtype=dtype, - ) - - def forward( - self, - hidden_states: Tensor, - attention_mask: Tensor = None, - position_ids: Tensor = None, # only used in ChatGLM-6B - use_cache: bool = False, - kv_cache_params: KeyValueCacheParams = None, - attention_params: AttentionParams = None, - ): - norm_output = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states=norm_output, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - encoder_output=None, - position_embedding=position_ids, - ) - - if use_cache: - attention_output, presents = attention_output - - if self.chatglm_version == 'chatglm': - residual = norm_output - - norm_input = residual * self.alpha + attention_output - - norm_output = self.post_layernorm(norm_input) - - mlp_output = self.mlp(norm_output) - - residual = norm_output - - output = residual * self.alpha + mlp_output - - else: - residual = norm_output if self.apply_residual_connection_post_layernorm else hidden_states - - norm_input = residual + attention_output - - norm_output = self.post_layernorm(norm_input) - - mlp_output = self.mlp(norm_output) - - residual = norm_output if self.apply_residual_connection_post_layernorm else norm_input - - output = residual + mlp_output - - if use_cache: - return (output, presents) - return output - - -class ChatGLMModel(Module): - - def __init__(self, config: ChatGLMConfig): - super().__init__() - self.chatglm_version = config.chatglm_version - norm_cls = RmsNorm if config.rmsnorm else LayerNorm - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - if config.chatglm_version == 'glm': - self.position_embedding = Embedding( - config.max_position_embeddings + 1, - config.hidden_size, - dtype=config.dtype, - ) - self.block_embedding = Embedding( - config.max_position_embeddings + 1, - config.hidden_size, - dtype=config.dtype, - ) - - self.layers = DecoderLayerList(ChatGLMDecoderLayer, config) - - self.ln_f = norm_cls( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - elementwise_affine=True, - dtype=config.dtype, - ) - - def forward( - self, - input_ids: Tensor = None, - position_ids: Tensor = None, # only used in ChatGLM-6B - use_cache: bool = False, - attention_mask: Tensor = None, - kv_cache_params: KeyValueCacheParams = None, - attention_params: AttentionParams = None, - ): - hidden_states = self.vocab_embedding(input_ids) - - if self.chatglm_version == 'glm': - if default_net().plugin_config.remove_input_padding: - position_ids_list = position_ids.split(1, dim=0) - else: - position_ids_list = position_ids.split(1, dim=1) - - position_embedding = self.position_embedding(position_ids_list[0]) - block_embedding = self.block_embedding(position_ids_list[1]) - position_embedding = position_embedding + block_embedding - - if default_net().plugin_config.remove_input_padding: - position_embedding = position_embedding.view( - concat([ - shape(position_embedding, 1), - shape(position_embedding, 2) - ])) - else: - position_embedding = position_embedding.view( - concat([ - shape(position_embedding, 0), - shape(position_embedding, 2), - shape(position_embedding, 3), - ])) - - hidden_states = hidden_states + position_embedding - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - position_ids=position_ids) - - 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 ChatGLMForCausalLM(DecoderModelForCausalLM): - config_class = ChatGLMConfig - - def __init__(self, config: ChatGLMConfig): - transformer = ChatGLMModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(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) - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - trust_remote_code = kwargs.pop('trust_remote_code', True) - - config = ChatGLMConfig.from_hugging_face(hf_model_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if config.chatglm_version == 'glm': - device_map = 'cuda' if not load_model_on_cpu else 'cpu' - else: - device_map = 'auto' if not load_model_on_cpu else 'cpu' - hf_model = AutoModel.from_pretrained( - hf_model_or_dir, - trust_remote_code=trust_remote_code, - dtype='auto' if config.chatglm_version != 'glm' else getattr( - torch, config.dtype), - device_map=device_map) - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = ChatGLMConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - convert.quantize(hf_model_dir, - output_dir, - config=config, - calib_dataset=calib_dataset, - device=device) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - def prepare_inputs(self, *args, **kwargs): - """See `PretrainedModel.prepare_inputs` for the detailed parameter list. - """ - if self.transformer.chatglm_version in GLM_ARCH1_VERSIONS: - position_encoding_2d = True - else: - position_encoding_2d = False - return super().prepare_inputs(*args, - **kwargs, - position_encoding_2d=position_encoding_2d) diff --git a/tensorrt_llm/models/clip/__init__.py b/tensorrt_llm/models/clip/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/clip/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/clip/model.py b/tensorrt_llm/models/clip/model.py deleted file mode 100644 index df131434b122..000000000000 --- a/tensorrt_llm/models/clip/model.py +++ /dev/null @@ -1,213 +0,0 @@ -# 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 ...functional import arange, concat, expand, expand_dims, shape -from ...layers import MLP, BertAttention, Conv2d, Embedding, LayerNorm -from ...mapping import Mapping -from ...module import Module, ModuleList -from ...parameter import Parameter - - -# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py#L164 -class CLIPVisionEmbeddings(Module): - - def __init__(self, image_size, num_channels, patch_size, hidden_size, - dtype): - super().__init__() - self.image_size = image_size - self.num_channels = num_channels - self.patch_size = patch_size - self.embed_dim = hidden_size - self.dtype = dtype - - self.class_embedding = Parameter(shape=[ - self.embed_dim, - ], - dtype=self.dtype) - - self.patch_embedding = Conv2d(in_channels=self.num_channels, - out_channels=self.embed_dim, - kernel_size=(self.patch_size, - self.patch_size), - stride=(self.patch_size, self.patch_size), - bias=False, - dtype=self.dtype) - - self.num_patches = (self.image_size // self.patch_size)**2 - self.num_positions = self.num_patches + 1 - self.position_embedding = Embedding(self.num_positions, - self.embed_dim, - dtype=self.dtype) - - def forward(self, pixel_values): - batch_size = shape(pixel_values, 0) - target_dtype = self.patch_embedding.weight.dtype - patch_embeds = self.patch_embedding( - pixel_values.cast( - dtype=target_dtype)) # shape = [*, width, grid, grid] - patch_embeds = patch_embeds.flatten(2).transpose(1, 2) - class_embeds = expand_dims(expand_dims(self.class_embedding.value, 0), - 0) - expand_shape = concat( - [batch_size, - shape(class_embeds, -2), - shape(class_embeds, -1)]) - class_embeds = expand(class_embeds, - expand_shape) # shape = [*, 1, grid, grid] - embeddings = concat([class_embeds, patch_embeds], - dim=1) # shape = [*, width + 1, grid, grid] - position_ids = arange(0, self.num_positions, dtype='int32') - position_embeds = self.position_embedding(position_ids) - position_embeds = expand_dims(position_embeds, 0) - expand_shape = concat([ - batch_size, - shape(position_embeds, -2), - shape(position_embeds, -1) - ]) - position_embeds = expand( - position_embeds, expand_shape) # shape = [*, width + 1, grid, grid] - embeddings = embeddings + position_embeds - return embeddings - - -class CLIPEncoderLayer(Module): - - def __init__(self, hidden_size, num_attention_heads, - max_position_embeddings, norm_epsilon, intermediate_size, - hidden_act, mapping: Mapping, dtype): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.mapping = mapping - - self.input_layernorm = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - self.attention = BertAttention( - hidden_size=self.hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - attention_head_size=self.hidden_size // num_attention_heads, - num_kv_heads=num_attention_heads, - dtype=self.dtype, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - tp_rank=self.mapping.tp_rank, - cp_group=self.mapping.cp_group, - cp_size=self.mapping.cp_size) - - self.post_layernorm = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - self.mlp = MLP(hidden_size=self.hidden_size, - ffn_hidden_size=intermediate_size, - hidden_act=hidden_act, - dtype=self.dtype, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size) - - def forward(self, hidden_states): - - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.attention(hidden_states) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -class CLIPEncoder(Module): - - def __init__(self, hidden_size, num_attention_heads, - max_position_embeddings, norm_epsilon, intermediate_size, - hidden_act, num_hidden_layers, mapping: Mapping, dtype): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.mapping = mapping - - self.layers = ModuleList([ - CLIPEncoderLayer(hidden_size=self.hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - norm_epsilon=norm_epsilon, - intermediate_size=intermediate_size, - hidden_act=hidden_act, - mapping=self.mapping, - dtype=self.dtype) for _ in range(num_hidden_layers) - ]) - - def forward(self, inputs_embeds): - - hidden_states = inputs_embeds - for layer in self.layers: - hidden_states = layer(hidden_states) - - return hidden_states - - -class CLIPVisionTransformer(Module): - - def __init__(self, image_size, num_channels, patch_size, hidden_size, - num_attention_heads, max_position_embeddings, norm_epsilon, - intermediate_size, hidden_act, num_hidden_layers, require_ln_f, - mapping: Mapping, dtype) -> None: - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.mapping = mapping - - self.embeddings = CLIPVisionEmbeddings(image_size=image_size, - num_channels=num_channels, - patch_size=patch_size, - hidden_size=hidden_size, - dtype=self.dtype) - self.pre_layernorm = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - self.encoder = CLIPEncoder( - hidden_size=self.hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - norm_epsilon=norm_epsilon, - intermediate_size=intermediate_size, - hidden_act=hidden_act, - num_hidden_layers=num_hidden_layers, - mapping=self.mapping, - dtype=self.dtype) - - self.ln_f = None - - if require_ln_f: - self.ln_f = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - def forward(self, pixel_values): - hidden_states = self.embeddings(pixel_values) - hidden_states = self.pre_layernorm(hidden_states) - hidden_states = self.encoder(inputs_embeds=hidden_states) - - if self.ln_f is None: - return hidden_states - - return self.ln_f(hidden_states) diff --git a/tensorrt_llm/models/cogvlm/__init__.py b/tensorrt_llm/models/cogvlm/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/cogvlm/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/cogvlm/config.py b/tensorrt_llm/models/cogvlm/config.py deleted file mode 100644 index 1588a95f6a67..000000000000 --- a/tensorrt_llm/models/cogvlm/config.py +++ /dev/null @@ -1,44 +0,0 @@ -# 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 - -from ..modeling_utils import PretrainedConfig - - -class CogVLMConfig(PretrainedConfig): - - def __init__(self, - *, - mlp_bias: bool = False, - attn_bias: bool = False, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - self.mlp_bias = mlp_bias - self.attn_bias = attn_bias - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.position_embedding_type = 'rope_gpt_neox' - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in CogVLMConfig - output['mlp_bias'] = self.mlp_bias - output['attn_bias'] = self.attn_bias - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - return output diff --git a/tensorrt_llm/models/cogvlm/convert.py b/tensorrt_llm/models/cogvlm/convert.py deleted file mode 100644 index 68413e4e224a..000000000000 --- a/tensorrt_llm/models/cogvlm/convert.py +++ /dev/null @@ -1,240 +0,0 @@ -import time - -import numpy as np -import torch - -from tensorrt_llm.logger import logger - -from ..._utils import pad_vocab_size -from ..llama.convert import (get_tllm_linear_weight, get_weight, split, - split_matrix_tp, split_qkv_tp) - - -def convert_hf_cogvlm(hf_model, - mapping, - vocab_size=32000, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - use_gemm_woq_plugin=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[]): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - if hasattr(hf_model.config, "num_key_value_heads"): - num_key_value_heads = hf_model.config.num_key_value_heads - else: - num_key_value_heads = num_attention_heads - mha_mode = (num_key_value_heads == num_attention_heads) - layers_range = mapping.pp_layers(hf_model.config.num_hidden_layers) - assert mha_mode, "CogVLM only supports mha mode" - assert not use_smooth_quant, "CogVLM currently doesn't support smooth quant" - assert not int8_kv_cache, "CogVLM currently doesn't support int8 kv cache" - - for l in layers_range: - prefix = f'model.layers.{l}.' - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - - qkv_weight = get_weight( - model_params, prefix + 'self_attn.language_expert_query_key_value', - dtype) - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vis_qkv_weight = get_weight( - model_params, prefix + 'self_attn.vision_expert_query_key_value', - dtype) - split_v = split_qkv_tp(vis_qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.vis_qkv.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - attn_dense_weight = get_weight( - model_params, prefix + 'self_attn.language_expert_dense', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - attn_vision_dense_weight = get_weight( - model_params, prefix + 'self_attn.vision_expert_dense', dtype) - split_v = split_matrix_tp(attn_vision_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.vis_dense.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_gate_weight = get_weight(model_params, - prefix + 'mlp.language_mlp.up_proj', dtype) - split_v = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.gate.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vision_mlp_gate_weight = get_weight(model_params, - prefix + 'mlp.vision_mlp.up_proj', - dtype) - split_v = split_matrix_tp(vision_mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'vis_mlp.gate.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_fc_weight = get_weight(model_params, - prefix + 'mlp.language_mlp.gate_proj', dtype) - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vision_mlp_fc_weight = get_weight(model_params, - prefix + 'mlp.vision_mlp.gate_proj', - dtype) - split_v = split_matrix_tp(vision_mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'vis_mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_proj_weight = get_weight(model_params, - prefix + 'mlp.language_mlp.down_proj', - dtype) - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vision_mlp_proj_weight = get_weight(model_params, - prefix + 'mlp.vision_mlp.down_proj', - dtype) - split_v = split_matrix_tp(vision_mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'vis_mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - cur_block_weights = [ - weight_name for weight_name in model_params - if weight_name.find(prefix) != -1 - ] - for weight_name in cur_block_weights: - model_params[weight_name] = None - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - - v = get_weight(model_params, 'model.embed_tokens', dtype) - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - v = torch.from_numpy( - np.pad(v.detach().cpu().numpy(), ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/cogvlm/model.py b/tensorrt_llm/models/cogvlm/model.py deleted file mode 100644 index c5c8ea1d02d1..000000000000 --- a/tensorrt_llm/models/cogvlm/model.py +++ /dev/null @@ -1,291 +0,0 @@ -# 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 - -import numpy as np - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import (Tensor, concat, constant, expand, op_and, recv, send, - shape, slice, unsqueeze, where) -from ...layers import (AttentionMaskType, CogVLMAttention, ColumnLinear, - Embedding, GatedMLP, PromptTuningEmbedding, RmsNorm) -from ...mapping import Mapping -from ...module import Module -# this is to use to module global algo string with a quant_algo prefix -from ...quantization import QuantMode -from ...top_model_mixin import TopModelMixin -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import CogVLMConfig - - -class CogvlmDecoderLayer(Module): - - def __init__(self, config: CogVLMConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_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] - self.attention = CogVLMAttention( - 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, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.hidden_size = config.hidden_size - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - self.vis_mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward( - self, - hidden_states, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - vision_token_mask=None, - position_ids=None, - ): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - vision_token_mask=vision_token_mask, - position_embedding=position_ids) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - vision_mlp_out = self.vis_mlp(hidden_states) - language_mlp_out = self.mlp(hidden_states) - hidden_states = where(vision_token_mask, vision_mlp_out, - language_mlp_out) - - # hidden_states = self.mlp(hidden_states, - # lora_layer_params=lora_layer_params) - - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class CogvlmModel(Module): - - def __init__(self, config: CogVLMConfig) -> None: - super().__init__() - - self.mapping = config.mapping - self.use_prompt_tuning = config.use_prompt_tuning - self.vocab_size = config.vocab_size - EmbeddingCls = PromptTuningEmbedding if config.use_prompt_tuning else Embedding - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = EmbeddingCls( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - dtype=config.dtype, - tp_size=self.mapping.tp_size - if config.use_parallel_embedding else 1, - tp_group=self.mapping.tp_group - if config.use_parallel_embedding else None, - sharding_dim=config.embedding_sharding_dim, - tp_rank=self.mapping.tp_rank, - ) - - self.layers = DecoderLayerList(CogvlmDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - kv_cache_params.fill_none_tensor_list(len(self.layers)) - - if use_cache: - presents = [] - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if self.use_prompt_tuning else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - vision_mask = input_ids > (self.vocab_size - 1) - - if default_net().plugin_config.remove_input_padding: - seq_length = shape(vision_mask, 0) # lvvvvvvllvvvvlll - zero = constant(np.ascontiguousarray(np.zeros([1], dtype=bool))) - one = constant(np.ascontiguousarray(np.ones([1], dtype=bool))) - - t1 = slice(vision_mask, [0], seq_length - 1) - t2 = slice(vision_mask, [1], seq_length - 1) - vision_token_mask = concat([op_and(t1 == one, t2 == one), - zero]) # 0111110001110000 - vision_token_mask = unsqueeze(vision_token_mask, - -1) # [num_tokens, 1] - else: - seq_length = shape(vision_mask, - 1) # lvvvvvvllvvvvlll, lvvvvvvllvvvvlll - batch_size = shape(vision_mask, 0) - t1 = slice(vision_mask, [0, 0], concat([batch_size, - seq_length - 1])) - t2 = slice(vision_mask, [0, 1], concat([batch_size, - seq_length - 1])) - zero = expand( - constant(np.ascontiguousarray(np.zeros([1, 1], dtype=bool))), - concat([batch_size, 1])) - one = constant(np.ascontiguousarray(np.ones([1, 1], dtype=bool))) - - vision_token_mask = concat([op_and(t1 == one, t2 == one), zero], - dim=1) # 0111110001110000 [bs, seqlen] - vision_token_mask = unsqueeze(vision_token_mask, - -1) # [bs, seqlen, 1] - - hidden_states = self.layers.forward(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - vision_token_mask=vision_token_mask, - position_ids=position_ids) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class CogVLMForCausalLM(DecoderModelForCausalLM, TopModelMixin): - config_class = CogVLMConfig - - def __init__(self, config: CogVLMConfig): - transformer = CogvlmModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(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) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face(cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None, - quant_mode: Optional[QuantMode] = None, - **kwargs): - pass - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.weight_only_groupwise_quant_matmul_plugin = 'auto' - return plugin_config - - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantConfig, - *, - dtype='float16', - mapping: Optional[Mapping] = None, - calib_batches=512, - calib_batch_size=1, - random_seed=1234, - tokenizer_max_seq_length=2048, - **kwargs, - ): - pass diff --git a/tensorrt_llm/models/commandr/__init__.py b/tensorrt_llm/models/commandr/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/commandr/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/commandr/config.py b/tensorrt_llm/models/commandr/config.py deleted file mode 100644 index 511640c22494..000000000000 --- a/tensorrt_llm/models/commandr/config.py +++ /dev/null @@ -1,88 +0,0 @@ -# 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 - -import transformers - -from ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class CohereConfig(PretrainedConfig): - - def __init__(self, - *, - output_multiplier_scale: float = 0.0625, - rotary_base: float = 10000.0, - attn_bias: bool = False, - **kwargs): - self.output_multiplier_scale = output_multiplier_scale - self.rotary_base = rotary_base - self.attn_bias = attn_bias - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in CohereConfig - output['output_multiplier_scale'] = self.output_multiplier_scale - output['rotary_base'] = self.rotary_base - output['attn_bias'] = self.attn_bias - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_or_dir, trust_remote_code=True) - - head_size = hf_config.hidden_size // hf_config.num_attention_heads - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - if hf_config.tie_word_embeddings: - kwargs['use_parallel_embedding'] = True - kwargs['embedding_sharding_dim'] = 0 - - return CohereConfig( - 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, - head_size=head_size, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gptj', # different rope type - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act=hf_config.hidden_act, - norm_epsilon=hf_config.layer_norm_eps, - output_multiplier_scale=hf_config.logit_scale, - rotary_base=get_hf_rope_theta(hf_config, 10000.0), - attn_bias=hf_config.attention_bias, - qk_layernorm=hf_config.use_qk_norm, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/commandr/model.py b/tensorrt_llm/models/commandr/model.py deleted file mode 100644 index ac0fe67dd754..000000000000 --- a/tensorrt_llm/models/commandr/model.py +++ /dev/null @@ -1,195 +0,0 @@ -from typing import Optional - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import recv, send -from ...layers import (Attention, AttentionMaskType, ColumnLinear, Embedding, - GatedMLP, LayerNorm, PositionEmbeddingType) -from ...mapping import Mapping -from ...module import Module -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import CohereConfig - - -class CohereDecoderLayer(Module): - - def __init__(self, config: CohereConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=self.local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gptj, - rotary_embedding_base=config.rotary_base, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - qk_layernorm=config.qk_layernorm, - layernorm_share=False, - eps=config.norm_epsilon, - quant_mode=config.quant_mode) - - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=False, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - assert not ( - default_net().plugin_config.reduce_fusion - ), "Custom all reduce and residual mlp can't be enabled at the same time." - - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - mlp_output = self.mlp(hidden_states) - hidden_states = residual + attention_output + mlp_output - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class CohereModel(Module): - - def __init__(self, config: CohereConfig) -> None: - super().__init__() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank) - - self.layers = DecoderLayerList(CohereDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - def forward( - self, - input_ids=None, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - ): - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class CohereForCausalLM(DecoderModelForCausalLM): - config_class = CohereConfig - - def __init__(self, config: CohereConfig): - transformer = CohereModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(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) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face(cls, - hf_model_or_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a CohereForCausalLM object from give parameters - ''' - - config = CohereConfig.from_hugging_face(hf_model_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - model = cls(config) - custom_dict = { - 'q_layernorm': 'q_norm', - 'k_layernorm': 'k_norm', - } - loader = ModelWeightsLoader(hf_model_or_dir, custom_dict) - loader.generate_tllm_weights(model) - - return model diff --git a/tensorrt_llm/models/dbrx/__init__.py b/tensorrt_llm/models/dbrx/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/dbrx/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/dbrx/config.py b/tensorrt_llm/models/dbrx/config.py deleted file mode 100644 index d97311ff88ca..000000000000 --- a/tensorrt_llm/models/dbrx/config.py +++ /dev/null @@ -1,59 +0,0 @@ -# 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 ...layers import MoeConfig -from ..modeling_utils import PretrainedConfig - - -class DbrxConfig(PretrainedConfig): - - def __init__(self, - *, - bias: bool = False, - clip_qkv: Optional[float] = None, - rotary_base: float = 500000.0, - rotary_scaling: Optional[dict] = None, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - self.bias = bias - self.clip_qkv = clip_qkv - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - if moe is None: - # Legacy MOE config fields - moe = MoeConfig( - num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE)) - elif isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in DbrxConfig - output['bias'] = self.bias - output['clip_qkv'] = self.clip_qkv - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['moe'] = self.moe.to_dict() - return output diff --git a/tensorrt_llm/models/dbrx/model.py b/tensorrt_llm/models/dbrx/model.py deleted file mode 100644 index 14617e0188aa..000000000000 --- a/tensorrt_llm/models/dbrx/model.py +++ /dev/null @@ -1,188 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import Tensor, recv, send -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, GatedMLP, LayerNorm) -from ...module import Module -from ..modeling_utils import DecoderLayerList, DecoderModelForCausalLM -from .config import DbrxConfig - - -class DbrxDecoderLayer(Module): - - def __init__(self, config: DbrxConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - 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, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - clip_qkv=config.clip_qkv) - - ClsMLP = GatedMLP - mlp_kwargs = {} - if config.moe.has_moe(): - ClsMLP = MOE - mlp_kwargs = { - "moe_config": config.moe, - "mapping": config.mapping, - } - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class DbrxModel(Module): - - def __init__(self, config: DbrxConfig): - super().__init__() - self.config = config - - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(DbrxDecoderLayer, config) - - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None): - - if self.config.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids) - else: - hidden_states = recv(hidden_states, - self.config.mapping.prev_pp_rank()) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.config.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, - self.config.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class DbrxForCausalLM(DecoderModelForCausalLM): - config_class = DbrxConfig - - def __init__(self, config: DbrxConfig): - transformer = DbrxModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=config.bias, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) diff --git a/tensorrt_llm/models/deepseek_v1/__init__.py b/tensorrt_llm/models/deepseek_v1/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/deepseek_v1/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/deepseek_v1/config.py b/tensorrt_llm/models/deepseek_v1/config.py deleted file mode 100755 index e7bff0d9aab8..000000000000 --- a/tensorrt_llm/models/deepseek_v1/config.py +++ /dev/null @@ -1,105 +0,0 @@ -# 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 ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class DeepSeekV1Config(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - if isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in DeepSeekV1Config - - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['moe'] = self.moe.to_dict() - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - rotary_scaling = getattr(hf_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - moe_config = MoeConfig( - num_experts=getattr(hf_config, 'n_routed_experts', 0), - shared_expert_intermediate_size=getattr(hf_config, - 'n_shared_experts', 0) * - getattr(hf_config, "moe_intermediate_size", 0), - top_k=getattr(hf_config, 'num_experts_per_tok', 0), - normalization_mode=getattr( - hf_config, 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.NONE), - ) - 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=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', - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - norm_epsilon=hf_config.rms_norm_eps, - mapping=mapping, - quantization=quant_config, - moe=moe_config, - moe_intermediate_size=hf_config.moe_intermediate_size, - **kwargs) diff --git a/tensorrt_llm/models/deepseek_v1/convert.py b/tensorrt_llm/models/deepseek_v1/convert.py deleted file mode 100755 index 6871ebbdad2f..000000000000 --- a/tensorrt_llm/models/deepseek_v1/convert.py +++ /dev/null @@ -1,290 +0,0 @@ -# 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 time - -import torch -from transformers import AutoModelForCausalLM - -from ..._utils import pad_vocab_size, release_gc - - -## Get HF model -def load_hf_deepseek(model_dir): - model = AutoModelForCausalLM.from_pretrained(model_dir, - device_map='auto', - dtype='auto', - trust_remote_code=True) - return model - - -## Prepare weights for TP -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() - - -def split_qkv_tp(v, n_head, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - v = v.reshape(3, n_hidden, n_hidden) - split_v = split(v, tensor_parallel, rank, dim=1) - split_v = split_v.reshape(3 * (n_hidden // tensor_parallel), n_hidden) - return split_v.contiguous() - - -def split_matrix_tp(v, tensor_parallel, rank, dim): - return split(v, tensor_parallel, rank, dim=dim) - - -def get_weight(config, prefix, dtype, postfix='.weight'): - if config[prefix + postfix].dtype != dtype: - config[prefix + postfix].data = config[prefix + postfix].to(dtype) - return config[prefix + postfix].detach().cpu() - - -def get_trtllm_linear_weight(weight, prefix, postfix='weight'): - results = {} - results[prefix + postfix] = weight - - return results - - -def convert_deepseek(hf_model, - config, - mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0): - - weights = {} - tik = time.time() - mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - moe_config = config.moe - layers_range = mapping.pp_layers(config.num_hidden_layers) - - def convert_layer(l): - prefix = f'model.layers.{l}.' - print(prefix) - trtllm_prex = f'transformer.layers.{l - layers_range[0]}.' - q_weight = get_weight(model_params, prefix + 'self_attn.q_proj', dtype) - k_weight = get_weight(model_params, prefix + 'self_attn.k_proj', dtype) - v_weight = get_weight(model_params, prefix + 'self_attn.v_proj', dtype) - - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, config.num_attention_heads, - config.hidden_size, mapping.tp_size, - mapping.tp_rank) - - weights.update( - get_trtllm_linear_weight(split_v, trtllm_prex + 'attention.qkv.')) - - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - split_v = split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - weights.update( - get_trtllm_linear_weight(split_v, trtllm_prex + 'attention.dense.')) - - if moe_config.has_moe() and l > 0: - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - for suffix in ["gate_proj", "down_proj", "up_proj"]: - model_params[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.mlp.experts.{expert}.{suffix}.weight'].detach().cpu() - for expert in rank_experts]) - - gate_proj = model_params[ - f'model.layers.{l}.mlp.experts.gate_proj.weight'] - down_proj = model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - up_proj = model_params[ - f'model.layers.{l}.mlp.experts.up_proj.weight'] - if mapping.has_moe_tp(): - gate_proj = split(gate_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - down_proj = split(down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=2) - up_proj = split(up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - model_params[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] = torch.concat( - [up_proj, gate_proj], dim=-2) - model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] = down_proj - - ## mlp.experts.down_proj.weight - moe_experts_down_proj_weights = get_weight( - model_params, prefix + 'mlp.experts.down_proj', dtype) - weights.update( - get_trtllm_linear_weight(moe_experts_down_proj_weights, - trtllm_prex + 'mlp.proj.')) - ##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_trtllm_linear_weight(moe_experts_up_gate_proj_weights, - trtllm_prex + 'mlp.fc.')) - ## 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', - torch.float32) - weights.update( - get_trtllm_linear_weight(moe_experts_gate_weights, - trtllm_prex + 'mlp.router.')) - - if moe_config.shared_expert_intermediate_size > 0: - shared_moe_up_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.up_proj', dtype) - shared_moe_up_proj_weights = split_matrix_tp( - shared_moe_up_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_down_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.down_proj', - dtype) - shared_moe_down_proj_weights = split_matrix_tp( - shared_moe_down_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_moe_gate_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.gate_proj', - dtype) - shared_moe_gate_proj_weights = split_matrix_tp( - shared_moe_gate_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_gate_up_proj_weights = torch.concat( - [shared_moe_up_proj_weights, shared_moe_gate_proj_weights], - dim=-2) - - ## mlp.shared_experts.gate_up_proj.weight - weights.update( - get_trtllm_linear_weight( - shared_moe_gate_up_proj_weights, - trtllm_prex + 'mlp.shared_expert.fc.')) - - ## mlp.shared_experts.down_proj.weight - weights.update( - get_trtllm_linear_weight( - shared_moe_down_proj_weights, - trtllm_prex + 'mlp.shared_expert.proj.')) - - else: - ## Current deepseek model has one MLP layer only, if it goes large consider to do fuse - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - split_gate = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_trtllm_linear_weight(split_gate, trtllm_prex + 'mlp.gate.')) - - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - split_fc = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_trtllm_linear_weight(split_fc, trtllm_prex + 'mlp.fc.')) - - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - split_proj = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_trtllm_linear_weight(split_proj, trtllm_prex + 'mlp.proj.')) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[trtllm_prex + 'input_layernorm.weight'] = input_ln_weight - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[trtllm_prex + 'post_layernorm.weight'] = post_ln_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = get_weight(model_params, 'model.embed_tokens', dtype) - 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: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_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, - mapping.tp_rank) - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_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', - value=0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - #print(set(weights.keys())) - return weights diff --git a/tensorrt_llm/models/deepseek_v1/model.py b/tensorrt_llm/models/deepseek_v1/model.py deleted file mode 100755 index 164f90387a92..000000000000 --- a/tensorrt_llm/models/deepseek_v1/model.py +++ /dev/null @@ -1,279 +0,0 @@ -# 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 os -from typing import Optional - -from ..._utils import pad_vocab_size -from ...functional import Tensor, non_gated_version, recv, send -from ...layers import (Attention, AttentionMaskType, ColumnLinear, Embedding, - GatedMLP, PositionEmbeddingType, RmsNorm, SharedMoE) -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 .config import DeepSeekV1Config -from .convert import convert_deepseek, load_hf_deepseek - - -class DeepseekDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - ### Input layernorm in Deepseek v1 is same as Llama - self.input_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] - ### Deepseek v1 model with standard attention - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=False, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode) - - ClsMLP = GatedMLP - moe_config = config.moe - - if moe_config.num_experts > 0 and layer_idx > 0: - mlp_hidden_size = config.moe_intermediate_size - hidden_act = config.hidden_act - mlp_kwargs = {'moe_config': moe_config, 'mapping': config.mapping} - if moe_config.shared_expert_intermediate_size > 0: - ClsMLP = SharedMoE - mlp_kwargs['use_shared_gate'] = False - mlp_kwargs['use_side_stream'] = False - else: - ClsMLP = MOE - else: - ClsMLP = GatedMLP - mlp_hidden_size = config.intermediate_size - hidden_act = non_gated_version( - config.hidden_act) # back to non gated for dense layers - mlp_kwargs = {} - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=hidden_act, - dtype=config.dtype, - 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 v1 is same as Llama ) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class DeepseekModel(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - init_all_reduce_helper() # enable use_customer_all_reduce - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(DeepseekDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class DeepseekForCausalLM(DecoderModelForCausalLM): - config_class = DeepSeekV1Config - - def __init__(self, config: PretrainedConfig): - transformer = DeepseekModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(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) - else: - lm_head = None - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face(cls, - model_dir, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - override_fields={}, - **kwargs): - if mapping is None: - mapping = Mapping() - - pretrained_config = DeepSeekV1Config.from_hugging_face(model_dir, - dtype=dtype, - mapping=mapping, - **kwargs) - deepseek = cls.from_config(pretrained_config) - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - - custom_dict = {} - - rank_experts = mapping.ep_experts(pretrained_config.moe.num_experts) - for index, module in enumerate(deepseek.transformer.layers): - if index > 0: - - module.mlp.shared_expert.fc.tllm_to_externel_key_dict = { - "fc": ["up_proj", "gate_proj"], - "shared_expert": "shared_experts" - } - module.mlp.shared_expert.proj.tllm_to_externel_key_dict = { - "shared_expert": "shared_experts" - } - module.mlp.fc.tllm_to_externel_key_dict = { - "fc": [ - f"experts.{expert}.up_proj" - for expert in rank_experts - ] + [ - f"experts.{expert}.gate_proj" - for expert in rank_experts - ] - } - module.mlp.proj.tllm_to_externel_key_dict = { - "proj": [ - f"experts.{expert}.down_proj" - for expert in rank_experts - ] - } - module.mlp.router.tllm_to_externel_key_dict = { - "mlp": "mlp", - "router": "gate" - } - - loader = ModelWeightsLoader(model_dir, custom_dict) - loader.generate_tllm_weights(deepseek) - return deepseek - else: - - hf_model = load_hf_deepseek(model_dir) - weights = convert_deepseek( - hf_model, - pretrained_config, - mapping=pretrained_config.mapping, - dtype=pretrained_config.dtype, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - sharding_dim=pretrained_config.embedding_sharding_dim) - deepseek.load(weights) - return deepseek diff --git a/tensorrt_llm/models/deepseek_v2/__init__.py b/tensorrt_llm/models/deepseek_v2/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/deepseek_v2/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/deepseek_v2/config.py b/tensorrt_llm/models/deepseek_v2/config.py deleted file mode 100644 index c110df0d53f1..000000000000 --- a/tensorrt_llm/models/deepseek_v2/config.py +++ /dev/null @@ -1,148 +0,0 @@ -# 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 ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class DeepSeekV2Config(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - if isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, - MoeConfig), "moe must be a MoeConfig or a dictionary" - self.moe = moe.validate() - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in DeepSeekV2Config - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['moe'] = self.moe.to_dict() - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = AutoConfig.from_pretrained( - hf_config_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': - 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( - f'Unsupported topk_method in hf_config: {hf_config.topk_method}' - ) - - rotary_scaling = None - if hf_config.rope_scaling is not None: - 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': - hf_config.rope_scaling['type'] - } - - 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=hf_config.num_experts_per_tok, - 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=hf_config.routed_scaling_factor - ) - - 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=get_hf_rope_theta(hf_config, 10000.0), - rotary_scaling=rotary_scaling, - 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, - scoring_func=hf_config.scoring_func, - **kwargs) diff --git a/tensorrt_llm/models/deepseek_v2/convert.py b/tensorrt_llm/models/deepseek_v2/convert.py deleted file mode 100755 index 5a23130fc528..000000000000 --- a/tensorrt_llm/models/deepseek_v2/convert.py +++ /dev/null @@ -1,929 +0,0 @@ -# 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 time - -import torch -from transformers import AutoConfig, AutoModelForCausalLM - -from tensorrt_llm.layers import MoeConfig - -from ..._utils import (get_hf_rope_theta, pad_vocab_size, release_gc, - str_dtype_to_torch) -from ...logger import logger -from ...mapping import Mapping -from ..convert_utils import get_tllm_linear_weight - -# `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 TODO: remove this function and change CI test -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 = get_hf_rope_theta(hf_config, 10000.0) - 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) - 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' - ) - - if load_model_on_cpu: - # Skip setting max_memory when loading on CPU, you might have OOM. - model = AutoModelForCausalLM.from_pretrained(model_dir, - config=hf_config, - device_map='cpu', - dtype='auto', - trust_remote_code=True) - else: - # Deepseek-v2 236B parameters with FP16 dtype need at least 472G GPU memory - # (official suggest at least 8x80G GPUs, see https://huggingface.co/deepseek-ai/DeepSeek-V2) - - max_memory = None - device_map = 'auto' - - gpu_memory = torch.cuda.get_device_properties(0).total_memory - gpu_count = torch.cuda.device_count() - - if gpu_memory < 90_000_000_000 and gpu_count == 8: - # WAR OOM loading on 8*80G GPUs - max_memory = {i: "76GB" for i in range(8)} - device_map = 'sequential' - elif gpu_memory < 180_000_000_000 and gpu_count == 4: - # WAR OOM loading on 4*141G GPUs - max_memory = {i: "128GB" for i in range(4)} - device_map = 'sequential' - elif gpu_memory < 180_000_000_000 and gpu_count == 8: - # WAR OOM loading on 8*141G GPUs - max_memory = {i: "128GB" for i in range(8)} - device_map = 'sequential' - - model = AutoModelForCausalLM.from_pretrained(model_dir, - config=hf_config, - device_map=device_map, - max_memory=max_memory, - dtype='auto', - trust_remote_code=True) - - return model - - -# Prepare weights for TP -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() - - -def split_matrix_tp(v, tensor_parallel, rank, dim): - return split(v, tensor_parallel, rank, dim=dim) - - -def get_weight(config, prefix, dtype, postfix='.weight'): - if config[prefix + postfix].dtype != dtype: - config[prefix + postfix].data = config[prefix + postfix].to(dtype) - return config[prefix + postfix].detach().cpu() - - -def get_param_weight(weight, prefix): - results = {} - results[prefix] = weight - - return results - - -def convert_deepseekv2(hf_model, - config, - mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0): - - weights = {} - tik = time.time() - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - moe_config = config.moe - - layers_range = mapping.pp_layers(config.num_hidden_layers) - - def convert_layer(l): - prefix = f'model.layers.{l}.' - trtllm_prefix = 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 - - if q_lora_rank is not None: - q_a_proj_weight = get_weight(model_params, - prefix + 'self_attn.q_a_proj', dtype) - # Layer normalization - q_a_layernorm_weight = get_weight( - model_params, - prefix + 'self_attn.q_a_layernorm', - dtype, - ) - - kv_a_proj_with_mqa_weight = get_weight( - model_params, prefix + 'self_attn.kv_a_proj_with_mqa', dtype) - - kv_a_layernorm_weight = get_weight( - model_params, - prefix + 'self_attn.kv_a_layernorm', - dtype, - ) - - if q_lora_rank is not None: - fused_a_weight = torch.cat( - [q_a_proj_weight, kv_a_proj_with_mqa_weight], - dim=0, - ) - - q_b_proj_weight = get_weight( - model_params, prefix + 'self_attn.q_b_proj', dtype).unflatten( - 0, - [ - num_heads, - qk_nope_head_dim + qk_rope_head_dim, - ], - ) - else: - fused_a_weight = kv_a_proj_with_mqa_weight - - q_b_proj_weight = get_weight( - model_params, prefix + 'self_attn.q_proj', dtype).unflatten( - 0, - [ - num_heads, - qk_nope_head_dim + qk_rope_head_dim, - ], - ) - - kv_b_proj_weight = get_weight(model_params, - prefix + 'self_attn.kv_b_proj', - dtype).unflatten( - 0, - [ - num_heads, - qk_nope_head_dim + v_head_dim, - ], - ) - - o_proj_weight = get_weight(model_params, prefix + 'self_attn.o_proj', - dtype) - - q_b_proj_weight = split_matrix_tp( - q_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0, - ) - kv_b_proj_weight = split_matrix_tp( - kv_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0, - ) - o_proj_weight = split_matrix_tp( - o_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1, - ) - - q_nope_weight, q_pe_weight = q_b_proj_weight.split( - [qk_nope_head_dim, qk_rope_head_dim], - dim=1, - ) - k_nope_weight, v_weight = kv_b_proj_weight.split( - [qk_nope_head_dim, v_head_dim], - dim=1, - ) - - if q_lora_rank is None: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, hidden_size) - else: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, q_lora_rank) - - kv_b_proj_weight = torch.concat([ - k_nope_weight.reshape( - num_heads * qk_nope_head_dim // mapping.tp_size, kv_lora_rank), - v_weight.reshape(num_heads * v_head_dim // mapping.tp_size, - kv_lora_rank) - ], - dim=0) - - # Fuse matrices for decompression - 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) - - weights.update( - get_tllm_linear_weight(fused_a_weight, - trtllm_prefix + 'attention.fused_a.')) - weights.update( - get_tllm_linear_weight(kv_a_layernorm_weight, - trtllm_prefix + 'attention.kv_a_layernorm.')) - weights.update( - get_param_weight(fused_q_weight, - trtllm_prefix + 'attention.fused_q_proj')) - weights.update( - get_param_weight(q_b_proj_weight, - trtllm_prefix + 'attention.q_b_proj')) - weights.update( - get_param_weight(kv_b_proj_weight, - trtllm_prefix + 'attention.kv_b_proj')) - weights.update( - get_tllm_linear_weight(o_proj_weight, - trtllm_prefix + 'attention.dense.')) - - if q_lora_rank is not None: - weights.update( - get_tllm_linear_weight( - q_a_layernorm_weight, - trtllm_prefix + 'attention.q_a_layernorm.')) - - if moe_config.has_moe() and l > 0: - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - for suffix in ["gate_proj", "down_proj", "up_proj"]: - model_params[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.mlp.experts.{expert}.{suffix}.weight'].detach().cpu() - for expert in rank_experts]) - - gate_proj = model_params[ - f'model.layers.{l}.mlp.experts.gate_proj.weight'] - down_proj = model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - up_proj = model_params[ - f'model.layers.{l}.mlp.experts.up_proj.weight'] - if mapping.has_moe_tp(): - gate_proj = split(gate_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - down_proj = split(down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=2) - up_proj = split(up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - model_params[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] = torch.concat( - [up_proj, gate_proj], dim=-2) - model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] = down_proj - - # mlp.experts.down_proj.weight - 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_prefix + 'mlp.proj.')) - # 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_prefix + 'mlp.fc.')) - # 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', - torch.float32) - weights.update( - get_tllm_linear_weight(moe_experts_gate_weights, - trtllm_prefix + 'mlp.router.')) - - if moe_config.shared_expert_intermediate_size > 0: - shared_moe_up_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.up_proj', dtype) - shared_moe_up_proj_weights = split_matrix_tp( - shared_moe_up_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_down_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.down_proj', - dtype) - shared_moe_down_proj_weights = split_matrix_tp( - shared_moe_down_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_moe_gate_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.gate_proj', - dtype) - shared_moe_gate_proj_weights = split_matrix_tp( - shared_moe_gate_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_gate_up_proj_weights = torch.concat( - [shared_moe_up_proj_weights, shared_moe_gate_proj_weights], - dim=-2) - - # mlp.shared_experts.gate_up_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_gate_up_proj_weights, - trtllm_prefix + 'mlp.shared_expert.fc.')) - - # mlp.shared_experts.down_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_down_proj_weights, - trtllm_prefix + 'mlp.shared_expert.proj.')) - - else: - # Current MLP layer is only one, if it goes large consider to do fuse - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - split_gate = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_gate, trtllm_prefix + 'mlp.gate.')) - - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - split_fc = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_fc, trtllm_prefix + 'mlp.fc.')) - - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - split_proj = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_proj, trtllm_prefix + 'mlp.proj.')) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[trtllm_prefix + 'input_layernorm.weight'] = input_ln_weight - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[trtllm_prefix + 'post_layernorm.weight'] = post_ln_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = get_weight(model_params, 'model.embed_tokens', dtype) - 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: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_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, - mapping.tp_rank) - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_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', - value=0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - #print(set(weights.keys())) - return weights - - -def load_weights_from_hf_safetensors(model_dir, - config, - mapping, - use_parallel_embedding=False, - sharding_dim=0): - logger.info('Loading weights from Huggingface safetensors...') - weights = {} - tik = time.time() - import json - import os - - import safetensors - - model_dir = model_dir if model_dir.endswith("/") else model_dir + "/" - safetensors_map = {} - has_safetensor_index_json = True - try: - with open(model_dir + "model.safetensors.index.json", 'r') as fr: - sharding_map = json.load(fr) - # safetensors_map structure: - # key: e.g., model.layers.0.self_attn.q_a_proj.weight - # value: e.g., model-00001-of-000055.safetensors - # int(value[6:11]) -> safetensors_idx - for k, v in sharding_map['weight_map'].items(): - safetensors_map[k] = int(v[6:11]) - 1 - except FileNotFoundError: - has_safetensor_index_json = False - # shard_files purpose to verify all .safetensors is aligned with .index.json - shard_files = [] - for name in os.listdir(model_dir): - if name.endswith(".safetensors"): - if has_safetensor_index_json and name not in sharding_map[ - 'weight_map'].values(): - continue - shard_files.append(name) - shard_files.sort() - # Create all .safetensors memory address - safetensors_ptrs = [ - safetensors.safe_open(model_dir + shard_file, - framework="pt", - device="cpu") for shard_file in shard_files - ] - - moe_config = MoeConfig( - num_experts=config.moe.num_experts, - shared_expert_intermediate_size=config.moe. - shared_expert_intermediate_size, - top_k=config.moe.top_k, - normalization_mode=config.moe.normalization_mode, - device_limited_n_group=config.moe.device_limited_n_group, - device_limited_topk_group=config.moe.device_limited_topk_group, - device_limited_routed_scaling_factor=config.moe. - device_limited_routed_scaling_factor) - - torch_dtype = str_dtype_to_torch(config.dtype) - - def load_weights(key, dtype): - assert key in safetensors_map, f"'{key}' not found in safetensors_map" - ptr_idx = safetensors_map[key] - - assert key in safetensors_ptrs[ptr_idx].keys( - ), f"'{key}' not found in safetensors file {ptr_idx}" - tensor_slice = safetensors_ptrs[ptr_idx].get_slice(key) - tensor_slice.get_shape() - res = tensor_slice[:] - if res.dtype != dtype: - res = res.to(dtype) - return res.contiguous().detach().cpu() - - layers_range = mapping.pp_layers(config.num_hidden_layers) - - def convert_layer(l): - prefix = f'model.layers.{l}.' - trtllm_prefix = 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 - - # MLA: - # q_lora_rank used for different deepseek-v2 and deepseek-v2-lite - if q_lora_rank is not None: - q_a_proj_weight = load_weights(prefix + 'self_attn.q_a_proj.weight', - torch_dtype) - q_a_layernorm_weight = load_weights( - prefix + 'self_attn.q_a_layernorm.weight', torch_dtype) - - kv_a_proj_with_mqa_weight = load_weights( - prefix + 'self_attn.kv_a_proj_with_mqa.weight', torch_dtype) - kv_a_layernorm_weight = load_weights( - prefix + 'self_attn.kv_a_layernorm.weight', torch_dtype) - - if q_lora_rank is not None: - fused_a_weight = torch.cat( - [q_a_proj_weight, kv_a_proj_with_mqa_weight], dim=0) - q_b_proj_weight = load_weights( - prefix + 'self_attn.q_b_proj.weight', torch_dtype).unflatten( - 0, [num_heads, qk_nope_head_dim + qk_rope_head_dim]) - else: - fused_a_weight = kv_a_proj_with_mqa_weight - q_b_proj_weight = load_weights( - prefix + 'self_attn.q_proj.weight', torch_dtype).unflatten( - 0, [num_heads, qk_nope_head_dim + qk_rope_head_dim]) - - kv_b_proj_weight = load_weights( - prefix + 'self_attn.kv_b_proj.weight', - torch_dtype).unflatten(0, - [num_heads, qk_nope_head_dim + v_head_dim]) - o_proj_weight = load_weights(prefix + 'self_attn.o_proj.weight', - torch_dtype) - - q_b_proj_weight = split_matrix_tp(q_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - kv_b_proj_weight = split_matrix_tp(kv_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - o_proj_weight = split_matrix_tp(o_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - q_nope_weight, q_pe_weight = q_b_proj_weight.split( - [qk_nope_head_dim, qk_rope_head_dim], dim=1) - k_nope_weight, v_weight = kv_b_proj_weight.split( - [qk_nope_head_dim, v_head_dim], dim=1) - - if q_lora_rank is None: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, hidden_size) - else: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, q_lora_rank) - - kv_b_proj_weight = torch.concat([ - k_nope_weight.reshape( - num_heads * qk_nope_head_dim // mapping.tp_size, kv_lora_rank), - v_weight.reshape(num_heads * v_head_dim // mapping.tp_size, - kv_lora_rank) - ], - dim=0) - - # Fuse matrices for decompression - 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) - - weights.update( - get_tllm_linear_weight(fused_a_weight, - trtllm_prefix + 'attention.fused_a.')) - weights.update( - get_tllm_linear_weight(kv_a_layernorm_weight, - trtllm_prefix + 'attention.kv_a_layernorm.')) - weights.update( - get_param_weight(fused_q_weight, - trtllm_prefix + 'attention.fused_q_proj')) - weights.update( - get_param_weight(q_b_proj_weight, - trtllm_prefix + 'attention.q_b_proj')) - weights.update( - get_param_weight(kv_b_proj_weight, - trtllm_prefix + 'attention.kv_b_proj')) - weights.update( - get_tllm_linear_weight(o_proj_weight, - trtllm_prefix + 'attention.dense.')) - - if q_lora_rank is not None: - weights.update( - get_tllm_linear_weight( - q_a_layernorm_weight, - trtllm_prefix + 'attention.q_a_layernorm.')) - - # MOE: - if moe_config.has_moe() and l > 0: - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - expert_weight_dict = {} - for suffix in ["gate_proj", "down_proj", "up_proj"]: - expert_weight_dict[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([load_weights(prefix + f'mlp.experts.{expert}.{suffix}.weight', torch_dtype) for expert in rank_experts]) - - gate_proj = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.gate_proj.weight'] - down_proj = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - up_proj = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.up_proj.weight'] - - if mapping.has_moe_tp(): - gate_proj = split(gate_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - down_proj = split(down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=2) - up_proj = split(up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - expert_weight_dict[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] = torch.concat( - [up_proj, gate_proj], dim=-2) - expert_weight_dict[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] = down_proj - - # mlp.experts.down_proj.weight - moe_experts_down_proj_weights = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - weights.update( - get_tllm_linear_weight(moe_experts_down_proj_weights, - trtllm_prefix + 'mlp.proj.')) - - # mlp.experts.up_gate.weight - moe_experts_up_gate_proj_weights = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] - weights.update( - get_tllm_linear_weight(moe_experts_up_gate_proj_weights, - trtllm_prefix + 'mlp.fc.')) - - # MOE hardcoded routing_input into trt.float32, please refer to moe.py line 397 - moe_experts_gate_weights = load_weights(prefix + 'mlp.gate.weight', - torch.float32) - weights.update( - get_tllm_linear_weight(moe_experts_gate_weights, - trtllm_prefix + 'mlp.router.')) - - if moe_config.shared_expert_intermediate_size > 0: - shared_moe_up_proj_weights = load_weights( - prefix + 'mlp.shared_experts.up_proj.weight', torch_dtype) - shared_moe_up_proj_weights = split_matrix_tp( - shared_moe_up_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_down_proj_weights = load_weights( - prefix + 'mlp.shared_experts.down_proj.weight', torch_dtype) - shared_moe_down_proj_weights = split_matrix_tp( - shared_moe_down_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_moe_gate_proj_weights = load_weights( - prefix + 'mlp.shared_experts.gate_proj.weight', torch_dtype) - shared_moe_gate_proj_weights = split_matrix_tp( - shared_moe_gate_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_gate_up_proj_weights = torch.concat( - [shared_moe_up_proj_weights, shared_moe_gate_proj_weights], - dim=-2) - - # mlp.shared_experts.gate_up_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_gate_up_proj_weights, - trtllm_prefix + 'mlp.shared_expert.fc.')) - - # mlp.shared_experts.down_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_down_proj_weights, - trtllm_prefix + 'mlp.shared_expert.proj.')) - - else: - # Current MLP layer is only one (layer 0), if it goes large consider to do fuse - mlp_gate_weight = load_weights(prefix + 'mlp.up_proj.weight', - torch_dtype) - split_gate = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_gate, trtllm_prefix + 'mlp.gate.')) - - mlp_fc_weight = load_weights(prefix + 'mlp.gate_proj.weight', - torch_dtype) - split_fc = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_fc, trtllm_prefix + 'mlp.fc.')) - - mlp_proj_weight = load_weights(prefix + 'mlp.down_proj.weight', - torch_dtype) - split_proj = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_proj, trtllm_prefix + 'mlp.proj.')) - - # Layer norms do not use tensor parallelism - input_ln_weight = load_weights(prefix + 'input_layernorm.weight', - torch_dtype) - weights[trtllm_prefix + 'input_layernorm.weight'] = input_ln_weight - post_ln_weight = load_weights( - prefix + 'post_attention_layernorm.weight', torch_dtype) - weights[trtllm_prefix + 'post_layernorm.weight'] = post_ln_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = load_weights('model.embed_tokens.weight', torch_dtype) - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - lm_head_weights = load_weights('lm_head.weight', torch_dtype) - - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_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', 0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = load_weights('model.norm.weight', torch_dtype) - weights['transformer.ln_f.weight'] = ln_f_w - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/deepseek_v2/model.py b/tensorrt_llm/models/deepseek_v2/model.py deleted file mode 100755 index c075d6fc7a95..000000000000 --- a/tensorrt_llm/models/deepseek_v2/model.py +++ /dev/null @@ -1,361 +0,0 @@ -# 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 os -from typing import Optional - -import torch -import transformers - -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, - PositionEmbeddingType, RmsNorm, SharedMoE) -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 .config import DeepSeekV2Config -from .convert import convert_deepseekv2, load_weights_from_hf_safetensors - - -class DeepseekV2DecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - # Input layernorm in Deepseek v2 is same as Llama - self.input_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] - - self.attention = DeepseekV2Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - q_lora_rank=config.q_lora_rank, - kv_lora_rank=config.kv_lora_rank, - 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, - max_position_embeddings=config.max_position_embeddings, - eps=config.norm_epsilon, - attention_mask_type=AttentionMaskType.causal, - dtype=config.dtype, - 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, - tp_rank=config.mapping.tp_rank) - - # Added deepseek MoE and shared_experts - # First decoder layer: MLA + dense MLP + input_layernorm(RMSNorm) + post_attention_layernorm(RMSNorm) - # Rest decoder layer: MLA + MoE MLP + MoE Gate + shared_experts(MLP) + input_layernorm(RMSNorm) + post_attention_layernorm(RMSNorm) - # Added MLA in co-testing phase, use standard attention for MoE testing - - # Distinguish dense MLP and MoE MLP - # dense_config = DenseConfig(intermediate_size=config.intermediate_size) - moe_config = config.moe - # In case of moe_config is a dict - if isinstance(moe_config, dict): - moe_config = MoeConfig.from_dict(moe_config) - - if moe_config.num_experts > 0 and layer_idx > 0: - hidden_act = config.hidden_act - mlp_hidden_size = config.moe_inter_size - mlp_kwargs = {'moe_config': moe_config, 'mapping': config.mapping} - if moe_config.shared_expert_intermediate_size > 0: - ClsMLP = SharedMoE - mlp_kwargs['use_shared_gate'] = False - mlp_kwargs['use_side_stream'] = False - else: - ClsMLP = MOE - else: - ClsMLP = GatedMLP - mlp_hidden_size = config.intermediate_size - hidden_act = non_gated_version( - config.hidden_act) # back to non gated for dense layers - mlp_kwargs = {} - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=hidden_act, - dtype=config.dtype, - 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 - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states=hidden_states, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual_attn = hidden_states - - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual_attn + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class DeepseekV2Model(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - init_all_reduce_helper() # enable use_customer_all_reduce - self.dtype = config.dtype - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(DeepseekV2DecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - self.head_num = config.num_attention_heads - self.head_size = config.qk_nope_head_dim + config.qk_rope_head_dim - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class DeepseekV2ForCausalLM(DecoderModelForCausalLM): - config_class = DeepSeekV2Config - - def __init__(self, config: PretrainedConfig): - transformer = DeepseekV2Model(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(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) - else: - lm_head = None - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - model_dir, - dtype: str = 'auto', - hf_model: Optional[transformers.PreTrainedModel] = None, - use_preloading: bool = False, - use_safetensors_loading: bool = False, - mapping: Optional[Mapping] = None, - override_fields={}, - **kwargs): - - if mapping is None: - mapping = Mapping() - pretrained_config = DeepSeekV2Config.from_hugging_face(model_dir, - dtype=dtype, - mapping=mapping, - **kwargs) - if dtype == 'auto': - dtype = getattr(pretrained_config, 'torch_dtype', None) - if dtype is None: - dtype = 'float16' - if isinstance(dtype, torch.dtype): - dtype = torch_dtype_to_str(dtype) - if dtype == 'float32': # should remove "float32" - dtype = 'float16' - if dtype == 'bfloat16' and torch.cuda.get_device_properties( - 0).major < 8: - logger.warning( - "Pre SM 80 GPUs do not support bfloat16, fallback to float16") - dtype = 'float16' - - deepseek = cls.from_config(pretrained_config) - - # If use_preloading is True, load the model from hf_model - # If use_safetensors_loading is True, load the model from safetensors - # if TRTLLM_DISABLE_UNIFIED_CONVERTER is not set, load the model use unified converter (recommended and default) - if use_preloading: - weights = convert_deepseekv2( - hf_model, - pretrained_config, - mapping, - dtype=dtype, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - sharding_dim=pretrained_config.embedding_sharding_dim) - deepseek.load(weights) - return deepseek - - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - - custom_dict = {} - rank_experts = mapping.ep_experts(pretrained_config.moe.num_experts) - for index, module in enumerate(deepseek.transformer.layers): - - if pretrained_config.q_lora_rank is not None: - module.attention.tllm_to_externel_key_dict = { - "fused_q_proj": ["q_b_proj.weight", "kv_b_proj.weight"], - "q_b_proj": "q_b_proj.weight", #v2 - "q_a_proj": "q_a_proj.weight", #v2 - "kv_b_proj": "kv_b_proj.weight", - "q_a_layernorm": "q_a_layernorm" - } - module.attention.fused_a.tllm_to_externel_key_dict = { - "fused_a": ["q_a_proj", "kv_a_proj_with_mqa"] - } #v2 - else: - module.attention.tllm_to_externel_key_dict = { - "fused_q_proj": ["q_proj.weight", - "kv_b_proj.weight"], #v2 lite - "q_b_proj": "q_proj.weight", #v2 lite - "kv_b_proj": "kv_b_proj.weight", - "q_a_layernorm": "q_a_layernorm" - } - module.attention.fused_a.tllm_to_externel_key_dict = { - "fused_a": "kv_a_proj_with_mqa" - } # v2 lite - - module.attention.kv_a_layernorm.tllm_to_externel_key_dict = { - 'kv_a_layernorm': 'kv_a_layernorm' - } - - if index > 0: - - module.mlp.shared_expert.fc.tllm_to_externel_key_dict = { - "fc": ["up_proj", "gate_proj"], - "shared_expert": "shared_experts" - } - module.mlp.shared_expert.proj.tllm_to_externel_key_dict = { - "shared_expert": "shared_experts" - } - module.mlp.fc.tllm_to_externel_key_dict = { - "fc": [ - f"experts.{expert}.up_proj" - for expert in rank_experts - ] + [ - f"experts.{expert}.gate_proj" - for expert in rank_experts - ] - } - module.mlp.proj.tllm_to_externel_key_dict = { - "proj": [ - f"experts.{expert}.down_proj" - for expert in rank_experts - ] - } - module.mlp.router.tllm_to_externel_key_dict = { - "mlp": "mlp", - "router": "gate" - } - - loader = ModelWeightsLoader(model_dir, custom_dict) - loader.generate_tllm_weights(deepseek) - return deepseek - - if use_safetensors_loading: - weights = load_weights_from_hf_safetensors( - model_dir, - pretrained_config, - mapping, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - sharding_dim=pretrained_config.embedding_sharding_dim) - deepseek.load(weights) - return deepseek diff --git a/tensorrt_llm/models/dit/__init__.py b/tensorrt_llm/models/dit/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/dit/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/dit/model.py b/tensorrt_llm/models/dit/model.py deleted file mode 100644 index 227262f99405..000000000000 --- a/tensorrt_llm/models/dit/model.py +++ /dev/null @@ -1,382 +0,0 @@ -# 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 math -from collections import OrderedDict - -import numpy as np -import tensorrt as trt - -from ..._utils import str_dtype_to_trt, trt_dtype_to_str -from ...functional import (Tensor, allgather, arange, chunk, concat, constant, - cos, exp, expand, shape, silu, sin, slice, split, - unsqueeze) -from ...layers import MLP, BertAttention, Conv2d, Embedding, LayerNorm, Linear -from ...mapping import Mapping -from ...module import Module, ModuleList -from ...parameter import Parameter -from ...plugin import current_all_reduce_helper -from ...quantization import QuantMode -from ..modeling_utils import PretrainedConfig, PretrainedModel - - -def modulate(x, shift, scale, dtype): - ones = 1.0 - if dtype is not None: - ones = constant(np.ones(1, dtype=np.float32)).cast(dtype) - return x * (ones + unsqueeze(scale, 1)) + unsqueeze(shift, 1) - - -class TimestepEmbedder(Module): - - def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None): - super().__init__() - self.dtype = dtype - self.mlp1 = Linear(frequency_embedding_size, - hidden_size, - bias=True, - dtype=dtype) - self.mlp2 = Linear(hidden_size, hidden_size, bias=True, dtype=dtype) - self.frequency_embedding_size = frequency_embedding_size - - def timestep_embedding(self, t, dim, max_period=10000): - half = dim // 2 - freqs = exp( - -math.log(max_period) * - arange(start=0, end=half, dtype=trt_dtype_to_str(trt.float32)) / - constant(np.array([half], dtype=np.float32))) - args = unsqueeze(t, -1).cast(trt.float32) * unsqueeze(freqs, 0) - embedding = concat([cos(args), sin(args)], dim=-1) - if self.dtype is not None: embedding = embedding.cast(self.dtype) - assert dim % 2 == 0 - return embedding - - def forward(self, t): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) - t_emb = self.mlp2(silu(self.mlp1(t_freq))) - - return t_emb - - -class LabelEmbedder(Module): - - def __init__(self, num_classes, hidden_size, dropout_prob, dtype=None): - super().__init__() - use_cfg_embedding = dropout_prob > 0 - self.embedding_table = Embedding(num_classes + use_cfg_embedding, - hidden_size, - dtype=dtype) - self.num_classes = num_classes - self.dropout_prob = dropout_prob - - def forward(self, labels, force_drop_ids=None): - assert force_drop_ids is None - embeddings = self.embedding_table(labels) - return embeddings - - -class PatchEmbed(Module): - - def __init__(self, - img_size: int, - patch_size: int, - input_c: int, - output_c: int, - bias: bool = True, - dtype: trt.DataType = None): - super().__init__() - self.img_size = img_size - self.patch_size = patch_size - self.num_patches = (img_size // patch_size)**2 - self.proj = Conv2d(input_c, - output_c, - kernel_size=(patch_size, patch_size), - stride=(patch_size, patch_size), - bias=bias, - dtype=dtype) - - def forward(self, x): - assert x.shape[2] == self.img_size - assert x.shape[3] == self.img_size - x = self.proj(x) - x = x.flatten(2).transpose(1, 2) # NCHW -> NLC - return x - - -class DiTBlock(Module): - - def __init__(self, - hidden_size, - num_heads, - mapping=Mapping(), - mlp_ratio=4.0, - dtype=None, - quant_mode=QuantMode(0)): - super().__init__() - self.dtype = dtype - self.norm1 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.attn = BertAttention(hidden_size, - num_heads, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - cp_group=mapping.cp_group, - cp_size=mapping.cp_size, - cp_rank=mapping.cp_rank, - dtype=dtype, - quant_mode=quant_mode) - self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=int(hidden_size * mlp_ratio), - hidden_act='gelu', - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - quant_mode=quant_mode) - self.adaLN_modulation = Linear(hidden_size, - 6 * hidden_size, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - bias=True, - dtype=dtype) - - def forward(self, x, c, input_lengths): - c = self.adaLN_modulation(silu(c)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = chunk( - c, 6, dim=1) - - x = x + unsqueeze(gate_msa, 1) * self.attn(modulate( - self.norm1(x), shift_msa, scale_msa, self.dtype), - input_lengths=input_lengths) - x = x + unsqueeze(gate_mlp, 1) * self.mlp( - modulate(self.norm2(x), shift_mlp, scale_mlp, self.dtype)) - return x - - -class FinalLayer(Module): - - def __init__(self, - hidden_size, - patch_size, - out_channels, - mapping=Mapping(), - dtype=None): - super().__init__() - self.dtype = dtype - self.norm_final = LayerNorm(hidden_size, - elementwise_affine=False, - eps=1e-6) - self.linear = Linear(hidden_size, - patch_size * patch_size * out_channels, - bias=True, - dtype=dtype) - self.adaLN_modulation = Linear(hidden_size, - 2 * hidden_size, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - bias=True, - dtype=dtype) - - def forward(self, x, c): - shift, scale = chunk(self.adaLN_modulation(silu(c)), 2, dim=1) - - x = modulate(self.norm_final(x), shift, scale, self.dtype) - x = self.linear(x) - - return x - - -class DiT(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - super().__init__(config) - self.learn_sigma = config.learn_sigma - self.in_channels = config.in_channels - self.out_channels = config.in_channels * 2 if config.learn_sigma else config.in_channels - self.input_size = config.input_size - self.patch_size = config.patch_size - self.num_heads = config.num_attention_heads - self.dtype = str_dtype_to_trt(config.dtype) - self.cfg_scale = config.cfg_scale - self.mapping = config.mapping - - self.x_embedder = PatchEmbed(config.input_size, - config.patch_size, - config.in_channels, - config.hidden_size, - bias=True, - dtype=self.dtype) - self.t_embedder = TimestepEmbedder(config.hidden_size, dtype=self.dtype) - self.y_embedder = LabelEmbedder(config.num_classes, - config.hidden_size, - config.class_dropout_prob, - dtype=self.dtype) - num_patches = self.x_embedder.num_patches - - self.pos_embed = Parameter(shape=(1, num_patches, config.hidden_size), - dtype=self.dtype) - self.blocks = ModuleList([ - DiTBlock(config.hidden_size, - config.num_attention_heads, - mlp_ratio=config.mlp_ratio, - mapping=config.mapping, - dtype=self.dtype, - quant_mode=config.quant_mode) - for _ in range(config.num_hidden_layers) - ]) - self.final_layer = FinalLayer(config.hidden_size, - config.patch_size, - self.out_channels, - mapping=config.mapping, - dtype=self.dtype) - - # We need to invoke default `__post_init__()` for quantized layers. - # def __post_init__(self): - # return - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('input_size', 32) - config.set_if_not_exist('patch_size', 2) - config.set_if_not_exist('in_channels', 4) - config.set_if_not_exist('mlp_ratio', 4.0) - config.set_if_not_exist('class_dropout_prob', 0.1) - config.set_if_not_exist('num_classes', 1000) - config.set_if_not_exist('learn_sigma', True) - config.set_if_not_exist('dtype', None) - config.set_if_not_exist('cfg_scale', None) - - def unpatchify(self, x: Tensor): - c = self.out_channels - p = self.x_embedder.patch_size - h = w = int(x.shape[1]**0.5) - assert h * w == x.shape[1] - - x = x.view(shape=(x.shape[0], h, w, p, p, c)) - x = x.permute((0, 5, 1, 3, 2, 4)) - imgs = x.view(shape=(x.shape[0], c, h * p, h * p)) - return imgs - - def forward(self, latent, timestep, label): - """ - Forward pass of DiT. - latent: (N, C, H, W) - timestep: (N,) - label: (N,) - """ - if self.cfg_scale is not None: - output = self.forward_with_cfg(latent, timestep, label) - else: - output = self.forward_without_cfg(latent, timestep, label) - output.mark_output('output', self.dtype) - return output - - def forward_without_cfg(self, x, t, y): - """ - Forward pass without classifier-free guidance. - """ - x = self.x_embedder(x) + self.pos_embed.value - t = self.t_embedder(t) - y = self.y_embedder(y) - self.register_network_output('t_embedder', t) - self.register_network_output('x_embedder', x) - self.register_network_output('y_embedder', y) - c = t + y - input_length = constant(np.array([x.shape[1]], dtype=np.int32)) - input_lengths = expand(input_length, unsqueeze(shape(x, 0), 0)) - # Split squeence for CP here - if self.mapping.cp_size > 1: - assert x.shape[1] % self.mapping.cp_size == 0 - x = chunk(x, self.mapping.cp_size, dim=1)[self.mapping.cp_rank] - input_lengths = input_lengths // self.mapping.cp_size - for block in self.blocks: - x = block(x, c, input_lengths) # (N, T, D) - self.register_network_output('before_final_layer', x) - x = self.final_layer(x, c) # (N, T, patch_size ** 2 * out_channels) - self.register_network_output('final_layer', x) - - # All gather after CP - if self.mapping.cp_size > 1: - x = allgather(x, self.mapping.cp_group, gather_dim=1) - x = self.unpatchify(x) # (N, out_channels, H, W) - self.register_network_output('unpatchify', x) - return x - - def forward_with_cfg(self, x, t, y): - """ - Forward pass with classifier-free guidance. - """ - batch_size = shape(x, 0) - half = slice( - x, [0, 0, 0, 0], - concat([batch_size / 2, x.shape[1], x.shape[2], x.shape[3]])) - combined = concat([half, half], dim=0) - self.register_network_output('combined', combined) - model_out = self.forward_without_cfg(combined, t, y) - - _, d, h, w = model_out.shape - eps, rest = split(model_out, [3, d - 3], dim=1) - cond_eps = slice(eps, [0, 0, 0, 0], concat([batch_size / 2, 3, h, w])) - uncond_eps = slice(eps, concat([batch_size / 2, 0, 0, 0]), - concat([batch_size / 2, 3, h, w])) - self.register_network_output('cond_eps', cond_eps) - self.register_network_output('uncond_eps', uncond_eps) - - half_eps = uncond_eps + self.cfg_scale * (cond_eps - uncond_eps) - eps = concat([half_eps, half_eps], dim=0) - self.register_network_output('eps', eps) - - return concat([eps, rest], dim=1) - - def prepare_inputs(self, max_batch_size, **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - mapping = self.config.mapping - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor(mapping, 1) - - def dit_default_range(max_batch_size): - return [2, max(2, (max_batch_size + 1) // 2), max_batch_size] - - default_range = dit_default_range - if self.cfg_scale is not None: - max_batch_size *= 2 - - latent = Tensor( - name='latent', - dtype=self.dtype, - shape=[-1, self.in_channels, self.input_size, self.input_size], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('in_channels', [[self.in_channels] * 3]), - ('latent_height', [[self.input_size] * 3]), - ('latent_width', [[self.input_size] * 3]), - ])) - timestep = Tensor(name='timestep', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - label = Tensor(name='label', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - return {'latent': latent, 'timestep': timestep, 'label': label} diff --git a/tensorrt_llm/models/eagle/__init__.py b/tensorrt_llm/models/eagle/__init__.py deleted file mode 100644 index a08b2c2049a3..000000000000 --- a/tensorrt_llm/models/eagle/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 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. diff --git a/tensorrt_llm/models/eagle/config.py b/tensorrt_llm/models/eagle/config.py deleted file mode 100644 index e7a559f34690..000000000000 --- a/tensorrt_llm/models/eagle/config.py +++ /dev/null @@ -1,222 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 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 json -from typing import Optional, Union - -from transformers import LlamaConfig - -from ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..llama.config import LLaMAConfig -from ..modeling_utils import QuantAlgo, QuantConfig - - -class EagleConfig(LLaMAConfig): - - def __init__(self, - *, - num_eagle_layers: int = 1, - max_draft_len: int = 63, - max_non_leaves_per_layer: int = 10, - **kwargs): - self.num_eagle_layers = num_eagle_layers - self.max_non_leaves_per_layer = max_non_leaves_per_layer - self.max_draft_len = max_draft_len - self.eagle_net_config = LLaMAConfig.from_dict( - kwargs["eagle_net_config"]) - del kwargs["eagle_net_config"] - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in EagleConfig - output['num_eagle_layers'] = self.num_eagle_layers - output['max_non_leaves_per_layer'] = self.max_non_leaves_per_layer - output['max_draft_len'] = self.max_draft_len - output['eagle_net_config'] = self.eagle_net_config.to_dict() - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - speculative_config_or_dir = kwargs.pop('speculative_model_dir', None) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - hf_config = None - hf_config_or_dir if speculative_config_or_dir is None else speculative_config_or_dir - if hf_config_or_dir is not None: - hf_config = LlamaConfig.from_pretrained(hf_config_or_dir) - - hf_config.model_type - n_head = hf_config.num_attention_heads - inter_size = hf_config.intermediate_size - n_layer = hf_config.num_hidden_layers - n_embd = hf_config.hidden_size - n_kv_head = hf_config.num_key_value_heads - rms_norm_eps = hf_config.rms_norm_eps - vocab_size = hf_config.vocab_size - rotary_scaling = hf_config.rope_scaling - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - n_positions = hf_config.max_position_embeddings - hidden_act = hf_config.hidden_act - dtype = str(hf_config.torch_dtype)[6:] if dtype == 'auto' else dtype - if hasattr(hf_config, 'head_dim'): - head_dim = hf_config.head_dim - else: - head_dim = hf_config.n_embd // hf_config.n_head - if hasattr(hf_config, 'head_size'): - head_size = hf_config.head_size - else: - head_size = head_dim - - if speculative_config_or_dir is None: - hf_config_eagle = hf_config.eagle - n_head_eagle = hf_config_eagle['num_attention_heads'] - inter_size_eagle = hf_config_eagle['intermediate_size'] - n_layer_eagle = hf_config_eagle['num_hidden_layers'] - n_embd_eagle = hf_config_eagle['hidden_size'] - n_kv_head_eagle = hf_config_eagle['num_key_value_heads'] - rms_norm_eps_eagle = hf_config_eagle['rms_norm_eps'] - n_positions_eagle = hf_config_eagle['max_position_embeddings'] - else: - hf_config_eagle = LlamaConfig.from_pretrained( - speculative_config_or_dir) - n_head_eagle = hf_config_eagle.num_attention_heads - inter_size_eagle = hf_config_eagle.intermediate_size - n_layer_eagle = hf_config_eagle.num_hidden_layers - n_embd_eagle = hf_config_eagle.hidden_size - n_kv_head_eagle = hf_config_eagle.num_key_value_heads - rms_norm_eps_eagle = hf_config_eagle.rms_norm_eps - n_positions_eagle = hf_config_eagle.max_position_embeddings - - if rotary_scaling is not None: - # assert use_gpt_attention_plugin, "RoPE scaling is only supported through GPT attention plugin." - rotary_scaling = { - "type": rotary_scaling["rope_type"], - } - rotary_scaling = rotary_scaling - - eagle_net_config = { - 'architecture': "LlamaForCausalLM", - 'dtype': dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': n_layer_eagle, - 'num_attention_heads': n_head_eagle, - 'hidden_size': n_embd_eagle, - 'intermediate_size': inter_size_eagle, - 'num_key_value_heads': n_kv_head_eagle, - 'vocab_size': vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': n_positions_eagle, - 'hidden_act': hidden_act, - 'rotary_base': rotary_base, - 'rotary_scaling': rotary_scaling, - 'norm_epsilon': rms_norm_eps_eagle, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': mapping.world_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - }, - 'use_parallel_embedding': kwargs['use_parallel_embedding'], - 'embedding_sharding_dim': kwargs['embedding_sharding_dim'], - 'head_dim': head_dim, - 'head_size': head_size - } - - config = { - 'architecture': 'EagleForCausalLM', - 'dtype': dtype, - 'logits_dtype': '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, - 'rotary_scaling': rotary_scaling, - 'norm_epsilon': rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': mapping.world_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - }, - 'use_parallel_embedding': kwargs['use_parallel_embedding'], - 'embedding_sharding_dim': kwargs['embedding_sharding_dim'], - 'num_eagle_layers': kwargs['speculative_config'].num_eagle_layers, - 'max_non_leaves_per_layer': - kwargs['speculative_config'].max_non_leaves_per_layer, - 'eagle_net_config': eagle_net_config - } - if quant_config: - config['quantization']['quant_algo'] = quant_config.quant_algo - config['quantization'][ - 'kv_cache_quant_algo'] = quant_config.kv_cache_quant_algo - - if quant_config.quant_algo == QuantAlgo.W4A16_GPTQ: - config['quantization'].update({ - "group_size": quant_config.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': QuantAlgo.W4A16_GPTQ - }) - eagle_quant_config = {} - try: - with open( - str(speculative_config_or_dir) + '/' + - 'hf_quant_config.json') as f: - eagle_quant_config = json.load(f) - if "lm_head" in eagle_quant_config['quantization'][ - 'exclude_modules']: - eagle_quant_config['quantization']['exclude_modules'] += [ - f"eagle_nets.{i}.lm_head" for i in range( - kwargs['speculative_config'].num_eagle_layers) - ] - config['quantization'].update( - eagle_quant_config['quantization']) - config['eagle_net_config']['quantization'].update( - eagle_quant_config['quantization']) - except IOError: - pass - - return cls.from_dict(config) diff --git a/tensorrt_llm/models/eagle/model.py b/tensorrt_llm/models/eagle/model.py deleted file mode 100644 index 779f397d70ed..000000000000 --- a/tensorrt_llm/models/eagle/model.py +++ /dev/null @@ -1,1327 +0,0 @@ -# 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 collections import OrderedDict -from typing import Optional, Union - -import numpy as np -import tensorrt as trt -from tqdm import tqdm - -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.generation_mixin import GenerationMixin -from tensorrt_llm.models.llama.model import LLaMAForCausalLM, LLaMAModel -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader - -from ..._common import default_net, default_trtnet -from ..._utils import pad_vocab_size -from ...functional import (Tensor, _create_tensor, cast, concat, - gather_last_token_logits, index_select, shape) -from ...layers import AttentionParams, ColumnLinear, SpecDecodingParams -from ...llmapi.kv_cache_type import KVCacheType -from ...module import Module, ModuleList -from ...plugin import TRT_LLM_PLUGIN_NAMESPACE -from ..modeling_utils import QuantConfig -from .config import EagleConfig - - -class TreeParams(object): - - def __init__(self, paths: Tensor = None): - self.paths = paths # on GPU - - -def eagle_sample_and_accept_draft_plugin(lm_logits: Tensor = None, - draft_tokens: Tensor = None, - draft_lens: Tensor = None, - eagle_temperature: Tensor = None, - rand_data_validation: Tensor = None, - posterior_alpha: Tensor = None, - posterior_threshold: Tensor = None, - tree_params: TreeParams = None, - greedy_sampling: Tensor = None, - use_dynamic_tree: Tensor = None): - ''' - Takes input logits and samples golden token + predictions from draft tokens. - Runs acceptance algorithm to accept draft tokens. - When greedy_sampling is True, all decoding is done using Top1 and token equality is used - for acceptance. Otherwise, typical acceptance and multinomial samplings are used. - - Visit tests/model/eagle/test_sample_accept_draft_tokens.py for input/output examples. - - Parameters: - lm_logits : Tensor - [num_tokens, vocab_size] - Logits produced by the base model. - - draft_tokens : Tensor - [batch_size, max_decoding_draft_tokens] - Input draft tokens. Only the first draft_lens[bi] tokens are relevant for bi'th row. - - draft_lens : Tensor - [batch_size] - Lengths of the draft_tokens. 0 for context request. Actual draft length for generation requests. - - eagle_temperature : Tensor - [batch_size] - Temperature of the decoding. - - rand_data_validation : Tensor - [batch_size, max_decoding_tokens] - Random data for multinomial sampling. - - posterior_alpha : Tensor - [batch_size] - Delta in typical acceptance in https://arxiv.org/pdf/2401.10774. - - posterior_threshold : Tensor - [batch_size] - Minimum probability threshold. - Epsilon in typical acceptance in https://arxiv.org/pdf/2401.10774. - - tree_params : TreeParams - Tree params of the input draft tokens. - - greedy_sampling : Tensor - Whether to do greedy or non-greedy sampling. - - use_dynamic_tree: Tensor - Whether to use dynamic tree (i.e., Eagle-2) - - - Return: - accepted_tokens : Tensor - [batch_size, max_path_len] - Accepted token ids. Only the first num_accepted_tokens[bi] tokens are relevant for bi'th row, - - num_accepted_tokens : Tensor - [batch_size] - Number of accepted tokens per request. Each entry is >= 1. - - accepted_paths : Tensor - [batch_size] - Indices of the accepted path per request of the input paths in tree_params.paths. - - next_draft_tokens : Tensor - [batch_size, max_decoding_draft_tokens] - Empty tensor used to allocate space for the next draft tokens. - - next_draft_lens : Tensor - [batch_size] - Empty tensor used to allocate space for lens of the next draft tokens. - - next_draft_paths : Tensor - [batch_size, max_decoding_len, max_path_len] - For EAGLE-1 just a copy of input path. - - hidden_size_batch_level_starts : Tensor - [max_draft_path_len * batch_size + 1] - Empty tensor used to allocate space for eagle_prepare_drafter_inputs_plugin. - ''' - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'EagleSampleAndAcceptDraftTokens', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pf_type = trt.PluginField("type_id", - np.array([int(lm_logits.dtype)], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type]) - plugin = plg_creator.create_plugin("eagle_sample_and_accept_draft_plugin", - pfc) - - plug_inputs = [ - lm_logits, draft_tokens, draft_lens, eagle_temperature, - rand_data_validation, posterior_alpha, posterior_threshold, - tree_params.paths, greedy_sampling, use_dynamic_tree - ] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, plugin) - - accepted_tokens = _create_tensor(layer.get_output(0), layer) - num_accepted_tokens = _create_tensor(layer.get_output(1), layer) - accepted_paths = _create_tensor(layer.get_output(2), layer) - next_draft_tokens = _create_tensor(layer.get_output(3), layer) - next_draft_lens = _create_tensor(layer.get_output(4), layer) - next_draft_paths = _create_tensor(layer.get_output(5), layer) - hidden_size_batch_level_starts = _create_tensor(layer.get_output(6), layer) - return tuple([ - accepted_tokens, num_accepted_tokens, accepted_paths, next_draft_tokens, - next_draft_lens, next_draft_paths, hidden_size_batch_level_starts - ]) - - -def eagle_draft_decoder_plugin( - layer_idx: int, num_eagle_layers: int, top_k_sampling: bool, - logits: Tensor, num_last_token_indices: Tensor, input_paths: Tensor, - use_dynamic_tree: Tensor, dynamic_tree_max_topK: Tensor, - input_draft_token_ids: Tensor, input_draft_lens: Tensor, - input_prev_scores: Tensor, input_current_expand_indices: Tensor, - input_all_layers_scores: Tensor, - input_all_layers_draft_token_ids: Tensor, - input_all_layers_draft_token_ids_predecessor: Tensor): - ''' - Parameters: - layer_idx : int - The index of the EagleNet. - - num_eagle_layers: int - The total number of eagle layers. - - top_k_sampling: bool - Whether to use top K sampling. Otherwise, use multinomial sampling. - - logits : Tensor - [num_logits, vocab_size] - Input logits. - - num_last_token_indices : Tensor - [1] - Number of valid logits in logits. - - input_paths: Tensor - [batch_size, max_decoding_tokens, max_path_len] - Input paths - - use_dynamic_tree: Tensor - [1] - Whether use dynamic tree (i.e., Eagle-2) - - dynamic_tree_max_topK: Tensor - [1] - Number of draft tokens expand in Eagle-2. - - input_draft_token_ids: Tensor - [batch_size, max_decoding_draft_tokens] - Draft tokens generated by previous EagleNets. - - input_draft_lens: Tensor - [batch_size] - Number of draft tokens for each request. - - input_prev_scores: Tensor - [batch_size, max_decoding_draft_tokens] - Last layer's scores - - input_current_expand_indices: Tensor - [batch_size, max_decoding_draft_tokens] - The indices of the nodes that expand in this layer. - The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - - input_all_layers_scores: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record scores from all EagleNets - - input_all_layers_draft_token_ids: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens from all EagleNets - - input_all_layers_draft_token_ids_predecessor: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens' predecessor - - Return: - output_draft_token_ids: Tensor - [batch_size, max_decoding_draft_tokens] - Draft tokens generated by this EagleNets, also include the previous draft tokens. - - output_draft_draft_lens: Tensor - [batch_size] - Number of draft tokens for each request. - - output_paths: Tensor - [batch_size, max_decoding_draft_tokens, max_path_len] - The latest path. - - output_current_scores: Tensor - [batch_size, max_decoding_draft_tokens] - This layer's scores, which will be used in next layer. - - output_next_expand_indices: - [batch_size, max_decoding_draft_tokens] - The indices of the nodes that expand in next layer. - The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - - output_all_layers_scores: - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record scores from all EagleNets - - output_all_layers_draft_token_ids: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens from all EagleNets - - output_all_layers_draft_token_ids_predecessor: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens' predecessor - - ''' - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'EagleDecodeDraftTokens', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pf_type = trt.PluginField("type_id", np.array([int(logits.dtype)], - np.int32), - trt.PluginFieldType.INT32) - - layer_idx_t = trt.PluginField("layer_idx", - np.array(layer_idx, dtype=np.int32), - trt.PluginFieldType.INT32) - - num_eagle_layers_t = trt.PluginField( - "num_eagle_layers", np.array(num_eagle_layers, dtype=np.int32), - trt.PluginFieldType.INT32) - - top_k_sampling_t = 1 if top_k_sampling else 0 - top_k_sampling_t = trt.PluginField( - "top_k_sampling", np.array(top_k_sampling_t, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [pf_type, layer_idx_t, num_eagle_layers_t, top_k_sampling_t]) - plugin = plg_creator.create_plugin("eagle_draft_decoder_plugin", pfc) - - plug_inputs = [ - logits, input_paths, num_last_token_indices, use_dynamic_tree, - dynamic_tree_max_topK, input_draft_token_ids, input_draft_lens, - input_prev_scores, input_current_expand_indices, - input_all_layers_scores, input_all_layers_draft_token_ids, - input_all_layers_draft_token_ids_predecessor - ] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, plugin) - - output_draft_token_ids = _create_tensor(layer.get_output(0), layer) - output_draft_lens = _create_tensor(layer.get_output(1), layer) - output_paths = _create_tensor(layer.get_output(2), layer) - output_current_scores = _create_tensor(layer.get_output(3), layer) - output_next_expand_indices = _create_tensor(layer.get_output(4), layer) - output_all_layers_scores = _create_tensor(layer.get_output(5), layer) - output_all_layers_draft_token_ids = _create_tensor(layer.get_output(6), - layer) - output_all_layers_draft_token_ids_predecessor = _create_tensor( - layer.get_output(7), layer) - return tuple([ - output_draft_token_ids, output_draft_lens, output_paths, - output_current_scores, output_next_expand_indices, - output_all_layers_scores, output_all_layers_draft_token_ids, - output_all_layers_draft_token_ids_predecessor - ]) - - -def eagle_prepare_drafter_inputs_plugin( - layer_idx: int, num_layers: int, max_non_leaves_per_layer: int, - attention_params: AttentionParams, input_ids: Tensor, - chunked_context_next_tokens: Tensor, accepted_token_ids: Tensor, - accepted_lens: Tensor, accepted_path_ids: Tensor, - next_draft_tokens: Tensor, next_draft_lens: Tensor, - next_draft_paths: Tensor, prev_draft_lens: Tensor, - prev_draft_paths: Tensor, hidden_size_batch_level_starts: Tensor, - input_gen_tokens: Tensor, - input_spec_decoding_generation_lengths: Tensor): - ''' - Prepares inputs for the EagleNet inference. - - Visit tests/model/eagle/test_prepare_drafter_inputs.py for input/output examples. - - Parameters: - layer_idx : int - Index of the EagleNet. 0 means context phase EagleNet or EagleNet0, - > 0 means EagleNetX or generation phase of EagleNet - - num_layers : int - Number of Eagle layers. - - max_non_leaves_per_layer : int - Number of nodes that can be non leaf in the tree at each level of the tree. - - attention_params : AttentionParams - - input_ids : Tensor - [num_tokens] - Tokens ids, inputs to the base model. - - chunked_context_next_tokens : Tensor - [batch_size] - The first token of the next chunk in chunked context. - -1 if current chunk is the last chunk or requests is in the gen phase. - - accepted_token_ids : Tensor - [batch_size, max_path_len] - Accepted tokens ids. - - accepted_lens : Tensor - [batch_size] - Number of accepted tokens. - - accepted_path_ids : Tensor - [batch_size] - Idx of the accepted path in prev_draft_paths. - - next_draft_tokens : Tensor - [batch_size, max_decoding_draft_tokens] - Tokens ids of the draft tokens being drafted by EagleNet - - next_draft_lens : Tensor - [batch_size] - Number of drafted tokens in next_draft_tokens - - next_draft_paths : Tensor - [batch_size, max_decoding_tokens, max_path_len] - Paths of the draft tokens for the next iteration. In EAGLE-1 is the same as prev_draft_paths - - prev_draft_lens : Tensor - [batch_size] - Number of draft tokens, inputs to the base model. - 0 for ctx requests and actual draft len for gen requests. - - prev_draft_paths : Tensor - [batch_size, max_decoding_tokens, max_path_len] - Paths of the draft tokens inputs to the base model. - - hidden_size_batch_level_starts : Tensor - [max_draft_path_len * batch_size + 1] - Exclusive sum of the starts of the segments of the hidden states in the concatenated array. - Hidden states shape is (flattened and w/o padding) - [max_draft_path_len, batch_size, num_output_tokens_i_j], where num_output_tokens_i_j - depends on the path of request j at level i. - - input_gen_tokens : Tensor - [num_gen_tokens] - Only needed to infer number of generation tokens from its shape. The content is irrelevant - - input_spec_decoding_generation_lengths : Tensor - [num_gen_requests] - Number of tokens for the base model. Only used to infer num_gen_requests from its shape, the content is irrelevant. - - Return: - sequence_length : Tensor - [batch_size] - Sequence length of the next EagleNet iteration. - For EagleNet0 equals to the (prompt_len + num_generated_tokens + accepted_lens). - For EagleNetX (X > 0) (seq_len_eagle_net_0 + spec_decoding_generation_lengths). - - context_length : Tensor - [batch_size] - Context length of the next EagleNet iteration. - For EagleNet0 it is either the actual context length of the request (for ctx requests) - or the number of accepted tokens in this iteration. EagleNet0's attn does chunked context attn. - For EagleNetX (X > 0), context length equals to the sequence length of the EagleNet0. - - spec_decoding_generation_lengths : Tensor - [batch_size] - Only relevant for EagleNetX (X > 0). - Number of draft tokens. - - spec_decoding_position_offsets : Tensor - [batch_size, max_decoding_tokens] - Only relevant for EagleNetX (X > 0). - Position offsets of the selected tokens from output_ids. - - spec_decoding_packed_mask : Tensor - [batch_size, max_decoding_tokens, ceil(max_decoding_tokens / 32)] - Only relevant for EagleNetX (X > 0). - uint32_t packed masks. - - output_ids : Tensor - [batch_size * max_non_leaves_per_layer * layer_idx] for layer_idx > 0 - [num_tokens - num_gen_tokens + num_gen_requests * (num_layers + 1)] for layer_idx == 0 - Token ids selected for the EagleNet iteration. - Tensor's actual size is larger than num_output_tokens. - - position_ids : Tensor - [batch_size] for layer_idx > 0 - [num_tokens - num_gen_tokens + num_gen_requests * (num_layers + 1)] for layer_idx == 0 - Position ids of the tokens selected for the EagleNet iteration. - Tensor's actual size is larger than num_output_tokens. - - hidden_states_indices : Tensor - [batch_size * max_non_leaves_per_layer * layer_idx] for layer_idx > 0 - [num_tokens - num_gen_tokens + num_gen_requests * (num_layers + 1)] for layer_idx == 0 - Indices of the hidden states to be selected from aggregated hidden states for the next iteration. - Tensor's actual size is larger than num_output_tokens. - - last_token_indices : Tensor - [batch_size * max_non_leaves_per_layer] - Indices of the hidden states to be converted to logits after the next EagleNet iteration. - Tensor's actual size is larger than num_output_tokens. - - num_last_token_indices : Tensor - [] - Number of logits selected after the next EagleNet iteration. - Tensors containing size of the outputs of V3 plugins. 0-D tensor. - - out_hidden_size_batch_level_starts : Tensor - [max_draft_path_len * batch_size + 1] - Same as hidden_size_batch_level_starts, but with updated path lens for the next level. - ''' - - plg_creator = trt.get_plugin_registry().get_creator( - 'EaglePrepareDrafterInputs', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - layer_idx = trt.PluginField("layer_idx", np.array(layer_idx, - dtype=np.int32), - trt.PluginFieldType.INT32) - - num_layers = trt.PluginField("num_layers", - np.array(num_layers, dtype=np.int32), - trt.PluginFieldType.INT32) - - max_non_leaves_per_layer = trt.PluginField( - "max_non_leaves_per_layer", - np.array(max_non_leaves_per_layer, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [layer_idx, num_layers, max_non_leaves_per_layer]) - plugin = plg_creator.create_plugin("eagle_prepare_drafter_inputs_plugin", - pfc, trt.TensorRTPhase.BUILD) - - plug_inputs = [ - attention_params.sequence_length, attention_params.context_lengths, - input_ids, chunked_context_next_tokens, accepted_token_ids, - accepted_lens, accepted_path_ids, next_draft_tokens, next_draft_lens, - next_draft_paths, prev_draft_lens, prev_draft_paths, - hidden_size_batch_level_starts, input_gen_tokens, - input_spec_decoding_generation_lengths - ] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - shape_inputs = [] - layer = default_trtnet().add_plugin_v3(plug_inputs, shape_inputs, plugin) - - sequence_length = _create_tensor(layer.get_output(0), layer) - context_length = _create_tensor(layer.get_output(1), layer) - spec_decoding_generation_lengths = _create_tensor(layer.get_output(2), - layer) - spec_decoding_position_offsets = _create_tensor(layer.get_output(3), layer) - spec_decoding_packed_mask = _create_tensor(layer.get_output(4), layer) - output_ids = _create_tensor(layer.get_output(5), layer) - position_ids = _create_tensor(layer.get_output(6), layer) - hidden_states_indices = _create_tensor(layer.get_output(7), layer) - last_token_indices = _create_tensor(layer.get_output(8), layer) - num_last_token_indices = _create_tensor(layer.get_output(9), layer) - out_hidden_size_batch_level_starts = _create_tensor(layer.get_output(10), - layer) - return tuple([ - sequence_length, context_length, spec_decoding_generation_lengths, - spec_decoding_position_offsets, spec_decoding_packed_mask, output_ids, - position_ids, hidden_states_indices, last_token_indices, - num_last_token_indices, out_hidden_size_batch_level_starts - ]) - - -class EagleNet(Module): - - def __init__(self, config, logits_dtype): - super().__init__() - self.drafter = LLaMAModel(config) - self.config = config - self.logits_dtype = logits_dtype - - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - self.lm_head = ColumnLinear(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) - else: - self.lm_head = None - - def forward(self, - input_ids, - position_ids=None, - hidden_states=None, - last_token_indices=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - hidden_states, cache = self.drafter( - input_ids, - position_ids=position_ids, - use_cache=True, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - hidden_states_for_embed=hidden_states) - - if self.config.mapping.is_last_pp_rank(): - hidden_states = gather_last_token_logits( - hidden_states, last_token_indices, - default_net().plugin_config.remove_input_padding) - return cast(self.lm_head(hidden_states), - dtype=self.logits_dtype), hidden_states, cache - - return None, hidden_states, cache - - -class EagleForCausalLM(LLaMAForCausalLM): - config_class = EagleConfig - - def __init__(self, config: EagleConfig): - - super().__init__(config) - - self.num_eagle_layers = config.num_eagle_layers - self.max_non_leaves_per_layer = config.max_non_leaves_per_layer - self.hidden_size = config.hidden_size - self.vocab_size = config.vocab_size - vocab_size_padded = pad_vocab_size(self.vocab_size, - config.mapping.tp_size) - eagle_net_config = config.eagle_net_config - eagle_net_config.mapping = Mapping(world_size=config.mapping.world_size, - rank=config.mapping.rank, - cp_size=1, - tp_size=config.mapping.world_size, - pp_size=1) - - eagle_net_config.fc_after_embed = True - eagle_net_config.use_input_layernorm_in_first_layer = False - eagle_net_config.use_last_layernorm = False - eagle_net_config.layer_idx_offset = config.num_hidden_layers - if self.mapping.is_last_pp_rank(): - self.eagle_nets = ModuleList([ - EagleNet(config=eagle_net_config, - logits_dtype=config.logits_dtype) - for _ in range(self.num_eagle_layers) - ]) - self.max_draft_len = config.max_draft_len - - def _prepare_drafter_inputs( - self, layer_idx, input_ids, chunked_context_next_tokens, - accepted_token_ids, accepted_lens, accepted_path_ids, - next_draft_tokens, next_draft_lens, next_draft_paths, - prev_draft_lens, prev_draft_paths, input_attention_params, - input_kv_cache_params, hidden_states, - host_ctx_eagle_net_request_types, - host_ctx_eagle_net_context_lengths, - host_ctx_eagle_net_past_key_value_lengths, - host_gen_eagle_net_request_types, - host_gen_eagle_net_context_lengths, - host_gen_eagle_net_past_key_value_lengths, - hidden_size_batch_level_starts, input_gen_tokens, - input_spec_decoding_generation_lengths, spec_decoding_use): - - drafter_inputs = eagle_prepare_drafter_inputs_plugin( - layer_idx, self.num_eagle_layers, self.max_non_leaves_per_layer, - input_attention_params, input_ids, chunked_context_next_tokens, - accepted_token_ids, accepted_lens, accepted_path_ids, - next_draft_tokens, next_draft_lens, next_draft_paths, - prev_draft_lens, prev_draft_paths, hidden_size_batch_level_starts, - input_gen_tokens, input_spec_decoding_generation_lengths) - - sequence_length, context_lengths, \ - spec_decoding_generation_lengths, spec_decoding_position_offsets, \ - spec_decoding_packed_mask, output_ids, position_ids, hidden_states_indices, \ - last_token_indices, num_last_token_indices, out_hidden_size_batch_level_starts \ - = drafter_inputs - - attention_params = input_attention_params - kv_cache_params = input_kv_cache_params - attention_params.sequence_length = sequence_length - attention_params.context_lengths = context_lengths - - if layer_idx == 0: - attention_params.host_request_types = host_ctx_eagle_net_request_types - attention_params.host_context_lengths = host_ctx_eagle_net_context_lengths - kv_cache_params.host_past_key_value_lengths = host_ctx_eagle_net_past_key_value_lengths - else: - attention_params.host_request_types = host_gen_eagle_net_request_types - attention_params.host_context_lengths = host_gen_eagle_net_context_lengths - kv_cache_params.host_past_key_value_lengths = host_gen_eagle_net_past_key_value_lengths - - spec_decoding_params = None - if layer_idx > 0: - spec_decoding_params = SpecDecodingParams( - True, self.max_draft_len, spec_decoding_generation_lengths, - spec_decoding_position_offsets, spec_decoding_packed_mask, - spec_decoding_use) - - # Get hidden states for accepted ids - hidden_states = self._slice_hidden_states(hidden_states, - hidden_states_indices) - - eagle_net_inputs = {} - eagle_net_inputs["input_ids"] = output_ids - eagle_net_inputs["position_ids"] = position_ids - eagle_net_inputs["last_token_indices"] = last_token_indices - eagle_net_inputs["attention_params"] = attention_params - eagle_net_inputs["kv_cache_params"] = kv_cache_params - eagle_net_inputs["spec_decoding_params"] = spec_decoding_params - eagle_net_inputs["hidden_states"] = hidden_states - return eagle_net_inputs, out_hidden_size_batch_level_starts, num_last_token_indices - - def _slice_hidden_states(self, hidden_states, indices): - hidden_states = index_select(hidden_states, 0, indices) - - hidden_states = hidden_states.view(concat( - [shape(indices, 0), shape(hidden_states, 1)]), - zero_is_placeholder=False) - return hidden_states - - def _eagle_fwd_helper(self, lm_logits, hidden_states, *args, **kwargs): - ''' - EAGLE inference can be viewed as - TRT_Engine(Target -> Draft0 -> Draft1 -> .. -> DraftK-1) -> Runtime -> TRT_Engine(..) -> .. - Target is Base model and Draft is EagleNet. - - Each EagleNet call can be viewed as call to Draft LLM in TensorRT-LLM. - We have to - 1. prepare input tensors before the EagleNet call (like in the the runtime), - 2. call EagleNet, - 3. decode draft tokens after the EagleNet. - The only difference with normal execution of the Draft model is that in EAGLE, - all these 3 things happen inside of the TensorRT engine execution. - We do 1 and 3 inside of the plugins. - For 1. We call eagle_prepare_drafter_inputs_plugin and for 3. eagle_draft_decoder_plugin. - - The first call to the EagleNet (Draft0 == EagleNet0) is the context phase. - For context request we populate the KV cache of the EagleNet. - For generation request that have accepted tokens we emulate KV cache reuse by doing chunked attention, - where chunk is the newly accepted tokens -- all previous tokens are already in the KV cache. - - The following calls to the EagleNet (EagleNetX (X > 0)) are generation phase. - For each EagleNetX we select tokens based on the current path which are going to be used for the generation. - - Let's consider an example: prompt ABCD. EAGLE-1, i.e tree is fixed for the iteration. - Tree: - ┌───┐ - │ 0 │ - └─┬─┘ - ┌─────┴─────┐ - ┌─┴─┐ ┌─┴─┐ ┌─┴─┐ - │ 1 │ │ 2 │ │ 3 │ - └─┬─┘ └─┬─┘ └───┘ - ┌─┴─┐ ┌─┴─┐ - │ 4 │ │ 5 │ - └─┬─┘ └─┬─┘ - ┌─┴─┐ ┌─┴─┐ - │ 6 │ │ 7 │ - └───┘ └───┘ - - First iteration of the TRT engine. Request is context request: - 1. Base model is called for [ABCD] tokens produces token E. - 2. Draft0 is called for tokens [BCDE] and produces - three possibilities F, G and H for positions 1, 2 and 3 respectively. - 3. Since H (position 3) is a leaf, it is not chosen as the input to Draft1 inference. - 4. Draft1 is called for tokens [FG] with appropriate mask of: - |F|G - F|1|0 - G|0|1 - It produces tokens I and J for positions 4 and 5. - 6. Draft2 is called for inputs [FGIJ] with mask of - |F|G|I|J - F|1|0|0|0 - G|0|1|0|0 - I|1|0|1|0 - J|0|1|0|1 - Note that we could've stored FG in KV cache and provide only IJ tokens here - with mask for past KV cache, but it is not supported in TensorRT LLM attention at the moment. - - Draft2 produces tokens K and L at positions 6 and 7. - 7. Resulting outputs are: - 7.1 accepted_ids [E] - 7.2 next_draft_tokens [FGHIJKL] - - Second iteration of the TRT engine. Request is the generation request. - 1. Base model is called for [EFGHIJKL]. Let's assume that it accepts [FIKM], i.e. the left-most path in the tree. - 2. Draft0 is called as context phase for [FIKM] -- to append to kv cache of the existing [BCDE]. - It produces tokens N, O and P for positions 1, 2 and 3. - 3. Draft1 is called as generation phase for [NO] tokens. - etc. - - ''' - input_tree_params = kwargs["tree_params"] - - draft_tokens = kwargs['draft_tokens'] - draft_lens = kwargs['draft_lens'] - eagle_temperature = kwargs['eagle_temperature'] - rand_data_validation = kwargs['rand_data_validation'] - posterior_alpha = kwargs['posterior_alpha'] - posterior_threshold = kwargs['posterior_threshold'] - input_ids = kwargs['input_ids'] - chunked_context_next_tokens = kwargs['chunked_context_next_tokens'] - host_ctx_eagle_net_request_types = kwargs[ - 'host_ctx_eagle_net_request_types'] - host_ctx_eagle_net_context_lengths = kwargs[ - 'host_ctx_eagle_net_context_lengths'] - host_ctx_eagle_net_past_key_value_lengths = kwargs[ - 'host_ctx_eagle_net_past_key_value_lengths'] - host_gen_eagle_net_request_types = kwargs[ - 'host_gen_eagle_net_request_types'] - host_gen_eagle_net_context_lengths = kwargs[ - 'host_gen_eagle_net_context_lengths'] - host_gen_eagle_net_past_key_value_lengths = kwargs[ - 'host_gen_eagle_net_past_key_value_lengths'] - input_gen_tokens = kwargs["input_gen_tokens"] - greedy_sampling = kwargs["greedy_sampling"] - - # Eagle-2 - use_dynamic_tree = kwargs['use_dynamic_tree'] - dynamic_tree_max_topK = kwargs['dynamic_tree_max_topK'] - prev_scores = kwargs['prev_scores'] - current_expand_indices = kwargs['current_expand_indices'] - all_layers_scores = kwargs['all_layers_scores'] - all_layers_draft_token_ids = kwargs['all_layers_draft_token_ids'] - all_layers_draft_token_ids_predecessor = kwargs[ - 'all_layers_draft_token_ids_predecessor'] - - # Sample target tokens and accept them - # next_draft_tokens, next_draft_lens, hidden_size_batch_level_starts are outputted here just to - # reserve the tensor with max size, which eagle_draft_decoder_plugin and - # eagle_prepare_drafter_inputs_plugin are going to directly write to - output = eagle_sample_and_accept_draft_plugin( - lm_logits, draft_tokens, draft_lens, eagle_temperature, - rand_data_validation, posterior_alpha, posterior_threshold, - input_tree_params, greedy_sampling, use_dynamic_tree) - accepted_tokens, num_accepted_tokens, accepted_paths, next_draft_tokens, \ - next_draft_lens, next_draft_paths, hidden_size_batch_level_starts = output - - attention_params = kwargs["attention_params"] - kv_cache_params = kwargs["kv_cache_params"] - spec_decoding_params = kwargs["spec_decoding_params"] - - input_hidden_states = hidden_states - - # Run EAGLE nets - for li in range(self.num_eagle_layers): - # Prepare EAGLE Net inputs. - eagle_net_inputs, hidden_size_batch_level_starts, num_last_token_indices = self._prepare_drafter_inputs( - layer_idx=li, - input_ids=input_ids, - chunked_context_next_tokens=chunked_context_next_tokens, - accepted_token_ids=accepted_tokens, - accepted_lens=num_accepted_tokens, - accepted_path_ids=accepted_paths, - next_draft_tokens=next_draft_tokens, - next_draft_lens=next_draft_lens, - next_draft_paths=next_draft_paths, - prev_draft_lens=draft_lens, - prev_draft_paths=input_tree_params.paths, - input_attention_params=attention_params, - input_kv_cache_params=kv_cache_params, - hidden_states=input_hidden_states, - host_ctx_eagle_net_request_types= - host_ctx_eagle_net_request_types, - host_ctx_eagle_net_context_lengths= - host_ctx_eagle_net_context_lengths, - host_ctx_eagle_net_past_key_value_lengths= - host_ctx_eagle_net_past_key_value_lengths, - host_gen_eagle_net_request_types= - host_gen_eagle_net_request_types, - host_gen_eagle_net_context_lengths= - host_gen_eagle_net_context_lengths, - host_gen_eagle_net_past_key_value_lengths= - host_gen_eagle_net_past_key_value_lengths, - hidden_size_batch_level_starts=hidden_size_batch_level_starts, - input_gen_tokens=input_gen_tokens, - input_spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_use=spec_decoding_params.spec_decoding_use) - - def single_eagle_net_iter(next_draft_tokens, next_draft_lens, - next_draft_paths, prev_scores, - current_expand_indices, all_layers_scores, - all_layers_draft_token_ids, - all_layers_draft_token_ids_predecessor): - # Run EAGLE Net - # NOTE: handle base net kv cache and eagle net kv cache are in the same tensor. - # EagleNet's kv cache is located starting at numBaseNetHiddenLayers in the kv tensor. - logits, hidden_states, _ = self.eagle_nets[li]( - **eagle_net_inputs) - - # FIXME We need to take top_k_sampling as an input - top_k_sampling = True - - # Decode draft tokens - next_draft_tokens, next_draft_lens, next_draft_paths, prev_scores, current_expand_indices, all_layers_scores, all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor = eagle_draft_decoder_plugin( - li, self.num_eagle_layers, top_k_sampling, logits, - num_last_token_indices, next_draft_paths, use_dynamic_tree, - dynamic_tree_max_topK, next_draft_tokens, next_draft_lens, - prev_scores, current_expand_indices, all_layers_scores, - all_layers_draft_token_ids, - all_layers_draft_token_ids_predecessor) - - return next_draft_tokens, next_draft_lens, hidden_states, next_draft_paths, prev_scores, current_expand_indices, all_layers_scores, all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor - - - next_draft_tokens, next_draft_lens, hidden_states, next_draft_paths, prev_scores, \ - current_expand_indices, all_layers_scores, all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor \ - = single_eagle_net_iter(next_draft_tokens, next_draft_lens, next_draft_paths, \ - prev_scores, current_expand_indices, all_layers_scores, \ - all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor) - - # Update params - if li == 0: - eagle_net_0_sequence_length = eagle_net_inputs[ - "attention_params"].sequence_length - input_hidden_states = hidden_states - else: - input_hidden_states = concat( - [input_hidden_states, hidden_states]) - - kv_cache_params = eagle_net_inputs["kv_cache_params"] - attention_params = eagle_net_inputs["attention_params"] - attention_params.context_lengths = eagle_net_0_sequence_length - attention_params.sequence_length = eagle_net_0_sequence_length - - # Mark tensors as output - accepted_tokens.mark_output('accepted_tokens') - num_accepted_tokens.mark_output('num_accepted_tokens') - accepted_paths.mark_output('accepted_paths') - next_draft_tokens.mark_output('next_draft_tokens') - next_draft_lens.mark_output('next_draft_lens') - next_draft_paths.mark_output('next_draft_paths') - - return next_draft_tokens - - def forward(self, *args, **kwargs): - extra_args = [ - "draft_tokens", "draft_lens", "eagle_temperature", - "rand_data_validation", "tree_params", - "host_ctx_eagle_net_request_types", - "host_ctx_eagle_net_context_lengths", - "host_ctx_eagle_net_past_key_value_lengths", - "host_gen_eagle_net_request_types", - "host_gen_eagle_net_context_lengths", - "host_gen_eagle_net_past_key_value_lengths", "input_gen_tokens", - "chunked_context_next_tokens", "posterior_alpha", - "posterior_threshold", "greedy_sampling", "use_dynamic_tree", - "dynamic_tree_max_topK", "prev_scores", "current_expand_indices", - "all_layers_scores", "all_layers_draft_token_ids", - "all_layers_draft_token_ids_predecessor" - ] - - base_kwargs = {k: v for k, v in kwargs.items() if k not in extra_args} - - # Base model forward - hidden_states = super().forward(*args, **base_kwargs) - - if self.mapping.is_last_pp_rank(): - extra_args = ["hidden_states"] - kwargs = {k: v for k, v in kwargs.items() if k not in extra_args} - - assert kwargs['use_cache'] and default_net( - ).plugin_config.paged_kv_cache - - if self.mapping.is_last_pp_rank(): - lm_logits, hidden_states, all_hidden_states = hidden_states - lm_logits = cast(lm_logits, self.config.logits_dtype) - # Call eagle logic to accept prev draft tokens and predict next draft tokens - next_draft_tokens = self._eagle_fwd_helper(lm_logits, - all_hidden_states, *args, - **kwargs) - else: - hidden_states.mark_output('hidden_states_output', self.config.dtype) - - if self.mapping.is_last_pp_rank(): - return next_draft_tokens - return hidden_states - - def prepare_inputs(self, *args, **kwargs): - """ - Inputs needed: - device_request_types: [bs] - draft_tokens: [bs, max_draft_len] - draft_lens: [bs] - spec_decoding_generation_lengths: [bs] - spec_decoding_position_offsets: [bs, max_gen_tokens] - spec_decoding_packed_mask: [bs, max_draft_len, packed_length] ** - eagle_temperature: [bs] - rand_data_validation: [bs, max_draft_len] - - ** The mask is tricky since the boolean mask will need to be - packed in runtime. So, the last dim will be: - packed_length = ceil((max_draft_len+1)/32) - """ - default_range = GenerationMixin.default_range - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - multiple_profiles = default_net().plugin_config.multiple_profiles - max_batch_size = kwargs['max_batch_size'] - assert max_batch_size is not None - gt_range = default_range(max_batch_size * (self.max_draft_len + 1), - min_range=0, - opt_offset=1) - - kwargs['speculative_decoding_draft_tokens_external'] = False - kwargs['max_draft_len'] = self.max_draft_len - kwargs['spec_decoding_is_generation_length_variable'] = True - kwargs[ - 'num_hidden_layers'] = self.config.num_hidden_layers + self.config.eagle_net_config.num_hidden_layers - - # Call base class prepare inputs - inputs = super().prepare_inputs(*args, **kwargs) - - assert inputs['spec_decoding_params'] is not None - - kv_cache_type = KVCacheType.PAGED if paged_kv_cache else KVCacheType.CONTINUOUS - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type) - - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=kwargs['max_beam_width'], - max_input_len=kwargs['max_input_len'], - max_num_tokens=kwargs['max_num_tokens'], - max_draft_len=self.max_draft_len, - opt_batch_size=None - if 'opt_batch_size' not in kwargs else kwargs['opt_batch_size'], - opt_num_tokens=None - if 'opt_num_tokens' not in kwargs else kwargs['opt_num_tokens'], - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles, - kv_cache_type=kv_cache_type) - - bb_range = ranges['bb_range'] - - draft_len_range = [self.max_draft_len for _ in range(len(bb_range))] - decoding_len_range = [(self.max_draft_len + 1) - for _ in range(len(bb_range))] - path_len_range = [(self.num_eagle_layers + 1) - for _ in range(len(bb_range))] - gen_tokens_range = [gt_range for _ in range(len(bb_range))] - num_eagle_layers_range = [ - self.num_eagle_layers for _ in range(len(bb_range)) - ] - draft_len_square_range = [(self.max_draft_len * self.max_draft_len) - for _ in range(len(bb_range))] - - draft_tokens = Tensor(name='draft_tokens', - dtype=trt.int32, - shape=[-1, self.max_draft_len], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('draft_len', draft_len_range), - ])) - draft_lens = Tensor(name='draft_lens', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - eagle_temperature = Tensor(name='eagle_temperature', - dtype=trt.float32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - rand_data_validation = Tensor(name='rand_data_validation', - dtype=trt.float32, - shape=[-1, self.max_draft_len + 1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('decoding_len', decoding_len_range), - ])) - posterior_alpha = Tensor(name='posterior_alpha', - dtype=trt.float32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - posterior_threshold = Tensor(name='posterior_threshold', - dtype=trt.float32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - draft_paths = Tensor( - name='draft_paths', - dtype=trt.int32, - shape=[-1, self.max_draft_len + 1, self.num_eagle_layers + 1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('decoding_len', decoding_len_range), - ('path_len', path_len_range), - ])) - - host_ctx_eagle_net_request_types = Tensor( - name='host_ctx_eagle_net_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_ctx_eagle_net_context_lengths = Tensor( - name='host_ctx_eagle_net_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_ctx_eagle_net_past_key_value_lengths = Tensor( - name='host_ctx_eagle_net_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_gen_eagle_net_request_types = Tensor( - name='host_gen_eagle_net_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_gen_eagle_net_context_lengths = Tensor( - name='host_gen_eagle_net_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_gen_eagle_net_past_key_value_lengths = Tensor( - name='host_gen_eagle_net_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - - input_gen_tokens = Tensor(name='input_gen_tokens', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('gen_tokens', gen_tokens_range), - ])) - chunked_context_next_tokens = Tensor(name='chunked_context_next_tokens', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - greedy_sampling = Tensor(name='greedy_sampling', - dtype=trt.int32, - shape=[1]) - - use_dynamic_tree = Tensor(name='use_dynamic_tree', - dtype=trt.int32, - shape=[1]) - - dynamic_tree_max_topK = Tensor(name='dynamic_tree_max_topK', - dtype=trt.int32, - shape=[1]) - - prev_scores = Tensor(name='prev_scores', - dtype=trt.float32, - shape=[-1, self.max_draft_len], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('draft_len', draft_len_range), - ])) - - current_expand_indices = Tensor(name='current_expand_indices', - dtype=trt.int32, - shape=[-1, self.max_draft_len], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('draft_len', draft_len_range), - ])) - - all_layers_scores = Tensor(name='all_layers_scores', - dtype=trt.float32, - shape=[ - -1, self.num_eagle_layers, - self.max_draft_len * self.max_draft_len - ], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('num_eagle_layers', - num_eagle_layers_range), - ('draft_len_square', - draft_len_square_range), - ])) - - all_layers_draft_token_ids = Tensor( - name='all_layers_draft_token_ids', - dtype=trt.int32, - shape=[ - -1, self.num_eagle_layers, - self.max_draft_len * self.max_draft_len - ], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('num_eagle_layers', num_eagle_layers_range), - ('draft_len_square', draft_len_square_range), - ])) - - all_layers_draft_token_ids_predecessor = Tensor( - name='all_layers_draft_token_ids_predecessor', - dtype=trt.int32, - shape=[ - -1, self.num_eagle_layers, - self.max_draft_len * self.max_draft_len - ], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('num_eagle_layers', num_eagle_layers_range), - ('draft_len_square', draft_len_square_range), - ])) - - tree_params = TreeParams(paths=draft_paths) - - inputs['draft_tokens'] = draft_tokens - inputs['draft_lens'] = draft_lens - inputs['eagle_temperature'] = eagle_temperature - inputs['posterior_alpha'] = posterior_alpha - inputs['posterior_threshold'] = posterior_threshold - inputs['rand_data_validation'] = rand_data_validation - inputs['tree_params'] = tree_params - inputs[ - 'host_ctx_eagle_net_request_types'] = host_ctx_eagle_net_request_types - inputs[ - 'host_ctx_eagle_net_context_lengths'] = host_ctx_eagle_net_context_lengths - inputs[ - 'host_ctx_eagle_net_past_key_value_lengths'] = host_ctx_eagle_net_past_key_value_lengths - inputs[ - 'host_gen_eagle_net_request_types'] = host_gen_eagle_net_request_types - inputs[ - 'host_gen_eagle_net_context_lengths'] = host_gen_eagle_net_context_lengths - inputs[ - 'host_gen_eagle_net_past_key_value_lengths'] = host_gen_eagle_net_past_key_value_lengths - inputs['input_gen_tokens'] = input_gen_tokens - inputs['chunked_context_next_tokens'] = chunked_context_next_tokens - inputs['greedy_sampling'] = greedy_sampling - inputs['use_dynamic_tree'] = use_dynamic_tree - inputs['dynamic_tree_max_topK'] = dynamic_tree_max_topK - inputs['prev_scores'] = prev_scores - inputs['current_expand_indices'] = current_expand_indices - inputs['all_layers_scores'] = all_layers_scores - inputs['all_layers_draft_token_ids'] = all_layers_draft_token_ids - inputs[ - 'all_layers_draft_token_ids_predecessor'] = all_layers_draft_token_ids_predecessor - - return inputs - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - assert hf_model_or_dir is not None - speculative_model_dir = kwargs.get('speculative_model_dir', None) - tllm_config = EagleConfig.from_hugging_face(hf_model_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - # for rank in range(mapping.world_size): - tllm_config.mapping = Mapping(world_size=mapping.world_size, - rank=mapping.rank, - cp_size=1, - tp_size=mapping.tp_size, - pp_size=mapping.pp_size) - - model = EagleForCausalLM(tllm_config) - - def check_and_update(module, dict): - if hasattr(module, 'tllm_to_externel_key_dict'): - module.tllm_to_externel_key_dict.update(dict) - else: - module.tllm_to_externel_key_dict = dict - - def copy(tensors): - if isinstance(tensors, list): - if None in tensors: - return tensors - else: - return [tensor.clone() for tensor in tensors] - elif tensors is None: - return tensors - else: - return tensors.clone() - - shared_weight_prefixs = [] - tllm_weights = {} - customized_dict = {"drafter": ""} - if speculative_model_dir is None: - # Single checkpoint for ModelOpt - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": "model"}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = 'eagle_module' - loader = ModelWeightsLoader(speculative_model_dir, customized_dict) - loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.named_parameters()): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update(loader.load(tllm_key, preprocess=copy)) - else: - tllm_weights.update(loader.load(tllm_key)) - loader.fill(tllm_weights) - else: - # Double checkpoint for HF - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": ""}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = '' - - # Load base model - base_loader = ModelWeightsLoader(hf_model_or_dir) - base_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.transformer.named_parameters()): - tllm_weights.update(base_loader.load("transformer." + tllm_key)) - tllm_weights.update(base_loader.load("lm_head.weight")) - # for idx in range(args.num_eagle_layers): - for idx in range(4): - tllm_weights.update( - base_loader.load(f"eagle_nets.{idx}.lm_head.weight", - preprocess=copy)) - - # Load eagle model - eagle_loader = ModelWeightsLoader(str(speculative_model_dir), - customized_dict) - eagle_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.eagle_nets.named_parameters()): - if not tllm_key.endswith("lm_head.weight"): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key, - preprocess=copy)) - else: - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key)) - base_loader.fill(tllm_weights) - - return model diff --git a/tensorrt_llm/models/enc_dec/__init__.py b/tensorrt_llm/models/enc_dec/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/enc_dec/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/enc_dec/model.py b/tensorrt_llm/models/enc_dec/model.py deleted file mode 100644 index be3c5afc49f6..000000000000 --- a/tensorrt_llm/models/enc_dec/model.py +++ /dev/null @@ -1,2195 +0,0 @@ -# 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 math -from collections import OrderedDict -from typing import List, Optional - -import tensorrt as trt -import torch - -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import (numpy_to_torch, pad_vocab_size, - str_dtype_to_torch) -from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType, - MLPType, PositionEmbeddingType, Tensor, - assertion, cast, gather_last_token_logits, - gelu, maximum, minimum, recv, send, shape, - transpose, unsqueeze) -# yapf: disable -from tensorrt_llm.layers import (MLP, Attention, AttentionMaskParams, - AttentionMaskType, AttentionParams, - BertAttention, ColumnLinear, Conv1d, Embedding, - FusedGatedMLP, GatedMLP, GroupNorm, - KeyValueCacheParams, LanguageAdapter, - LanguageAdapterConfig, LayerNorm, LoraParams, - PromptTuningEmbedding, RmsNorm) -# yapf: enable -from tensorrt_llm.lora_helper import (LoraConfig, - get_default_trtllm_modules_to_hf_modules, - use_lora) -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import PretrainedConfig, PretrainedModel -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.parameter import Parameter -from tensorrt_llm.plugin.plugin import current_all_reduce_helper -from tensorrt_llm.quantization import QuantMode - -layernorm_map = { - LayerNormType.LayerNorm: LayerNorm, - LayerNormType.RmsNorm: RmsNorm, - LayerNormType.GroupNorm: GroupNorm, -} - -mlp_map = { - MLPType.MLP: MLP, - MLPType.GatedMLP: GatedMLP, - MLPType.FusedGatedMLP: FusedGatedMLP, -} - - -class EncDecTransformer(Module): - - def __init__(self, - vocab_size, - hidden_size, - has_vocab_embeddings=True, - max_position_embeddings=None, - has_position_embedding=False, - type_vocab_size=None, - has_embedding_layernorm=False, - has_embedding_scale=False, - layernorm_eps=1e-5, - layernorm_type=LayerNormType.LayerNorm, - dtype=None, - use_parallel_embedding=False, - embedding_sharding_dim=0, - mapping=Mapping(), - has_model_final_layernorm=False, - norm_epsilon=1e-05, - is_decoder=False): - super().__init__() - - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - self.mapping = mapping - - if self.mapping.is_first_pp_rank(): - if has_vocab_embeddings: - self.vocab_embedding = Embedding( - vocab_size, - hidden_size, - dtype=dtype, - tp_size=mapping.tp_size if use_parallel_embedding else 1, - tp_group=mapping.tp_group - if use_parallel_embedding else None, - sharding_dim=embedding_sharding_dim, - tp_rank=mapping.tp_rank) - - self.position_embedding = None - self.max_position_embeddings = max_position_embeddings - if has_position_embedding: - self.position_embedding = Embedding( - max_position_embeddings, - hidden_size, - dtype=dtype, - tp_size=mapping.tp_size if use_parallel_embedding else 1, - tp_group=mapping.tp_group - if use_parallel_embedding else None, - sharding_dim=embedding_sharding_dim, - tp_rank=mapping.tp_rank) - - self.token_type_embedding = None - if type_vocab_size: - self.token_type_embedding = Embedding( - type_vocab_size, - hidden_size, - dtype=dtype, - tp_size=mapping.tp_size if use_parallel_embedding else 1, - tp_group=mapping.tp_group - if use_parallel_embedding else None, - sharding_dim=embedding_sharding_dim, - tp_rank=mapping.tp_rank) - - # e.g. BART true, T5 false - self.ln_embed = None - if has_embedding_layernorm: - self.ln_embed = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # e.g. BART true, T5 false - self.embedding_scale = 1.0 - if has_embedding_scale: - self.embedding_scale = math.sqrt(hidden_size) - - # Note: embedding offset in BART is not considered as a standard. For the specific case, - # we just need to shrink its position embedding table by [offset:] during weight loading - - self.layers = None - self.is_decoder = is_decoder - self.has_model_final_layernorm = has_model_final_layernorm - if self.mapping.is_last_pp_rank(): - if self.has_model_final_layernorm: - self.ln_f = ln_type(normalized_shape=hidden_size, - eps=norm_epsilon, - dtype=dtype) - - def assign_module(self, module, module_name): - setattr(self, module_name, module) - - def embedding(self, - input_ids, - position_ids=None, - token_type_ids=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None): - # position_ids and token_type_ids are provided inputs - # and should not be formulated deterministically - - args = [prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - x = self.vocab_embedding(input_ids, *args) * self.embedding_scale - self.register_network_output('word_embeddings', x) - - if self.position_embedding: - pos_emb = self.position_embedding(position_ids) - self.register_network_output('position_embeddings', pos_emb) - x = x + pos_emb - if self.token_type_embedding: - x = x + self.token_type_embedding(token_type_ids) - - if self.ln_embed: - x = self.ln_embed(x) - - return x - - -class EncoderLayer(Module): - - def __init__(self, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.LayerNorm, - layernorm_eps=1e-5, - hidden_act="relu", - mlp_type=MLPType.MLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - quant_mode=QuantMode(0), - language_adapter_config: LanguageAdapterConfig = None): - super().__init__() - - # e.g. BART regular, T5 RMS - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - # e.g. BART post, T5 pre - self.layernorm_position = layernorm_position - - # e.g. BART q_scaling = 1.f, T5 q_scaling = 1.f/sqrt(head_size) - self.attention = BertAttention( - hidden_size, - num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - relative_attention=relative_attention, - max_distance=max_distance, - num_buckets=num_buckets, - quant_mode=quant_mode) - - self.attention_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # T5/BART MLP, Flan-T5 GatedMLP - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - quant_mode=quant_mode, - ) - - self.mlp_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.residual_scaling = residual_scaling - - # T5-series model(e.g. t5-large, t5-3b, flan-t5-small) has accuracy issue due to fp16 overflow - # after residual add. We add workaround for clamping fp16 range [-64000, 64000] after every - # residual add to avoid accuracy drop. - self.fp16_clamping = fp16_clamping - - self.adapter = None - self.adapter_layer_norm = None - - if language_adapter_config: - self.adapter_layer_norm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.adapter = LanguageAdapter( - language_adapter_config=language_adapter_config, - hidden_size=hidden_size, - hidden_act=hidden_act, - mapping=mapping, - has_mlp_bias=has_mlp_bias, - dtype=dtype, - quant_mode=quant_mode) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - input_lengths=None, - max_input_length=None, - lora_layer_params=None, - language_adapter_routings: Optional[Tensor] = None): - assert isinstance(hidden_states, Tensor) - - # self attention - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.attention_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - input_lengths=input_lengths, - max_input_length=max_input_length, - lora_layer_params=lora_layer_params) - - self.register_network_output('attention_output', attention_output) - - hidden_states = residual + attention_output - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.attention_layernorm(hidden_states) - - # MLP - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - self.register_network_output('mlp_output', hidden_states) - - hidden_states = residual + hidden_states - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - # MT Specific: adapters - if self.adapter: - residual = hidden_states - hidden_states = self.adapter_layer_norm(hidden_states) - hidden_states = self.adapter.layers( - hidden_states, static_routing_input=language_adapter_routings) - hidden_states = residual + hidden_states - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - return hidden_states - - -class DecoderLayer(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.LayerNorm, - layernorm_eps=1e-5, - hidden_act="relu", - mlp_type=MLPType.MLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - skip_cross_kv=False, - use_implicit_relative_attention=False, - quant_mode=QuantMode(0), - language_adapter_config: LanguageAdapterConfig = None): - super().__init__() - - # e.g. BART regular, T5 RMS - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - # e.g. BART post, T5 pre - self.layernorm_position = layernorm_position - - # e.g. BART q_scaling = 1.f, T5 q_scaling = 1.f/sqrt(head_size) - self.self_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=False, - relative_attention=relative_attention, - max_distance=max_distance if use_implicit_relative_attention else 0, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType.relative - if relative_attention else PositionEmbeddingType.learned_absolute, - use_implicit_relative_attention=use_implicit_relative_attention, - quant_mode=quant_mode) - - self.self_attention_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # Note: self attn uses MMHA, mask is always causal triangular - # cross attn has two scenarios: - # - in context phase, all ones mask, same as padding type - # - in generation phase, same causal triangular mask as MMHA - # - context phase special handling is done in plugin by resetting mask type - # - # e.g. BART q_scaling = 1.f, T5 q_scaling = 1.f/sqrt(head_size) - self.cross_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=True, - relative_attention= - False, # Cross attention has no relative attention bias - max_distance=max_distance, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType.learned_absolute, - skip_cross_kv=skip_cross_kv, - quant_mode=quant_mode) - - self.cross_attention_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # T5/BART MLP, Flan-T5 GatedMLP - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - quant_mode=quant_mode, - ) - - self.mlp_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.residual_scaling = residual_scaling - - # T5-series model(e.g. t5-large, t5-3b, flan-t5-small) has accuracy issue due to fp16 overflow - # after residual add. We add workaround for clamping fp16 range [-64000, 64000] after every - # residual add to avoid accuracy drop. - self.fp16_clamping = fp16_clamping - - self.adapter = None - self.adapter_layer_norm = None - - if language_adapter_config: - self.adapter_layer_norm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.adapter = LanguageAdapter( - language_adapter_config=language_adapter_config, - hidden_size=hidden_size, - hidden_act=hidden_act, - mapping=mapping, - has_mlp_bias=has_mlp_bias, - dtype=dtype, - quant_mode=quant_mode) - - def forward(self, - hidden_states: Tensor, - encoder_output: Optional[Tensor] = None, - attention_mask_params=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - language_adapter_routings: Optional[Tensor] = None): - assert isinstance(hidden_states, Tensor) - - if encoder_output: - assert isinstance(encoder_output, Tensor) - - # self-attention - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.self_attention_layernorm(hidden_states) - - attention_output = self.self_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.self_attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params) - - if use_cache: - attention_output, presents_self = attention_output - - self.register_network_output('self_attention_output', attention_output) - - hidden_states = residual + attention_output - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.self_attention_layernorm(hidden_states) - - # cross attention - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.cross_attention_layernorm(hidden_states) - - attention_output = self.cross_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.cross_attention_mask, - attention_packed_mask=attention_mask_params. - cross_attention_packed_mask, - encoder_output=encoder_output, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse) - - if use_cache: - attention_output, presents_cross = attention_output - - self.register_network_output('cross_attention_output', attention_output) - - hidden_states = residual + attention_output - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.cross_attention_layernorm(hidden_states) - - # MLP - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - self.register_network_output('mlp_output', hidden_states) - - hidden_states = residual + hidden_states - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - # MT Specific: adapters - if self.adapter: - residual = hidden_states - hidden_states = self.adapter_layer_norm(hidden_states) - hidden_states = self.adapter.layers( - hidden_states, static_routing_input=language_adapter_routings) - hidden_states = residual + hidden_states - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if use_cache: - return (hidden_states, presents_self, presents_cross) - return hidden_states - - -class EncoderModel(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - super().__init__(config) - self.mapping = self.config.mapping - - self.has_position_embedding = self.config.has_position_embedding - type_vocab_size = self.config.type_vocab_size - self.has_token_type_embedding = False if type_vocab_size is None else True - - # e.g. BART regular, T5 RMS - self.layernorm_type = self.config.layernorm_type - - # e.g. BART true, T5 false - self.has_attention_qkvo_bias = self.config.has_attention_qkvo_bias - self.has_mlp_bias = self.config.has_mlp_bias - - # e.g. BART false, T5 true - self.has_model_final_layernorm = self.config.has_model_final_layernorm - - self._dtype = self.config.dtype - - self.total_num_layers = self.config.num_hidden_layers - self.num_layers = self.config.num_hidden_layers // self.mapping.pp_size - - self.hidden_size = self.config.hidden_size - self.num_heads = self.config.num_attention_heads - num_kv_heads = self.num_heads - if num_kv_heads is None or num_kv_heads <= 0: - num_kv_heads = self.config.num_attention_heads - self.num_kv_heads = num_kv_heads - self.head_size = self.hidden_size // self.num_heads if self.config.head_size is None else self.config.head_size - - self.fp16_clamping = (self.config.dtype - == 'float16') and (self.config.model_type == 't5') - self.mlp_type = MLPType.MLP if not hasattr( - self.config, "mlp_type") else self.config.mlp_type - self.language_adapter_config = None if not hasattr( - self.config, - 'language_adapter_config') else LanguageAdapterConfig.from_dict( - self.config.language_adapter_config) - - self.transformer = EncDecTransformer( - self.config.vocab_size, - self.config.hidden_size, - max_position_embeddings=self.config.max_position_embeddings, - has_position_embedding=self.has_position_embedding, - type_vocab_size=type_vocab_size, - has_embedding_layernorm=self.config.has_embedding_layernorm, - has_embedding_scale=self.config.has_embedding_scale, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.layernorm_type, - dtype=self.config.dtype, - use_parallel_embedding=self.config.use_parallel_embedding, - embedding_sharding_dim=self.config.embedding_sharding_dim, - mapping=self.mapping, - has_model_final_layernorm=self.has_model_final_layernorm, - norm_epsilon=self.config.norm_epsilon, - is_decoder=False) - - encoder_layers = ModuleList([ - EncoderLayer( - hidden_size=self.hidden_size, - ffn_hidden_size=self.config.intermediate_size, - num_attention_heads=self.num_heads, - num_kv_heads=num_kv_heads, - head_size=self.head_size, - max_position_embeddings=self.config.max_position_embeddings, - q_scaling=self.config.q_scaling, - has_attention_qkvo_bias=self.has_attention_qkvo_bias, - has_mlp_bias=self.has_mlp_bias, - layernorm_position=self.config.layernorm_position, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.layernorm_type, - hidden_act=self.config.hidden_act, - mlp_type=self.mlp_type, - mapping=self.mapping, - dtype=self.config.dtype, - residual_scaling=1.0 - if not hasattr(self.config, "residual_scaling") else - self.config.residual_scaling, - relative_attention=self.config.relative_attention, - max_distance=self.config.max_distance, - num_buckets=self.config.num_buckets, - fp16_clamping=self.fp16_clamping, - quant_mode=self.config.quant_mode, - language_adapter_config=self.language_adapter_config) - for _ in self.mapping.pp_layers(self.total_num_layers) - ]) - self.transformer.assign_module(encoder_layers, "layers") - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('has_position_embedding', False) - config.set_if_not_exist('type_vocab_size', None) - config.set_if_not_exist('rescale_before_lm_head', False) - config.set_if_not_exist('layernorm_type', LayerNormType.LayerNorm) - config.set_if_not_exist('layernorm_position', - LayerNormPositionType.pre_layernorm) - config.set_if_not_exist('has_attention_qkvo_bias', False) - config.set_if_not_exist('has_mlp_bias', False) - config.set_if_not_exist('has_model_final_layernorm', False) - config.set_if_not_exist('encoder_hidden_size', None) - config.set_if_not_exist('encoder_num_heads', None) - config.set_if_not_exist('encoder_num_kv_heads', None) - config.set_if_not_exist('encoder_head_size', None) - config.set_if_not_exist('model_type', 't5') - config.set_if_not_exist('skip_cross_kv', False) - config.set_if_not_exist('mlp_type', MLPType.MLP) - config.set_if_not_exist('has_embedding_scale', False) - config.set_if_not_exist('residual_scaling', 1.0) - config.set_if_not_exist('has_lm_head_bias', False) - config.set_if_not_exist('num_buckets', None) - config.set_if_not_exist('max_distance', None) - config.set_if_not_exist('relative_attention', False) - config.set_if_not_exist('residual_scaling', 1.0) - - def forward(self, - input_ids: Tensor, - input_lengths=None, - position_ids=None, - token_type_ids=None, - hidden_states=None, - max_input_length=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_mask=None, - lora_params: LoraParams = None, - language_adapter_routings: Optional[Tensor] = None): - # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs - if self.mapping.is_first_pp_rank(): - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - hidden_states = self.transformer.embedding(input_ids, position_ids, - token_type_ids, - *ptuning_args) - self.register_network_output('embedding_layer_output', - hidden_states) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - for layer_idx, encoder_layer in enumerate(self.transformer.layers): - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(layer_idx) - hidden_states = encoder_layer( - hidden_states=hidden_states, - attention_mask=attention_mask, - input_lengths=input_lengths, - max_input_length=max_input_length, - lora_layer_params=lora_layer_params, - language_adapter_routings=language_adapter_routings) - - if self.mapping.is_last_pp_rank(): - if self.has_model_final_layernorm: - hidden_states = self.transformer.ln_f(hidden_states) - hidden_states.mark_output('encoder_output', self._dtype) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - hidden_states.mark_output('hidden_states_output', self._dtype) - - return hidden_states - - def prepare_inputs(self, - max_batch_size, - max_input_len, - prompt_embedding_table_size: int = 0, - lora_target_modules: List[str] = None, - *args, - **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - hidden_size = self.hidden_size - - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - inlen_range = [1, (max_input_len + 1) // 2, max_input_len] - num_tokens_range = [ - 1, - (max_input_len * max_batch_size + 1) // 2, - max_input_len * max_batch_size, - ] - - input_ids, position_ids, token_type_ids, hidden_states = None, None, None, None - remove_input_padding = default_net().plugin_config.remove_input_padding - use_lora_plugin = default_net().plugin_config.lora_plugin - - attention_mask = None - if remove_input_padding: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("num_tokens", [num_tokens_range])]), - ) - if self.has_position_embedding: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', - [num_tokens_range])]), - ) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', - [num_tokens_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, hidden_size], - dim_range=OrderedDict([ - ('num_tokens', [num_tokens_range]), - ('hidden_size', [hidden_size]), - ])) - else: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([("batch_size", [bs_range]), - ("input_len", [inlen_range])]), - ) - if self.has_position_embedding: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, -1, hidden_size], - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('input_len', [inlen_range]), - ('hidden_size', [hidden_size]), - ])) - - if not default_net().plugin_config.bert_attention_plugin: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('input_len', [inlen_range]), - ]), - ) - - # if self.mapping.tp_size > 1: - # current_all_reduce_helper().set_workspace_tensor(self.mapping, 1) - # FIXME(TRTLLM-996): Support custom allreduce for encoder models on C++ runtime - - input_lengths = Tensor( - name="input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size", [bs_range])]), - ) - max_input_length = Tensor( - name="max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("max_input_length", [inlen_range])]), - ) - - prompt_embedding_table = None - tasks = None - prompt_vocab_size = None - - if self.mapping.is_first_pp_rank() and prompt_embedding_table_size > 0: - p_embedding_range = [[ - 1, prompt_embedding_table_size // 2, prompt_embedding_table_size - ]] - - prompt_embedding_table = Tensor(name='prompt_embedding_table', - dtype=self._dtype, - shape=[-1, hidden_size], - dim_range=OrderedDict([ - ('prompt_embedding_table_size', - p_embedding_range), - ('hidden_size', [hidden_size]), - ])) - if remove_input_padding: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('input_len_task', - [num_tokens_range])])) - else: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([ - ('batch_size', bs_range), - ('broadcast_dim', [1]), - ])) - prompt_vocab_size = Tensor(name='prompt_vocab_size', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('size', [1])])) - ''' - LoRA plugin related inputs: - lora_target_modules for BART-encoder: - ['attn_q', 'attn_v'] - For BART-decoder, the lora_target_modules is different. - See comments in the DecoderModel.prepare_inputs() for more details. - ''' - lora_weights_pointers = None - lora_ranks = None - lora_params = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - # In current design, q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time. - # However, BART lora modules only contain two of them, so we use zero tensor to fill the missing ones. - missing_qkv_modules = [] - if any(x in lora_target_modules - for x in ["attn_q", "attn_k", "attn_v"]): - for lora_module in ["attn_q", "attn_k", "attn_v"]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in (lora_target_modules + missing_qkv_modules): - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('in_out_scales', [3])])) - lora_weight_pointer_dict.update({ - f'{lora_module}_lora_weights_pointers': - lora_weight_pointer - }) - - lora_rank = Tensor(name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', - [bs_range])])) - lora_rank_dict.update( - {f'{lora_module}_lora_ranks': lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - host_request_types = Tensor(name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', - [bs_range])])) - - host_context_lengths = None - if remove_input_padding: - host_context_lengths = Tensor(name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [bs_range]) - ])) - - lora_params = LoraParams( - lora_ranks=lora_ranks, - lora_weights_pointers=lora_weights_pointers, - host_request_types=host_request_types, - host_context_lengths=host_context_lengths, - ) - - language_adapter_routings = None - if self.language_adapter_config: - language_adapter_routings = Tensor( - name="language_adapter_routings", - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([('num_tokens', [num_tokens_range]), - ('routing_dim', [1])]), - ) - - result = { - 'input_ids': input_ids, - 'input_lengths': input_lengths, - 'position_ids': position_ids, - 'token_type_ids': token_type_ids, - 'hidden_states': hidden_states, - 'max_input_length': max_input_length, - 'prompt_embedding_table': prompt_embedding_table, - 'prompt_tasks': tasks, - 'prompt_vocab_size': prompt_vocab_size, - 'attention_mask': attention_mask, - 'lora_params': lora_params, - 'language_adapter_routings': language_adapter_routings, - } - - return result - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config) - - def use_prompt_tuning(self): - embedding = self.transformer.vocab_embedding - self.transformer.vocab_embedding = PromptTuningEmbedding( - num_embeddings=embedding.num_embeddings, - embedding_dim=embedding.embedding_dim, - dtype=embedding.dtype, - tp_size=embedding.tp_size, - tp_group=embedding.tp_group, - sharding_dim=embedding.sharding_dim, - tp_rank=embedding.tp_rank) - - self.transformer.vocab_embedding.weight.value = embedding.weight.raw_value - - def precompute_relative_attention_bias(self, build_config): - pass - - -class DecoderModel(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - super().__init__(config) - - self.mapping = self.config.mapping - - self.has_position_embedding = self.config.has_position_embedding - type_vocab_size = self.config.type_vocab_size - self.has_token_type_embedding = (type_vocab_size is not None) - self.rescale_before_lm_head = self.config.rescale_before_lm_head - - # e.g. BART regular, T5 RMS - self.layernorm_type = self.config.layernorm_type - - # e.g. BART true, T5 false - self.has_attention_qkvo_bias = self.config.has_attention_qkvo_bias - self.has_mlp_bias = self.config.has_mlp_bias - - # e.g. BART false, T5 true - self.has_model_final_layernorm = self.config.has_model_final_layernorm - self._dtype = self.config.dtype - # no quantization considered for now - self._kv_dtype = self._dtype - self._logits_dtype = self.config.logits_dtype - - self.total_num_layers = self.config.num_hidden_layers - self.num_layers = self.config.num_hidden_layers // self.mapping.pp_size - - self.hidden_size = self.config.hidden_size - self.num_heads = self.config.num_attention_heads - num_kv_heads = self.num_heads - if num_kv_heads is None or num_kv_heads <= 0: - num_kv_heads = self.num_heads - self.num_kv_heads = num_kv_heads - self.head_size = self.hidden_size // self.num_heads if self.config.head_size is None else self.config.head_size - - self.encoder_hidden_size = self.config.encoder_hidden_size - self.encoder_num_heads = self.config.encoder_num_heads - encoder_num_kv_heads = None if not hasattr( - self.config, - "encoder_num_kv_heads") else self.config.encoder_num_kv_heads - if encoder_num_kv_heads is None or encoder_num_kv_heads <= 0: - encoder_num_kv_heads = self.encoder_num_heads - self.encoder_num_kv_heads = encoder_num_kv_heads - self.encoder_head_size = self.encoder_hidden_size // self.num_heads if self.config.encoder_head_size is None else self.config.encoder_head_size - - self.has_position_embedding = self.config.has_position_embedding - self.has_token_type_embedding = type_vocab_size is not None - - self.fp16_clamping = (self.config.dtype - == 'float16') and (self.config.model_type - in ['t5', 'pix2struct']) - - self.skip_cross_kv = self.config.skip_cross_kv - self.mlp_type = MLPType.MLP if not hasattr( - self.config, "mlp_type") else self.config.mlp_type - self.use_implicit_relative_attention = self.config.use_implicit_relative_attention if hasattr( - self.config, "use_implicit_relative_attention") else False - - self.language_adapter_config = None if not hasattr( - self.config, - 'language_adapter_config') else LanguageAdapterConfig.from_dict( - self.config.language_adapter_config) - - self.transformer = EncDecTransformer( - self.config.vocab_size, - self.config.hidden_size, - max_position_embeddings=self.config.max_position_embeddings, - has_position_embedding=self.config.has_position_embedding, - type_vocab_size=type_vocab_size, - has_embedding_layernorm=self.config.has_embedding_layernorm, - has_embedding_scale=self.config.has_embedding_scale, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.config.layernorm_type, - dtype=self._dtype, - use_parallel_embedding=self.config.use_parallel_embedding, - embedding_sharding_dim=self.config.embedding_sharding_dim, - mapping=self.mapping, - has_model_final_layernorm=self.has_model_final_layernorm, - norm_epsilon=self.config.norm_epsilon, - is_decoder=True) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - decoder_layers = ModuleList([ - DecoderLayer( - local_layer_idx=layer_idx - layers_range[0], - hidden_size=self.config.hidden_size, - ffn_hidden_size=self.config.intermediate_size, - num_attention_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_size=self.head_size, - max_position_embeddings=self.config.max_position_embeddings, - q_scaling=self.config.q_scaling, - has_attention_qkvo_bias=self.config.has_attention_qkvo_bias, - has_mlp_bias=self.config.has_mlp_bias, - layernorm_position=self.config.layernorm_position, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.config.layernorm_type, - hidden_act=self.config.hidden_act, - mlp_type=self.mlp_type, - mapping=self.mapping, - dtype=self._dtype, - residual_scaling=self.config.residual_scaling, - relative_attention=self.config.relative_attention, - max_distance=self.config.max_distance, - num_buckets=self.config.num_buckets, - fp16_clamping=self.fp16_clamping, - skip_cross_kv=self.skip_cross_kv, - use_implicit_relative_attention=self. - use_implicit_relative_attention, - quant_mode=self.config.quant_mode, - language_adapter_config=self.language_adapter_config) - for layer_idx in layers_range - ]) - self.transformer.assign_module(decoder_layers, "layers") - - if self.mapping.is_last_pp_rank(): - vocab_size_padded = pad_vocab_size(self.config.vocab_size, - self.config.mapping.tp_size) - self.lm_head = ColumnLinear( - self.config.hidden_size, - vocab_size_padded, - bias=False if not hasattr(self.config, "has_lm_head_bias") else - self.config.has_lm_head_bias, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True, - ) - - self.trtllm_modules_to_hf_modules = { - **get_default_trtllm_modules_to_hf_modules(), - "attn_q": "self_attn.q_proj", - "attn_k": "self_attn.k_proj", - "attn_v": "self_attn.v_proj", - "attn_dense": "self_attn.o_proj", - "cross_attn_q": "encoder_attn.q_proj", - "cross_attn_k": "encoder_attn.k_proj", - "cross_attn_v": "encoder_attn.v_proj", - "cross_attn_dense": "encoder_attn.o_proj", - } - - if self.config.relative_attention and not self.use_implicit_relative_attention: - self.rel_attn_table = Parameter( - shape=(self.config.num_attention_heads // self.mapping.tp_size, - self.config.num_buckets), - dtype=self._dtype) - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('has_position_embedding', False) - config.set_if_not_exist('type_vocab_size', None) - config.set_if_not_exist('rescale_before_lm_head', False) - config.set_if_not_exist('layernorm_type', LayerNormType.LayerNorm) - config.set_if_not_exist('layernorm_position', - LayerNormPositionType.pre_layernorm) - config.set_if_not_exist('has_attention_qkvo_bias', False) - config.set_if_not_exist('has_mlp_bias', False) - config.set_if_not_exist('has_model_final_layernorm', False) - config.set_if_not_exist('encoder_hidden_size', None) - config.set_if_not_exist('encoder_num_heads', None) - config.set_if_not_exist('encoder_num_kv_heads', None) - config.set_if_not_exist('encoder_head_size', None) - config.set_if_not_exist('model_type', 't5') - config.set_if_not_exist('skip_cross_kv', False) - config.set_if_not_exist('mlp_type', MLPType.MLP) - config.set_if_not_exist('has_embedding_scale', False) - config.set_if_not_exist('residual_scaling', 1.0) - config.set_if_not_exist('has_lm_head_bias', False) - config.set_if_not_exist('num_buckets', None) - config.set_if_not_exist('max_distance', None) - config.set_if_not_exist('relative_attention', False) - - def forward(self, - decoder_input_ids: Tensor, - encoder_output: Tensor, - position_ids=None, - token_type_ids=None, - use_cache=False, - attention_mask_params=None, - last_token_ids=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - lora_params: LoraParams = None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - language_adapter_routings: Optional[Tensor] = None): - if self.mapping.is_first_pp_rank(): - assert isinstance(decoder_input_ids, Tensor) - else: - assert isinstance(hidden_states, Tensor) - - # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs - if self.mapping.is_first_pp_rank(): - hidden_states = self.transformer.embedding(decoder_input_ids, - position_ids, - token_type_ids) - self.register_network_output('embedding_layer_output', - hidden_states) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - kv_cache_params.fill_none_tensor_list(len(self.transformer.layers)) - - if use_cache: - presents = [] - - for i, (decoder_layer, past) in enumerate( - zip(self.transformer.layers, kv_cache_params.past_key_value)): - - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(i) - - hidden_states = decoder_layer( - hidden_states, - encoder_output=encoder_output, - attention_mask_params=attention_mask_params, - use_cache=use_cache, - kv_cache_params=KeyValueCacheParams( - past_key_value=past, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - cache_indirection=kv_cache_params.cache_indirection, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=kv_cache_params. - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers=kv_cache_params. - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping=kv_cache_params. - host_cross_kv_cache_pool_mapping), - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - language_adapter_routings=language_adapter_routings) - - if use_cache: - presents_self, presents_cross = hidden_states[1], hidden_states[ - 2] - presents.append((presents_self, presents_cross)) - hidden_states = hidden_states[0] - self.register_network_output(f'decoder_layer_{i}_output', - hidden_states) - - if self.mapping.is_last_pp_rank(): - if self.has_model_final_layernorm: - hidden_states = self.transformer.ln_f(hidden_states) - - # [bs, seq, hidden_size] or [num_tokens, hidden_size] -> [bs, hidden_size] - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) - self.register_network_output('logits_before_lmhead', hidden_states) - - # Rescale output before projecting on vocab (for T5) - # See https://github.com/huggingface/transformers/blob/0b192de1f353b0e04dad4813e02e2c672de077be/src/transformers/models/t5/modeling_t5.py#L1769-L1772 - # Note: this is specific for T5, to make it more generic, one can pass in a config: - # self.config.tie_word_embeddings - default to be True for T5 - # openai whisper model didn't use this rescale - if self.rescale_before_lm_head: - hidden_states = hidden_states * (self.hidden_size**-0.5) - - # [bs, hidden_size] -> [bs, vocab_size] - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output('logits', self._logits_dtype) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - hidden_states.mark_output('hidden_states_output', self._dtype) - - if use_cache and default_net().plugin_config.paged_kv_cache == False: - for i, present in zip(self.mapping.pp_layers(self.total_num_layers), - presents): - present[0].mark_output(f'present_key_value_{i}', self._kv_dtype) - if default_net().plugin_config.gpt_attention_plugin: - present[1].mark_output(f'cross_present_key_value_{i}', - self._kv_dtype) - if self.mapping.is_last_pp_rank(): - return (lm_logits, tuple(presents)) - return (hidden_states, tuple(presents)) - else: - if self.mapping.is_last_pp_rank(): - return lm_logits - return hidden_states - - def prepare_inputs(self, - max_batch_size, - max_beam_width, - max_decoder_input_len, - max_seq_len, - max_encoder_input_len, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - use_cache=True, - *args, - **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - # Prepare inputs - max_output_len = max_decoder_input_len + max_seq_len - - head_size = self.head_size - num_kv_heads = (self.num_kv_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - - encoder_head_size = self.encoder_head_size - encoder_num_kv_heads = (self.encoder_num_kv_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - - bb_range = [ - 1, (max_batch_size * max_beam_width + 1) // 2, - max_batch_size * max_beam_width - ] - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - beam_width_range = [1, (max_beam_width + 1) // 2, max_beam_width] - inlen_range = [ - 1, 1, max_decoder_input_len - ] # context phase >= 1 (if forced_input_ids), generation phase = 1 - encoder_inlen_range = [ - 1, (max_encoder_input_len + 1) // 2, max_encoder_input_len - ] - mask_len_range = [1, (max_output_len + 1) // 2 + 1, max_output_len + 1] - max_output_len_range = [0, (max_output_len + 1) // 2, max_output_len] - - encoder_num_tokens_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len * max_batch_size + 1) // 2, - max_encoder_input_len * max_batch_size, - ] - decoder_num_tokens_range = [ - 1, - max_batch_size * max_beam_width, - max(max_decoder_input_len * max_batch_size, - max_beam_width * max_batch_size), - ] - - # No enable_two_optimization_profiles support yet - - encoder_input_len_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len + 1) // 2, - max_encoder_input_len - ] - max_cross_packed_mask_dim0 = max_batch_size * ( - (max_decoder_input_len + 128 - 1) // 128) * 128 - max_cross_packed_mask_dim1 = ( - (max_encoder_input_len + 256 - 1) // 256) * 256 // 32 - cross_packed_mask_dim0_range = [ - 1, (max_cross_packed_mask_dim0 + 1) // 2, max_cross_packed_mask_dim0 - ] - cross_packed_mask_dim1_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_cross_packed_mask_dim1 + 1) // 2, - max_cross_packed_mask_dim1 - ] - - past_key_value = [] - sequence_length = None - host_past_key_value_lengths = None - runtime_perf_knobs = None - context_progress = None - attention_mask = None - cross_attention_mask = None - cross_attention_packed_mask = None - attention_mask_params = AttentionMaskParams() - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - remove_input_padding = default_net().plugin_config.remove_input_padding - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - use_lora_plugin = default_net().plugin_config.lora_plugin - - input_ids, position_ids, token_type_ids, hidden_states = None, None, None, None - if remove_input_padding: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - if self.has_position_embedding: - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('decoder_num_tokens', - [decoder_num_tokens_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, self.hidden_size], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ('hidden_size', [self.hidden_size]), - ])) - else: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('input_len', [inlen_range]), - ])) - if self.has_position_embedding: - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]), - ('input_len', [inlen_range]), - ])) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size_beam_width', - [bb_range]), - ('input_len', [inlen_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, -1, self.hidden_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range - ]), - ('input_len', [inlen_range]), - ('hidden_size', [self.hidden_size]), - ])) - - encoder_input_lengths = Tensor( - name="encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range])]), - ) - encoder_max_input_length = Tensor( - name="encoder_max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("encoder_max_input_length", - [encoder_inlen_range])]), - ) - encoder_output = None - if remove_input_padding: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, self.encoder_hidden_size], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("encoder_hidden_size", [self.encoder_hidden_size]), - ]), - ) - else: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, -1, self.encoder_hidden_size], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("encoder_hidden_size", [self.encoder_hidden_size]), - ]), - ) - - if use_gpt_attention_plugin: - host_past_key_value_lengths = Tensor( - name='host_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range])]), - ) - - context_lengths = None - host_context_lengths = None - host_request_types = None - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = Tensor(name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - - if use_gpt_attention_plugin: - sequence_length = Tensor( - name='sequence_length', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range])]), - ) - - context_lengths = Tensor(name='context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - host_request_types = Tensor(name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - runtime_perf_knobs = Tensor(name='host_runtime_perf_knobs', - dtype=trt.int64, - shape=[16], - dim_range=OrderedDict([ - ('perf_knob_size', [16]) - ])) - context_progress = Tensor(name='host_context_progress', - dtype=trt.int64, - shape=[1], - dim_range=OrderedDict([ - ('context_progress_size', [1]) - ])) - - last_token_ids = None - if self.mapping.is_last_pp_rank() and not gather_context_logits: - last_token_ids = Tensor( - name="last_token_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_last_token_ids", [bb_range]) - ]), - ) - - if not use_gpt_attention_plugin: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('mask_len', [mask_len_range]), - ]), - ) - - cross_attention_mask = Tensor( - name='cross_attention_mask', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('query_len', [inlen_range]), - ('encoder_input_len_2', [encoder_input_len_range]), - ]), - ) - else: - cross_attention_mask = Tensor( - name='cross_attention_mask', - dtype=trt.bool, - shape=[-1, -1], - dim_range=OrderedDict([ - ('decoder_num_tokens_2', - [decoder_num_tokens_range - ]), # TODO should use same name as input_ids - ('encoder_input_len_2', [encoder_input_len_range]), - ]), - ) - - cross_attention_packed_mask = Tensor( - name='cross_attention_packed_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('cross_packed_mask_dim0', [cross_packed_mask_dim0_range]), - ('cross_packed_mask_dim1', [cross_packed_mask_dim1_range]), - ]), - ) - - # create the attention_mask_params. - attention_mask_params = AttentionMaskParams( - attention_mask, None, cross_attention_mask, - cross_attention_packed_mask) - - cache_indirection = Tensor( - name='cache_indirection', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_cache', [bs_range]), - ('beam_width', [beam_width_range]), - ('max_seq_len', [max_output_len_range]), - ]), - ) - - if self.mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor(self.mapping, 1) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - num_pp_layers = len(layers_range) - - host_max_attention_window_sizes = None - host_sink_token_length = None - if use_gpt_attention_plugin: - host_max_attention_window_sizes = Tensor( - name=f'host_max_attention_window_sizes', - dtype=trt.int32, - shape=[num_pp_layers], - dim_range=OrderedDict([('num_layers', [num_pp_layers])])) - host_sink_token_length = Tensor(name='host_sink_token_length', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('scalar', - [1])])) - ''' - LoRA plugin related inputs: - lora_target_modules for BART-decoder: - ['attn_q', 'cross_attn_q', - 'attn_v', 'cross_attn_v'] - This is NOT directly loaded from the adapter-config file - We make it this way because BART has LoRA weights for both self-attention and cross-attention in decoder - ''' - lora_weights_pointers = None - lora_ranks = None - lora_params = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - # In current design, q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time. - # However, BART lora modules only contain two of them, so we use zero tensor to fill the missing ones. - missing_qkv_modules = [] - if any(x in lora_target_modules - for x in ["attn_q", "attn_k", "attn_v"]): - for lora_module in [ - "attn_q", - "attn_k", - "attn_v", - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - if any(x in lora_target_modules - for x in ["cross_attn_q", "cross_attn_k", "cross_attn_v"]): - for lora_module in [ - "cross_attn_q", "cross_attn_k", "cross_attn_v" - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in (lora_target_modules + missing_qkv_modules): - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([('batch_size_beam_width', - [bb_range]), - ('in_out_scales', [3])])) - lora_weight_pointer_dict.update({ - f'{lora_module}_lora_weights_pointers': - lora_weight_pointer - }) - - lora_rank = Tensor(name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - lora_rank_dict.update( - {f'{lora_module}_lora_ranks': lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - # For cross attention, we need to use encoder_input_lengths (in CPU) to pass - # as the host_context_lengths to the lora_plugin. But for self attention, we - # should keep using the original host_context_lengths. Therefore, we keep both - # of them in the lora_params. - host_encoder_input_lengths = None - if remove_input_padding: - host_encoder_input_lengths = Tensor( - name="host_encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range]) - ]), - ) - - lora_params = LoraParams( - lora_ranks=lora_ranks, - lora_weights_pointers=lora_weights_pointers, - host_context_lengths=host_context_lengths, - max_encoder_context_length=max_encoder_input_len, - host_request_types=host_request_types, - host_encoder_input_lengths=host_encoder_input_lengths, - ) - - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - host_kv_cache_pool_pointers = None - host_kv_cache_pool_mapping = None - - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - host_cross_kv_cache_pool_pointers = None - host_cross_kv_cache_pool_mapping = None - - if use_cache: - if not paged_kv_cache: - for i in layers_range: - kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('num_heads', [num_kv_heads]), - ('past_key_len', [max_output_len_range]), - ('head_size', [head_size]), - ]) - kv = Tensor(name=f'past_key_value_{i}', - dtype=self._kv_dtype, - shape=[-1, 2, num_kv_heads, -1, head_size], - dim_range=kv_dim_range) - - if use_gpt_attention_plugin: - cross_kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('cross_num_heads', [encoder_num_kv_heads]), - ('cross_past_key_len', [encoder_input_len_range]), - ('cross_head_size', [encoder_head_size]), - ]) - cross_kv = Tensor(name=f'cross_past_key_value_{i}', - dtype=self._kv_dtype, - shape=[ - -1, 2, encoder_num_kv_heads, -1, - encoder_head_size - ], - dim_range=cross_kv_dim_range) - past_key_value.append((kv, cross_kv)) - else: - # use encoder_output directly, no need to save cross_past_key_value - past_key_value.append((kv, )) - - # TODO: Remove this when TRT fix the named dimension - if not remove_input_padding: - assertion( - shape( - input_ids if self.mapping.is_first_pp_rank() else - hidden_states, 0) == shape(kv, 0), 'batch size') - - else: # paged_kv_cache == True - # PagedKV setup for KV cache of self-attention - max_blocks_per_seq_range = [[ - math.ceil(max_output_len_range[0] / tokens_per_block), - math.ceil(max_output_len_range[1] / tokens_per_block), - math.ceil(max_output_len_range[2] / tokens_per_block) - ]] - max_blocks_per_seq_range = [[ - x for x in max_blocks_per_seq_range[0] - ]] - - # PagedKV setup for KV cache of cross-attention - max_cross_blocks_per_seq_range = [[ - math.ceil(encoder_input_len_range[0] / tokens_per_block), - math.ceil(encoder_input_len_range[1] / tokens_per_block), - math.ceil(encoder_input_len_range[2] / tokens_per_block) - ]] - max_cross_blocks_per_seq_range = [[ - x for x in max_cross_blocks_per_seq_range[0] - ]] - - # TODO(oargov): add support for vgqa + vwindow, meanwhile assume a single kv cache pool - num_kv_cache_pools = 1 - - kv_cache_block_offsets = Tensor( - name=f'kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_block_offsets = Tensor( - name=f'host_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_pool_pointers = Tensor( - name=f'host_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_pools_layers', [num_kv_cache_pools]), - ('num_pools_kv', [2]), - ])) - host_kv_cache_pool_mapping = Tensor( - name=f"host_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - # paged blocks for cross kv - cross_kv_cache_block_offsets = Tensor( - name=f'cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_block_offsets = Tensor( - name=f'host_cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_pool_pointers = Tensor( - name=f'host_cross_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('num_pools', [2]), - ])) - host_cross_kv_cache_pool_mapping = Tensor( - name=f"host_cross_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - for i in layers_range: - past_key_value.append(None) - - kv_cache_params = KeyValueCacheParams( - past_key_value=past_key_value, - host_past_key_value_lengths=host_past_key_value_lengths, - host_max_attention_window_sizes=host_max_attention_window_sizes, - host_sink_token_length=host_sink_token_length, - cache_indirection=cache_indirection, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers= - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping= - host_cross_kv_cache_pool_mapping, - ) - - attention_params = AttentionParams( - sequence_length=sequence_length, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_decoder_input_len, - host_request_types=host_request_types, - encoder_input_lengths=encoder_input_lengths, - encoder_max_input_length=encoder_max_input_length, - host_runtime_perf_knobs=runtime_perf_knobs, - host_context_progress=context_progress) - - cross_kv_cache_gen = Tensor(name='cross_kv_cache_gen', - dtype=trt.bool, - shape=[1], - dim_range=OrderedDict([ - ('boolean', [1]), - ])) - cross_kv_reuse = None - num_heads = (self.num_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - cross_kv_out_dim = 2 * num_kv_heads * self.head_size - if self.skip_cross_kv: - if remove_input_padding: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - else: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, -1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - - language_adapter_routings = None - if self.language_adapter_config: - language_adapter_routings = Tensor( - name="language_adapter_routings", - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([('decoder_num_tokens_range', - [decoder_num_tokens_range]), - ('routing_dim', [1])]), - ) - - result = { - 'decoder_input_ids': input_ids, - 'encoder_output': encoder_output, - 'position_ids': position_ids, - 'token_type_ids': token_type_ids, - 'use_cache': True, - 'attention_mask_params': attention_mask_params, - 'last_token_ids': last_token_ids, - 'kv_cache_params': kv_cache_params, - 'attention_params': attention_params, - 'hidden_states': hidden_states, - 'lora_params': lora_params, - 'cross_kv_cache_gen': cross_kv_cache_gen, - 'cross_kv_reuse': cross_kv_reuse, - 'language_adapter_routings': language_adapter_routings, - } - - return result - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) - - def precompute_relative_attention_bias(self, build_config): - if self.config.relative_attention and not self.use_implicit_relative_attention: - relative_attention_bias_builder = torch.ops.tensorrt_llm.relative_attention_bias - rel_attn_precomputed = torch.zeros( - (self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len + 1, build_config.max_seq_len + 1), - dtype=str_dtype_to_torch(self.config.dtype), - device='cuda') - rel_attn_table = numpy_to_torch( - self.rel_attn_table.raw_value).to('cuda') - relative_attention_bias_builder( - rel_attn_precomputed, - rel_attn_table, - self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len, - self.config.num_buckets, - False, - self.config.max_distance, - ) - for layer_idx in range(self.num_layers): - self.transformer.layers[ - layer_idx].self_attention.set_rel_attn_table( - build_config.max_seq_len, rel_attn_precomputed) - - -class WhisperEncoder(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - super().__init__(config) - self._dtype = self.config.dtype - conv1 = Conv1d(config.n_mels, - config.hidden_size, - kernel_size=3, - padding=1, - dtype=self._dtype) - conv2 = Conv1d(config.hidden_size, - config.hidden_size, - kernel_size=3, - stride=2, - padding=1, - dtype=self._dtype) - self.transformer = EncDecTransformer( - 0, - self.config.hidden_size, - has_vocab_embeddings=False, - max_position_embeddings=self.config.max_position_embeddings, - has_position_embedding=self.config.has_position_embedding, - layernorm_type=LayerNormType.LayerNorm, - dtype=self.config.dtype, - has_model_final_layernorm=True, - is_decoder=False) - encoder_layers = ModuleList([ - EncoderLayer( - hidden_size=config.hidden_size, - ffn_hidden_size=config.hidden_size * 4, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_attention_heads, - head_size=config.hidden_size // config.num_attention_heads, - q_scaling=1.0, - has_attention_qkvo_bias=True, - has_mlp_bias=True, - hidden_act='gelu', - dtype=self._dtype) for _ in range(config.num_hidden_layers) - ]) - self.transformer.assign_module(encoder_layers, "layers") - self.transformer.assign_module(conv1, "conv1") - self.transformer.assign_module(conv2, "conv2") - self.max_audio_feature_seq_len = 3000 - self.downsample_factor = 2 - - def forward(self, - input_features: Tensor, - input_lengths=None, - position_ids=None): - if default_net().plugin_config.remove_input_padding: - # BXT,D -> 1,BxT,D -> 1,D,BxT - input_features = unsqueeze(input_features, 0) - input_features = transpose(input_features, 1, 2) - x_type = input_features.dtype - input_features = cast(input_features, self._dtype) - x = self.transformer.conv1(input_features) - x = gelu(x) - x = self.transformer.conv2(x) - x = cast(x, x_type) - x = gelu(x) - x = transpose(x, 2, 1) - x = x + cast(self.transformer.position_embedding(position_ids), x.dtype) - - if default_net().plugin_config.remove_input_padding: - #B,T,D -> BxT,D - x = x.view([-1, self.config.hidden_size]) - hidden_states = x - input_lengths = input_lengths // self.downsample_factor - for encoder_layer in self.transformer.layers: - hidden_states = encoder_layer(hidden_states, - input_lengths=input_lengths) - - x = hidden_states - x = self.transformer.ln_f(x) - x.mark_output('encoder_output', self._dtype) - return x - - def prepare_inputs(self, max_batch_size=16): - - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - min_feat_len, optimal_feat_len = 10, 1000 # 100ms, 10s - inlen_range = [ - min_feat_len, optimal_feat_len, self.max_audio_feature_seq_len - ] - inlen_range_after_downsample = [ - min_feat_len // self.downsample_factor, - optimal_feat_len // self.downsample_factor, - self.max_audio_feature_seq_len // self.downsample_factor - ] - if not default_net().plugin_config.remove_input_padding: - x = Tensor(name="input_features", - dtype=self._dtype, - shape=[-1, self.config.n_mels, -1], - dim_range=OrderedDict([ - ("batch_size", [bs_range]), - ("feature_dim", [self.config.n_mels]), - ("feature_len_range", [inlen_range]), - ])) - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('feature_len_downsample_range', - [inlen_range_after_downsample])]), - ) - else: - batch_seqlen_range = [ - 1, - (self.max_audio_feature_seq_len * max_batch_size + 1) // 2, - self.max_audio_feature_seq_len * max_batch_size, - ] - batch_seqlen_downsample_range = [ - 1, - (self.max_audio_feature_seq_len // self.downsample_factor * - max_batch_size + 1) // 2, - self.max_audio_feature_seq_len // self.downsample_factor * - max_batch_size, - ] - x = Tensor(name="input_features", - dtype=self._dtype, - shape=[-1, self.config.n_mels], - dim_range=OrderedDict([ - ("batch_seqlen_range", [batch_seqlen_range]), - ("feature_dim", [self.config.n_mels]), - ])) - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_seqlen_downsample_range', - [batch_seqlen_downsample_range])]), - ) - input_lengths = Tensor( - name="input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size", [bs_range])]), - ) - - return { - 'input_features': x, - 'input_lengths': input_lengths, - 'position_ids': position_ids - } - - def precompute_relative_attention_bias(self, build_config): - pass diff --git a/tensorrt_llm/models/falcon/__init__.py b/tensorrt_llm/models/falcon/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/falcon/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/falcon/config.py b/tensorrt_llm/models/falcon/config.py deleted file mode 100644 index 1ff2ff0391c0..000000000000 --- a/tensorrt_llm/models/falcon/config.py +++ /dev/null @@ -1,118 +0,0 @@ -# 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 ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class FalconConfig(PretrainedConfig): - - def __init__(self, - *, - bias: bool = False, - parallel_attention: bool = False, - num_ln_in_parallel_attn: int | None = None, - new_decoder_architecture: bool = False, - rotary_base: float = 10000.0, - **kwargs): - self.bias = bias - self.parallel_attention = parallel_attention - self.num_ln_in_parallel_attn = num_ln_in_parallel_attn - self.new_decoder_architecture = new_decoder_architecture - self.rotary_base = rotary_base - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in LLaMAConfig - output['bias'] = self.bias - output['parallel_attention'] = self.parallel_attention - output['new_decoder_architecture'] = self.new_decoder_architecture - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - # Falcon-7B config may not have num_kv_heads or n_head_kv. - # Although Falcon-180B uses GQA (num_kv_heads=8), its config - # has multi_query=True. - if getattr(hf_config, 'multi_query', False) and not getattr( - hf_config, 'new_decoder_architecture', False): - hf_config.num_kv_heads = 1 - - if hf_config.model_type == 'RefinedWeb': - # Case 1. Falcon-40B / Falcon-40B-instruct - # https://huggingface.co/tiiuae/falcon-40b/blob/main/config.json - hf_config.num_hidden_layers = hf_config.n_layer - hf_config.num_attention_heads = hf_config.n_head - hf_config.num_kv_heads = hf_config.n_head_kv - hf_config.new_decoder_architecture = True - elif hf_config.model_type == 'RefinedWebModel': - # Case 2. Falcon-7B / Falcon-7B-instruct - # https://huggingface.co/tiiuae/falcon-7b/blob/main/config.json - hf_config.num_hidden_layers = hf_config.n_layer - hf_config.num_attention_heads = hf_config.n_head - hf_config.num_kv_heads = 1 if hf_config.multi_query else hf_config.n_head - hf_config.new_decoder_architecture = False - elif hf_config.model_type != 'falcon': - raise ValueError("Shouldn't reach here.") - hf_config.model_type = 'falcon' - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture='FalconForCausalLM', - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - num_key_value_heads=hf_config.num_kv_heads, - hidden_size=hf_config.hidden_size, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type='alibi_with_scale' - if hf_config.alibi else 'rope_gpt_neox', - hidden_act='gelu', - bias=hf_config.bias, - parallel_attention=hf_config.parallel_attn, - num_ln_in_parallel_attn=getattr(hf_config, - 'num_ln_in_parallel_attn', - None), - new_decoder_architecture=hf_config.new_decoder_architecture, - max_position_embeddings=getattr(hf_config, - 'max_position_embeddings', - 2048), - rotary_base=get_hf_rope_theta(hf_config, 10000.0), - intermediate_size=getattr(hf_config, 'ffn_hidden_size', - None), - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/falcon/convert.py b/tensorrt_llm/models/falcon/convert.py deleted file mode 100644 index ef36e9e7383f..000000000000 --- a/tensorrt_llm/models/falcon/convert.py +++ /dev/null @@ -1,499 +0,0 @@ -import time -from typing import Dict, Optional - -import torch - -from ...quantization import QuantAlgo -from ..convert_utils import (get_weight, get_weight_and_bias, - iterate_shard_files, load_state_dict, - retrieved_layer_index_from_name) -from .config import FalconConfig - - -def split(weight: torch.Tensor, - tp_size: int, - rank: int = 0, - dim: int = 0) -> torch.Tensor: - if tp_size == 1: - return weight - elif weight.ndim == 1: - return torch.chunk(weight, tp_size)[rank].clone() - else: - return torch.chunk(weight, tp_size, dim=dim)[rank].clone() - - -def reorder_qkv_weight_or_bias(weight: torch.Tensor, - head_dim: int, - num_heads: int, - num_kv_heads: Optional[int] = None, - tp_size: int = 1, - is_bias: bool = False) -> torch.Tensor: - """ Reorder the qkv weight for TRT-LLM use. - - The shape of the fused QKV weights in HF is different from the shape that - TRT-LLM requires. In particular, the weight of HF consists of interleaved - q, k, v head weights, while that of TRT-LLM is contiguous. - HF : [q1, k1, v1, ..., qh, kh, vh] - TRT-LLM: [q1, ..., qh, k1, ..., kh, v1, vh] - where qi, vi, ki are weight vectors corresponding to attention head i. - It's similar to multi/grouped query attention cases. - - We reorder and split the weight of an attention layer to fit into TRT-LLM. - The reordered weight and bias will be - weight: (T, Qh * D + 2 * KVh * D, H) - bias : (T, Qh * D + 2 * KVh * D) - where T=tp_size, Qh=local_num_q_heads, KVh=local_num_kv_heads, D=head_dim, - H=hidden_dim. In the multi/grouped query attention, the number of K/V - attention heads are less than that of Q attention, so that K/V attention - heads may be shared across different ranks if necessary. - - For tensor parallelism, we use the first dimension to select the - corresponding weights. - """ - - # Query types and expected kv heads. - # - Conventional MHA: num_heads = num_kv_heads - # - Multi-Query Attention: num_kv_heads = 1 - # - Grouped-Query Attention: num_heads % num_kv_heads = 0 - num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads - assert num_heads % num_kv_heads == 0, \ - f'num_heads({num_heads}) must be divisible by ' \ - f'num_kv_heads({num_kv_heads})).' - - # The number of attention heads per group: N q head + 1 k head + 1 v head. - num_group_heads = num_heads // num_kv_heads + 2 - assert weight.shape[0] == num_kv_heads * num_group_heads * head_dim, \ - f'{weight.shape[0]} != {num_kv_heads} * {num_group_heads} * {head_dim}' - - qkv_in = num_heads * head_dim if not is_bias else 1 - - # Split Q/K/V weights - weight = weight.reshape(num_kv_heads, num_heads // num_kv_heads + 2, - head_dim, qkv_in) - q_w = weight[:, :-2, ...] # (nKV, num_heads // nKV, head_dim, qkv_in) - k_w = weight[:, -2:-1, ...] # (nKV, 1, head_dim, qkv_in) - v_w = weight[:, -1:, ...] # (nKV, 1, head_dim, qkv_in) - - if num_kv_heads < num_heads and num_kv_heads < tp_size: - # Duplicate K/V heads to make sure that each rank has at least one - # K/V heads. For instance, num_heads=8, num_kv_heads=2, tp_size=4, - # we will make the qkv weight as below. - # Orig: [q0 q1 q2 q3 k0 v0 q4 q5 q6 q7 k1 v0 v1] - # >>>> [[q0 q1 k0 v0], [q2 q3 k0 v0], [q4 q5 k1 v1], [q6 q7 k1 v1]] - assert tp_size % num_kv_heads == 0 - num_dups = tp_size // num_kv_heads - - # k_w and v_w have the same shape. - new_shape = (num_kv_heads, num_dups) + k_w.shape[2:] - k_w = torch.broadcast_to(k_w, size=new_shape) - v_w = torch.broadcast_to(v_w, size=new_shape) - - # Update the number of kv heads. - num_kv_heads = tp_size - - reordered = torch.concat( - [ - q_w.reshape(tp_size, num_heads // tp_size, head_dim, qkv_in), - k_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - v_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - ], - dim=1, - ) - - qkv_out = (num_heads + 2 * num_kv_heads) // tp_size * head_dim - return reordered.reshape((tp_size, qkv_out, -1)) - - -def split_qkv_weight(weight: torch.Tensor, - hidden_size: int, - num_heads: int, - tp_size: int, - rank: int, - is_bias: bool, - num_kv_heads: Optional[int] = None) -> torch.Tensor: - """ Splits the QKV matrix according to tensor parallelism """ - head_dim = hidden_size // num_heads - weight = reorder_qkv_weight_or_bias(weight, - head_dim=head_dim, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - tp_size=tp_size, - is_bias=is_bias) - - # Copy a sliced tensor to prevent memory leak. A sliced tensor shares the - # memory buffer of the original tensor. So, returning without copying makes - # the buffer of a loaded "qkv" be referenced, resulting GC can't release - # those weights until the whole process ends. - if not is_bias: - return weight[rank, ...].clone() - else: - return weight[rank, ...].ravel().clone() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def get_tllm_param( - param: torch.Tensor, - name: str, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if name.endswith('.weight') and use_weight_only: - v = param.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[name] = processed_torch_weights - results[name.replace('weight', - 'per_channel_scale')] = torch_weight_scales - else: - results[name] = param - - return results - - -def load_weights_from_hf_model(model, config: FalconConfig): - weights = {} - tik = time.time() - - model_params = dict(model.named_parameters()) - dtype = getattr(torch, config.dtype) - mapping = config.mapping - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - vocab_size = config.vocab_size - num_kv_heads = getattr(config, 'num_key_value_heads', num_attention_heads) - num_hidden_layers = config.num_hidden_layers - parallel_attention = config.parallel_attention - new_decoder_architecture = config.new_decoder_architecture - use_parallel_embedding = config.use_parallel_embedding - sharding_dim = config.embedding_sharding_dim - 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 - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f'transformer.h.{l}' - tllm_prex = f'transformer.layers.{l - layers_range[0]}' - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, f'{prefix}.self_attention.query_key_value', dtype) - qkv_w = split_qkv_weight(qkv_weight, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_kv_heads) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = split_qkv_weight(qkv_bias, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', qkv_b, - use_weight_only, - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, f'{prefix}.self_attention.dense', dtype) - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, f'{tllm_prex}.attention.dense', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_h_to_4h', dtype) - mlp_fc_w = split_matrix(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - if mlp_fc_bias is None: - mlp_fc_b = None - else: - mlp_fc_b = split_matrix(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', mlp_fc_b, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_4h_to_h', dtype) - mlp_proj_w = split_matrix(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - if new_decoder_architecture: - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.ln_attn', dtype) - if input_ln_weight is None: - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - if input_ln_bias is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - mlp_ln_weight, mlp_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.ln_mlp', dtype) - if mlp_ln_weight is not None: - weights[f'{tllm_prex}.mlp_layernorm.weight'] = mlp_ln_weight - if mlp_ln_bias is not None: - weights[f'{tllm_prex}.mlp_layernorm.bias'] = mlp_ln_bias - else: - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - if input_ln_bias is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - if not parallel_attention: - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.post_attention_layernorm', dtype) - if post_ln_weight is not None: - weights[ - f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - if post_ln_bias is not None: - weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype) - if mapping.is_first_pp_rank(): - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - if sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix( - embed_w, mapping.tp_size, mapping.tp_rank, sharding_dim) - - if mapping.is_last_pp_rank(): - lm_head = get_weight(model_params, 'lm_head', dtype) - if lm_head is None: - # No lm_head in the checkpoint, cloning word_embedding. - lm_head = embed_w.clone() - weights['lm_head.weight'] = split_matrix(lm_head, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f', - dtype) - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def load_weights_from_hf_by_shard(model_dir: str, config: FalconConfig): - weights = {} - tik = time.time() - - dtype = getattr(torch, config.dtype) - mapping = config.mapping - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - vocab_size = config.vocab_size - num_kv_heads = getattr(config, 'num_key_value_heads', num_attention_heads) - num_hidden_layers = config.num_hidden_layers - use_weight_only = config.quantization.quant_algo in [ - QuantAlgo.W8A16, QuantAlgo.W4A16 - ] - if config.quantization.quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif config.quantization.quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - - layers_range = mapping.pp_layers(num_hidden_layers) - for model_file in iterate_shard_files(model_dir, mapping.tp_rank): - state_dict = load_state_dict(model_file, dtype) - for name, param in state_dict.items(): - l = retrieved_layer_index_from_name(name) - if l is not None: - if l not in layers_range: - continue - prefix = f'transformer.layers.{l-layers_range[0]}' - - if 'self_attention.query_key_value' in name: - if name.endswith('weight'): - qkv_w = split_qkv_weight(param, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_param(qkv_w, - f'{prefix}.attention.qkv.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - qkv_b = split_qkv_weight(param, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_param(qkv_b, - f'{prefix}.attention.qkv.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'self_attention.dense' in name: - if name.endswith('weight'): - attn_dense_w = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_param(attn_dense_w, - f'{prefix}.attention.dense.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - weights.update( - get_tllm_param(param, - f'{prefix}.attention.dense.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'mlp.dense_h_to_4h' in name: - if name.endswith('weight'): - mlp_fc_w = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_param(mlp_fc_w, f'{prefix}.mlp.fc.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - mlp_fc_b = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_param(mlp_fc_b, f'{prefix}.mlp.fc.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'mlp.dense_4h_to_h' in name: - if name.endswith('weight'): - mlp_proj_w = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_param(mlp_proj_w, - f'{prefix}.mlp.proj.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - weights.update( - get_tllm_param(param, f'{prefix}.mlp.proj.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'ln_attn' in name or 'input_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.input_layernorm.weight'] = param - else: - weights[f'{prefix}.input_layernorm.bias'] = param - elif 'ln_mlp' in name: - if name.endswith('weight'): - weights[f'{prefix}.mlp_layernorm.weight'] = param - else: - weights[f'{prefix}.mlp_layernorm.bias'] = param - elif 'post_attention_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.post_layernorm.weight'] = param - else: - weights[f'{prefix}.post_layernorm.bias'] = param - elif 'word_embeddings' in name: - if mapping.is_first_pp_rank(): - if not config.use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = param - else: - if config.embedding_sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights[ - 'transformer.vocab_embedding.weight'] = split_matrix( - param, mapping.tp_size, mapping.tp_rank, - config.embedding_sharding_dim) - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=0) - elif 'ln_f' in name: - if mapping.is_last_pp_rank(): - if name.endswith('weight'): - weights['transformer.ln_f.weight'] = param - else: - weights['transformer.ln_f.bias'] = param - del state_dict - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/falcon/model.py b/tensorrt_llm/models/falcon/model.py deleted file mode 100644 index 01e523fa66c9..000000000000 --- a/tensorrt_llm/models/falcon/model.py +++ /dev/null @@ -1,273 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import Tensor, allreduce, recv, send -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import FalconConfig -from .convert import load_weights_from_hf_by_shard, load_weights_from_hf_model - - -class FalconDecoderLayer(Module): - - def __init__(self, config: FalconConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - layernorm_epsilon = config.norm_epsilon - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - dtype=dtype) - - self.new_decoder_architecture = config.new_decoder_architecture - self.parallel_attn = config.parallel_attention - self.num_ln_in_parallel_attn = config.num_ln_in_parallel_attn - if self.num_ln_in_parallel_attn is None and self.new_decoder_architecture: - self.num_ln_in_parallel_attn = 2 - if self.is_parallel_attention: - # Not to apply allreduce inside the Attention/MLP layers. - # allreduce applies after those layer. - tp_group = None - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - bias=config.bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - quant_mode=config.quantization.quant_mode, - ) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - if self.new_decoder_architecture and self.num_ln_in_parallel_attn == 2: - # Layernorm before MLP. - self.mlp_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - dtype=dtype) - else: - self.mlp_layernorm = None - self.mlp = MLP( - hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=dtype, - bias=config.bias, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quantization.quant_mode, - ) - if self.is_parallel_attention: - self.post_layernorm = None - else: - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - @property - def is_parallel_attention(self): - return self.new_decoder_architecture or self.parallel_attn - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - if self.new_decoder_architecture and self.num_ln_in_parallel_attn == 2: - mlp_ln_output = self.mlp_layernorm(hidden_states) - hidden_states = self.input_layernorm(hidden_states) - input_ln_output = hidden_states - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - if not self.new_decoder_architecture: - if self.parallel_attn: - hidden_states = input_ln_output - else: - hidden_states = residual + attention_output - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - elif self.num_ln_in_parallel_attn == 2: - hidden_states = mlp_ln_output - - if (self.new_decoder_architecture and self.parallel_attn - and self.num_ln_in_parallel_attn == 1): - hidden_states = input_ln_output - - hidden_states = self.mlp(hidden_states) - - if self.is_parallel_attention: - hidden_states = hidden_states + attention_output - if self.config.mapping.tp_size > 1: - hidden_states = allreduce(hidden_states, - self.config.mapping.tp_group) - - hidden_states = residual + hidden_states - if use_cache: - return hidden_states, presents - return hidden_states - - -class FalconModel(Module): - - def __init__(self, config: FalconConfig): - super().__init__() - self.config = config - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(FalconDecoderLayer, config) - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - 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, - hidden_states=None): - if self.config.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids) - else: - hidden_states = recv(hidden_states, - self.config.mapping.prev_pp_rank()) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.config.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, - self.config.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class FalconForCausalLM(DecoderModelForCausalLM): - config_class = FalconConfig - - def __init__(self, config: FalconConfig): - self.check_config(config) - transformer = FalconModel(config) - - if config.mapping.is_last_pp_rank(): - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - lm_head = ColumnLinear(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) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('bias', True) - config.set_if_not_exist('new_decoder_architecture', False) - config.set_if_not_exist('parallel_attention', False) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a FalconForCausalLM object from give parameters - ''' - import transformers - - load_by_shard = kwargs.pop('load_by_shard', False) - # load_model_on_cpu is ignored here, since specify target device_map will fail when workers > 1. - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = FalconConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if use_preloading: - assert not load_by_shard - weights = load_weights_from_hf_model(hf_model, config) - elif load_by_shard: - weights = load_weights_from_hf_by_shard(hf_model_dir, config) - else: - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype='auto') - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model diff --git a/tensorrt_llm/models/gemma/__init__.py b/tensorrt_llm/models/gemma/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gemma/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/gemma/config.py b/tensorrt_llm/models/gemma/config.py deleted file mode 100644 index 5755bce17aee..000000000000 --- a/tensorrt_llm/models/gemma/config.py +++ /dev/null @@ -1,224 +0,0 @@ -# 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 json import loads -from pathlib import Path -from typing import TYPE_CHECKING, Optional, Union - -from tensorrt_llm._utils import get_hf_rope_theta -from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import infer_dtype -from tensorrt_llm.models.modeling_utils import (Gemma2ConfigGroup, - Gemma3ConfigGroup, - PretrainedConfig, QuantConfig) - -if TYPE_CHECKING: - from os import PathLike - - import transformers - - HfConfigOrDir = Union[str, PathLike, transformers.PretrainedConfig] - -GEMMA_ARCHITECTURE = "GemmaForCausalLM" -GEMMA2_ARCHITECTURE = "Gemma2ForCausalLM" -GEMMA3_ARCHITECTURE = "Gemma3ForCausalLM" - - -class GemmaConfig(PretrainedConfig): - - def __init__( - self, - *, - architecture: str, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - attn_bias: bool = False, - mlp_bias: bool = False, - position_embedding_type: PositionEmbeddingType = PositionEmbeddingType. - rope_gpt_neox, - query_pre_attn_scalar: Optional[int] = None, - final_logit_softcapping: Optional[float] = None, - attn_logit_softcapping: Optional[float] = None, - mapping: Optional[Union[Mapping, dict]] = None, - _sliding_window_pattern: int = None, - rope_local_base_freq: int = None, - sliding_window: int = None, - **kwargs, - ): - use_parallel_embedding = False - if mapping: - use_parallel_embedding = mapping.tp_size > 1 if isinstance( - mapping, Mapping) else mapping["tp_size"] > 1 - if use_parallel_embedding != kwargs.pop("use_parallel_embedding", None): - """ - We always pass `bool(mapping.tp_size > 1)` - the passed value is `False` by default, and ignored either way. - We can't just raise an exception here, because this will force the user to explicitly pass `LLM(use_parallel_embedding=True)`. - """ - logger.debug( - f"Using `use_parallel_embedding={use_parallel_embedding}` for Gemma" - ) - - super().__init__( - architecture=architecture, - use_parallel_embedding=use_parallel_embedding, - rotary_base=rotary_base, - attn_bias=attn_bias, - mlp_bias=mlp_bias, - position_embedding_type=position_embedding_type, - mapping=mapping, - **kwargs, - ) - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.attn_bias = attn_bias - self.mlp_bias = mlp_bias - - self.inter_layernorms = False - if self.is_gemma_2 or self.is_gemma_3: - self.inter_layernorms = True - assert query_pre_attn_scalar is not None, "Gemma2 / Gemma3 models must configure `query_pre_attn_scalar`" - self.query_pre_attn_scalar = query_pre_attn_scalar - self.final_logit_softcapping = final_logit_softcapping - if self.is_gemma_2: - self.attn_logit_softcapping = attn_logit_softcapping - if self.is_gemma_3: - self._sliding_window_pattern = _sliding_window_pattern - self.rope_local_base_freq = rope_local_base_freq - self.sliding_window = sliding_window - - GEMMA_ADDED_FIELDS = { - "rotary_base", "rotary_scaling", "attn_bias", "mlp_bias", - "inter_layernorms" - } - GEMMA2_ADDED_FIELDS = Gemma2ConfigGroup.keys() - GEMMA3_ADDED_FIELDS = Gemma3ConfigGroup.keys() - VERBATIM = { - "num_hidden_layers", "num_attention_heads", "hidden_size", - "intermediate_size", "vocab_size", "max_position_embeddings", - "hidden_act", "use_parallel_embedding" - } | GEMMA2_ADDED_FIELDS | GEMMA3_ADDED_FIELDS - - @property - def is_gemma_2(self) -> bool: - return self.architecture == GEMMA2_ARCHITECTURE - - def gemma2_config(self): - if self.is_gemma_2: - return self.get_config_group(Gemma2ConfigGroup) - return None - - @property - def is_gemma_3(self) -> bool: - return self.architecture == GEMMA3_ARCHITECTURE - - def gemma3_config(self): - if self.is_gemma_3: - return self.get_config_group(Gemma3ConfigGroup) - return None - - def to_dict(self): - """Serialize the fields added in GemmaConfig""" - - return { - **super().to_dict(), - **{ - f: getattr(self, f) - for f in self.GEMMA_ADDED_FIELDS - }, - **({ - f: getattr(self, f) - for f in self.GEMMA2_ADDED_FIELDS - } if self.is_gemma_2 else {}), - **({ - f: getattr(self, f) - for f in self.GEMMA3_ADDED_FIELDS - } if self.is_gemma_3 else {}) - } - - @staticmethod - def get_hf_config(config_dir: "Union[str, PathLike]"): - import transformers - model_type = loads( - (Path(config_dir) / "config.json").read_text())["model_type"] - HFConfigClass = { - "gemma": transformers.GemmaConfig, - "gemma2": transformers.Gemma2Config, - "gemma3_text": transformers.Gemma3TextConfig, - }[model_type] - hf_config = HFConfigClass.from_pretrained(config_dir) - return hf_config - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: "HfConfigOrDir", - dtype: str = "auto", - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs, - ) -> "GemmaConfig": - import transformers - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = cls.get_hf_config(hf_config_or_dir) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - assert isinstance(quant_config, QuantConfig) or quant_config is None - assert isinstance(mapping, Mapping) or mapping is None - - # transformers 5.x moved Gemma 3 RoPE fields into a nested - # rope_parameters dict. Resolve the global and sliding-window thetas - # from either the legacy top-level attributes or the nested form. - rope_parameters = getattr(hf_config, "rope_parameters", None) or {} - full_rope = rope_parameters.get("full_attention") or {} - sliding_rope = rope_parameters.get("sliding_attention") or {} - - rotary_base = getattr(hf_config, "rope_theta", None) - if rotary_base is None: - rotary_base = full_rope.get("rope_theta") - if rotary_base is None: - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - - verbatim_fields = { - k: v - for k, v in hf_config.to_dict().items() if k in cls.VERBATIM - } - # Gemma 3's local RoPE base is no longer a top-level attribute in - # transformers 5.x; pull it from rope_parameters.sliding_attention. - if ("rope_local_base_freq" in cls.VERBATIM - and verbatim_fields.get("rope_local_base_freq") is None): - local_freq = getattr(hf_config, "rope_local_base_freq", None) - if local_freq is None: - local_freq = sliding_rope.get("rope_theta") - if local_freq is not None: - verbatim_fields["rope_local_base_freq"] = local_freq - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - head_size=hf_config.head_dim, - norm_epsilon=hf_config.rms_norm_eps, - num_key_value_heads=getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads), - rotary_base=rotary_base, - rotary_scaling=getattr(hf_config, "rotary_scaling", None), - quantization=quant_config, - mapping=mapping, - **verbatim_fields, - **kwargs, - ) diff --git a/tensorrt_llm/models/gemma/convert.py b/tensorrt_llm/models/gemma/convert.py deleted file mode 100644 index 5aa40d9fd9c5..000000000000 --- a/tensorrt_llm/models/gemma/convert.py +++ /dev/null @@ -1,926 +0,0 @@ -# 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 json -import math -import os -import re -import time -from argparse import Namespace -from dataclasses import dataclass -from pathlib import Path -from types import SimpleNamespace -from typing import TYPE_CHECKING, Dict, Optional, Union - -import h5py -import numpy -import numpy as np -import sentencepiece as sp -import torch - -import tensorrt_llm -import tensorrt_llm.models.modeling_utils as tllm_utils -from tensorrt_llm._utils import (np_bfloat16, numpy_to_torch, - str_dtype_to_torch, torch_to_numpy) -from tensorrt_llm.logger import logger -from tensorrt_llm.models.convert_utils import load_calib_dataset -from tensorrt_llm.models.gemma.config import GemmaConfig -from tensorrt_llm.models.gemma.smoothquant import (capture_activation_range, - convert_hf_model, - create_model_from_config, - smooth_model) -from tensorrt_llm.models.gemma.utils import params as gemma_params -from tensorrt_llm.models.gemma.weight import (dummy_weights_awq, - load_from_fp8_gemma, - quantize_fp8_weights) -from tensorrt_llm.quantization.mode import QuantAlgo - -Weights = Dict[str, torch.Tensor] -Flattened = Dict[str, np.ndarray] - -if TYPE_CHECKING: - import transformers - - -class JAXParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> Weights: - checkpoint_path = checkpoint_path.absolute() - os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" - return gemma_params.nest_params( - gemma_params.param_remapper( - gemma_params.load_params(checkpoint_path))) - - def embedding_weights(self, ckpt_params) -> np.ndarray: - return ckpt_params["transformer"]["embedder"]["input_embedding"] - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - from tensorrt_llm.models.gemma.utils.transformer import \ - TransformerConfig - - return TransformerConfig.from_params(ckpt_params, num_embed=num_embed) - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix, name = name.split(".", maxsplit=1) - assert prefix == "transformer" - sub_patterns = ( - (r"embedder.input_embedding", r"vocab_embedding.weight"), - (r"layer_(\d+).pre_attention_norm.scale", - r"layers.\1.input_layernorm.weight"), - (r"layer_(\d+).attn.q_einsum.w", r"layers.\1.attention.qkv.weight"), - (r"layer_(\d+).attn.kv_einsum.w", - None), # drop as kv will be concatenated with q - (r"layer_(\d+).attn.qkv_einsum.w", - r"layers.\1.attention.qkv.weight"), - (r"layer_(\d+).attn.attn_vec_einsum.w", - r"layers.\1.attention.dense.weight"), - (r"layer_(\d+).mlp.gating_einsum", r"layers.\1.mlp.fc.weight"), - (r"layer_(\d+).mlp.linear", r"layers.\1.mlp.proj.weight"), - (r"layer_(\d+).pre_ffw_norm.scale", - r"layers.\1.post_layernorm.weight"), - (r"final_norm.scale", r"ln_f.weight"), - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {prefix}.{name}") - - def flatten_params(self, params) -> Flattened: - import flax.traverse_util - - new_params = flax.traverse_util.flatten_dict(params, sep=".") - # if the dtype is bfloat16, cast to float32 - for k in new_params: - if new_params[k].dtype != np.float32 and new_params[ - k].dtype != np.float16: - new_params[k] = new_params[k].astype(np.float32) - return new_params - - -class KerasParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> h5py.File: - checkpoint_path = checkpoint_path.absolute() - config_file = "config.json" - weights_file = json.load(open(checkpoint_path / config_file))["weights"] - h5_path = checkpoint_path / weights_file - return h5py.File(h5_path, "r") - - def embedding_weights(self, ckpt_params) -> np.ndarray: - return np.array(ckpt_params["layers/reversible_embedding/vars/0"]) - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - checkpoint_path = checkpoint_path.absolute() - config_file = "config.json" - config_old = json.load(open(checkpoint_path / config_file))["config"] - config_new = SimpleNamespace() - config_new.num_layers = config_old["num_layers"] - config_new.num_embed = config_old["vocabulary_size"] - config_new.embed_dim = config_old["hidden_dim"] - config_new.hidden_dim = config_old["intermediate_dim"] // 2 - config_new.num_heads = config_old["num_query_heads"] - config_new.head_dim = config_old["head_dim"] - config_new.num_kv_heads = config_old["num_key_value_heads"] - return config_new - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix = "transformer" - name = name.replace("/gemma_decoder_block/", "/gemma_decoder_block_0/") - sub_patterns = ( - (r"layers/reversible_embedding/vars/0", r"vocab_embedding.weight"), - (r"layers/gemma_decoder_block_(\d+)/pre_attention_norm/vars/0", - r"layers.\1.input_layernorm.weight"), - (r"layers/gemma_decoder_block_(\d+)/attention/query_dense/vars/0", - r"layers.\1.attention.qkv.weight"), - (r"layers/gemma_decoder_block_(\d+)/attention/key_dense/vars/0", - None), # drop as k will be concatenated with q - (r"layers/gemma_decoder_block_(\d+)/attention/value_dense/vars/0", - None), # drop as v will be concatenated with q - (r"layers/gemma_decoder_block_(\d+)/attention/output_dense/vars/0", - r"layers.\1.attention.dense.weight"), - (r"layers/gemma_decoder_block_(\d+)/gating_ffw/vars/0", - r"layers.\1.mlp.fc.weight"), - (r"layers/gemma_decoder_block_(\d+)/gating_ffw_2/vars/0", - None), # merged with above - (r"layers/gemma_decoder_block_(\d+)/ffw_linear/vars/0", - r"layers.\1.mlp.proj.weight"), - (r"layers/gemma_decoder_block_(\d+)/pre_ffw_norm/vars/0", - r"layers.\1.post_layernorm.weight"), - (r"layers/rms_normalization/vars/0", r"ln_f.weight"), - (r"optimizer/vars/(\d+)", None), # Not used - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {prefix}.{name}") - - def flatten_params(self, params) -> Flattened: - f_params = {} - - def walk(name, obj): - if isinstance(obj, h5py.Dataset): - if obj.dtype == "|V2": - # bfloat16 case - f_params[name] = torch_to_numpy( - numpy_to_torch(np.array(obj).astype(np_bfloat16)).to( - torch.float32)) - else: - f_params[name] = np.array(obj) - - params.visititems(walk) - return f_params - - -class TorchParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> Weights: - ckpt_path = list(checkpoint_path.glob('*.ckpt'))[0] - model_params = torch.load(ckpt_path)['model_state_dict'] - model_params.pop('freqs_cis') - return model_params - - def embedding_weights(self, ckpt_params) -> np.ndarray: - return ckpt_params["embedder.weight"] - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - checkpoint_path = checkpoint_path.absolute() - config_file = "config.json" - with open(checkpoint_path / config_file, 'r') as f: - json_str = f.read() - json_str = json_str.replace("'", "\"") - json_str = json_str.replace(",\n}", "\n}") - config_old = json.loads(json_str) - config_new = SimpleNamespace() - config_new.num_layers = config_old["num_hidden_layers"] - config_new.num_embed = config_old["vocab_size"] - config_new.embed_dim = config_old["hidden_size"] - config_new.hidden_dim = config_old["intermediate_size"] - config_new.num_heads = config_old["num_attention_heads"] - config_new.head_dim = config_old["head_dim"] - config_new.num_kv_heads = config_old["num_key_value_heads"] - return config_new - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix = "transformer" - sub_patterns = ( - (r"embedder.weight", r"vocab_embedding.weight"), - (r"model.layers.(\d+).input_layernorm.weight", - r"layers.\1.input_layernorm.weight"), - (r"model.layers.(\d+).self_attn.qkv_proj.weight", - r"layers.\1.attention.qkv.weight"), - (r"model.layers.(\d+).self_attn.o_proj.weight", - r"layers.\1.attention.dense.weight"), - (r"model.layers.(\d+).mlp.gate_proj.weight", - r"layers.\1.mlp.fc.weight"), - (r"model.layers.(\d+).mlp.up_proj.weight", - None), # merged with above - (r"model.layers.(\d+).mlp.down_proj.weight", - r"layers.\1.mlp.proj.weight"), - (r"model.layers.(\d+).post_attention_layernorm.weight", - r"layers.\1.post_layernorm.weight"), - (r"model.norm.weight", r"ln_f.weight"), - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {name}") - - def flatten_params(self, params) -> Flattened: - f_params = {} - for k, v in params.items(): - if v.dtype == torch.bfloat16: - v = v.float() - f_params[k] = torch_to_numpy(v) - return f_params - - -class HfParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> Weights: - """`AutoModelForCausalLM.from_pretrained` will parse the correct gemma, whether Gemma or Gemma2 or future versions.""" - from transformers import AutoModelForCausalLM - hf_model = AutoModelForCausalLM.from_pretrained( - checkpoint_path, - device_map="cpu" if load_model_on_cpu else "auto", - dtype='auto', - trust_remote_code=True, - ) - model_params = dict(hf_model.named_parameters()) - return model_params - - def embedding_weights(self, ckpt_params) -> np.ndarray: - raise RuntimeError( - "This method shouldn't be called - `GemmaConfig.from_hugging_face` takes care of this." - ) - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - raise RuntimeError( - "This method shouldn't be called - `GemmaConfig.from_hugging_face` takes care of this." - ) - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix = "transformer" - sub_patterns = ( - (r"model.embed_tokens.weight", r"vocab_embedding.weight"), - (r"model.layers.(\d+).input_layernorm.weight", - r"layers.\1.input_layernorm.weight"), - (r"model.layers.(\d+).self_attn.q_proj.weight", - r"layers.\1.attention.qkv.weight"), - (r"model.layers.(\d+).self_attn.k_proj.weight", - None), # merged with above - (r"model.layers.(\d+).self_attn.v_proj.weight", - None), # merged with above - (r"model.layers.(\d+).self_attn.o_proj.weight", - r"layers.\1.attention.dense.weight"), - (r"model.layers.(\d+).self_attn.q_norm.weight", - r"layers.\1.attention.q_layernorm.weight"), - (r"model.layers.(\d+).self_attn.k_norm.weight", - r"layers.\1.attention.k_layernorm.weight"), - (r"model.layers.(\d+).mlp.gate_proj.weight", - r"layers.\1.mlp.fc.weight"), - (r"model.layers.(\d+).mlp.up_proj.weight", - None), # merged with above - (r"model.layers.(\d+).mlp.down_proj.weight", - r"layers.\1.mlp.proj.weight"), - (r"model.layers.(\d+).post_attention_layernorm.weight", - r"layers.\1.post_layernorm.weight"), - (r"model.layers.(\d+).pre_feedforward_layernorm.weight", - r"layers.\1.pre_feedforward_layernorm.weight"), - (r"model.layers.(\d+).post_feedforward_layernorm.weight", - r"layers.\1.post_feedforward_layernorm.weight"), - (r"model.norm.weight", r"ln_f.weight"), - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {prefix}.{name}") - - def flatten_params(self, params: Weights) -> Dict[str, numpy.ndarray]: - f_params = {} - for k, v in params.items(): - if v.dtype == torch.bfloat16: - v = v.float() - f_params[k] = torch_to_numpy(v) - return f_params - - -def split(v: np.ndarray, tp_size: int, idx: int, dim: int = 0) -> np.ndarray: - if tp_size == 1: - return v - return np.split(v, tp_size, axis=dim)[idx] - - -def split_matrix_tp(v: np.ndarray, tensor_parallel: int, rank: int, - dim: int) -> np.ndarray: - return split(v, tensor_parallel, rank, dim=dim) - - -def add_trt_llm_weight(weights: Flattened, - name: str, - param: np.ndarray, - dtype: Optional[np.dtype] = None): - assert name not in weights, f"{name} is already added." - param = numpy_to_torch(param) - if dtype is not None: - assert isinstance(dtype, - str), f"dtype must be str, but get type {type(dtype)}" - param = param.to(str_dtype_to_torch(dtype)) - weights[name] = param.contiguous() - - -def quantize(param: np.ndarray, - quant_mode: tensorrt_llm.quantization.QuantMode): - if quant_mode.is_int8_weight_only(): - quant_dtype = torch.int8 - elif quant_mode.is_int4_weight_only(): - quant_dtype = torch.quint4x2 - else: - raise ValueError(f"Invalid configuration got quant_mode={quant_mode}") - - param = numpy_to_torch(param) - param = param.t().contiguous() - - # previously this fn was available in torch.ops.fastertransformer namespace - ( - quantized_weights, - scales, - ) = torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - param, quant_dtype) - - if scales.dtype == torch.bfloat16: - scales = scales.to(torch.float32).numpy().astype("bfloat16") - else: - scales = scales.numpy() - return quantized_weights.numpy(), scales - - -Parsers = Union[JAXParser, KerasParser, TorchParser, HfParser] - - -def load_gemma_weights( - *, - trt_llm_config: tllm_utils.PretrainedConfig, - parameters_or_model_dir: "Union[Weights, Path]", - ckpt_parser: Parsers, - load_model_on_cpu: bool = True, -): - print("Loading weights...") - tik = time.time() - - tp_rank = trt_llm_config.mapping.tp_rank - tp_size = trt_llm_config.mapping.tp_size - hidden_size = trt_llm_config.hidden_size - head_dim = trt_llm_config.head_size - - weights = {} - - if isinstance(parameters_or_model_dir, Path): - logger.debug(f"Loading directory {str(parameters_or_model_dir)}...") - model_params = ckpt_parser.load_parameters( - parameters_or_model_dir, - load_model_on_cpu=load_model_on_cpu, - ) - else: - model_params = parameters_or_model_dir - - model_params = ckpt_parser.flatten_params(model_params) - - for name, param in model_params.items(): - logger.debug(f"Converting weight {name}...") - trt_llm_name = ckpt_parser.rename_to_trt_llm(name) - if trt_llm_name is None: # omit as used with other params - continue - - if "attn.q_einsum" in name: - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - assert gqa_mode - - # initial shape: (num_q_heads, hidden_size, head_dim) - q_param = param.transpose(1, 0, 2) - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=1) - - # initial shape: (2, num_kv_heads, hidden_size, head_dim) - kv_name = name.replace("q_einsum", "kv_einsum") - kv_param = model_params[kv_name] - kv_param = kv_param.reshape( - trt_llm_config.num_key_value_heads * 2, - hidden_size, - head_dim, - ).transpose(1, 0, 2) - - # -> (hidden_size, num_q_heads / tp_size + 2, head_dim) - qkv_param = np.concatenate([q_param, kv_param], axis=1) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1) - qkv_param = qkv_param.transpose(1, 0) - - # If int8 kv enabled, weight-only quantization will be done later. - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "self_attn.qkv_proj" in name: - q_param, k_param, v_param = np.split(param, [ - trt_llm_config.num_attention_heads * trt_llm_config.head_size, - trt_llm_config.num_attention_heads * trt_llm_config.head_size + - trt_llm_config.num_key_value_heads * trt_llm_config.head_size - ], - axis=0) - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=0) - if not gqa_mode: - k_param = split_matrix_tp(k_param, tp_size, tp_rank, dim=0) - v_param = split_matrix_tp(v_param, tp_size, tp_rank, dim=0) - - qkv_param = np.concatenate([q_param, k_param, v_param], axis=0) - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_per_group_scaling(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "attn.qkv_einsum" in name: - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - assert not gqa_mode - # initial shape: [3, num_heads, hidden_size, head_dim] -> [3, num_heads, head_dim, hidden_size] - qkv_param = param.transpose(0, 1, 3, 2) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1, - qkv_param.shape[3]) - qkv_param = split_matrix_tp(qkv_param, tp_size, tp_rank, dim=1) - qkv_param = qkv_param.reshape(-1, qkv_param.shape[2]) - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() \ - and not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "attention/query_dense" in name: - # Keras specific KQV convert - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - if gqa_mode: - - # initial shape: (num_q_heads, hidden_size, head_dim) - q_param = param.transpose(1, 0, 2) - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=1) - - # initial shape: (2, num_kv_heads, hidden_size, head_dim) - k_name = name.replace("query", "key") - k_param = model_params[k_name] - v_name = name.replace("query", "value") - v_param = model_params[v_name] - kv_param = np.stack((k_param, v_param), axis=0) - - kv_param = kv_param.reshape( - trt_llm_config.num_key_value_heads * 2, - hidden_size, - head_dim, - ).transpose(1, 0, 2) - - # -> (hidden_size, num_q_heads / tp_size + 2, head_dim) - qkv_param = np.concatenate([q_param, kv_param], axis=1) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1) - qkv_param = qkv_param.transpose(1, 0) - - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, - qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - else: - q_param = param - k_name = name.replace("query", "key") - k_param = model_params[k_name] - v_name = name.replace("query", "value") - v_param = model_params[v_name] - # initial shape: [3, num_heads, hidden_size, head_dim] -> [3, num_heads, head_dim, hidden_size] - qkv_param = np.stack((q_param, k_param, v_param), axis=0) - qkv_param = qkv_param.transpose(0, 1, 3, 2) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1, - qkv_param.shape[3]) - qkv_param = split_matrix_tp(qkv_param, tp_size, tp_rank, dim=1) - qkv_param = qkv_param.reshape(-1, qkv_param.shape[2]) - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, - qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "q_proj" in name: - mqa_mode = trt_llm_config.num_key_value_heads == 1 - - if mqa_mode: - # initial shape: (num_heads * head_dim, hidden_size) - q_param = param - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=0) - - k_name = name.replace("q_proj", "k_proj") - k_param = model_params[k_name] - - v_name = name.replace("q_proj", "v_proj") - v_param = model_params[v_name] - else: - # initial shape: (num_heads * head_dim, hidden_size) - q_param = param - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=0) - - k_name = name.replace("q_proj", "k_proj") - k_param = model_params[k_name] - k_param = split_matrix_tp(k_param, tp_size, tp_rank, dim=0) - - v_name = name.replace("q_proj", "v_proj") - v_param = model_params[v_name] - v_param = split_matrix_tp(v_param, tp_size, tp_rank, dim=0) - - qkv_param = np.concatenate([q_param, k_param, v_param], axis=0) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1) - - # If int8 kv enabled, weight-only quantization will be done later. - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "attention.dense.weight" in trt_llm_name: - # initial shape: (num_heads, head_dim, hidden_size) - if len(param.shape) == 3: - param = param.reshape(-1, param.shape[2]) - param = param.transpose( - 1, 0) # (hidden_size, num_heads * head_dum) - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_int8_kv_cache(): - param_quantized, param_scales = quantize( - param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - elif "mlp.fc.weight" in trt_llm_name: - if isinstance(ckpt_parser, KerasParser): - # initial shape: (hidden_size, intermediate_size) - fc_param, gate_param = param, model_params[name.replace( - "gating_ffw", "gating_ffw_2")] - elif isinstance(ckpt_parser, TorchParser): - # initial shape: (intermediate_size, hidden_size) - fc_param, gate_param = param, model_params[name.replace( - "mlp.gate_proj", "mlp.up_proj")] - fc_param = fc_param.transpose(1, 0) - gate_param = gate_param.transpose(1, 0) - elif isinstance(ckpt_parser, HfParser): - # initial shape: (intermediate_size, hidden_size) - fc_param, gate_param = param, model_params[name.replace( - "mlp.gate_proj", "mlp.up_proj")] - fc_param = fc_param.transpose(1, 0) - gate_param = gate_param.transpose(1, 0) - else: - # initial shape: (2, hidden_size, intermediate_size) - fc_param, gate_param = param[0], param[1] - fc_param = fc_param.transpose(1, 0) - fc_param = split_matrix_tp(fc_param, tp_size, tp_rank, dim=0) - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - fc_param_quantized, fc_param_scales = quantize( - fc_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, fc_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - fc_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, fc_param, - trt_llm_config.dtype) - - gate_param = gate_param.transpose(1, 0) - gate_param = split_matrix_tp(gate_param, tp_size, tp_rank, dim=0) - trt_llm_name = trt_llm_name.replace("mlp.fc.weight", - "mlp.gate.weight") - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - gate_param_quantized, gate_param_scales = quantize( - gate_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, gate_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - gate_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, gate_param, - trt_llm_config.dtype) - elif "mlp.proj.weight" in trt_llm_name: - if not isinstance(ckpt_parser, TorchParser) and not isinstance( - ckpt_parser, HfParser): - # initial shape: (intermediate_size, hidden_size) - param = param.transpose(1, 0) - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - param_quantized, param_scales = quantize( - param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - elif "embedder.input_embedding" in name or "reversible_embedding" in name or "embedder.weight" in name \ - or "embed_tokens.weight" in name: - # TODO: safetensor doesn't allow to save a shared tensor. - # Currently, we clone the weight but to save the disk, it - # would be better to skip saving lm_head weights and - # handle it at the loading phase. - lm_head = split_matrix_tp(param, tp_size, tp_rank, dim=0) - add_trt_llm_weight(weights, "lm_head.weight", np.copy(lm_head), - trt_llm_config.dtype) - - if trt_llm_config.use_parallel_embedding: - assert trt_llm_config.vocab_size % tp_size == 0 - param = split_matrix_tp( - param, - tp_size, - tp_rank, - dim=trt_llm_config.embedding_sharding_dim, - ) - if trt_llm_config.quant_mode.is_int8_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache() and trt_llm_config.quantization.exclude_modules is not None: - - # shape of embedding table: [V, K], V: vocab size, K: embedding dim - - # quantize will do following work: - # 1. transpose the input from [V, K] to [K, V] - # 2. compute V scales across dimension K - # 3. quantize the input to int8_weight - # 4. transpose the int8_weight to [V, K], but there is a bug in 'quantize' that the claimed shape is [K, V] - param_quantized, param_scales = quantize( - param, trt_llm_config.quant_mode) - - # Reshape the [K, V] to [V, K] to match the real data layout. - param_quantized = np.ascontiguousarray( - param_quantized.reshape( - [param_quantized.shape[1], param_quantized.shape[0]])) - - add_trt_llm_weight(weights, trt_llm_name, param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_token_scale"), - param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - elif any(keyword in name for keyword in ( - "pre_attention_norm.scale", - "pre_ffw_norm.scale", - "final_norm.scale", - "pre_attention_norm/vars/0", - "pre_ffw_norm/vars/0", - "rms_normalization/vars/0", - "input_layernorm", - "post_attention_layernorm", - "pre_feedforward_layernorm", - "post_feedforward_layernorm", - "model.norm.weight", - "q_norm.weight", - "k_norm.weight", - )): - param = param + 1.0 # upcasted to float32 in case of bfloat16 - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - else: - raise RuntimeError(f"Unhandled {name} module weights") - del model_params - - print( - f"Weights loaded. Total time: {time.strftime('%H:%M:%S', time.gmtime(time.time() - tik))}" - ) - return weights - - -@dataclass(frozen=True, kw_only=True) -class QuantizeModifiers: - """ - Bag of additional conversion parameters, sourced from argparse or from defaults. - We may use this class to easily create wrappers of `convert_gemma` - (Otherwise, we would need to repeat these params many times or spread `**kwargs` in an untypesafe manner) - """ - - use_weight_only_with_precision: Optional[str] = None - per_channel: bool = False - per_token: bool = False - calib_dataset: str = "ccdv/cnn_dailymail" - tokenizer_dir: Optional[str] = None - enable_fp8: bool = False - fp8_kv_cache: bool = False - quant_ckpt_path: Optional[str] = None - - @classmethod - def from_args(cls, args: Namespace) -> "QuantizeModifiers": - return cls(**{ - k: v - for k, v in vars(args).items() if k in cls.__annotations__ - }) - - -def non_modelopt_quantize_if_needed( - weights: Weights, *, model_dir: Path, - quantize_modifiers: QuantizeModifiers, - trt_llm_config: Union[GemmaConfig, - tllm_utils.PretrainedConfig]) -> Weights: - tokenizer_dir = quantize_modifiers.tokenizer_dir or model_dir - - quant_cfg = trt_llm_config.quantization - - kv_cache_is_int8 = quant_cfg.kv_cache_quant_algo == QuantAlgo.INT8 - use_smooth_quant = quant_cfg._use_plugin_sq - - if use_smooth_quant or kv_cache_is_int8: - qkv_para = {} - smoother = {} - dataset = load_calib_dataset(quantize_modifiers.calib_dataset) - tokenizer = sp.SentencePieceProcessor(model_file=tokenizer_dir) - if "transformer.vocab_embedding.weight" in weights: - # To use the HF to do SmoothQuant, we need to scale the embedding. - weights["transformer.vocab_embedding.weight"] = torch.multiply( - weights["transformer.vocab_embedding.weight"].to(torch.float32), - math.sqrt(trt_llm_config.hidden_size), - ) - hf_model = create_model_from_config(trt_llm_config, weights) - act_range = capture_activation_range(hf_model, tokenizer, dataset) - if use_smooth_quant: - smooth_model(hf_model, act_range, quant_cfg.smoothquant_val, - qkv_para, smoother) - weights = convert_hf_model( - hf_model=hf_model, - mapping=trt_llm_config.mapping, - vocab_size=trt_llm_config.vocab_size or 32000, - dtype=trt_llm_config.dtype, - use_parallel_embedding=trt_llm_config.use_parallel_embedding, - sharding_dim=0, - use_weight_only=quantize_modifiers.use_weight_only_with_precision - is not None, - plugin_weight_only_quant_type=torch.int8 - if quantize_modifiers.use_weight_only_with_precision == 'int8' else - torch.quint4x2, - use_smooth_quant=use_smooth_quant, - per_channel=quantize_modifiers.per_channel, - per_token=quantize_modifiers.per_token, - int8_kv_cache=kv_cache_is_int8, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother) - if "transformer.vocab_embedding.weight" in weights: - # Revert the scaling of embedding - weights["transformer.vocab_embedding.weight"] = torch.divide( - weights["transformer.vocab_embedding.weight"].to(torch.float32), - math.sqrt(trt_llm_config.hidden_size), - ).to(str_dtype_to_torch(trt_llm_config.dtype)) - if "lm_head.weight" in weights: # Remove lm_head before saving it in unified checkpoint. - del weights["lm_head.weight"] - return weights - if quantize_modifiers.use_weight_only_with_precision and quantize_modifiers.use_weight_only_with_precision.endswith( - "awq"): - weights = dummy_weights_awq( - weights=weights, - precision=quantize_modifiers.use_weight_only_with_precision, - trt_llm_config=trt_llm_config, - group_size=128) - elif quantize_modifiers.enable_fp8 or quantize_modifiers.fp8_kv_cache: - weight_scales = quantize_fp8_weights(weights, - trt_llm_config.num_hidden_layers, - trt_llm_config.mapping) - scales = load_from_fp8_gemma(quantize_modifiers.quant_ckpt_path, - trt_llm_config.num_hidden_layers, - trt_llm_config.mapping, - quantize_modifiers.fp8_kv_cache, - weight_scales) - weights.update(scales) - - return weights - - -def load_gemma_weights_from_hf_model( - hf_model: "transformers.AutoModelForCausalLM", - config: GemmaConfig) -> Weights: - return load_gemma_weights(parameters_or_model_dir=dict( - hf_model.named_parameters()), - trt_llm_config=config, - ckpt_parser=HfParser()) diff --git a/tensorrt_llm/models/gemma/model.py b/tensorrt_llm/models/gemma/model.py deleted file mode 100644 index c8c5307c5d4c..000000000000 --- a/tensorrt_llm/models/gemma/model.py +++ /dev/null @@ -1,396 +0,0 @@ -# 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 math -from typing import TYPE_CHECKING, Any, Dict, Optional - -from tensorrt_llm.models.gemma.convert import (QuantizeModifiers, Weights, - load_gemma_weights_from_hf_model, - non_modelopt_quantize_if_needed) -from tensorrt_llm.quantization.mode import (MODELOPT_FLOW_QUANTIZATIONS, - QuantAlgo) - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import (AllReduceFusionOp, AllReduceParams, LayerNormType, - Tensor, cast, recv, send) -from ...layers import (Attention, AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, GatedMLP, KeyValueCacheParams, - LoraParams, PositionEmbeddingType, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig, save_checkpoint, save_config) -from .config import GemmaConfig - -if TYPE_CHECKING: - - from .config import HfConfigOrDir - - -class GemmaDecoderLayer(Module): - - def __init__(self, config: GemmaConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - - q_scaling = 1.0 - max_attn_value = 0.0 - qk_layernorm = False - is_sliding = False - rotary_base = config.rotary_base - rotary_base_local = None - - gemma2_config = config.gemma2_config() - gemma3_config = config.gemma3_config() - if gemma2_config: - q_scaling = math.sqrt( - gemma2_config.query_pre_attn_scalar) / math.sqrt( - config.head_size) - max_attn_value = config.attn_logit_softcapping or 0.0 - elif gemma3_config: - qk_layernorm = True - q_scaling = math.sqrt( - gemma3_config.query_pre_attn_scalar) / math.sqrt( - config.head_size) - is_sliding = bool( - (layer_idx + 1) % gemma3_config._sliding_window_pattern) - rotary_base_local = config.rope_local_base_freq - - self.attention = Attention( - local_layer_idx=self.local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - attention_head_size=config.head_size, - qk_layernorm=qk_layernorm, - layernorm_type=LayerNormType.RmsNorm, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=rotary_base, - rotary_embedding_base_local=rotary_base_local, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - q_scaling=q_scaling, - max_attn_value=max_attn_value, - is_local=is_sliding, - ) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - if self.config.inter_layernorms: - self.pre_feedforward_layernorm = RmsNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - self.post_feedforward_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) - - def forward(self, - hidden_states: Tensor, - attention_mask: Optional[Tensor] = None, - use_cache: bool = False, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - lora_layer_params: Optional[LoraParams] = None, - next_layer_input_layernorm_args=None): - # assert not ( - # default_net().plugin_config.reduce_fusion and self.has_residual_mlp - # ), "Custom all reduce and residual mlp can't be enabled at the same time." - if default_net( - ).plugin_config.reduce_fusion and self.local_layer_idx > 0: - hidden_states, residual = hidden_states #FIXME:AN need to check if appropriate residual value is hidden state is pulled out. - else: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - norm_before_bmm1=True, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_PREPOST_NORM - if default_net().plugin_config.reduce_fusion else - AllReduceFusionOp.NONE, - residual=residual, - norm_weight=self.pre_feedforward_layernorm.weight.value - if self.config.inter_layernorms else None, - norm_pre_residual_weight=self.post_layernorm.weight.value, - eps=self.pre_feedforward_layernorm.eps - if self.config.inter_layernorms else 1e-06, - ), - ) - - if use_cache: - attention_output, presents = attention_output - - if default_net().plugin_config.reduce_fusion: - hidden_states, residual = attention_output - else: - if self.config.inter_layernorms: - attention_output = self.post_layernorm(attention_output) - hidden_states = residual + attention_output - residual = hidden_states - if self.config.inter_layernorms: - hidden_states = self.pre_feedforward_layernorm(hidden_states) - else: - hidden_states = self.post_layernorm(hidden_states) - - if next_layer_input_layernorm_args is not None: - hidden_states = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_PREPOST_NORM - if default_net().plugin_config.reduce_fusion else - AllReduceFusionOp.NONE, - residual=residual, - norm_weight=next_layer_input_layernorm_args[0], - norm_pre_residual_weight=self.post_feedforward_layernorm. - weight.value, - eps=next_layer_input_layernorm_args[1])) - else: - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - if self.config.inter_layernorms: - hidden_states = self.post_feedforward_layernorm(hidden_states) - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GemmaModel(Module): - - def __init__(self, config: GemmaConfig) -> None: - super().__init__() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(GemmaDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - self.hidden_size = config.hidden_size - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - hidden_states = cast(hidden_states * math.sqrt(self.hidden_size), - hidden_states.dtype) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - hidden_states = self.layers.forward( - 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 - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GemmaForCausalLM(DecoderModelForCausalLM): - config_class = GemmaConfig - - def __init__(self, config: GemmaConfig): - transformer = GemmaModel(config) - - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(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) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - - super().__init__(config, transformer, lm_head) - - @staticmethod - def _load_gemma_weights_from_hf(hf_model_dir: "HfConfigOrDir", - trt_llm_config: GemmaConfig, *, - load_model_on_cpu: bool) -> Weights: - """`AutoModelForCausalLM.from_pretrained` will parse the correct gemma, whether Gemma or Gemma2 or future versions.""" - import transformers - hf_gemma = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map="cpu" if load_model_on_cpu else "auto", - dtype='auto', - ) - weights = load_gemma_weights_from_hf_model(hf_gemma, trt_llm_config) - del hf_gemma - return weights - - @classmethod - def from_hugging_face(cls, - hf_model_dir: "HfConfigOrDir", - dtype='float16', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - load_model_on_cpu: bool = True, - **kwargs): - config = GemmaConfig.from_hugging_face(hf_config_or_dir=hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - model = GemmaForCausalLM(config) - weights = cls._load_gemma_weights_from_hf( - hf_model_dir, config, load_model_on_cpu=load_model_on_cpu) - model.load(weights) - return model - - NATIVE_QUANT_FLOW = { - QuantAlgo.W8A16, QuantAlgo.W4A16, - QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN, - QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN, - QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN, - QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - } - - @classmethod - def assert_valid_quant_algo(cls, quant_algo: Optional[QuantAlgo]): - allowed_quant_values = { - None - } | cls.NATIVE_QUANT_FLOW | MODELOPT_FLOW_QUANTIZATIONS - assert quant_algo in allowed_quant_values, f"{quant_algo} isn't in the allowed `QuantAlgo` values for this model: {allowed_quant_values}" - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'float16', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - gemma_config_kwargs: Dict[str, Any] = None, - **quantize_kwargs: Dict[str, Any], - ): - config = GemmaConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **(gemma_config_kwargs or {})) - - quant_algo = config.quantization.quant_algo - if quant_algo is None and config.quantization.kv_cache_quant_algo is None: - raise ValueError( - "There is no point in calling `quantize()` if both `quant_algo` and `kv_cache_quant_algo` are `None`" - ) - elif quant_algo in MODELOPT_FLOW_QUANTIZATIONS: - super().quantize(hf_model_dir, - output_dir, - dtype=config.dtype, - mapping=config.mapping, - quant_config=config.quantization, - **quantize_kwargs) - elif quant_algo in cls.NATIVE_QUANT_FLOW: - save_config(config, output_dir=output_dir, log=True) - for config in config.for_each_rank(): - hf_weights = cls._load_gemma_weights_from_hf( - hf_model_dir, config) - ranked_weights = non_modelopt_quantize_if_needed( - hf_weights, - model_dir=hf_model_dir, - quantize_modifiers=QuantizeModifiers(), - trt_llm_config=config) - save_checkpoint( - output_dir=output_dir, - weights=ranked_weights, - rank=config.mapping.rank, - ) - del hf_weights - else: - cls.assert_valid_quant_algo(quant_algo) - - def use_lora(self, lora_config: LoraConfig) -> None: - return use_lora( - self, lora_config) # Use the default trtllm->hf module mapping diff --git a/tensorrt_llm/models/gemma/smoothquant.py b/tensorrt_llm/models/gemma/smoothquant.py deleted file mode 100644 index 63311a8eb11d..000000000000 --- a/tensorrt_llm/models/gemma/smoothquant.py +++ /dev/null @@ -1,848 +0,0 @@ -# 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 copy -import functools -import math -import time -from collections import defaultdict -from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from tqdm import tqdm -from transformers import LlamaConfig, LlamaForCausalLM -from transformers.models.llama.modeling_llama import (LlamaAttention, - LlamaDecoderLayer, - LlamaRotaryEmbedding, - apply_rotary_pos_emb, - repeat_kv) -from transformers.pytorch_utils import Conv1D - -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (dup_kv_weight, generate_int8, - get_weight, smooth_gemm, - smooth_gemm_fc1_gate, split, - split_matrix_tp) - -if TYPE_CHECKING: - from transformers import AutoModelForCausalLM, Cache - - # transformers included ⬆️ `Cache` in https://github.com/huggingface/transformers/commit/633215ba58fe5114d8c8d32e415a04600e010701 - transformers 4.33, which is installed in the tests, is before this. - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=1, - seq_len=512): - model.cuda().eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - # tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip( - 1e-8, None).max(dim=1)[0].float() - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - # input_ids = tokenizer(line, - # return_tensors="pt", - # max_length=seq_len, - # padding=True, - # truncation=True).input_ids.to(device) - inputs = tokenizer.EncodeAsIds(line[0]) - inputs = np.array([[tokenizer.bos_id()] + inputs], dtype=np.int32) - input_ids = torch.tensor(inputs, dtype=torch.int32).to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -@torch.no_grad() -def smooth_model(model, scales, alpha: Optional[float], qkv_para, - smoother_dict): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, LlamaDecoderLayer): - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - smoother_dict[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - smoother_dict[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.split(q, q.shape[-1] // tp_size, dim=-1) - k_split = torch.split(k, q.shape[-1] // tp_size, dim=-1) - v_split = torch.split(v, q.shape[-1] // tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), dim=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.split(original_weights, - original_weights.shape[-1] // - tensor_parallel, - dim=-1)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=-1)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.split( - vals["scale_w_quant_orig"], - vals["scale_w_quant_orig"].shape[-1] // tensor_parallel, - dim=-1)[rank] - - results[prefix + - 'per_channel_scale'] = cur_per_channel_value.reshape(col_shape) - else: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.split(original_weights, - original_weights.shape[-1] // - tensor_parallel, - dim=-1)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = torch.split(smoother_value, - smoother_value.shape[-1] // - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def split_qkv_tp(qkv, n_head, n_kv_heads, head_size, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - kv_head_size = n_kv_heads * head_size - q, k, v = torch.split(qkv, [n_head * head_size, kv_head_size, kv_head_size], - dim=0) - q = split(q, tensor_parallel, rank, dim=0) - k = split(k, tensor_parallel, rank, dim=0) - v = split(v, tensor_parallel, rank, dim=0) - return torch.concatenate([q, k, v], dim=0).contiguous() - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}weight'] = processed_torch_weights - results[f'{prefix}per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}bias'] = bias - - return results - - -class LlamaAttentionExtend(LlamaAttention): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.head_dim = self.config.head_size - self.q_proj = nn.Linear(self.config.hidden_size, - self.config.num_attention_heads * self.head_dim, - bias=False) - self.k_proj = nn.Linear(self.config.hidden_size, - self.config.num_key_value_heads * self.head_dim, - bias=False) - self.v_proj = nn.Linear(self.config.hidden_size, - self.config.num_key_value_heads * self.head_dim, - bias=False) - self.o_proj = nn.Linear(self.config.num_attention_heads * self.head_dim, - self.config.hidden_size, - bias=False) - self.config.head_dim = self.head_dim - self.rotary_emb = LlamaRotaryEmbedding(config=self.config) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.IntTensor] = None, - past_key_value: "Optional[Cache]" = None, - output_attentions: bool = False, - use_cache: bool = False, - cache_position: Optional[torch.IntTensor] = None, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], - Optional[Tuple[torch.Tensor]]]: - bsz, q_len, _ = hidden_states.size() - - if self.config.pretraining_tp > 1: - key_value_slicing = (self.config.num_key_value_heads * - self.head_dim) // self.config.pretraining_tp - query_slices = self.q_proj.weight.split( - (self.config.num_attention_heads * self.head_dim) // - self.config.pretraining_tp, - dim=0) - key_slices = self.k_proj.weight.split(key_value_slicing, dim=0) - value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) - - query_states = [ - F.linear(hidden_states, query_slices[i]) - for i in range(self.config.pretraining_tp) - ] - query_states = torch.cat(query_states, dim=-1) - - key_states = [ - F.linear(hidden_states, key_slices[i]) - for i in range(self.config.pretraining_tp) - ] - key_states = torch.cat(key_states, dim=-1) - - value_states = [ - F.linear(hidden_states, value_slices[i]) - for i in range(self.config.pretraining_tp) - ] - value_states = torch.cat(value_states, dim=-1) - - else: - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, - self.config.num_attention_heads, - self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, - self.config.num_key_value_heads, - self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, - self.config.num_key_value_heads, - self.head_dim).transpose(1, 2) - - past_key_value = getattr(self, "past_key_value", past_key_value) - cos, sin = self.rotary_emb(value_states, position_ids) - query_states, key_states = apply_rotary_pos_emb(query_states, - key_states, cos, sin) - - if past_key_value is not None: - # sin and cos are specific to RoPE models; position_ids needed for the static cache - cache_kwargs = { - "sin": sin, - "cos": cos, - "cache_position": cache_position - } - key_states, value_states = past_key_value.update( - key_states, value_states, self.layer_idx, cache_kwargs) - - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) - - attn_weights = torch.matmul(query_states, key_states.transpose( - 2, 3)) / math.sqrt(self.head_dim) - - if attention_mask is not None: # no matter the length, we just slice it - if cache_position is not None: - causal_mask = attention_mask[:, :, cache_position, :key_states. - shape[-2]] - attn_weights = attn_weights + causal_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, - dim=-1, - dtype=torch.float32).to( - query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, - p=self.attention_dropout, - training=self.training) - attn_output = torch.matmul(attn_weights, value_states) - - if attn_output.size() != (bsz, self.config.num_attention_heads, q_len, - self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.config.num_attention_heads, q_len, self.head_dim)}, but is" - f" {attn_output.size()}") - - attn_output = attn_output.transpose(1, 2).contiguous() - - # Here is what we extend. - attn_output = attn_output.reshape( - bsz, q_len, self.config.num_attention_heads * self.head_dim) - - if self.config.pretraining_tp > 1: - attn_output = attn_output.split(self.hidden_size // - self.config.pretraining_tp, - dim=2) - o_proj_slices = self.o_proj.weight.split(self.hidden_size // - self.config.pretraining_tp, - dim=1) - attn_output = sum([ - F.linear(attn_output[i], o_proj_slices[i]) - for i in range(self.config.pretraining_tp) - ]) - else: - attn_output = self.o_proj(attn_output) - - if not output_attentions: - attn_weights = None - - return attn_output, attn_weights - - -def create_model_from_config(trt_llm_config, weights): - model_config = LlamaConfig() - model_config.vocab_size = trt_llm_config.vocab_size - model_config.dtype = trt_llm_config.dtype - model_config.max_position_embeddings = trt_llm_config.max_position_embeddings - model_config.hidden_size = trt_llm_config.hidden_size - model_config.num_hidden_layers = trt_llm_config.num_hidden_layers - model_config.num_attention_heads = trt_llm_config.num_attention_heads - model_config.num_key_value_heads = trt_llm_config.num_key_value_heads - model_config.hidden_act = trt_llm_config.hidden_act - model_config.head_size = trt_llm_config.head_size - model_config.intermediate_size = trt_llm_config.intermediate_size - model = LlamaForCausalLM(model_config) - # Hack attention module since head_dim * num_heads > hidden_size for 7B. - for i in range(model_config.num_hidden_layers): - module = model.model.layers[i].self_attn - model.model.layers[i].self_attn = LlamaAttentionExtend( - module.config, module.layer_idx) - # Copy wegiht to LLAMA model. - replace_name_dict = { - 'attention.dense': 'self_attn.o_proj', - 'mlp.proj': 'mlp.down_proj', - 'mlp.gate': 'mlp.up_proj', - 'mlp.fc': 'mlp.gate_proj', - 'ln_f': 'norm', - 'post_layernorm': 'post_attention_layernorm', - 'vocab_embedding': 'embed_tokens', - } - for name in list(weights): - param = weights[name] - weights.pop(name) - new_name = name.replace('transformer', 'model') - for _name in replace_name_dict: - if _name in new_name: - new_name = new_name.replace(_name, replace_name_dict[_name]) - if 'attention.qkv' in name: - qw, kw, vw = torch.split(param, [ - model_config.num_attention_heads * model_config.head_size, - model_config.num_key_value_heads * model_config.head_size, - model_config.num_key_value_heads * model_config.head_size, - ], - dim=0) - weights[new_name.replace('attention.qkv', 'self_attn.q_proj')] = qw - weights[new_name.replace('attention.qkv', 'self_attn.k_proj')] = kw - weights[new_name.replace('attention.qkv', 'self_attn.v_proj')] = vw - else: - weights[new_name] = param - - if "lm_head.weight" not in weights: - weights["lm_head.weight"] = weights["model.embed_tokens.weight"].clone() - model.load_state_dict(weights) - return model - - -def convert_hf_model(*, hf_model: "AutoModelForCausalLM", mapping: Mapping, - vocab_size: int, dtype: str, use_parallel_embedding: bool, - sharding_dim: int, use_weight_only: bool, - plugin_weight_only_quant_type: torch.dtype, - use_smooth_quant: bool, per_channel: bool, per_token: bool, - int8_kv_cache: bool, - act_range: "defaultdict[Any, dict[str, None]]", - qkv_para: Dict, smoother: Dict): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - intermediate_size = hf_model.config.intermediate_size - head_size = hf_model.config.head_size - num_key_value_heads = hf_model.config.num_key_value_heads - mha_mode = (num_key_value_heads == num_attention_heads) - - num_hidden_layers = hf_model.config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - print("Processing layer", l) - prefix = f'model.layers.{l}.' - layer_idx = int(l) - layers_range[0] - tllm_prex = f'transformer.layers.{layer_idx}.' - - if use_smooth_quant: - qkv_weight = qkv_para[prefix + 'self_attn.qkv_proj'] - qkv_out_dim = qkv_weight.shape[1] - - if not mha_mode: - hidden_size = qkv_weight.shape[0] - local_dim = hidden_size - head_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(hidden_size, - local_dim + 2 * head_size) - else: - qkv_weight = qkv_weight.reshape(hidden_size, 3, - head_size * num_attention_heads) - - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.qkv_proj'), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - q_weight = get_weight(model_params, prefix + 'self_attn.q_proj', - dtype) - k_weight = get_weight(model_params, prefix + 'self_attn.k_proj', - dtype) - v_weight = get_weight(model_params, prefix + 'self_attn.v_proj', - dtype) - if not mha_mode: - if num_key_value_heads < tensor_parallel: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, num_key_value_heads, - tensor_parallel) - v_weight = dup_kv_weight(v_weight, num_key_value_heads, - tensor_parallel) - assert (k_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - assert (v_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - split_v = torch.concat((wq, wk, wv)) - - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, num_attention_heads, - num_key_value_heads, head_size, - tensor_parallel, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_y = torch.cat([ - act_range.get(prefix + 'self_attn.q_proj')["y"], - act_range.get(prefix + 'self_attn.k_proj')["y"], - act_range.get(prefix + 'self_attn.v_proj')["y"] - ], - dim=0) - int8_kv_scales = qkv_y.max() / 127. - kv_cache_weights = {} - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - weights.update(kv_cache_weights) - - # Attention dense. - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'self_attn.o_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'self_attn.o_proj')], - smoother_shape=[ - 1, head_size * num_attention_heads // tensor_parallel - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - attn_dense_weight = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_weight, - tllm_prex + 'attention.dense.', None, - use_weight_only, - plugin_weight_only_quant_type)) - # MLP hf up to trt gate - mlp_up_weight = get_weight(model_params, prefix + 'mlp.up_proj', dtype) - if use_smooth_quant: - mlp_up_weight = mlp_up_weight.t() - int8_weights = generate_int8(mlp_up_weight, - act_range.get(prefix + 'mlp.up_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_up_weight = split_matrix_tp(mlp_up_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_up_weight, tllm_prex + 'mlp.gate.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP trt Gate to mlp fc - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - if use_smooth_quant: - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8( - mlp_gate_weight, act_range.get(prefix + 'mlp.gate_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_gate_weight = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_gate_weight, tllm_prex + 'mlp.fc.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP down - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'mlp.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'mlp.down_proj'], - smoother_shape=[1, intermediate_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_proj_weight = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_weight, tllm_prex + 'mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - v = get_weight(model_params, 'model.embed_tokens', dtype) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/gemma/utils/__init__.py b/tensorrt_llm/models/gemma/utils/__init__.py deleted file mode 100644 index c506269a5304..000000000000 --- a/tensorrt_llm/models/gemma/utils/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# 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. -# ============================================================================ diff --git a/tensorrt_llm/models/gemma/utils/layers.py b/tensorrt_llm/models/gemma/utils/layers.py deleted file mode 100644 index 0c2f471298a8..000000000000 --- a/tensorrt_llm/models/gemma/utils/layers.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# 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. -# ============================================================================ -"""Base layers.""" - -import jax -import jax.numpy as jnp -from flax import linen as nn - - -class Einsum(nn.Module): - shape: tuple[int, ...] - - @nn.compact - def __call__(self, eqn: str, x: jax.Array) -> jax.Array: - w = self.param('w', nn.initializers.zeros_init(), self.shape) - return jnp.einsum(eqn, x, w) - - -class RMSNorm(nn.Module): - - @nn.compact - def __call__(self, x): - scale = self.param('scale', nn.initializers.zeros_init(), (x.shape[-1])) - var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) - normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06))) - normed_inputs = normed_inputs * (1 + scale) - return normed_inputs diff --git a/tensorrt_llm/models/gemma/utils/modules.py b/tensorrt_llm/models/gemma/utils/modules.py deleted file mode 100644 index 2ec10855e496..000000000000 --- a/tensorrt_llm/models/gemma/utils/modules.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# 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. -# ============================================================================ -"""Transformer sub-modules. -""" - -import jax -import jax.numpy as jnp -from flax import linen as nn - -from . import layers, positional_embeddings - -K_MASK = -2.3819763e38 # Set to a large negative number. -LayerCache = dict[str, jax.Array] - - -def init_layer_cache(cache_size: int, num_heads: int, head_dim: int, - batch_size: int) -> LayerCache: - return { - 'v': - jnp.zeros((batch_size, cache_size, num_heads, head_dim), - dtype=jnp.float32), - 'k': - jnp.zeros((batch_size, cache_size, num_heads, head_dim), - dtype=jnp.float32), - } - - -class Embedder(nn.Module): - """Embedder module.""" - - vocab_size: int - embed_dim: int - - def setup(self): - self.input_embedding_table = self.param( - 'input_embedding', - nn.initializers.zeros_init(), - (self.vocab_size, self.embed_dim), - ) - - def encode(self, x: jax.Array) -> jax.Array: - x = self.input_embedding_table[(x, )] - x *= jnp.sqrt(self.embed_dim).astype(x.dtype) - return x - - def decode(self, x: jax.Array) -> jax.Array: - return jnp.dot(x, self.input_embedding_table.T) - - -class Attention(nn.Module): - """Attention module.""" - - num_heads: int - num_kv_heads: int - features: int - head_dim: int - - @property - def use_qkv_einsum(self): - return self.num_kv_heads == self.num_heads - - def setup(self): - self.attn_vec_einsum = layers.Einsum(shape=(self.num_heads, - self.head_dim, - self.features), ) - - if self.use_qkv_einsum: - self.qkv_einsum = layers.Einsum(shape=(3, self.num_heads, - self.features, - self.head_dim), ) - else: - self.q_einsum = layers.Einsum(shape=(self.num_heads, self.features, - self.head_dim), ) - self.kv_einsum = layers.Einsum(shape=(2, self.num_kv_heads, - self.features, - self.head_dim), ) - - def __call__( - self, - x: jax.Array, - segment_pos: int, - cache: LayerCache, - attn_mask: jax.Array, - time_step: int, - ) -> tuple[LayerCache, jax.Array]: - - bsz = x.shape[0] - - if self.use_qkv_einsum: - query_proj, key_proj, value_proj = self.qkv_einsum( - 'BTD,SNDH->SBTNH', x) - else: - query_proj = self.q_einsum('BTD,NDH->BTNH', x) - key_proj, value_proj = self.kv_einsum('BSD,CKDH->CBSKH', x) - - query_proj = positional_embeddings.apply_rope( - query_proj, - segment_pos, - head_dim=self.head_dim, - ) - query_scaled = query_proj * self.head_dim**-0.5 - - key_proj = positional_embeddings.apply_rope( - key_proj, - segment_pos, - head_dim=self.head_dim, - ) - - # Cache is left aligned. - cache['v'] = (cache['v'].at[:bsz, [time_step], :, :].set(value_proj) - ) # values - cache['k'] = (cache['k'].at[:bsz, [time_step], :, :].set(key_proj) - ) # rotated_keys - - logits = jnp.einsum('BTNH,BSNH->BTNS', query_scaled, cache['k']) - logits = logits.astype(jnp.float32) - - padded_logits = jnp.where( - (jnp.expand_dims(attn_mask, -2) >= K_MASK * 0.5), logits, K_MASK) - probs = jax.nn.softmax(padded_logits, axis=-1).astype(cache['k'].dtype) - - encoded = jnp.einsum('BTNS,BSNH->BTNH', probs, cache['v']) - attn_output = self.attn_vec_einsum('BTNH,NHD->BTD', encoded) - - return cache, attn_output - - -class FeedForward(nn.Module): - """Feed forward module.""" - - features: int - hidden_dim: int - - @nn.compact - def __call__(self, x): - w_gating = self.param( - 'gating_einsum', - nn.initializers.zeros_init(), - ((2, self.features, self.hidden_dim)), - ) - ff_gate = jnp.dot(x, w_gating[0]) - gate_value = nn.gelu(ff_gate) - - ff1 = jnp.dot(x, w_gating[1]) - activations = gate_value * ff1 - - w_linear = self.param( - 'linear', - nn.initializers.zeros_init(), - (self.hidden_dim, self.features), - ) - outputs = jnp.dot(activations, w_linear) - - return outputs - - -class Block(nn.Module): - """Transformer block.""" - - num_heads: int - num_kv_heads: int - embed_dim: int - head_dim: int - hidden_dim: int - - def setup(self): - self.pre_attention_norm = layers.RMSNorm() - self.attn = Attention( - num_heads=self.num_heads, - features=self.embed_dim, - head_dim=self.head_dim, - num_kv_heads=self.num_kv_heads, - ) - self.pre_ffw_norm = layers.RMSNorm() - self.mlp = FeedForward(features=self.embed_dim, - hidden_dim=self.hidden_dim) - - def __call__( - self, - x: jax.Array, - segment_pos: int, - cache: LayerCache, - attn_mask: jax.Array, - time_step: int, - ): - inputs_normalized = self.pre_attention_norm(x) - cache, attn_output = self.attn(inputs_normalized, segment_pos, cache, - attn_mask, time_step) - attn_output += x - residual = attn_output - attn_output = self.pre_ffw_norm(attn_output) - outputs = self.mlp(attn_output) - outputs = residual + outputs - return cache, outputs diff --git a/tensorrt_llm/models/gemma/utils/params.py b/tensorrt_llm/models/gemma/utils/params.py deleted file mode 100644 index f0445f918e2e..000000000000 --- a/tensorrt_llm/models/gemma/utils/params.py +++ /dev/null @@ -1,72 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# 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. -# ============================================================================ -"""Utils for loading Gemma params. - -These utilities are just helpers for current development. They will not be -needed once Gemma switches to Orbax and changes checkpoint formats ahead of -open sourcing. -""" - -import functools -from typing import Any - -Params = dict[str, Any] - - -@functools.cache -def load_params(path: str) -> Params: - import orbax.checkpoint - """Loads parameters from a checkpoint path.""" - checkpointer = orbax.checkpoint.PyTreeCheckpointer() - params = checkpointer.restore(path) - return params - - -def param_remapper(orig_params: Params) -> Params: - """Remaps params to new module layout. - - This is needed here because the model definition does not have a separate - `mlp` module. For the real code release, we will just save the params in a - different format and this will not be needed. - - Args: - orig_params: original dict of parameters in Gemma format. - - Returns: - dict of params with different names. - """ - new_params = {} - for k, v in orig_params.items(): - if 'mlp/' in k: - layer_name, param = k.rsplit('/', maxsplit=1) - if layer_name not in new_params: - new_params[layer_name] = {} - if 'w' in v: - new_params[layer_name][param] = v['w'] - else: - new_params[k] = v - return new_params - - -def nest_params(params: Params) -> Params: - """Nests params as a dict of dicts rather than a flat dict.""" - nested_params = {} - for path, param in params.items(): - *path, leaf = path.split('/') - subdict = nested_params - for key in path: - subdict = subdict.setdefault(key, {}) - subdict[leaf] = param - return nested_params diff --git a/tensorrt_llm/models/gemma/utils/positional_embeddings.py b/tensorrt_llm/models/gemma/utils/positional_embeddings.py deleted file mode 100644 index 25707692d758..000000000000 --- a/tensorrt_llm/models/gemma/utils/positional_embeddings.py +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# 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. -# ============================================================================ -"""Utils for positional embeddings (including RoPE). -""" - -import jax -import jax.numpy as jnp - -_MAX_WAVELENGTH = 10_000 - - -def add_positional_embedding( - input_embedding: jax.Array, - position: int, - max_wavelength: int = _MAX_WAVELENGTH, -) -> jax.Array: - """Adds positional embeddings to input embeddings.""" - embed_dim = input_embedding.shape[-1] - num_timescales = embed_dim // 2 - log_timescale_increment = jnp.log(float(max_wavelength)) / jnp.maximum( - jnp.asarray(num_timescales, dtype=jnp.float32) - 1, 1) - inv_timescales = jnp.exp( - jnp.arange(num_timescales, dtype=jnp.float32) * - -log_timescale_increment) - scaled_time = position * inv_timescales - signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)]) - signal = jnp.pad(signal, [[0, jnp.mod(embed_dim, 2)]]) - position_embedding = signal.astype(jnp.float32) - - return input_embedding + position_embedding - - -def _rotary_embed( - inputs: jax.Array, # [B, 1, H, D] - position: jax.Array, # [B,] - head_dim: int, - max_wavelength: int = _MAX_WAVELENGTH, -) -> jax.Array: - """Helper for RoPE.""" - fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim - timescale = max_wavelength**fraction - timescale = timescale[jnp.newaxis, jnp.newaxis, jnp.newaxis, :] - - sinusoid_inp = position[:, jnp.newaxis, jnp.newaxis, - jnp.newaxis] / timescale - sin = jnp.sin(sinusoid_inp) - cos = jnp.cos(sinusoid_inp) - - first_half, second_half = jnp.split(inputs, 2, axis=-1) - first_part = first_half * cos - second_half * sin - second_part = second_half * cos + first_half * sin - - return jnp.concatenate([first_part, second_part], axis=-1) - - -def apply_rope( - inputs: jax.Array, - position: int, - head_dim: int, - max_wavelength: int = _MAX_WAVELENGTH, -) -> jax.Array: - """Applies RoPE.""" - batch_size, seq_length = inputs.shape[0:2] - - position = jnp.broadcast_to(position, [batch_size])[:, jnp.newaxis] - prefix_position = jnp.arange(seq_length, dtype=jnp.int32) - prefix_position = (position - jnp.flip(prefix_position)[jnp.newaxis, :] - ) # [B, seq_len] - prefix_position = jnp.where(prefix_position < 0, - jnp.zeros_like(prefix_position), - prefix_position).reshape((batch_size, )) - - output = _rotary_embed( - inputs, - position=prefix_position, - head_dim=head_dim, - max_wavelength=max_wavelength, - ) - - return output diff --git a/tensorrt_llm/models/gemma/utils/sampler.py b/tensorrt_llm/models/gemma/utils/sampler.py deleted file mode 100644 index 895a1557d475..000000000000 --- a/tensorrt_llm/models/gemma/utils/sampler.py +++ /dev/null @@ -1,191 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# 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. -# ============================================================================ -"""Sampler for Gemma transformer. - -An example of a sampling class for a Gemma model. -""" - -import chex -import jax -import jax.numpy as jnp -import sentencepiece as spm - -from . import modules -from . import params as params_lib -from . import transformer as transformer_lib - - -def _compute_attention_masks(time_step: jax.Array, seq_len: int, - input_mask: jax.Array) -> jax.Array: - """Computes causal attention mask.""" - bsz = input_mask.shape[0] - batch_time_step = jnp.full((bsz, 1), time_step, dtype=jnp.uint32) - causal_padding = jnp.greater(jnp.expand_dims(jnp.arange(seq_len), 0), - batch_time_step) - causal_padding = causal_padding * jnp.expand_dims(input_mask, axis=-1) - attention_mask = ( - causal_padding[:, jnp.newaxis, jnp.newaxis, :].astype(jnp.float32) * - modules.K_MASK) - attention_mask = jnp.squeeze(attention_mask, axis=1) - return attention_mask - - -@chex.dataclass -class _SamplingState: - - # Number of tokens in the prompt. - num_input_tokens: jnp.int32 # [B] - - # Fixed-size buffer for accumulating the output tokens. - token_buffer: jnp.ndarray # [B, L] - - # Model state for conditioning the model on autoregressively. - cache: dict[str, modules.LayerCache] - - -class Sampler: - """Sampler for Gemma transformer.""" - - def __init__( - self, - transformer_config: transformer_lib.TransformerConfig, - vocab: spm.SentencePieceProcessor, - params: params_lib.Params, - cache_size: int, - buffer_size: int, - max_decode_steps: int, - ): - self.transformer = transformer_lib.Transformer( - config=transformer_config) - self.vocab = vocab - self.params = params - self.cache_size = cache_size - self.buffer_size = buffer_size - self.max_decode_steps = max_decode_steps - self._compiled_sample_fn = jax.jit(self._sample_fn) - - def _sample_step(self, params, time_step, - sampler_state: _SamplingState) -> _SamplingState: - """Performs a single sampling step.""" - time_step = jnp.asarray(time_step, dtype=jnp.int32) - last_token = sampler_state.token_buffer[:, time_step] - input_mask = last_token != self.vocab.pad_id() - attention_mask = _compute_attention_masks( - time_step, self.cache_size, input_mask).astype(jnp.float32) - - logits, cache = self.transformer.apply( - {'params': params}, - last_token, - time_step, - sampler_state.cache, - attention_mask, - time_step, - ) - - next_token_candidate = jnp.argmax(logits, axis=-1) # [B, 1] - next_token_candidate = next_token_candidate[:, 0] # [B,] - - next_token_candidate = jnp.where( - time_step < sampler_state.num_input_tokens - 1, - sampler_state.token_buffer[:, time_step + 1], - next_token_candidate, - ) - - token_buffer = sampler_state.token_buffer.at[:, time_step + 1].set( - next_token_candidate) - - return _SamplingState( - num_input_tokens=sampler_state.num_input_tokens, - token_buffer=token_buffer, - cache=cache, - ) - - def init_cache(self, bsz) -> dict[str, modules.LayerCache]: - """Initializes the attention cache for each layer.""" - return { - f'layer_{i}': - modules.init_layer_cache( - self.cache_size, - self.transformer.config.num_heads, - self.transformer.config.head_dim, - bsz, - ) - for i in range(self.transformer.config.num_layers) - } - - def init_sample_state(self, - all_input_ids: list[jax.Array]) -> _SamplingState: - """Initializes the sampling state given input prompts.""" - bsz = len(all_input_ids) - num_input_tokens = [len(input_ids) for input_ids in all_input_ids] - - token_buffer = jnp.full( - ( - bsz, - self.buffer_size, - ), - self.vocab.pad_id(), - dtype=jnp.int32, - ) - for i, (input_ids, - num_tokens) in enumerate(zip(all_input_ids, num_input_tokens)): - token_buffer = token_buffer.at[i, :num_tokens].set(input_ids) - - return _SamplingState( - num_input_tokens=jnp.array(num_input_tokens, dtype=jnp.int32), - token_buffer=token_buffer, - cache=self.init_cache(bsz), - ) - - def tokenize(self, input_string: str) -> jax.Array: - """Tokenizes the input string.""" - input_ids = self.vocab.EncodeAsIds(input_string) - input_ids = jnp.array([self.vocab.bos_id()] + - jnp.array(input_ids).tolist(), - dtype=jnp.int32) - return input_ids - - def _sample_fn( - self, - params: params_lib.Params, - initial_sampling_state: _SamplingState, - ) -> _SamplingState: - - def sample_with_params(time_step: int, sampler_state: _SamplingState): - return self._sample_step(params, time_step, sampler_state) - - return jax.lax.fori_loop(0, self.max_decode_steps, sample_with_params, - initial_sampling_state) - - def __call__(self, input_strings: list[str] | str) -> list[str]: - """Samples a completion of the input string.""" - if isinstance(input_strings, str): - input_strings = [input_strings] - all_input_ids = [self.tokenize(x) for x in input_strings] - initial_sampling_state = self.init_sample_state(all_input_ids) - - sampling_state = self._compiled_sample_fn(self.params, - initial_sampling_state) - - out_tokens = [ - buffer[num_tokens:num_tokens + self.max_decode_steps] - for buffer, num_tokens in zip(sampling_state.token_buffer, - sampling_state.num_input_tokens) - ] - decoded_outputs = [ - self.vocab.DecodeIds(out_tokens.tolist()) - for out_tokens in out_tokens - ] - return decoded_outputs diff --git a/tensorrt_llm/models/gemma/utils/transformer.py b/tensorrt_llm/models/gemma/utils/transformer.py deleted file mode 100644 index de3a0da61b7b..000000000000 --- a/tensorrt_llm/models/gemma/utils/transformer.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Gemma transformer.""" - -import dataclasses - -import jax -import jax.numpy as jnp -from flax import linen as nn - -from . import layers, modules -from . import params as params_lib - -Cache = dict[str, modules.LayerCache] - - -@dataclasses.dataclass -class TransformerConfig: - """Configuration for the Gemma transformer.""" - - num_layers: int - num_embed: int - embed_dim: int - hidden_dim: int - num_heads: int - head_dim: int - num_kv_heads: int - - @classmethod - def from_params(cls, params: params_lib.Params, - num_embed: int) -> 'TransformerConfig': - """Creates a TransformerConfig from loaded parameters.""" - num_layers = (max([ - int(k.split('_')[1]) - for k in params['transformer'].keys() if 'layer_' in k - ]) + 1) - hidden_dim, embed_dim = ( - params['transformer']['layer_0']['mlp']['linear'].shape) - num_heads, head_dim, _ = (params['transformer']['layer_0']['attn'] - ['attn_vec_einsum']['w'].shape) - use_qkv_einsum = 'qkv_einsum' in params['transformer']['layer_0'][ - 'attn'] - if use_qkv_einsum: - num_kv_heads = num_heads - else: - num_kv_heads = params['transformer']['layer_0']['attn'][ - 'kv_einsum']['w'].shape[1] - return cls( - num_layers=num_layers, - num_embed=num_embed, - embed_dim=embed_dim, - hidden_dim=hidden_dim, - num_heads=num_heads, - head_dim=head_dim, - num_kv_heads=num_kv_heads, - ) - - -def init_cache(config: TransformerConfig, cache_size: int, - batch_size: int) -> Cache: - """Initializes a new Transformer cache.""" - return { - f'layer_{i}': - modules.init_layer_cache(cache_size, config.num_heads, config.head_dim, - batch_size) - for i in range(config.num_layers) - } - - -class Transformer(nn.Module): - """Gemma transformer.""" - - config: TransformerConfig - - def setup(self): - self.embedder = modules.Embedder( - vocab_size=self.config.num_embed, - embed_dim=self.config.embed_dim, - ) - self.blocks = [ - modules.Block( - name=f'layer_{i}', - num_heads=self.config.num_heads, - num_kv_heads=self.config.num_kv_heads, - embed_dim=self.config.embed_dim, - head_dim=self.config.head_dim, - hidden_dim=self.config.hidden_dim, - ) for i in range(self.config.num_layers) - ] - self.final_norm = layers.RMSNorm() - - def __call__( - self, - last_tokens: jax.Array, # [B,] - current_token_position: int, - cache: Cache, - attention_mask: jax.Array, # [B, 1, L] - time_step: int, - ) -> tuple[jax.Array, Cache]: - input_emb = self.embedder.encode(last_tokens) - x = jnp.expand_dims(input_emb, axis=1) # adding temporal dimension - - for i, block in enumerate(self.blocks): - layer_name = f'layer_{i}' - cache[layer_name], x = block( - x, - current_token_position, - cache[layer_name], - attention_mask, - time_step, - ) - - x = self.final_norm(x) - logits = self.embedder.decode(x) - - return logits, cache diff --git a/tensorrt_llm/models/gemma/weight.py b/tensorrt_llm/models/gemma/weight.py deleted file mode 100644 index 56b4a91e4048..000000000000 --- a/tensorrt_llm/models/gemma/weight.py +++ /dev/null @@ -1,181 +0,0 @@ -# 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 os -from typing import Dict, Optional - -import numpy as np -import torch - -from tensorrt_llm.models.modeling_utils import PretrainedConfig - -from ..._utils import numpy_to_torch -from ...logger import logger -from ...mapping import Mapping - -Weights = Dict[str, torch.Tensor] - - -def quantize_fp8_weights(weights: Weights, num_layers: int, - mapping: Mapping) -> Weights: - - def get_scaling_factor(weight): - amax = weight.max() - scale = 448.0 / amax - return scale - - layers_range = mapping.pp_layers(num_layers) - scaling_factors = {} - scaled_weights = {} - trt_llm_prefix = "transformer.layers" - for l in layers_range: - # attention.qkv.weight - for name in [ - "attention.qkv", "attention.dense", "mlp.fc", "mlp.gate", - "mlp.proj" - ]: - trt_llm_name = ".".join((trt_llm_prefix, str(l), name, "weight")) - scale_name = ".".join( - (trt_llm_prefix, str(l), name, "weights_scaling_factor")) - weight = weights[trt_llm_name].float() - dtype = weights[trt_llm_name].dtype - scale = get_scaling_factor(weight) - scaled_weights[trt_llm_name] = (weight * - scale).to(dtype).contiguous() - scaling_factors[scale_name] = numpy_to_torch( - np.asarray([1 / scale]).astype(np.float32)) - return scaling_factors - - -def load_from_fp8_gemma(quant_ckpt_path: Optional[str], num_layers: int, - mapping: Mapping, fp8_kv_cache: bool, - weight_scales: Weights): - """ - Get the fp8 scaling factors. - """ - fake_fp8_sf_dt = torch.float32 - - if quant_ckpt_path is not None and os.path.isfile(quant_ckpt_path): - fp8_gemma = np.load(quant_ckpt_path) - else: - fp8_gemma = None - logger.info( - f"There is not quantized checkpoint, use dummy fp8 scaling factors instead." - ) - weights = {} - - def get_fp8_gemma(name: str) -> np.ndarray: - if fp8_gemma is not None: - return fp8_gemma[name] - else: - return torch.tensor([1.0], dtype=fake_fp8_sf_dt).numpy() - - layers_range = mapping.pp_layers(num_layers) - for l in layers_range: - prefix = f'_np:layers:{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - - weights[f'{tllm_prex}.attention.qkv.activation_scaling_factor'] = max( - get_fp8_gemma( - f'{prefix}:attention:qkv:q:activation_scaling_factor'), - get_fp8_gemma( - f'{prefix}:attention:qkv:k:activation_scaling_factor'), - get_fp8_gemma( - f'{prefix}:attention:qkv:v:activation_scaling_factor')) - weights[f'{tllm_prex}.attention.qkv.weights_scaling_factor'] = max( - get_fp8_gemma(f'{prefix}:attention:qkv:q:weights_scaling_factor'), - get_fp8_gemma(f'{prefix}:attention:qkv:k:weights_scaling_factor'), - get_fp8_gemma(f'{prefix}:attention:qkv:v:weights_scaling_factor')) - weights[ - f'{tllm_prex}.attention.dense.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:attention:dense:activation_scaling_factor') - weights[ - f'{tllm_prex}.attention.dense.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:attention:dense:weights_scaling_factor') - - weights[ - f'{tllm_prex}.mlp.fc.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:fc:activation_scaling_factor') - weights[f'{tllm_prex}.mlp.fc.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:fc:weights_scaling_factor') - - weights[ - f'{tllm_prex}.mlp.gate.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:gate:activation_scaling_factor') - weights[f'{tllm_prex}.mlp.gate.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:gate:weights_scaling_factor') - - weights[ - f'{tllm_prex}.mlp.proj.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:proj:activation_scaling_factor') - weights[f'{tllm_prex}.mlp.proj.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:proj:weights_scaling_factor') - - if fp8_kv_cache: - # Not calibrating KV cache. - scaling_factor = 1.0 - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = torch.tensor( - [scaling_factor], dtype=fake_fp8_sf_dt).numpy() - if fp8_gemma is None: - weights.update(weight_scales) - - for key in weights: - if isinstance(weights[key], np.ndarray): - weights[key] = numpy_to_torch(weights[key]) - return weights - - -def dummy_weights_awq(weights: Weights, precision: str, - trt_llm_config: PretrainedConfig, group_size: int): - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - use_fp8_kv_cache = trt_llm_config.quant_mode.has_fp8_kv_cache() - use_int8_kv_cache = trt_llm_config.quant_mode.has_int8_kv_cache() - num_layers = trt_llm_config.num_hidden_layers - for name in list(weights): - if any([ - _name in name for _name in [ - 'mlp.proj.weight', 'mlp.gate.weight', 'mlp.fc.weight', - 'attention.qkv.weight', 'attention.dense.weight' - ] - ]): - print("Processing:", name) - weight = np.ascontiguousarray(weights[name].T) - in_dim, out_dim = weight.shape - scale = np.amax(weight) / 7 - weights_scaling_factor = np.ones([out_dim, in_dim // group_size - ]) * scale.astype(np.float32) - weight_smoothed = (weight.astype(np.float32) / scale).astype( - np.int8) - weight_smoothed[weight_smoothed < -8] = -8 - weight_smoothed[weight_smoothed > 7] = 7 - prequant_scaling_factor = np.ones([in_dim], dtype=weight.dtype) - weights[name] = packer( - torch.from_numpy(weight_smoothed)).T.contiguous().numpy() - weights[name.replace( - 'weight', 'prequant_scaling_factor')] = prequant_scaling_factor - weights[name.replace( - 'weight', - 'weights_scaling_factor')] = weights_scaling_factor.astype( - weight.dtype) - if precision == "w4a8_awq": - alpha = np.array([1], dtype=np.float32) - weights[name.replace('weight', 'alpha')] = alpha - if use_fp8_kv_cache or use_int8_kv_cache: - for l in range(num_layers): - t = np.array([1], dtype=np.float32) - weights[ - f"transformer.layers.{l}.attention.kv_cache_scaling_factor"] = t - - return weights diff --git a/tensorrt_llm/models/generation_mixin.py b/tensorrt_llm/models/generation_mixin.py deleted file mode 100644 index 5f590f1ee4bb..000000000000 --- a/tensorrt_llm/models/generation_mixin.py +++ /dev/null @@ -1,889 +0,0 @@ -# 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 math -from collections import OrderedDict -from typing import List, Optional - -import tensorrt as trt - -from ..functional import Tensor -from ..layers import MropeParams, SpecDecodingParams -from ..llmapi.kv_cache_type import KVCacheType -from ..mapping import Mapping -from ..plugin import current_all_reduce_helper - - -class GenerationMixin: - - @staticmethod - def has_ctx_gen_opt_profiles( - use_gpt_attention_plugin: bool = False, - use_gemm_plugin: bool = False, - use_mamba_conv1d_plugin: bool = False, - remove_input_padding: bool = False, - paged_state: bool = False, - kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS) -> bool: - res = False - - if not use_gpt_attention_plugin or not use_gemm_plugin: - use_in_flight_batching = False - # Refer to modelConfig.h: supportsInflightBatching(), this should be consistent its implementation. - # We skip check transformer or rnn arch for simplification. - if remove_input_padding and use_gpt_attention_plugin: - use_in_flight_batching = kv_cache_type in [ - KVCacheType.PAGED, KVCacheType.DISABLED - ] - elif remove_input_padding and use_mamba_conv1d_plugin: - use_in_flight_batching = paged_state == True - - res = not use_in_flight_batching - - return res - - @staticmethod - def default_range(max_range, offset=0, min_range=1, opt_offset=0): - result = [ - min_range, (max_range + min_range + opt_offset) // 2, max_range - ] - return [elem + offset for elem in result] - - @staticmethod - def split_num_tokens_range(max_num_tokens): - split_point = [64, 128, 256, 512, 1024] - num_tokens_ranges = [] - for i, p in enumerate(split_point): - if i == 0 and max_num_tokens <= p: - return [1, max_num_tokens, max_num_tokens] - elif max_num_tokens <= p: - num_tokens_ranges.append( - [split_point[i - 1], max_num_tokens, max_num_tokens]) - return num_tokens_ranges - elif i == 0 and max_num_tokens > p: - num_tokens_ranges = [[1, 64, 64]] - else: - num_tokens_ranges.append( - [split_point[i - 1], split_point[i], split_point[i]]) - num_tokens_ranges.append( - [split_point[-1], max_num_tokens, max_num_tokens]) - return num_tokens_ranges - - @staticmethod - def get_profiles_ranges( - *, - max_batch_size, - max_beam_width, - max_input_len, - max_num_tokens, - max_draft_len, - opt_batch_size, - opt_num_tokens, - enable_ctx_gen_opt_profiles, - multiple_profiles, - kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS): - default_range = GenerationMixin.default_range - if opt_batch_size: - bb_range_cxt = [1, opt_batch_size, max_batch_size] - bb_range_gen = [ - 1, opt_batch_size * max_beam_width, - max_batch_size * max_beam_width - ] - else: - bb_range_cxt = default_range(max_batch_size) - bb_range_gen = default_range(max_batch_size * max_beam_width) - tokens_per_engine_step = max_draft_len + 1 - tokens_per_engine_step_range = [ - 1, tokens_per_engine_step, tokens_per_engine_step - ] - bbd_range_ctx = [ - bb_range_cxt[i] * (tokens_per_engine_step if i != 0 else 1) - for i in range(len(bb_range_cxt)) - ] - bbd_range_gen = [ - bb_range_gen[i] * (tokens_per_engine_step if i != 0 else 1) - for i in range(len(bb_range_gen)) - ] - inlen_range_cxt = default_range(max_input_len) - inlen_range_gen = [1, 1, tokens_per_engine_step] - if enable_ctx_gen_opt_profiles: - num_profiles = 2 - bb_range = [bb_range_cxt, bb_range_gen] - bbd_range = [bbd_range_ctx, bbd_range_gen] - inlen_range = [inlen_range_cxt, inlen_range_gen] - position_ids_inlen_range = [inlen_range_cxt, [1, 1, 1]] - num_tokens_range_ctx = default_range(max_batch_size * max_input_len) - # Draft tokens cannot be combined with beam search - num_tokens_range_gen = default_range( - max_batch_size * max(tokens_per_engine_step, max_beam_width)) - num_tokens_range = [num_tokens_range_ctx, num_tokens_range_gen] - - # Only keep context range when kv cache is disabled. - if kv_cache_type == KVCacheType.DISABLED: - num_profiles = 1 - bb_range = [bb_range[0]] - bbd_range = [bbd_range[0]] - inlen_range = [inlen_range[0]] - position_ids_inlen_range = [position_ids_inlen_range[0]] - num_tokens_range_ctx = [num_tokens_range_ctx[0]] - # Draft tokens cannot be combined with beam search - num_tokens_range_gen = [num_tokens_range_gen[0]] - num_tokens_range = [num_tokens_range[0]] - - else: - if multiple_profiles: - num_tokens_range = GenerationMixin.split_num_tokens_range( - max_num_tokens) - else: - if opt_num_tokens is None: - opt_num_tokens = min(max_num_tokens, - max_batch_size * max_beam_width) - num_tokens_range = [[1, opt_num_tokens, max_num_tokens]] - num_profiles = len(num_tokens_range) - bb_range = [bb_range_gen] * num_profiles - bbd_range = [bbd_range_gen] * num_profiles - inlen_range = [[1, 1, max_input_len]] * num_profiles - position_ids_inlen_range = [[1, 1, max_input_len]] * num_profiles - tokens_per_engine_step_range = [tokens_per_engine_step_range - ] * num_profiles - - position_ids_num_tokens_range = num_tokens_range - # If max_draft_len != 0, the input_ids may include draft tokens. And the length of position_ids may be not the same as input_ids. - # In extreme cases, input_ids may contain (max_draft_token + 1) * N, and the actual position_ids value is only 1 * N. - # Therefore, we need to adjust the min value in the ranges of position_ids. - if max_draft_len != 0: - position_ids_num_tokens_range = list( - map( - lambda x: - [math.ceil(x[0] / (max_draft_len + 1)), x[1], x[2]], - num_tokens_range)) - - ranges = { - 'bb_range': bb_range, - 'bbd_range': bbd_range, - 'inlen_range': inlen_range, - 'position_ids_inlen_range': position_ids_inlen_range, - 'num_tokens_range': num_tokens_range, - 'tokens_per_engine_step_range': tokens_per_engine_step_range, - 'position_ids_num_tokens_range': position_ids_num_tokens_range, - } - return num_profiles, ranges - - def prepare_attention_inputs( - self, - *, - max_batch_size, - max_beam_width, - max_input_len, - max_seq_len, - num_kv_heads, - head_size, - num_layers, - kv_dtype, - kv_cache_type: KVCacheType, - num_profiles=1, - enable_ctx_gen_opt_profiles=False, - remove_input_padding=False, - use_gpt_attention_plugin=False, - tokens_per_block=32, - mapping=Mapping(), - streamingllm=False, - attn_layer_idx=None, - opt_batch_size=None, - num_kv_heads_per_layer: Optional[List[int]] = None): - - if attn_layer_idx is not None and num_kv_heads_per_layer is not None: - assert len(attn_layer_idx) == len(num_kv_heads_per_layer), ( - f"Expected len(attn_layer_idx) ({len(attn_layer_idx)})" - f" == len(num_kv_heads_per_layer) ({len(num_kv_heads_per_layer)})" - ) - default_range = GenerationMixin.default_range - - if opt_batch_size: - bb_range_cxt = [1, opt_batch_size, max_batch_size] - bb_range_gen = [ - 1, opt_batch_size * max_beam_width, - max_batch_size * max_beam_width - ] - else: - bb_range_cxt = default_range(max_batch_size) - bb_range_gen = default_range(max_batch_size * max_beam_width) - - _bs_range = default_range(max_batch_size) - _beam_width_range = default_range(max_beam_width) - _max_len_range = default_range(max_seq_len) - _mask_len_ctx = default_range(max_input_len) - _kv_cache_range_ctx = [0, 0, 0] - _kv_cache_range_gen = default_range(max_seq_len, -1) - if kv_cache_type == KVCacheType.DISABLED: - _kv_cache_range = default_range(max_seq_len) - else: - kv_max_seq_len = max_seq_len - if streamingllm: - # add the max bubble length - kv_max_seq_len += tokens_per_block - 1 - if max_beam_width > 1: - # support cyclic kv cache cases that use one more block - kv_max_seq_len += tokens_per_block - _kv_cache_range = default_range(kv_max_seq_len) - - if enable_ctx_gen_opt_profiles: - if kv_cache_type != KVCacheType.DISABLED: - assert num_profiles == 2 - bb_range = [bb_range_cxt, bb_range_gen] - mask_len_range = [_mask_len_ctx, _max_len_range] - if use_gpt_attention_plugin: - kv_cache_range = [_kv_cache_range, _kv_cache_range] - else: - kv_cache_range = [_kv_cache_range_ctx, _kv_cache_range_gen] - else: - assert num_profiles == 1 - bb_range = [bb_range_cxt] - mask_len_range = [_mask_len_ctx] - if use_gpt_attention_plugin: - kv_cache_range = [_kv_cache_range] - else: - kv_cache_range = [_kv_cache_range_ctx] - - else: - bb_range = [bb_range_gen] * num_profiles - mask_len_range = [_max_len_range] * num_profiles - kv_cache_range = [_kv_cache_range] * num_profiles - bs_range = [_bs_range] * num_profiles - beam_width_range = [_beam_width_range] * num_profiles - max_len_range = [_max_len_range] * num_profiles - - num_kv_heads = (num_kv_heads + mapping.tp_size - 1) // mapping.tp_size - if num_kv_heads_per_layer is not None: - num_kv_heads_per_layer = [ - (nheads + mapping.tp_size - 1) // mapping.tp_size - for nheads in num_kv_heads_per_layer - ] - - layers_range = mapping.pp_layers(num_layers) - if attn_layer_idx is None: - attn_layer_idx = [i for i in range(num_layers)] - # layer indices of attention layers local to the current pp rank - local_attn_layers = [i for i in layers_range if i in attn_layer_idx] - # number of attention layers local to previous pp ranks - num_attn_layers_lower_ranks = attn_layer_idx.index(local_attn_layers[0]) - num_attn_layers = len(local_attn_layers) - num_layers_prev_rank = layers_range[ - 0] // mapping.pp_rank if mapping.pp_rank != 0 else len(layers_range) - past_key_value = [] - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - host_kv_cache_pool_pointers = None - host_kv_cache_pool_mapping = None - if kv_cache_type == KVCacheType.DISABLED: - past_key_value = [None] * num_layers_prev_rank - else: - if kv_cache_type != KVCacheType.PAGED: - for layer_idx in layers_range: - if layer_idx not in local_attn_layers: - # not an attention layer ==> give it None pkv input - past_key_value.append(None) - continue - - attn_idx = local_attn_layers.index(layer_idx) - if num_kv_heads_per_layer is not None: - heads_dim_name = f"num_heads_{layer_idx}" - kv_heads = num_kv_heads_per_layer[ - num_attn_layers_lower_ranks + attn_idx] - else: - heads_dim_name = "num_heads" - kv_heads = num_kv_heads - - kv_dim_range = OrderedDict([ - ('batch_size_beam_width', bb_range), - ('kv', [2] * num_profiles), - (heads_dim_name, [kv_heads] * num_profiles), - ('past_key_len', kv_cache_range), - ('head_size', [head_size] * num_profiles), - ]) - - kv = Tensor(name=f'past_key_value_{layer_idx}', - dtype=kv_dtype, - shape=[-1, 2, kv_heads, -1, head_size], - dim_range=kv_dim_range) - past_key_value.append(kv) - else: - if enable_ctx_gen_opt_profiles: - max_blocks_per_seq_range = [ - [ - math.ceil(kv_cache_range[0][0] / tokens_per_block), - math.ceil(kv_cache_range[0][1] / tokens_per_block), - math.ceil(kv_cache_range[0][2] / tokens_per_block) - ], - [ - math.ceil(kv_cache_range[1][0] / tokens_per_block), - math.ceil(kv_cache_range[1][1] / tokens_per_block), - math.ceil(kv_cache_range[1][2] / tokens_per_block) - ] - ] - else: - max_blocks_per_seq_range = [[ - math.ceil(kv_cache_range[0][0] / tokens_per_block), - math.ceil(kv_cache_range[0][1] / tokens_per_block), - math.ceil(kv_cache_range[0][2] / tokens_per_block) - ]] * num_profiles - - NUM_KV_CACHE_POOLS = -1 # the number of unique variable window sizes, which is only known at runtime, affects the number of pools. - # dim_range=(min=1, opt=1 (this is the usual case - non vgqa, non vsliding_window), max=num_layers, - # TODO(nhaber): Benchmark if making NUM_KV_CACHE_POOLS dynamic has a significant performance hit? - - kv_pools_range = [[1, 1, len(local_attn_layers)]] * num_profiles - kv_cache_block_offsets = Tensor( - name=f'kv_cache_block_offsets', - dtype=trt.int32, - shape=[NUM_KV_CACHE_POOLS, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', kv_pools_range), - ('batch_size_beam_width', bb_range), - ('kv', [2] * num_profiles), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_block_offsets = Tensor( - name=f'host_kv_cache_block_offsets', - dtype=trt.int32, - shape=[NUM_KV_CACHE_POOLS, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', kv_pools_range), - ('batch_size_beam_width', bb_range), - ('kv', [2] * num_profiles), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_pool_pointers = Tensor( - name=f'host_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[NUM_KV_CACHE_POOLS, 2], - dim_range=OrderedDict([ - ('num_pools_layers', kv_pools_range), - ('num_pools_kv', [2] * num_profiles), - ])) - - host_kv_cache_pool_mapping = Tensor( - name=f'host_kv_cache_pool_mapping', - dtype=trt.int32, - shape=[num_attn_layers, - 2], # 2: (Index of pool, Index of layer within pool) - dim_range=OrderedDict([ - ('pools_mapping', [num_attn_layers] * num_profiles), - ('layer_cache_pool_locator', [2] * num_profiles) - ])) - - past_key_value = [None] * num_layers_prev_rank - - assert len(past_key_value) == num_layers_prev_rank - sequence_length = None - context_lengths = None - host_context_lengths = None - host_past_key_value_lengths = None - host_max_attention_window_sizes = None - host_sink_token_length = None - attention_mask = None - cache_indirection = None - host_request_types = None - runtime_perf_knobs = None - context_progress = None - - if use_gpt_attention_plugin: - if kv_cache_type != KVCacheType.DISABLED: - sequence_length = Tensor( - name='sequence_length', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range) - ]), - ) - - host_request_types = Tensor( - name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range)]), - ) - if kv_cache_type != KVCacheType.DISABLED: - host_past_key_value_lengths = Tensor( - name='host_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range) - ]), - ) - context_lengths = Tensor( - name='context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range)]), - ) - runtime_perf_knobs = Tensor(name='host_runtime_perf_knobs', - dtype=trt.int64, - shape=[16], - dim_range=OrderedDict([ - ('perf_knob_size', - [16] * num_profiles) - ])) - context_progress = Tensor(name='host_context_progress', - dtype=trt.int64, - shape=[1], - dim_range=OrderedDict([ - ('context_progress_size', - [1] * num_profiles) - ])) - else: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('mask_len', mask_len_range), - ]), - ) - - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = Tensor( - name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range)]), - ) - - if use_gpt_attention_plugin: - # TODO: change shape to [1] - host_max_attention_window_sizes = Tensor( - name=f'host_max_attention_window_sizes', - dtype=trt.int32, - shape=[num_attn_layers], - dim_range=OrderedDict([('num_layers', - [num_attn_layers] * num_profiles)])) - - host_sink_token_length = Tensor(name='host_sink_token_length', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([ - ('scalar', [1] * num_profiles) - ])) - - if kv_cache_type != KVCacheType.DISABLED: - cache_indirection = Tensor( - name='cache_indirection', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_cache', bs_range), - ('beam_width', beam_width_range), - ('max_seq_len', max_len_range), - ]), - ) - - return { - 'attention_mask': attention_mask, - 'sequence_length': sequence_length, - 'host_past_key_value_lengths': host_past_key_value_lengths, - 'host_max_attention_window_sizes': host_max_attention_window_sizes, - 'host_sink_token_length': host_sink_token_length, - 'past_key_value': past_key_value, - 'cache_indirection': cache_indirection, - 'kv_cache_block_offsets': kv_cache_block_offsets, - 'host_kv_cache_block_offsets': host_kv_cache_block_offsets, - 'host_kv_cache_pool_pointers': host_kv_cache_pool_pointers, - 'host_kv_cache_pool_mapping': host_kv_cache_pool_mapping, - 'context_lengths': context_lengths, - 'host_context_lengths': host_context_lengths, - 'host_request_types': host_request_types, - 'host_runtime_perf_knobs': runtime_perf_knobs, - 'host_context_progress': context_progress, - } - - def prepare_basic_inputs( - self, - *, - max_batch_size, - max_beam_width, - max_input_len, - max_seq_len, - max_num_tokens, - hidden_size, - num_kv_heads, - head_size, - num_layers, - kv_dtype, - kv_cache_type: KVCacheType, - remove_input_padding=False, - use_gpt_attention_plugin=False, - use_gemm_plugin=False, - tokens_per_block=32, - gather_context_logits=False, - dtype=None, - num_heads=None, - mapping=Mapping(), - opt_num_tokens=None, - prompt_embedding_table_size: int = 0, - position_encoding_2d=False, - use_lora_plugin: bool = False, - lora_target_modules: List[str] = None, - speculative_decoding_draft_tokens_external: bool = False, - spec_decoding_is_generation_length_variable: bool = False, - max_draft_len=0, - multiple_profiles: bool = False, - streamingllm: bool = False, - opt_batch_size=None, - pp_reduce_scatter: bool = False, - mrope_rotary_cos_sin_size: int = None, - ): - - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type) - - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_num_tokens=max_num_tokens, - max_draft_len=max_draft_len, - opt_batch_size=opt_batch_size, - opt_num_tokens=opt_num_tokens, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles, - kv_cache_type=kv_cache_type) - bb_range = ranges['bb_range'] - bbd_range = ranges['bbd_range'] - inlen_range = ranges['inlen_range'] - num_tokens_range = ranges['num_tokens_range'] - position_ids_inlen_range = ranges['position_ids_inlen_range'] - tokens_per_engine_step_range = ranges['tokens_per_engine_step_range'] - position_ids_num_tokens_range = ranges['position_ids_num_tokens_range'] - - input_ids = None - position_ids = None - hidden_states = None - if remove_input_padding: - if mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('num_tokens', num_tokens_range), - ])) - if position_encoding_2d: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[2, -1], - dim_range=OrderedDict([ - ('2', [2] * num_profiles), - ('position_ids_num_tokens_range', - position_ids_num_tokens_range), - ]), - ) - else: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('position_ids_num_tokens_range', - position_ids_num_tokens_range), - ]), - ) - else: - assert dtype is not None - assert num_heads is not None - pp_hidden_size = hidden_size // mapping.tp_size if pp_reduce_scatter else hidden_size - hidden_states = Tensor( - name='hidden_states_input', - dtype=dtype, - shape=[-1, pp_hidden_size], - dim_range=OrderedDict([ - ('num_tokens', num_tokens_range), - ('hidden_size', [pp_hidden_size] * num_profiles), - ]), - ) - - else: - if mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('input_len', inlen_range), - ])) - if position_encoding_2d: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, 2, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('2', [2] * num_profiles), - ('position_ids_inlen_range', - position_ids_inlen_range), - ]), - ) - else: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('position_ids_inlen_range', - position_ids_inlen_range), - ]), - ) - else: - assert dtype is not None - assert num_heads is not None - hidden_states = Tensor( - name='hidden_states_input', - dtype=dtype, - shape=[-1, -1, hidden_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('input_len', inlen_range), - ('hidden_size', [hidden_size] * num_profiles), - ]), - ) - - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor( - mapping, num_profiles) - - prompt_embedding_table = None - tasks = None - prompt_vocab_size = None - if prompt_embedding_table_size > 0: - _p_embedding_range = [ - 1, prompt_embedding_table_size // 2, prompt_embedding_table_size - ] - p_embedding_range = [_p_embedding_range] * num_profiles - - prompt_embedding_table = Tensor(name='prompt_embedding_table', - dtype=dtype, - shape=[-1, hidden_size], - dim_range=OrderedDict([ - ('prompt_embedding_table_size', - p_embedding_range), - ('hidden_size', - [hidden_size] * num_profiles), - ])) - if remove_input_padding: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('input_len_task', num_tokens_range), - ])) - else: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('broadcast_dim', [1] * num_profiles), - ])) - prompt_vocab_size = Tensor(name='prompt_vocab_size', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([ - ('size', [1] * num_profiles) - ])) - - lora_weights_pointers = None - lora_ranks = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - layers_range = mapping.pp_layers(num_layers) - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in lora_target_modules: - - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('in_out_scales', [3] * num_profiles), - ])) - lora_weight_pointer_dict.update({ - f"{lora_module}_lora_weights_pointers": - lora_weight_pointer - }) - - lora_rank = Tensor( - name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', - bb_range)]), - ) - lora_rank_dict.update( - {f"{lora_module}_lora_ranks": lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - last_token_ids = None - if mapping.is_last_pp_rank() and not gather_context_logits: - if not remove_input_padding and max_draft_len > 0: - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('last_token_ids', tokens_per_engine_step_range), - ]), - ) - else: - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_last_token_ids', bbd_range), - ]), - ) - - spec_decoding_params = None - # Use positional offsets and packed mask only when not in SpS spec decoding - if speculative_decoding_draft_tokens_external == False and max_draft_len > 0: - tokens_per_engine_step = max_draft_len + 1 - # 32 bits packed mask aligned. - num_packed_masks = (tokens_per_engine_step + 32 - 1) // 32 - packed_mask_len_range = [[0, 1, num_packed_masks]] * num_profiles - # total number of spec decoding tokens for all sequences (sequence length can be variable). - num_gen_tokens_range = [ - GenerationMixin.default_range( - max_batch_size * max_beam_width * tokens_per_engine_step, - min_range=0) - ] * num_profiles - bb_range_0 = [[0] + bbr[1:] for bbr in bb_range] - - # support variable sequence lengths for medusa. - spec_decoding_generation_lengths = Tensor( - name='spec_decoding_generation_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width_0', bb_range_0) - ]), - ) - - # position offsets that are fixed during the whole session. - # it will be shared among all sequences. - spec_decoding_position_offsets = Tensor( - name='spec_decoding_position_offsets', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width_0', bb_range_0), - ('spec_decoding_position_ids_dim0', - tokens_per_engine_step_range), - ]), - ) - - spec_decoding_packed_mask = Tensor( - name='spec_decoding_packed_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('spec_decoding_packed_mask_dim0', num_gen_tokens_range), - ('spec_decoding_packed_mask_dim1', packed_mask_len_range), - ]), - ) - spec_decoding_use = Tensor(name='spec_decoding_use', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([ - ('spec_decoding_use_dim', - [1] * num_profiles), - ])) - spec_decoding_params = SpecDecodingParams( - spec_decoding_is_generation_length_variable= - spec_decoding_is_generation_length_variable, - spec_decoding_max_generation_length=tokens_per_engine_step, - spec_decoding_generation_lengths= - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_packed_mask, - spec_decoding_use=spec_decoding_use) - - mrope_params = None - if mrope_rotary_cos_sin_size is not None: - mrope_rotary_cos_sin = Tensor( - name='mrope_rotary_cos_sin', - dtype=trt.float32, - shape=[-1, mrope_rotary_cos_sin_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('mult_dim', [mrope_rotary_cos_sin_size] * num_profiles), - ]), - ) - - mrope_position_deltas = Tensor( - name='mrope_position_deltas', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range), - ('mult_dim_delta', [1] * num_profiles)]), - ) - mrope_params = MropeParams( - mrope_rotary_cos_sin=mrope_rotary_cos_sin, - mrope_position_deltas=mrope_position_deltas, - ) - - basic_inputs = { - 'input_ids': input_ids, - 'hidden_states_input': hidden_states, - 'position_ids': position_ids, - 'last_token_ids': last_token_ids, - 'prompt_embedding_table': prompt_embedding_table, - 'tasks': tasks, - 'prompt_vocab_size': prompt_vocab_size, - 'lora_ranks': lora_ranks, - 'lora_weights_pointers': lora_weights_pointers, - 'spec_decoding_params': spec_decoding_params, - 'mrope_params': mrope_params, - } - - attention_inputs = self.prepare_attention_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - num_kv_heads=num_kv_heads, - head_size=head_size, - num_layers=num_layers, - kv_dtype=kv_dtype, - num_profiles=num_profiles, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - mapping=mapping, - streamingllm=streamingllm, - opt_batch_size=opt_batch_size) - for key, value in attention_inputs.items(): - basic_inputs[key] = value - - return basic_inputs diff --git a/tensorrt_llm/models/gpt/__init__.py b/tensorrt_llm/models/gpt/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gpt/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/gpt/config.py b/tensorrt_llm/models/gpt/config.py deleted file mode 100644 index ba09d1f86942..000000000000 --- a/tensorrt_llm/models/gpt/config.py +++ /dev/null @@ -1,324 +0,0 @@ -# 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 - -import torch - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...logger import logger -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class GPTConfig(PretrainedConfig): - - def __init__(self, - *, - gpt_variant: str = 'gpt2', - bias: bool = True, - q_scaling: float = 1.0, - embedding_scale: Optional[float] = None, - apply_query_key_layer_scaling: bool = False, - rotary_pct: float = 1.0, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - inner_layernorm: bool = False, - norm_before_bmm1: bool = False, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - self.gpt_variant = gpt_variant - self.bias = bias - self.q_scaling = q_scaling - self.embedding_scale = embedding_scale - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.rotary_pct = rotary_pct - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.inner_layernorm = inner_layernorm - self.norm_before_bmm1 = norm_before_bmm1 - if moe is None: - # Legacy MOE config fields - moe = MoeConfig( - num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE)) - elif isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in GPTConfig - output['gpt_variant'] = self.gpt_variant - output['bias'] = self.bias - output['q_scaling'] = self.q_scaling - output['embedding_scale'] = self.embedding_scale - output[ - 'apply_query_key_layer_scaling'] = self.apply_query_key_layer_scaling - output['rotary_pct'] = self.rotary_pct - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['inner_layernorm'] = self.inner_layernorm - output['norm_before_bmm1'] = self.norm_before_bmm1 - output['moe'] = self.moe.to_dict() - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - from .convert import get_needed_padding - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_or_dir, trust_remote_code=trust_remote_code) - - gpt_variant = kwargs.pop('gpt_variant', None) - if gpt_variant is None: - logger.info("Inferring gpt variant from path...") - for v in [ - 'starcoder2', 'starcoder', 'santacoder', 'gpt2', - 'persimmon', 'fuyu', 'kosmos-2', 'jais', 'nemotron' - ]: - if v in hf_config._name_or_path or v == hf_config.model_type: - gpt_variant = v - break - if gpt_variant == 'fuyu': - gpt_variant = 'persimmon' - - assert gpt_variant in [ - 'gpt2', 'santacoder', 'starcoder', 'starcoder2', 'persimmon', - 'kosmos-2', 'jais', 'nemotron' - ] - logger.info(f"Gpt variant: {gpt_variant}") - - if gpt_variant in ['starcoder2', 'nemotron', 'persimmon']: - hf_config.n_embd = hf_config.hidden_size - hf_config.n_inner = hf_config.intermediate_size - hf_config.n_head = hf_config.num_attention_heads - hf_config.n_kv_head = hf_config.num_key_value_heads if hasattr( - hf_config, 'num_key_value_heads') else hf_config.n_head - hf_config.n_layer = hf_config.num_hidden_layers - hf_config.n_positions = hf_config.max_position_embeddings - hf_config.activation_function = 'gelu' if gpt_variant == 'starcoder2' else 'squared-relu' - if gpt_variant == "nemotron": - hf_config.layer_norm_eps = hf_config.norm_eps - hf_config.layer_norm_epsilon = hf_config.norm_epsilon if gpt_variant == 'starcoder2' else hf_config.layer_norm_eps - hf_config.bias = hf_config.use_bias if gpt_variant == 'starcoder2' else gpt_variant != 'nemotron' - hf_config.position_embedding_type = 'rope_gpt_neox' - hf_config.rotary_base = get_hf_rope_theta(hf_config, 10000.0) - hf_config.rotary_pct = getattr( - hf_config, 'partial_rotary_factor', - getattr(hf_config, 'rope_percent', 1.0)) - try: - # only for persimmon, not starcoder2 - hf_config.vocab_size = hf_config.text_config.vocab_size - except AttributeError: - pass - elif gpt_variant == "kosmos-2": - hf_config.n_embd = hf_config.text_config.embed_dim - hf_config.n_inner = hf_config.text_config.ffn_dim - hf_config.n_head = hf_config.text_config.attention_heads - hf_config.n_kv_head = hf_config.n_head - hf_config.n_layer = hf_config.text_config.layers - hf_config.n_positions = hf_config.text_config.max_position_embeddings - hf_config.activation_function = hf_config.text_config.activation_function - hf_config.layer_norm_epsilon = hf_config.text_config.layer_norm_eps - hf_config.bias = True - hf_config.vocab_size = hf_config.text_config.vocab_size - else: - if hf_config.n_inner is None: - hf_config.n_inner = hf_config.n_embd * 4 - if gpt_variant in ['santacoder', 'starcoder']: - hf_config.n_kv_head = 1 - else: - hf_config.n_kv_head = hf_config.n_head - - if gpt_variant == 'jais': - hf_config.q_scaling = (hf_config.n_embd // hf_config.n_head)**0.5 - if hasattr(hf_config, 'width_scale'): - hf_config.logits_scale = hf_config.width_scale - else: - hf_config.logits_scale = hf_config.mup_output_alpha * hf_config.mup_width_scale - - if hasattr(hf_config, 'mup_embeddings_scale'): - hf_config.embeddings_scale = hf_config.mup_embeddings_scale - else: - assert hasattr(hf_config, 'embeddings_scale') - - hf_config.n_inner += get_needed_padding(hf_config.n_inner, - mapping.tp_size) - - if gpt_variant == 'kosmos-2': - if hf_config.text_config.scale_embedding: - hf_config.embeddings_scale = hf_config.n_embd**0.5 - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.n_layer, - num_attention_heads=hf_config.n_head, - num_key_value_heads=hf_config.n_kv_head, - hidden_size=hf_config.n_embd, - intermediate_size=hf_config.n_inner, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type=getattr(hf_config, - 'position_embedding_type', - 'learned_absolute'), - max_position_embeddings=hf_config.n_positions, - hidden_act=hf_config.activation_function, - gpt_variant=gpt_variant, - bias=getattr(hf_config, 'bias', True), - apply_query_key_layer_scaling=getattr( - hf_config, 'apply_query_key_layer_scaling', False), - rotary_pct=getattr(hf_config, 'rotary_pct', 1.0), - rotary_base=getattr(hf_config, 'rotary_base', 10000.0), - rotary_scaling=getattr(hf_config, 'rotary_scaling', None), - qk_layernorm=gpt_variant == 'persimmon', - inner_layernorm=gpt_variant == 'kosmos-2', - norm_before_bmm1=gpt_variant == 'kosmos-2', - q_scaling=getattr(hf_config, 'q_scaling', 1), - embedding_scale=getattr(hf_config, 'embeddings_scale', None), - mapping=mapping, - quantization=quant_config, - **kwargs) - - @classmethod - def from_nemo(cls, - nemo_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - from .convert import (UnpackedNemoCheckpointDir, cpu_map_location, - gpu_map_location, rename_keys) - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - nemo_rename_key = kwargs.pop('nemo_rename_key', []) - layer_rename_config = { - pattern.split(':')[0]: pattern.split(':')[1] - for pattern in nemo_rename_key - } - - unpacked_checkpoints_dir = UnpackedNemoCheckpointDir( - nemo_ckpt_dir, load_checkpoints_to_cpu=load_model_on_cpu) - nemo_model_config = unpacked_checkpoints_dir.model_config - - training_tp_size = nemo_model_config.get("tensor_model_parallel_size", - 1) - training_pp_size = nemo_model_config.get("pipeline_model_parallel_size", - 1) - - checkpoints_paths = unpacked_checkpoints_dir.get_checkpoints_paths( - training_tp_size, - training_pp_size, - ) - if unpacked_checkpoints_dir._load_checkpoints_to_cpu: - map_location_fn = cpu_map_location - else: - map_location_fn = gpu_map_location - model_00 = torch.load(checkpoints_paths[0][0], - map_location=map_location_fn) - model_00 = rename_keys(model_00, layer_rename_config) - vocab_size = model_00[ - "model.language_model.embedding.word_embeddings.weight"].shape[ - 0] * training_tp_size - del model_00 - - hf_config = transformers.GPT2Config( - vocab_size=vocab_size, - n_positions=nemo_model_config['max_position_embeddings'], - n_embd=nemo_model_config['hidden_size'], - n_layer=nemo_model_config['num_layers'], - n_head=nemo_model_config['num_attention_heads'], - n_inner=nemo_model_config['ffn_hidden_size'], - activation_function=nemo_model_config['activation'], - layer_norm_epsilon=nemo_model_config['layernorm_epsilon'], - ) - hf_config.n_kv_head = hf_config.n_head - hf_config.bias = nemo_model_config['bias'] - hf_config.apply_query_key_layer_scaling = False - - hf_config.position_embedding_type = nemo_model_config.get( - 'position_embedding_type', 'learned_absolute') - if hf_config.position_embedding_type == 'rope': - hf_config.position_embedding_type = 'rope_gpt_neox' - hf_config.rotary_base = nemo_model_config.get('rotary_base', 10000.0) - hf_config.rotary_pct = nemo_model_config.get('rotary_percentage', 1.0) - assert hf_config.rotary_pct >= 0 and hf_config.rotary_pct <= 1 - - rotary_scaling_factor = nemo_model_config.get( - 'seq_len_interpolation_factor', None) - if rotary_scaling_factor is None: - hf_config.rotary_scaling = None - else: - assert rotary_scaling_factor > 1 - hf_config.rotary_scaling = { - 'type': 'linear', - 'factor': rotary_scaling_factor - } - - if dtype == 'auto': - dtype = nemo_model_config.get('precision', None) - if dtype is None: - dtype = 'float16' - elif 'bf16' in dtype or 'bfloat16' in dtype: - dtype = 'bfloat16' - else: - dtype = 'float16' - logger.info(f"Specified dtype 'auto'; inferred dtype {dtype!r}.") - - return cls(architecture='GPTForCausalLM', - dtype=dtype, - num_hidden_layers=hf_config.n_layer, - num_attention_heads=hf_config.n_head, - num_key_value_heads=hf_config.n_kv_head, - hidden_size=hf_config.n_embd, - intermediate_size=hf_config.n_inner, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type=hf_config.position_embedding_type, - max_position_embeddings=hf_config.n_positions, - hidden_act=hf_config.activation_function, - bias=hf_config.bias, - apply_query_key_layer_scaling=hf_config. - apply_query_key_layer_scaling, - rotary_pct=hf_config.rotary_pct, - rotary_base=hf_config.rotary_base, - rotary_scaling=hf_config.rotary_scaling, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/gpt/convert.py b/tensorrt_llm/models/gpt/convert.py deleted file mode 100644 index 799fac22d40f..000000000000 --- a/tensorrt_llm/models/gpt/convert.py +++ /dev/null @@ -1,1455 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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 copy -import functools -import os -import shutil -import tarfile -import time -from collections import defaultdict -from pathlib import Path -from typing import Dict, Optional, Union - -import numpy as np -import safetensors -import torch -import torch.nn as nn -import yaml -from tqdm import tqdm -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.models.gpt2.modeling_gpt2 import GPT2Block -from transformers.pytorch_utils import Conv1D - -from ..._utils import pad_vocab_size, str_dtype_to_torch -from ...logger import logger -from ...quantization import QuantAlgo -from ..convert_utils import (generate_int8, get_weight, get_weight_and_bias, - load_calib_dataset, - retrieved_layer_index_from_name, smooth_gemm) -from .config import GPTConfig - - -def rename_keys(model_state, layer_rename_config: Dict[str, str]): - if not layer_rename_config: - return model_state - - new_state_dict = {} - for key, value in model_state.items(): - for old, new in layer_rename_config.items(): - key = key.replace(old, new) - assert key not in new_state_dict, f"Key already exists: {key}" - new_state_dict[key] = value - - return new_state_dict - - -def get_needed_padding(value: int, multiple: int) -> int: - return (multiple - value % multiple) % multiple - - -def pad_array_up_to(v: torch.Tensor, axis: int, multiple: int) -> torch.Tensor: - a = [0 for i in range(len(v.shape) * 2)] - a[axis * 2 - 1] = get_needed_padding(v.shape[axis], multiple) - return torch.nn.functional.pad(v, a) - - -def split(param: torch.Tensor, - tp_rank: int, - tp_size: int, - is_column: bool = True) -> torch.Tensor: - """Split linear layer's weight, bias or scaling factors for tensor parallelism.""" - if param is None: - return None - assert param.ndim in [1, 2] - if tp_size == 1: - return param - if param.numel() == 1: - return param - if param.ndim == 1 and not is_column: - return param - split_dim = 0 if (param.ndim == 1 or is_column) else 1 - return torch.chunk(param, tp_size, dim=split_dim)[tp_rank].contiguous() - - -def split_qkv( - param: torch.Tensor, - tp_rank: int, - tp_size: int, - hidden_size: int, - num_heads: int, - num_kv_heads: Optional[int] = None, -) -> torch.Tensor: - """Split qkv layer's weight, bias or scaling factors for tensor parallelism. - - param: (num_heads*head_dim + 2*num_kv_heads*head_dim, in_dim) - """ - if param is None: - return None - assert hidden_size % num_heads == 0 - head_dim = hidden_size // num_heads - num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads - assert num_heads % num_kv_heads == 0 - assert num_heads % tp_size == 0 - - q_param, k_param, v_param = torch.split( - param, [hidden_size, num_kv_heads * head_dim, num_kv_heads * head_dim], - dim=0) - - if num_kv_heads < tp_size: - assert tp_size % num_kv_heads == 0 - num_dups = tp_size // num_kv_heads - remain_shape = k_param.shape[1:] - k_param = k_param.view( - num_kv_heads, head_dim, - *remain_shape).repeat_interleave(num_dups, dim=0).view( - num_kv_heads * head_dim * num_dups, *remain_shape) - v_param = v_param.view( - num_kv_heads, head_dim, - *remain_shape).repeat_interleave(num_dups, dim=0).view( - num_kv_heads * head_dim * num_dups, *remain_shape) - else: - assert num_kv_heads % tp_size == 0 - - q_param = split(q_param, tp_rank, tp_size, is_column=True) - k_param = split(k_param, tp_rank, tp_size, is_column=True) - v_param = split(v_param, tp_rank, tp_size, is_column=True) - return torch.cat([q_param, k_param, v_param], dim=0) - - -def split_embedding( - param: torch.Tensor, - tp_rank: int, - tp_size: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_rank, tp_size, is_column=(sharding_dim == 0)) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=0)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - input_ids = tokenizer(dataset[i], - return_tensors="pt", - max_length=seq_len, - truncation=True).input_ids.to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -@torch.no_grad() -def smooth_gpt_model(model, scales, alpha): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, GPT2Block): - continue - - # qkv_proj - layer_name = name + ".attn.c_attn" - smoother = smooth_gemm(module.attn.c_attn.weight.T, - scales[layer_name]["x"], module.ln_1.weight, - module.ln_1.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.c_attn.weight.abs().max(dim=0)[0] - - # fc1 - layer_name = name + ".mlp.c_fc" - smoother = smooth_gemm(module.mlp.c_fc.weight.T, - scales[layer_name]["x"], module.ln_2.weight, - module.ln_2.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.c_fc.weight.abs().max(dim=0)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.split(q, q.shape[-1] // tp_size, dim=-1) - k_split = torch.split(k, q.shape[-1] // tp_size, dim=-1) - v_split = torch.split(v, q.shape[-1] // tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), dim=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def load_weights_from_hf_model(hf_model, - config: GPTConfig, - act_range: Optional[dict] = None): - 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_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - if use_smooth_quant or int8_kv_cache: - assert act_range is not None - - weights = {} - tik = time.time() - - hf_config = hf_model.config - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - gpt_variant = config.gpt_variant - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - vocab_size = config.vocab_size - num_kv_heads = config.num_key_value_heads - num_hidden_layers = config.num_hidden_layers - multi_query_mode = (num_kv_heads != num_attention_heads) - - mapping = config.mapping - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - if gpt_variant in ['starcoder2', 'nemotron']: - prefix = f'model.layers.{l}' - elif gpt_variant == 'persimmon': - is_fuyu = f'language_model.model.embed_tokens.weight' in model_params - prefix = f'language_model.model.layers.{l}' if is_fuyu else f'model.layers.{l}' - elif gpt_variant == 'kosmos-2': - prefix = f'text_model.model.layers.{l}' - else: - prefix = f'transformer.h.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - - # (1) Attention QKV Linear - if gpt_variant == 'santacoder': - q_w, q_b = get_weight_and_bias(model_params, - f'{prefix}.attn.q_attn', dtype) - kv_w, kv_b = get_weight_and_bias(model_params, - f'{prefix}.attn.kv_attn', dtype) - qkv_w = torch.cat([q_w, kv_w], dim=-1) - qkv_b = torch.cat([q_b, kv_b], dim=-1) - elif gpt_variant in ['starcoder2', 'nemotron', 'kosmos-2']: - q_w, q_b = get_weight_and_bias(model_params, - f'{prefix}.self_attn.q_proj', dtype) - k_w, k_b = get_weight_and_bias(model_params, - f'{prefix}.self_attn.k_proj', dtype) - v_w, v_b = get_weight_and_bias(model_params, - f'{prefix}.self_attn.v_proj', dtype) - qkv_w = torch.cat([q_w.cuda(), k_w.cuda(), v_w.cuda()], dim=0) - qkv_b = torch.cat([q_b.cuda(), k_b.cuda(), - v_b.cuda()], dim=0) if q_b is not None else None - elif gpt_variant == 'persimmon': - qkv_w, qkv_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.query_key_value', dtype) - else: - qkv_w, qkv_b = get_weight_and_bias(model_params, - f'{prefix}.attn.c_attn', dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - qkv_w = qkv_w.t().contiguous() # transpose for Conv1D - - if use_smooth_quant: - qkv_out_dim = qkv_w.shape[0] - qkv_w_t = qkv_w.t() - if not multi_query_mode: - qkv_w_t = qkv_w_t.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_w_t, - act_range.get(f'{prefix}.attn.c_attn'), - is_qkv=True, - multi_query_mode=multi_query_mode) - qkv_b = split_qkv(qkv_b, mapping.tp_rank, mapping.tp_size, - hidden_size, num_attention_heads, num_kv_heads) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.attention.qkv.', - [1, qkv_out_dim // mapping.tp_size], - mapping.tp_size, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.input_layernorm.scale_to_int', - bias=qkv_b, - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=multi_query_mode)) - else: - if gpt_variant == 'persimmon': - qkv_w = split(qkv_w, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - - qkv_b = split(qkv_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - else: - qkv_w = split_qkv(qkv_w, mapping.tp_rank, mapping.tp_size, - hidden_size, num_attention_heads, - num_kv_heads) - qkv_b = split_qkv(qkv_b, mapping.tp_rank, mapping.tp_size, - hidden_size, num_attention_heads, - num_kv_heads) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', - qkv_b, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_w_t = qkv_w.t() - if not multi_query_mode: - qkv_w_t = qkv_w_t.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_w_t, - act_range.get(f'{prefix}.attn.c_attn'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = int8_weights[ - 'scale_y_quant_orig'].contiguous() - - # (2) Attention Dense Linear - if gpt_variant in ['starcoder2', 'nemotron']: - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.o_proj', dtype) - elif gpt_variant == 'persimmon': - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.dense', dtype) - elif gpt_variant == 'kosmos-2': - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.out_proj', dtype) - else: - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.attn.c_proj', dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - attn_dense_w = attn_dense_w.t().contiguous() # transpose for Conv1D - - if use_smooth_quant: - attn_dense_w_t = attn_dense_w.t() - int8_weights = generate_int8(attn_dense_w_t, - act_range.get(f'{prefix}.attn.c_proj')) - # change it to the real smoother if dense layer is applied smooth quant - fake_smoother_value = torch.ones([1, hidden_size], - dtype=torch.float32) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.attention.dense.', [1, hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix= - f'{tllm_prex}.attention.quantization_scaling_factor', - bias=attn_dense_b, - smoother_value=fake_smoother_value, - smoother_shape=[1, hidden_size // mapping.tp_size], - rank=mapping.tp_rank, - cat_dim=0)) - else: - attn_dense_w = split(attn_dense_w, - mapping.tp_rank, - mapping.tp_size, - is_column=False) - weights.update( - get_tllm_linear_weight(attn_dense_w, - f'{tllm_prex}.attention.dense', - attn_dense_b, use_weight_only, - plugin_weight_only_quant_type)) - - # (3) MLP FC Linear - if gpt_variant == 'persimmon': - suffix = "mlp.dense_h_to_4h" - elif gpt_variant == 'kosmos-2': - suffix = "ffn.fc1" - elif gpt_variant == 'nemotron': - suffix = "mlp.up_proj" - else: - suffix = "mlp.c_fc" - mlp_fc_w, mlp_fc_b = get_weight_and_bias(model_params, - f'{prefix}.{suffix}', dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - mlp_fc_w = mlp_fc_w.t().contiguous() # transpose for Conv1D - if gpt_variant in ['jais']: - mlp_fc_w = pad_array_up_to(mlp_fc_w, 0, mapping.tp_size) - mlp_fc_b = pad_array_up_to(mlp_fc_b, 0, mapping.tp_size) - if use_smooth_quant: - mlp_fc_w_t = mlp_fc_w.t() - int8_weights = generate_int8(mlp_fc_w_t, - act_range.get(f'{prefix}.mlp.c_fc')) - mlp_fc_b = split(mlp_fc_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.mlp.fc.', - [1, 4 * hidden_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.post_layernorm.scale_to_int', - bias=mlp_fc_b, - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_fc_w = split(mlp_fc_w, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - mlp_fc_b = split(mlp_fc_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - if gpt_variant in ['jais']: - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.gate', - mlp_fc_b, use_weight_only, - plugin_weight_only_quant_type)) - else: - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', - mlp_fc_b, use_weight_only, - plugin_weight_only_quant_type)) - if gpt_variant in ['jais']: - mlp_fc2_w, mlp_fc2_b = get_weight_and_bias( - model_params, f'{prefix}.mlp.c_fc2', dtype) - mlp_fc2_w = mlp_fc2_w.t().contiguous() - mlp_fc2_w = pad_array_up_to(mlp_fc2_w, 0, mapping.tp_size) - mlp_fc2_b = pad_array_up_to(mlp_fc2_b, 0, mapping.tp_size) - mlp_fc2_w = split(mlp_fc2_w, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - mlp_fc2_b = split(mlp_fc2_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - weights.update( - get_tllm_linear_weight(mlp_fc2_w, f'{tllm_prex}.mlp.fc', - mlp_fc2_b, use_weight_only, - plugin_weight_only_quant_type)) - - # (4) MLP Proj Layer - if gpt_variant == 'persimmon': - suffix = "mlp.dense_4h_to_h" - elif gpt_variant == 'kosmos-2': - suffix = "ffn.fc2" - elif gpt_variant == 'nemotron': - suffix = "mlp.down_proj" - else: - suffix = "mlp.c_proj" - mlp_proj_w, mlp_proj_b = get_weight_and_bias(model_params, - f'{prefix}.{suffix}', - dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - mlp_proj_w = mlp_proj_w.t().contiguous() # transpose for Conv1D - if gpt_variant in ['jais']: - mlp_proj_w = pad_array_up_to(mlp_proj_w, 1, mapping.tp_size) - if use_smooth_quant: - mlp_proj_w_t = mlp_proj_w.t() - int8_weights = generate_int8(mlp_proj_w_t, - act_range.get(f'{prefix}.mlp.c_proj')) - # change it to the real smoother if proj layer is applied smooth quant - fake_smoother_value = torch.ones([1, 4 * hidden_size], - dtype=torch.float32) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.mlp.proj.', [1, hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.mlp.quantization_scaling_factor', - bias=mlp_proj_b, - smoother_value=fake_smoother_value, - smoother_shape=[1, 4 * hidden_size // mapping.tp_size], - rank=mapping.tp_rank, - cat_dim=0)) - else: - mlp_proj_w = split(mlp_proj_w, - mapping.tp_rank, - mapping.tp_size, - is_column=False) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_b, use_weight_only, - plugin_weight_only_quant_type)) - - # (5) Input layernorm - apply_layernorm_1p = gpt_variant == 'nemotron' - if gpt_variant in ['starcoder2', 'nemotron', 'persimmon']: - input_ln_w, input_ln_b = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - elif gpt_variant == 'kosmos-2': - input_ln_w, input_ln_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn_layer_norm', dtype) - else: - input_ln_w, input_ln_b = get_weight_and_bias( - model_params, f'{prefix}.ln_1', dtype) - if apply_layernorm_1p: - input_ln_w += 1.0 - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_w - if input_ln_b is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_b - - # (6) Post layernorm - if gpt_variant in ['starcoder2', 'nemotron', 'persimmon']: - post_ln_w, post_ln_b = get_weight_and_bias( - model_params, f'{prefix}.post_attention_layernorm', dtype) - elif gpt_variant == 'kosmos-2': - post_ln_w, post_ln_b = get_weight_and_bias( - model_params, f'{prefix}.final_layer_norm', dtype) - else: - post_ln_w, post_ln_b = get_weight_and_bias(model_params, - f'{prefix}.ln_2', dtype) - if apply_layernorm_1p: - post_ln_w += 1.0 - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_w - if post_ln_b is not None: - weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_b - - if gpt_variant == 'persimmon': - q_layernorm_w, q_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.q_layernorm', dtype) - - weights[f'{tllm_prex}.attention.q_layernorm.weight'] = q_layernorm_w - weights[f'{tllm_prex}.attention.q_layernorm.bias'] = q_layernorm_b - - k_layernorm_w, k_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.k_layernorm', dtype) - - weights[f'{tllm_prex}.attention.k_layernorm.weight'] = k_layernorm_w - weights[f'{tllm_prex}.attention.k_layernorm.bias'] = k_layernorm_b - - if gpt_variant == 'kosmos-2': - q_layernorm_w, q_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.inner_attn_ln', dtype) - - weights[ - f'{tllm_prex}.attention.inner_layernorm.weight'] = q_layernorm_w - weights[ - f'{tllm_prex}.attention.inner_layernorm.bias'] = q_layernorm_b - - k_layernorm_w, k_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.ffn.ffn_layernorm', dtype) - - weights[f'{tllm_prex}.mlp.inner_layernorm.weight'] = k_layernorm_w - weights[f'{tllm_prex}.mlp.inner_layernorm.bias'] = k_layernorm_b - - if mapping.is_first_pp_rank(): - if gpt_variant in ['starcoder2', 'nemotron']: - embed_w = get_weight(model_params, 'model.embed_tokens', dtype) - elif gpt_variant == 'kosmos-2': - embed_w = get_weight(model_params, 'text_model.model.embed_tokens', - dtype) - elif gpt_variant == 'persimmon': - embed_w = get_weight(model_params, - ('language_model.' if is_fuyu else '') + - 'model.embed_tokens', dtype) - else: - embed_w = get_weight(model_params, 'transformer.wte', dtype) - weights['transformer.vocab_embedding.weight'] = split_embedding( - embed_w, - mapping.tp_rank, - mapping.tp_size, - use_parallel_embedding=config.use_parallel_embedding, - sharding_dim=config.embedding_sharding_dim) - - if gpt_variant == 'kosmos-2': - padding_idx = hf_config.text_config.pad_token_id - sin_pos_embedding = hf_model.text_model.model.embed_positions.get_embedding( - padding_idx + 1 + hf_config.text_config.max_position_embeddings, - hf_config.text_config.embed_dim, - padding_idx=padding_idx) # [2 + num_embeddings, embed_dim] - pos_embed_w = sin_pos_embedding[2:].to(dtype).detach().cpu() - else: - pos_embed_w = get_weight(model_params, 'transformer.wpe', dtype) - if pos_embed_w is not None: - weights['transformer.position_embedding.weight'] = split_embedding( - pos_embed_w, - mapping.tp_rank, - mapping.tp_size, - use_parallel_embedding=config.use_parallel_embedding, - sharding_dim=config.embedding_sharding_dim) - - if mapping.is_last_pp_rank(): - if gpt_variant in ['starcoder2', 'nemotron']: - embed_w = get_weight(model_params, 'lm_head', dtype) - if embed_w is None: - embed_w = get_weight(model_params, 'model.embed_tokens', dtype) - elif gpt_variant == 'persimmon': - embed_w = get_weight(model_params, - ('language_model.' if is_fuyu else '') + - 'lm_head', dtype) - elif gpt_variant == 'kosmos-2': - embed_w = get_weight(model_params, 'text_model.model.embed_tokens', - dtype) - else: - embed_w = get_weight(model_params, 'transformer.wte', dtype) - - if vocab_size % mapping.tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - embed_w = torch.nn.functional.pad(embed_w, (0, 0, 0, pad_width), - value=0) - if hasattr(hf_config, 'logits_scale'): - embed_w *= hf_config.logits_scale - weights['lm_head.weight'] = split(embed_w.clone(), - mapping.tp_rank, - mapping.tp_size, - is_column=True) - - if gpt_variant in ['starcoder2', 'nemotron']: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'model.norm', - dtype) - elif gpt_variant == 'persimmon': - ln_f_w, ln_f_b = get_weight_and_bias( - model_params, ('language_model.' if is_fuyu else '') + - 'model.final_layernorm', dtype) - elif gpt_variant == 'kosmos-2': - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'text_model.model.layer_norm', - dtype) - else: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'transformer.ln_f', dtype) - if apply_layernorm_1p: - ln_f_w += 1.0 - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: GPTConfig, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - trust_remote_code: bool = True): - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map='auto' if device != 'cpu' else 'cpu', - dtype='auto' if not use_smooth_quant else torch.float16, - trust_remote_code=trust_remote_code) - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - use_fast=False, - padding_side='left') - - dataset = load_calib_dataset(calib_dataset) - act_range = capture_activation_range(hf_model, tokenizer, dataset) - if use_smooth_quant: - smooth_gpt_model(hf_model, act_range, quant_config.smoothquant_val) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model( - hf_model, - config=config, - act_range=act_range, - ) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -def load_hf_gpt(model_dir: str, load_model_on_cpu: bool = False): - if 'kosmos-2' in model_dir: - # AutoModelForVision2Seq was removed in transformers 5.x; import lazily - # so `import tensorrt_llm` works under newer transformers versions even - # though kosmos-2 conversion itself still requires transformers<5. - from transformers import AutoModelForVision2Seq - hf_model = AutoModelForVision2Seq.from_pretrained( - model_dir, trust_remote_code=True) - else: - hf_model = AutoModelForCausalLM.from_pretrained( - model_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype='auto', - trust_remote_code=True, - ) - return hf_model - - -def cpu_map_location(storage, loc): - return storage.cpu() - - -def gpu_map_location(storage, loc): - if loc.startswith("cuda"): - training_gpu_idx = int(loc.split(":")[1]) - inference_gpu_idx = training_gpu_idx % torch.cuda.device_count() - return storage.cuda(inference_gpu_idx) - elif loc.startswith("cpu"): - return storage.cpu() - else: - raise ValueError(f"Not handled {loc}") - - -def copy_tokenizer_files(config, out_dir): - basenames = { - "model": "tokenizer", - "vocab_file": "vocab", - "merge_file": "merges", - } - - for key in basenames.keys(): - if config[key] is None: - continue - path = Path(config[key]) - if not path.exists(): - logger.debug(f"Tokenizer {key}: {path} file not found") - continue - - dst_path = out_dir / f"{basenames[key]}{path.suffix}" - logger.debug(f"Copy tokenizer {key}: {path}->{dst_path}") - shutil.copy(path.as_posix(), dst_path.as_posix()) - - -def update_tokenizer_paths(tokenizer_config: Dict, - tokenizer_file_paths: Dict[str, Optional[str]]): - for key, new_path in tokenizer_file_paths.items(): - old_path = tokenizer_config[key] - if old_path is None: - continue - old_path = Path(old_path) - if new_path: - logger.debug(f"Update tokenizer {key} {old_path} -> {new_path}") - tokenizer_config[key] = new_path.as_posix() - elif not old_path.exists(): - logger.warning( - f"Tokenizer {key}'s path {old_path} does not exists: set it to None" - ) - tokenizer_config[key] = None - return tokenizer_config - - -def unpack_nemo_ckpt(nemo_archive_path: Union[str, Path], - out_dir_path: Union[str, Path]): - nemo_archive_path = Path(nemo_archive_path) - if not nemo_archive_path.exists(): - raise FileNotFoundError(f"{nemo_archive_path} does not exist") - - for tar_mode in ["r:", "r:gz"]: - try: - with tarfile.open(nemo_archive_path, mode=tar_mode) as tar_file: - - def is_within_directory(directory, target): - - abs_directory = os.path.abspath(directory) - abs_target = os.path.abspath(target) - - prefix = os.path.commonprefix([abs_directory, abs_target]) - - return prefix == abs_directory - - def safe_members(tar_file): - members = [] - for member in tar_file.getmembers(): - member_path = os.path.join(out_dir_path, member.name) - if not is_within_directory(out_dir_path, member_path): - raise Exception( - "Attempted Path Traversal in Tar File") - members.append(member) - return members - - for member in safe_members(tar_file): - tar_file.extract(member, - path=out_dir_path, - numeric_owner=False, - filter=tarfile.data_filter) - - return out_dir_path - except tarfile.ReadError: - pass - - raise RuntimeError(f"Could not unpack {nemo_archive_path}") - - -def extract_layers_with_prefix(model_, prefix): - length_to_trim = len(prefix) - model_state = model_.get("state_dict", model_) - return { - key[length_to_trim:]: model_state[key] - for key in model_state.keys() if prefix in key - } - - -class UnpackedNemoCheckpointDir: - - def __init__(self, - checkpoints_dir: Union[str, Path], - load_checkpoints_to_cpu: bool = False): - self._checkpoints_dir = Path(checkpoints_dir) - self._load_checkpoints_to_cpu = load_checkpoints_to_cpu - - @property - @functools.lru_cache - def model_config(self): - model_config = None - - model_config_filename = "model_config.yaml" - model_configs_paths = list( - self._checkpoints_dir.rglob(model_config_filename)) - if model_configs_paths: - if len(model_configs_paths) > 1: - raise RuntimeError( - f"There are more than single {model_config_filename} " - f"in {self._checkpoints_dir}: {', '.join(map(lambda p: p.as_posix(), model_configs_paths))}" - ) - model_config_path = model_configs_paths[0] - logger.debug(f"Loading model config from {model_config_path}") - with model_config_path.open("r") as model_config_file: - model_config = yaml.load(model_config_file, - Loader=yaml.SafeLoader) - else: - logger.debug("Searching model config in checkpoints") - # try to obtain from checkpoint - checkpoint_name = self.checkpoint_name - checkpoints_paths = sorted( - self._checkpoints_dir.rglob(checkpoint_name)) - if checkpoints_paths: - # assume that parallel ranks 0 checkpoint should have model config embedded - checkpoint_path = checkpoints_paths[0] - - map_location_fn = cpu_map_location if self._load_checkpoints_to_cpu else gpu_map_location - - model_00 = torch.load(checkpoint_path, - map_location=map_location_fn) - if "hyper_parameters" in model_00 and "cfg" in model_00[ - "hyper_parameters"]: - model_config = model_00["hyper_parameters"]["cfg"] - logger.debug( - f"Loaded model config from checkpoint {checkpoint_path}" - ) - else: - logger.debug( - f"Could not find model config in checkpoint {checkpoint_path}" - ) - del model_00 - - if model_config is None: - logger.warning( - f"Could not find checkpoint with NeMo model config in {self._checkpoints_dir}" - ) - - logger.debug(f"Loaded model config {model_config}") - - return model_config - - @property - def checkpoints_dir(self): - return self._checkpoints_dir - - def get_checkpoints_paths(self, - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1): - """ - Injects tensor/pipeline model parallel ranks into the filepath. - Does nothing if not using model parallelism. - """ - - checkpoint_path_without_rank = self.checkpoints_dir / self.checkpoint_name - - def _inject_parallel_ranks(tp_rank, pp_rank): - if tensor_model_parallel_size > 1 or pipeline_model_parallel_size > 1: - if pipeline_model_parallel_size is None or pipeline_model_parallel_size == 1: - checkpoint_path = (checkpoint_path_without_rank.parent / - f"mp_rank_{tp_rank:02d}" / - checkpoint_path_without_rank.name) - else: - checkpoint_path = ( - checkpoint_path_without_rank.parent / - f"tp_rank_{tp_rank:02d}_pp_rank_{pp_rank:03d}" / - checkpoint_path_without_rank.name) - return checkpoint_path - else: - return checkpoint_path_without_rank - - return [[ - _inject_parallel_ranks(tp_rank=tp_rank, pp_rank=pp_rank) - for pp_rank in range(pipeline_model_parallel_size) - ] for tp_rank in range(tensor_model_parallel_size)] - - @property - @functools.lru_cache - def checkpoint_name(self): - patterns = [ - "model_weights.ckpt", # older megatron checkpoints - "*last.ckpt", # newer format of checkpoints - ] - for pattern in patterns: - model_files = sorted(list(self._checkpoints_dir.rglob(pattern))) - if model_files: - return model_files[0].name - - raise ValueError( - f"Could not find checkpoint files in {self._checkpoints_dir}") - - @functools.lru_cache - def get_tokenizer_file_path(self, tokenizer_key, file_key, - default_filename_pattern): - model_config = self.model_config - file_property = None - if tokenizer_key in model_config and file_key in model_config[ - tokenizer_key]: - file_property = model_config[tokenizer_key][file_key] - elif file_key in model_config: - file_property = model_config[file_key] - - logger.debug( - f"model_config[{tokenizer_key}][{file_key}]={file_property}") - - if file_property and file_property.startswith("nemo:"): - filename = file_property.split("nemo:")[1] - filename_pattern = f"*{filename}" - elif file_property and file_property.startswith("/artifacts/"): - filename = Path(file_property).name - filename_pattern = f"*{filename}" - elif file_property is None or file_property == "None": - filename_pattern = None - else: - filename_pattern = default_filename_pattern - logger.warning( - f"Tokenizer file from config: {tokenizer_key}.{file_key}={file_property} " - f"looks like unsupported path. Pattern {filename_pattern} will be used." - ) - - file_path = None - if filename_pattern is not None: - files_paths = list(self._checkpoints_dir.glob(filename_pattern)) - if files_paths: - assert len(files_paths) == 1 - file_path = files_paths[0] - - return file_path - - @functools.lru_cache - def get_all_tokenizer_file_paths(self): - return { - "model": - self.get_tokenizer_file_path("tokenizer", "model", "*.model"), - "vocab_file": - self.get_tokenizer_file_path("tokenizer", "vocab_file", "*vocab*"), - "merge_file": - self.get_tokenizer_file_path("tokenizer", "merge_file", - "*merge*.txt"), - } - - -@torch.no_grad() -def load_torch_checkpoints(checkpoints_paths, - merge_factor, - tp_rank, - pp_rank, - map_location_fn, - handle_model_level_weights, - layer_rename_config: Dict[str, str] = {}): - models = [] - for k in range(merge_factor): - rank_weights = checkpoints_paths[tp_rank * merge_factor + k][pp_rank] - model = torch.load(rank_weights, map_location=map_location_fn) - model = rename_keys(model, layer_rename_config) - handle_model_level_weights(model, tp_rank * merge_factor + k, pp_rank) - layers = extract_layers_with_prefix(model, - "model.language_model.encoder.") - models.append(layers) - return models - - -@torch.no_grad() -def load_weights_from_nemo(nemo_ckpt_dir: str, config: GPTConfig, **kwargs): - assert config.mapping.pp_size == 1, \ - "Pipeline parallelism is not supported." - assert not config.quantization.quant_mode.has_any_quant(), \ - "Quantization is not supported." - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - nemo_rename_key = kwargs.pop('nemo_rename_key', []) - layer_rename_config = { - pattern.split(':')[0]: pattern.split(':')[1] - for pattern in nemo_rename_key - } - - unpacked_checkpoints_dir = UnpackedNemoCheckpointDir( - nemo_ckpt_dir, load_checkpoints_to_cpu=load_model_on_cpu) - nemo_model_config = unpacked_checkpoints_dir.model_config - - checkpoints_paths = unpacked_checkpoints_dir.get_checkpoints_paths( - nemo_model_config.get("tensor_model_parallel_size", 1), - nemo_model_config.get("pipeline_model_parallel_size", 1), - ) - - if unpacked_checkpoints_dir._load_checkpoints_to_cpu: - map_location_fn = cpu_map_location - else: - map_location_fn = gpu_map_location - dtype = str_dtype_to_torch(config.dtype) - - # load position_embedding from rank 0 - model_00 = torch.load(checkpoints_paths[0][0], map_location=map_location_fn) - model_00 = model_00.get("state_dict", model_00) - model_00 = rename_keys(model_00, layer_rename_config) - has_position_embedding = "model.language_model.embedding.position_embeddings.weight" in model_00 - has_lm_head = "model.language_model.output_layer.weight" in model_00 - del model_00 - - num_layers = nemo_model_config["num_layers"] - training_tp_size = nemo_model_config.get("tensor_model_parallel_size", 1) - training_pp_size = nemo_model_config.get("pipeline_model_parallel_size", 1) - inference_tp_size = config.mapping.tp_size - inference_tp_rank = config.mapping.tp_rank - - apply_layernorm_1p = (nemo_model_config.get('normalization', - '') == "layernorm1p") - split_gated_activation = ("swiglu" - in nemo_model_config.get('activation', "gelu")) - num_attention_heads = nemo_model_config["num_attention_heads"] - # use_attention_nemo_shape = True - transpose_weights = True - # multi_query_mode = False - local_dim = None - - # merge_factor: how many TP training nodes are merged into an inference TP node - # split_factor: in how many parts a TP training node is split - gcd = np.gcd(training_tp_size, inference_tp_size) - merge_factor = training_tp_size // gcd - split_factor = inference_tp_size // gcd - - model_level_weights = defaultdict(list) - - def handle_model_level_weights(model, tp_idx: int, pp_idx: int): - if tp_idx == 0 and pp_idx == 0: - if has_position_embedding: - val = model[ - "model.language_model.embedding.position_embeddings.weight"].detach( - ).cpu() - model_level_weights[ - "transformer.position_embedding.weight"].append(val) - if pp_idx == 0: - val = model.get( - "state_dict", model - )["model.language_model.embedding.word_embeddings.weight"].detach( - ).cpu() - model_level_weights["transformer.vocab_embedding.weight"].append( - val) - if has_lm_head and pp_idx == training_pp_size - 1: - val = model.get( - "state_dict", - model)["model.language_model.output_layer.weight"].detach().cpu( - ) - model_level_weights["lm_head.weight"].append(val) - - weights = {} - tik = time.time() - tp_rank = inference_tp_rank // split_factor - # for tp_rank in range(training_tp_size // merge_factor): - for pp_rank in range(training_pp_size): - models = load_torch_checkpoints(checkpoints_paths, merge_factor, - tp_rank, pp_rank, map_location_fn, - handle_model_level_weights, - layer_rename_config) - for name in list(models[0].keys()): - params = [model[name].detach().cpu() for model in models] - if transpose_weights and params[0].ndim == 2: - params = [p.T for p in params] - if "layernorm.weight" in name and apply_layernorm_1p: - params = [p + 1.0 for p in params] - - l = retrieved_layer_index_from_name(name) - if l is not None: - new_l = l + pp_rank * num_layers // training_pp_size - prefix = f'transformer.layers.{new_l}' - - if 'attention.query_key_value' in name: - if name.endswith('weight'): - hidden_dim = params[0].shape[0] - if local_dim is None: - local_dim = params[0].shape[-1] // 3 - - # multi_query_mode = False; use_attention_nemo_shape = True - head_num = num_attention_heads // training_tp_size - size_per_head = hidden_dim // num_attention_heads - params = [ - param.reshape(hidden_dim, head_num, 3, - size_per_head) for param in params - ] - params = [param.permute(0, 2, 1, 3) for param in params] - params = [ - param.reshape(hidden_dim, 3, local_dim) - for param in params - ] - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[ - f'{prefix}.attention.qkv.weight'] = param.reshape( - hidden_dim, -1).t() - else: - if local_dim is None: - local_dim = params[0].shape[-1] // 3 - - # multi_query_mode = False; use_attention_nemo_shape = True - head_num = num_attention_heads // training_tp_size - size_per_head = local_dim // head_num - params = [ - param.reshape(head_num, 3, size_per_head) - for param in params - ] - params = [param.permute(1, 0, 2) for param in params] - params = [ - param.reshape(3, local_dim) for param in params - ] - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.attention.qkv.bias'] = param.reshape( - -1) - - elif 'attention.dense' in name: - if name.endswith('weight'): - cat_dim = 0 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.attention.dense.weight'] = param.t() - else: - weights[f'{prefix}.attention.dense.bias'] = params[0] - - elif 'mlp.dense_h_to_4h' in name: - if name.endswith('weight'): - if split_gated_activation: - params = [torch.chunk(p, 2, dim=-1) for p in params] - params, gate_params = list(zip(*params)) - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.mlp.fc.weight'] = param.t() - if split_gated_activation: - gate_param = torch.concat(gate_params, dim=cat_dim) - gate_param = torch.chunk( - gate_param, split_factor, - dim=cat_dim)[inference_tp_rank % split_factor] - weights[f'{prefix}.mlp.gate.weight'] = gate_param.t( - ) - else: - if split_gated_activation: - params = [torch.chunk(p, 2, dim=-1) for p in params] - params, gate_params = list(zip(*params)) - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.mlp.fc.bias'] = param - if split_gated_activation: - gate_param = torch.concat(gate_params, dim=cat_dim) - gate_param = torch.chunk( - gate_param, split_factor, - dim=cat_dim)[inference_tp_rank % split_factor] - weights[f'{prefix}.mlp.gate.bias'] = gate_param - - elif 'mlp.dense_4h_to_h' in name: - if name.endswith('weight'): - cat_dim = 0 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.mlp.proj.weight'] = param.t() - else: - weights[f'{prefix}.mlp.proj.bias'] = params[0] - - elif 'input_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.input_layernorm.weight'] = params[0] - else: - weights[f'{prefix}.input_layernorm.bias'] = params[0] - elif 'post_attention_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.post_layernorm.weight'] = params[0] - else: - weights[f'{prefix}.post_layernorm.bias'] = params[0] - - elif 'final_layernorm' in name: - if name.endswith('weight'): - weights['transformer.ln_f.weight'] = params[0] - else: - weights['transformer.ln_f.bias'] = params[0] - for model in models: - del model[name] - del models - - for key in list(model_level_weights.keys()): - weights[key] = torch.concat(model_level_weights[key], dim=0) - weights[key] = torch.chunk(weights[key], split_factor, - dim=0)[inference_tp_rank % split_factor] - del model_level_weights[key] - for key, param in weights.items(): - weights[key] = weights[key].to(dtype).contiguous() - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/gpt/model.py b/tensorrt_llm/models/gpt/model.py deleted file mode 100644 index 89267a90f71e..000000000000 --- a/tensorrt_llm/models/gpt/model.py +++ /dev/null @@ -1,417 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import (Tensor, is_gated_activation, non_gated_version, recv, - send) -from ...layers import (MLP, MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, GatedMLP, LayerNorm, MoeConfig, - PositionEmbeddingType) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ...quantization import QuantMode -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import GPTConfig -from .convert import (load_hf_gpt, load_weights_from_hf_model, - load_weights_from_nemo) - - -def MLPFactory(hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - moe_config: MoeConfig = MoeConfig(), - tp_group=None, - tp_size=1, - mapping=Mapping(), - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05): - if moe_config.has_moe(): - return MOE(moe_config, - hidden_size, - ffn_hidden_size, - hidden_act, - mapping=mapping, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - MLPClass = GatedMLP if is_gated_activation(hidden_act) else MLP - hidden_act = non_gated_version(hidden_act) - return MLPClass( - hidden_size, - ffn_hidden_size, - hidden_act, - bias, - dtype, - tp_group, - tp_size, - quant_mode, - inner_layernorm=inner_layernorm, - eps=eps, - ) - - -class GPTDecoderLayer(Module): - - def __init__(self, config: GPTConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - - self.input_layernorm = LayerNorm(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] - inner_layernorm = config.inner_layernorm if hasattr( - config, "inner_layernorm") else False - attention_head_size = config.head_size if hasattr(config, - "head_size") else None - 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, - max_position_embeddings=config.max_position_embeddings, - num_layers=config.num_hidden_layers, - q_scaling=config.q_scaling, - apply_query_key_layer_scaling=config.apply_query_key_layer_scaling, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - attention_head_size=attention_head_size, - position_embedding_type=config.position_embedding_type, - rotary_embedding_percentage=config.rotary_pct, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - bias=config.bias, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=config.quant_mode, - qk_layernorm=config.qk_layernorm, - inner_layernorm=inner_layernorm, - eps=config.norm_epsilon) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - self.norm_before_bmm1 = config.norm_before_bmm1 if hasattr( - config, "norm_before_bmm1") else False - - self.mlp = MLPFactory(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.bias, - moe_config=config.moe, - tp_group=tp_group, - tp_size=tp_size, - mapping=config.mapping, - quant_mode=config.quant_mode, - inner_layernorm=inner_layernorm, - eps=config.norm_epsilon) - - self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - spec_decoding_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - norm_before_bmm1=self.norm_before_bmm1) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GPTModel(Module): - - def __init__(self, config: GPTConfig): - super().__init__() - self.mapping = config.mapping - self.position_embedding_type = config.position_embedding_type - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.embedding_scale = config.embedding_scale - - if config.position_embedding_type == PositionEmbeddingType.learned_absolute: - self.position_embedding = Embedding( - num_embeddings=config.max_position_embeddings, - embedding_dim=config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(GPTDecoderLayer, config) - - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - lora_params=None, - spec_decoding_params=None): - if self.mapping.is_first_pp_rank(): - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - if self.embedding_scale is not None: - hidden_states *= self.embedding_scale - if self.position_embedding_type == PositionEmbeddingType.learned_absolute: - hidden_states = hidden_states + self.position_embedding( - position_ids) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - 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, - spec_decoding_params=spec_decoding_params) - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GPTForCausalLM(DecoderModelForCausalLM): - config_class = GPTConfig - - def __init__(self, config: GPTConfig): - transformer = GPTModel(config) - - if config.mapping.is_last_pp_rank(): - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - lm_head = ColumnLinear(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) - else: - lm_head = None - self.trtllm_modules_to_hf_modules = { - "attn_q": "q_proj", - "attn_k": "k_proj", - "attn_v": "v_proj", - "attn_dense": "o_proj", - "mlp_h_to_4h": "c_fc", - "mlp_4h_to_h": "c_proj", - } - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - import transformers - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - is_prequantized_to_fp8 = kwargs.pop('is_prequantized_to_fp8', False) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = GPTConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if is_prequantized_to_fp8: - custom_dict = {'fc': 'up_proj'} - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - - # This is to account for all layernorms in nemotron variants being NemotronLayerNorm-1P. - def apply_layernorm_1p(weights): - return weights + 1.0 - - loader.update_key_mapping(model) - tllm_weights = {} - for tllm_key, _ in model.named_parameters(): - if config.gpt_variant == "nemotron" and ( - 'layernorm.weight' in tllm_key - or 'ln_f.weight' in tllm_key): - tllm_weights.update( - loader.load(tllm_key, - preprocess=apply_layernorm_1p, - custom_postprocess_kwargs={})) - else: - tllm_weights.update( - loader.load(tllm_key, custom_postprocess_kwargs={})) - loader.fill(tllm_weights) - else: - if not use_preloading: - hf_model = load_hf_gpt(hf_model_dir, load_model_on_cpu) - weights = load_weights_from_hf_model(hf_model, config) - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = GPTConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - convert.quantize(hf_model_dir, - output_dir, - config=config, - device=device, - calib_dataset=calib_dataset) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - @classmethod - def from_nemo(cls, - nemo_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - config = GPTConfig.from_nemo(nemo_ckpt_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - weights = load_weights_from_nemo(nemo_ckpt_dir, config, **kwargs) - - model = cls(config) - model.load(weights) - return model - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/gptj/__init__.py b/tensorrt_llm/models/gptj/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gptj/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/gptj/config.py b/tensorrt_llm/models/gptj/config.py deleted file mode 100644 index e4a7d3d99e9b..000000000000 --- a/tensorrt_llm/models/gptj/config.py +++ /dev/null @@ -1,55 +0,0 @@ -from typing import Mapping, Optional, Union - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class GPTJConfig(PretrainedConfig): - """ - This is the configuration class to store the configuration of GPTJ model. - """ - - def __init__(self, *, rotary_dim: int = 64, **kwargs): - self.rotary_dim = rotary_dim - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - output.update(rotary_dim=self.rotary_dim) - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - 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, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gptj', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act='gelu', - rotary_dim=hf_config.rotary_dim, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/gptj/convert.py b/tensorrt_llm/models/gptj/convert.py deleted file mode 100644 index 1f10b3b260b0..000000000000 --- a/tensorrt_llm/models/gptj/convert.py +++ /dev/null @@ -1,170 +0,0 @@ -import time -from typing import Dict, Optional - -import torch - -from tensorrt_llm.quantization import QuantAlgo - -from ..convert_utils import get_weight, get_weight_and_bias, split_matrix_tp -from .config import GPTJConfig - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def get_tllm_param( - param: torch.Tensor, - name: str, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if name.endswith('.weight') and use_weight_only: - v = param.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[name] = processed_torch_weights - results[name.replace('weight', - 'per_channel_scale')] = torch_weight_scales - else: - results[name] = param - - return results - - -def load_weights_from_hf_model(hf_model, config: GPTJConfig): - 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 - - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_hidden_layers = config.num_hidden_layers - mapping = config.mapping - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f'transformer.h.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # Attention QKV (no bias) - q_weight = get_weight(model_params, f'{prefix}.attn.q_proj', dtype) - k_weight = get_weight(model_params, f'{prefix}.attn.k_proj', dtype) - v_weight = get_weight(model_params, f'{prefix}.attn.v_proj', dtype) - q_w = split_matrix_tp(q_weight, mapping.tp_size, mapping.tp_rank, dim=0) - k_w = split_matrix_tp(k_weight, mapping.tp_size, mapping.tp_rank, dim=0) - v_w = split_matrix_tp(v_weight, mapping.tp_size, mapping.tp_rank, dim=0) - qkv_w = torch.concatenate([q_w, k_w, v_w], dim=0) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Attention dense (not bias) - attn_dense_weight = get_weight(model_params, f'{prefix}.attn.out_proj', - dtype) - attn_dense_w = split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, f'{tllm_prex}.attention.dense', - None, use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_in (with bias) - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.fc_in', dtype) - mlp_fc_w = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - mlp_fc_b = split_matrix_tp(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', mlp_fc_b, - use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_out (with bias) - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.fc_out', dtype) - mlp_proj_w = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - # Only rank0 will get bias - if mapping.tp_size > 1 and mapping.tp_rank > 0: - mlp_proj_bias = torch.zeros(mlp_proj_weight.shape[0], - dtype=mlp_proj_weight.dtype) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.ln_1', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - if mapping.is_first_pp_rank(): - # Embedding - embed_w = get_weight(model_params, 'transformer.wte', dtype) - if config.use_parallel_embedding: - embed_w = split_matrix_tp(embed_w, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights['transformer.vocab_embedding.weight'] = embed_w - - if mapping.is_last_pp_rank(): - # lm_head weight and bias - lm_head_w, ln_head_bias = get_weight_and_bias(model_params, 'lm_head', - dtype) - weights['lm_head.weight'] = split_matrix_tp(lm_head_w, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights['lm_head.bias'] = split_matrix_tp(ln_head_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f', - dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/gptj/model.py b/tensorrt_llm/models/gptj/model.py deleted file mode 100644 index ff2d1deab24c..000000000000 --- a/tensorrt_llm/models/gptj/model.py +++ /dev/null @@ -1,202 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor, allreduce -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import DecoderLayerList, DecoderModelForCausalLM -from .config import GPTJConfig -from .convert import load_weights_from_hf_model - - -class GPTJDecoderLayer(Module): - - def __init__(self, config: GPTJConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - num_attention_heads = config.num_attention_heads - rotary_dim = config.rotary_dim - dtype = config.dtype - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - layernorm_epsilon = config.norm_epsilon - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - rotary_embedding_percentage=rotary_dim / - (hidden_size // num_attention_heads), - max_position_embeddings=config.max_position_embeddings, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=None, - tp_size=tp_size, - tp_rank=tp_rank, - bias=False, - position_embedding_type=PositionEmbeddingType.rope_gptj, - quant_mode=config.quant_mode) - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=config.hidden_act, - dtype=dtype, - bias=True, - tp_group=None, - tp_size=tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - attention_output = attention_output - - feed_forward_hidden_states = self.mlp(hidden_states) - hidden_states = attention_output + feed_forward_hidden_states - if self.config.mapping.tp_size > 1: - hidden_states = allreduce(hidden_states, - self.config.mapping.tp_group) - hidden_states = hidden_states + residual - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GPTJModel(Module): - - def __init__(self, config: GPTJConfig): - super().__init__() - self.config = config - - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(GPTJDecoderLayer, config) - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - 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): - - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_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 GPTJForCausalLM(DecoderModelForCausalLM): - config_class = GPTJConfig - - def __init__(self, config: GPTJConfig): - transformer = GPTJModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=True, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config=None, - **kwargs): - import transformers - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = GPTJConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype='auto', trust_remote_code=trust_remote_code) - weights = load_weights_from_hf_model(hf_model, config) - - model = GPTJForCausalLM(config) - model.load(weights) - return model diff --git a/tensorrt_llm/models/gptneox/__init__.py b/tensorrt_llm/models/gptneox/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gptneox/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/gptneox/model.py b/tensorrt_llm/models/gptneox/model.py deleted file mode 100644 index 0ac5d356336c..000000000000 --- a/tensorrt_llm/models/gptneox/model.py +++ /dev/null @@ -1,147 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class GPTNeoXDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - self.post_attention_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - rotary_embedding_percentage=config.rotary_pct, - rotary_embedding_base=config.rotary_emb_base, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - max_position_embeddings=config.max_position_embeddings, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=config.hidden_act, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - - input_layernorm_output = self.input_layernorm(hidden_states) - post_attention_layernorm_output = self.post_attention_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=True) - - if use_cache: - attention_output, presents = attention_output - - feed_forward_hidden_states = self.mlp(post_attention_layernorm_output) - hidden_states = attention_output + feed_forward_hidden_states + residual - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GPTNeoXModel(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(GPTNeoXDecoderLayer, config) - - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - 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): - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_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 GPTNeoXForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - transformer = GPTNeoXModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(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) - super().__init__(config, transformer, lm_head) diff --git a/tensorrt_llm/models/grok/__init__.py b/tensorrt_llm/models/grok/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/grok/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/grok/convert.py b/tensorrt_llm/models/grok/convert.py deleted file mode 100644 index c0ef4b630b9f..000000000000 --- a/tensorrt_llm/models/grok/convert.py +++ /dev/null @@ -1,518 +0,0 @@ -# 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 time -from pathlib import Path -from typing import Optional - -import jax -import numpy as np -import torch -from jax import dlpack as jax_dlpack -from torch.utils import dlpack as torch_dlpack - -from ..._utils import pad_vocab_size, release_gc -from ...layers import MoeConfig -from ...quantization import QuantAlgo -from ..convert_utils import split -from ..modeling_utils import PretrainedConfig, QuantConfig, optimize_model - - -def get_jax_weight(config, prefix, dtype, postfix='.weight', key_name='scale'): - - return torch.as_tensor((config[prefix + postfix][key_name])._value, - dtype=dtype).T - - -def get_jax_weight_scale_tp(params, key, rank): - jax_obj = params[key]['w'] - jax_scales = jax.device_put(jax_obj.scales, device=jax.devices('gpu')[rank]) - torch_scales = torch_dlpack.from_dlpack(jax_dlpack.to_dlpack(jax_scales)) - return torch.as_tensor( - np.asarray(jax_obj.weight.addressable_shards[rank].data)), torch_scales - - -def get_jax_weight_scale(params, key): - jax_obj = params[key]['w'] - - jax_scales = jax.device_put(jax_obj.scales, device=jax.devices('cpu')[0]) - - torch_scales = torch_dlpack.from_dlpack( - jax_dlpack.to_dlpack(jax_scales, copy=False)) - return torch.as_tensor(np.asarray(jax_obj.weight), - dtype=torch.int8), torch_scales - - -def get_tllm_linear_weight( - weight, - torch_weight_scales, - prefix, - plugin_weight_only_quant_type=torch.int8, - postfix='weight', -): - results = {} - processed_weight = torch.ops.trtllm.preprocess_weights_for_mixed_gemm( - weight if weight.is_contiguous() else weight.contiguous(), - plugin_weight_only_quant_type, torch.bfloat16) - results[prefix + postfix] = processed_weight - - results[prefix + 'per_channel_scale'] = torch_weight_scales.contiguous() - - return results - - -def convert_grok(hf_model, - config, - mapping, - vocab_size=32000, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - use_gemm_woq_plugin=False, - plugin_weight_only_quant_type=torch.int8, - moe_config=None): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = hf_model - dtype = getattr(torch, dtype) - - config['num_attention_heads'] - config['hidden_size'] - - layers_range = mapping.pp_layers(config['num_hidden_layers']) - - def convert_layer(l): - prefix = f'transformer/decoder_layer_{l}/' - print(prefix) - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - - wq, q_scale = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/query', - mapping.tp_rank) - wk, k_scale = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/key', mapping.tp_rank) - wv, v_scale = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/value', - mapping.tp_rank) - - qs = split(q_scale, mapping.tp_size, mapping.tp_rank, dim=1) - ks = split(k_scale, mapping.tp_size, mapping.tp_rank, dim=1) - vs = split(v_scale, mapping.tp_size, mapping.tp_rank, dim=1) - split_v = torch.concat((wq, wk, wv), dim=1) - scale_v = torch.concat((qs, ks, vs), dim=1) - - weights.update( - get_tllm_linear_weight(split_v, scale_v.squeeze(), - tllm_prex + 'attention.qkv.', - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_scales = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/linear', - mapping.tp_rank) - - split_scales = split(attn_dense_scales, - tensor_parallel, - mapping.tp_rank, - dim=0) - - weights.update( - get_tllm_linear_weight(attn_dense_weight, split_scales.squeeze(), - tllm_prex + 'attention.dense.', - plugin_weight_only_quant_type)) - if mapping.moe_ep_size > 1: - w3, s3 = get_jax_weight_scale( - model_params, f'transformer/decoder_layer_{l}/moe/linear_v') - - w2, s2 = get_jax_weight_scale( - model_params, f'transformer/decoder_layer_{l}/moe/linear_1') - - w1, s1 = get_jax_weight_scale( - model_params, f'transformer/decoder_layer_{l}/moe/linear') - - # moe expert parallel - w3_split = split(w3, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - w2_split = split(w2, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - w1_split = split(w1, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - - s3_split = split(s3, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - s2_split = split(s2, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - s1_split = split(s1, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - - # moe tensor parallel - w3_split = split(w3_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - w2_split = split(w2_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - w1_split = split(w1_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - - s3_split = split(s3_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - s2_split = split(s2_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - s1_split = split(s1_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - else: - w3_split, s3 = get_jax_weight_scale_tp( - model_params, f'transformer/decoder_layer_{l}/moe/linear_v', - mapping.tp_rank) - - w2_split, s2 = get_jax_weight_scale_tp( - model_params, f'transformer/decoder_layer_{l}/moe/linear_1', - mapping.tp_rank) - - w1_split, s1 = get_jax_weight_scale_tp( - model_params, f'transformer/decoder_layer_{l}/moe/linear', - mapping.tp_rank) - - s3_split = split(s3, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - s2_split = split(s2, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - s1_split = split(s1, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - weights.update( - get_tllm_linear_weight(w2_split, - s2_split.reshape(moe_config.num_experts, -1), - tllm_prex + 'mlp.proj.', - plugin_weight_only_quant_type)) - - weights.update( - get_tllm_linear_weight( - torch.concat([w3_split, w1_split], dim=-1), - torch.concat([s3_split, s1_split], - dim=-1).reshape(moe_config.num_experts, -1), - tllm_prex + 'mlp.fc.', - plugin_weight_only_quant_type, - )) - - moe_experts_gate_weights = get_jax_weight(model_params, - prefix + 'router', - torch.float32, - postfix='', - key_name='w').contiguous() - - weights[tllm_prex + 'mlp.router.weight'] = moe_experts_gate_weights - - # Layer norms do not use tensor parallelism - input_ln_weight = get_jax_weight(model_params, - prefix + 'rms_norm', - dtype, - postfix='') - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_attn_weight = get_jax_weight(model_params, - prefix + 'rms_norm_1', - dtype, - postfix='') - weights[tllm_prex + 'post_attn_layernorm.weight'] = post_attn_weight - - post_ln_weight = get_jax_weight(model_params, - prefix + 'rms_norm_2', - dtype, - postfix='') - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - post_mlp_weight = get_jax_weight(model_params, - prefix + 'rms_norm_3', - dtype, - postfix='') - weights[tllm_prex + 'post_mlp_layernorm.weight'] = post_mlp_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = get_jax_weight(model_params, - 'language_model/in_out_embed', - dtype, - postfix='', - key_name='embeddings').T - tie_word_embeddings = config['tie_word_embeddings'] - if tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - v = torch.nn.functional.pad(v, (0, pad_width, 0, 0), 'constant', - 0) - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - - if use_parallel_embedding: - v = split(v, mapping.tp_size, mapping.tp_rank, dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - ln_f_w = get_jax_weight(model_params, - 'language_model/rms_norm', - dtype, - postfix='') - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def create_config_from_xai(dtype, - mapping, - quantization: QuantConfig = None, - override_fields: dict = {}): - config = {} - hf_config = { - "architectures": ["Grok1ModelForCausalLM"], - "vocab_size": 131072, - "hidden_size": 6144, - "intermediate_size": 32768, - "num_hidden_layers": 64, - "num_attention_heads": 48, - "num_key_value_heads": 8, - "attn_output_multiplier": 0.08838834764831845, - "embedding_multiplier_scale": 78.38367176906169, - "output_multiplier_scale": 0.5773502691896257, - "max_attn_value": 30.0, - "max_position_embeddings": 8192, - "rms_norm_eps": 1e-5, - "use_cache": True, - "pad_token_id": 0, - "bos_token_id": 1, - "eos_token_id": 2, - "tie_word_embeddings": True, - "num_experts_per_tok": 2, - "num_experts": 8, - "output_router_logits": False, - "router_aux_loss_coef": 0.001, - "torch_dtype": "bfloat16", - "transformers_version": "4.35.0" - } - # same for from_meta and from_cli_args - n_head = hf_config['num_attention_heads'] - inter_size = hf_config['intermediate_size'] - n_layer = hf_config['num_hidden_layers'] - n_embd = hf_config['hidden_size'] - n_kv_head = hf_config['num_key_value_heads'] - rms_norm_eps = hf_config['rms_norm_eps'] - vocab_size = hf_config['vocab_size'] - n_positions = hf_config['max_position_embeddings'] - hidden_act = 'geglu' - config['rotary_scaling'] = None - rotary_base = 10000.0 - - config[ - 'moe_normalization_mode'] = MoeConfig.ExpertScaleNormalizationMode.NONE - - moe_num_experts = hf_config['num_experts'] - - moe_top_k = hf_config['num_experts_per_tok'] - - attn_output_multiplier = hf_config['attn_output_multiplier'] - embedding_multiplier_scale = hf_config['embedding_multiplier_scale'] - - output_multiplier_scale = hf_config['output_multiplier_scale'] - max_attn_value = hf_config['max_attn_value'] - - architecture = hf_config['architectures'][0] - - attn_bias = False - - config.update({ - 'architecture': - architecture, - 'dtype': - dtype, - 'logits_dtype': - '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, - 'moe_num_experts': - moe_num_experts, - 'moe_top_k': - moe_top_k, - 'moe_normalization_mode': - MoeConfig.ExpertScaleNormalizationMode.NONE, - #TODO: should have directly map from the Mapping object to the TRT-LLM checkpoint fields - '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, - }, - 'attn_bias': - attn_bias, - "attn_output_multiplier": - attn_output_multiplier, - "embedding_multiplier_scale": - embedding_multiplier_scale, - "output_multiplier_scale": - output_multiplier_scale, - "max_attn_value": - max_attn_value, - "tie_word_embeddings": - True, - }) - - config['quantization'] = quantization.model_dump() - config.update(override_fields) - - return config - - -def from_hugging_face(cls, - model_dir, - dtype, - *, - mapping, - quantization: QuantConfig = None, - override_fields={}, - skip_loading_weights=False, - preloaded_model=None): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - assert model_dir is not None - if isinstance(model_dir, Path): # some code relies on this as string - model_dir = str(model_dir) - - config = create_config_from_xai(dtype, - mapping, - quantization, - override_fields=override_fields) - - pretrained_config = PretrainedConfig.from_dict(config) - pretrained_config.set_rank(mapping.rank) # TODO:remove this hack - - grok = cls.from_config(pretrained_config) - grok = optimize_model( - grok, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - ) - - if skip_loading_weights: - return grok - - weights = load_weights_from_xai(config=config, - mapping=mapping, - model=preloaded_model) - - grok.load(weights) - return grok - - -def quantize(dtype, - model_dir, - output_dir, - mapping, - quantization: QuantConfig, - *, - override_fields, - dataset_cache_dir: Optional[str] = None): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - pass #The official grok-1 model is published under int8 wo format, we don't need to quantize again. - - -def load_weights_from_xai(*, config, mapping, model): - assert model is not None - plugin_weight_only_quant_type = None # the value does not matter when use_weight_only is False - quant_algo = config['quantization']['quant_algo'] - assert quant_algo == QuantAlgo.W8A16 - plugin_weight_only_quant_type = torch.int8 - - moe_config = MoeConfig( - num_experts=config['moe_num_experts'], - top_k=config['moe_top_k'], - normalization_mode=config['moe_normalization_mode']).validate() - - use_weight_only = quant_algo in [QuantAlgo.W8A16] - - weights = convert_grok( - model, - config, - mapping, - vocab_size=config['vocab_size'], - dtype=config['dtype'], - use_weight_only=use_weight_only, - use_gemm_woq_plugin=not config.get('disable_weight_only_quant_plugin', - False), - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=config.get('use_parallel_embedding', False), - sharding_dim=config.get('embedding_sharding_dim', 0), - moe_config=moe_config) - return weights diff --git a/tensorrt_llm/models/grok/model.py b/tensorrt_llm/models/grok/model.py deleted file mode 100644 index 9ff22cd71c71..000000000000 --- a/tensorrt_llm/models/grok/model.py +++ /dev/null @@ -1,273 +0,0 @@ -# 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 - -from ..._utils import pad_vocab_size -from ...functional import Tensor, recv, send -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, MoeConfig, PositionEmbeddingType, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) - - -class GrokDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_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] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode, - max_attn_value=config.max_attn_value) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - self.post_attn_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - self.post_mlp_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - mlp_kwargs = {} - assert config.moe_num_experts > 1, "Grok model is a MoE model." - ClsMLP = MOE - moe_config = MoeConfig( - num_experts=config.moe_num_experts, - top_k=config.moe_top_k, - normalization_mode=config.moe_normalization_mode).validate() - mlp_kwargs = { - "moe_config": moe_config, - "mapping": config.mapping, - } - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params) - - if use_cache: - attention_output, presents = attention_output - - attention_output = self.post_attn_layernorm(attention_output) - hidden_states = residual + attention_output - - residual_attn = hidden_states - - # regular llama/mixtral layers - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - hidden_states = self.post_mlp_layernorm(hidden_states) - hidden_states = residual_attn + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GrokModel(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(GrokDecoderLayer, config) - - self.embedding_multiplier_scale = config.embedding_multiplier_scale - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - hidden_states *= 78.38367176906169 - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GrokForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = GrokModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(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) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('mlp_bias', False) - config.set_if_not_exist('attn_bias', False) - config.set_if_not_exist('rotary_base', 10000.0) - config.set_if_not_exist('rotary_scaling', None) - config.set_if_not_exist('moe_num_experts', 0) - config.set_if_not_exist('moe_top_k', 0) - config.set_if_not_exist('moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.NONE) - - @classmethod - def from_hugging_face(cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None, - **kwargs): - from . import convert - if mapping is None: - mapping = Mapping() - grok = convert.from_hugging_face( - cls, - hf_model_dir, - dtype, - mapping=mapping, - quantization=kwargs.get('quantization', QuantConfig()), - override_fields=kwargs.get('override_fields', {}), - skip_loading_weights=kwargs.get('skip_loading_weights', False), - preloaded_model=kwargs.get('preloaded_model', None)) - return grok - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.set_weight_only_groupwise_quant_matmul_plugin() - return plugin_config - - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantConfig, - *, - dtype='float16', - mapping: Optional[Mapping] = None, - calib_batches=512, - calib_batch_size=1, - random_seed=1234, - tokenizer_max_seq_length=2048, - **kwargs, - ): - pass - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config) diff --git a/tensorrt_llm/models/grok/weight.py b/tensorrt_llm/models/grok/weight.py deleted file mode 100644 index 0c412305a388..000000000000 --- a/tensorrt_llm/models/grok/weight.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 Union - -import numpy as np -import torch - - -def split(v: Union[np.ndarray, torch.Tensor], - tp_size: int, - tp_rank: int, - dim=0): - if tp_size == 1: - return v - assert len(v.shape) > 1 or dim == 0 - if isinstance(v, np.ndarray): - return np.ascontiguousarray( - np.split(v, tp_size, axis=dim)[tp_rank].copy()) - else: - assert v.shape[dim] % tp_size == 0, \ - 'Unable to split: shape={v.shape} (dim={dim}) tp_size={tp_size}.' - split_size = v.shape[dim] // tp_size - return v.split(split_size, dim=dim)[tp_rank].clone().detach() diff --git a/tensorrt_llm/models/llama/__init__.py b/tensorrt_llm/models/llama/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/llama/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/llama/config.py b/tensorrt_llm/models/llama/config.py deleted file mode 100644 index 54038e32c4ff..000000000000 --- a/tensorrt_llm/models/llama/config.py +++ /dev/null @@ -1,280 +0,0 @@ -# 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 json -import math -import sys -from pathlib import Path -from typing import Optional, Union - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class LLaMAConfig(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, - embedding_multiplier: float = 1.0, - attention_multiplier: float = 1.0, - residual_multiplier: float = 1.0, - output_multiplier_scale: float = 1.0, - **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 moe is None: - # Legacy MOE config fields - moe = MoeConfig( - num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE)) - elif 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 - self.embedding_multiplier = embedding_multiplier - self.attention_multiplier = attention_multiplier - self.residual_multiplier = residual_multiplier - self.output_multiplier_scale = output_multiplier_scale - self.has_partial_lora_mask = False - - 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, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - trust_remote_code = kwargs.pop('trust_remote_code', True) - has_partial_lora_mask = False - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - if "vila" in hf_config_dir.lower(): - sys.path.append(hf_config_dir + "/../VILA") - from llava.model import LlavaLlamaConfig # noqa - from llava.model import LlavaLlamaModel - transformers.AutoConfig.register("llava_llama", - LlavaLlamaConfig, - exist_ok=True) - transformers.AutoModelForCausalLM.register( - LlavaLlamaConfig, LlavaLlamaModel) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - if hf_config.model_type == "llava": - # LLaVA = Vision model + Llama LLM - # We load a llava config and use its' text config as llama config - from transformers import LlavaConfig - hf_config = LlavaConfig.from_pretrained( - hf_config_dir).text_config - if hf_config.model_type == "llava_next": - from transformers import LlavaNextConfig - hf_config = LlavaNextConfig.from_pretrained( - hf_config_dir).text_config - if hf_config.model_type == "llava_llama": - hf_config.llm_cfg["architecture"] = hf_config.llm_cfg[ - "architectures"][0] - hf_config.llm_cfg["dtype"] = hf_config.llm_cfg["torch_dtype"] - hf_config = PretrainedConfig.from_dict(hf_config.llm_cfg) - if hf_config.model_type == 'internlmxcomposer2': - # InternLM-XComposer2 has a mask for partial lora - # Therefore we need an additional flag for this mask - has_partial_lora_mask = True - if hf_config.model_type == 'mistral3': - from transformers import Mistral3Config - hf_config = Mistral3Config.from_pretrained( - hf_config_dir).text_config - hf_config.architectures = ["MistralForCausalLM"] - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - if hf_config.model_type == "exaone": - hidden_act = hf_config.activation_function - # NOTE - # EXAONE also uses RMS norm but they represent as layer_norm_epsilon. - norm_epsilon = getattr(hf_config, "layer_norm_epsilon", 1e-5) - else: - hidden_act = hf_config.hidden_act - norm_epsilon = hf_config.rms_norm_eps - head_dim = getattr( - hf_config, "head_dim", - hf_config.hidden_size // hf_config.num_attention_heads) - head_size = getattr(hf_config, "kv_channels", head_dim) - attn_bias = getattr(hf_config, 'bias', False) or getattr( - hf_config, 'attention_bias', False) - rotary_scaling = getattr(hf_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - residual_mlp = getattr(hf_config, "parallel_attn_mlp_res", False) - disable_weight_only_quant_plugin = kwargs.pop( - 'disable_weight_only_quant_plugin', False) - remove_duplicated_kv_heads = kwargs.pop('remove_duplicated_kv_heads', - False) - embedding_multiplier = getattr(hf_config, "embedding_multiplier", 1.0) - attention_multiplier = getattr(hf_config, "attention_multiplier", 1.0) - if attention_multiplier != 1.0: - attention_multiplier *= math.sqrt(head_size) - residual_multiplier = getattr(hf_config, "residual_multiplier", 1.0) - output_multiplier_scale = 1.0 / getattr(hf_config, "logits_scaling", - 1.0) - if hf_config.model_type in ["mixtral", "arctic", "granitemoe"]: - # HF LLaMA-type models are implicitly using gated activation. - # With our MoE implementation, we must make it explicit - hidden_act = "swiglu" - moe_normalization_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE - else: - moe_normalization_mode = None - moe_num_experts = getattr(hf_config, "num_local_experts", 0) - moe_top_k = getattr(hf_config, "num_experts_per_tok", 0) - moe_config = MoeConfig(num_experts=moe_num_experts, - top_k=moe_top_k, - normalization_mode=moe_normalization_mode) - moe_config.validate() - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - tie_word_embeddings = getattr(hf_config, 'tie_word_embeddings', False) - - 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=num_key_value_heads, - head_size=head_size, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act=hidden_act, - norm_epsilon=norm_epsilon, - attn_bias=attn_bias, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - residual_mlp=residual_mlp, - disable_weight_only_quant_plugin=disable_weight_only_quant_plugin, - moe=moe_config, - mapping=mapping, - quantization=quant_config, - has_partial_lora_mask=has_partial_lora_mask, - remove_duplicated_kv_heads=remove_duplicated_kv_heads, - tie_word_embeddings=tie_word_embeddings, - embedding_multiplier=embedding_multiplier, - attention_multiplier=attention_multiplier, - residual_multiplier=residual_multiplier, - output_multiplier_scale=output_multiplier_scale, - **kwargs) - - @classmethod - def from_meta_ckpt(cls, - meta_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - - with open(Path(meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - - n_embd = meta_config["dim"] - n_head = meta_config["n_heads"] - n_kv_head = meta_config.get("n_kv_heads", n_head) - vocab_size = meta_config.get("vocab_size", 32000) - - # Reset vocab_size to 32000 for LLama v2 checkpoint. - if vocab_size == -1: - vocab_size = 32000 - - if "hidden_dim" in meta_config: - inter_size = meta_config["hidden_dim"] - else: - multiple_of = meta_config.get("multiple_of", 1) - n_embd_ = int(4 * n_embd * 2 / 3) - ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - inter_size = multiple_of * ( - (int(n_embd_ * ffn_dim_multiplier) + multiple_of - 1) // - multiple_of) - - dtype = infer_dtype(dtype, 'bfloat16') - - if meta_config.get('use_scaled_rope'): - rotary_scaling = {"type": "llama3"} - else: - rotary_scaling = meta_config.get("rope_scaling") - - # meta checkpoint don't have vocab_size|hidden_act|rotary_base specified, use same default value as HF - return cls(architecture="LlamaForCausalLM", - dtype=dtype, - num_hidden_layers=meta_config["n_layers"], - 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=2048, - hidden_act='silu', - rotary_scaling=rotary_scaling, - rotary_base=meta_config.get('rope_theta', 10000), - norm_epsilon=meta_config["norm_eps"], - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/llama/convert.py b/tensorrt_llm/models/llama/convert.py deleted file mode 100644 index c36526f8283c..000000000000 --- a/tensorrt_llm/models/llama/convert.py +++ /dev/null @@ -1,2449 +0,0 @@ -# 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 copy -import functools -import os -import sys -import time -from collections import defaultdict -from pathlib import Path -from typing import List, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer -from transformers.models.llama.modeling_llama import LlamaDecoderLayer -from transformers.pytorch_utils import Conv1D - -from ..._utils import pad_vocab_size, release_gc, str_dtype_to_torch -from ...logger import logger -from ...quantization import QuantAlgo -from ...quantization.quantize import (qserve_pack_reorder_per_channel, - qserve_pack_reorder_per_group, - qserve_quantize_weight_per_channel, - qserve_quantize_weight_per_group) -from ..convert_utils import (dup_kv_weight, generate_int8, - get_tllm_linear_weight, iterate_shard_files, - load_calib_dataset, load_state_dict, - retrieved_layer_index_from_name, smooth_gemm, - smooth_gemm_fc1_gate, split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) -from ..modeling_utils import PretrainedConfig -from .config import LLaMAConfig - - -@torch.no_grad() -def smooth_llama_model(model, scales, alpha, llama_qkv_para, llama_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance( - module, - LlamaDecoderLayer) and module.__class__.__name__ not in [ - "InternLMDecoderLayer", "MistralDecoderLayer" - ]: - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - llama_qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - if hasattr(module, 'residual_mlp'): - fc1_layer_name = name + ".residual_mlp.w1" - gate_layer_name = name + ".residual_mlp.w3" - - smoother = smooth_gemm_fc1_gate(module.residual_mlp.w1.weight, - module.residual_mlp.w3.weight, - scales[fc1_layer_name]["x"], - module.residual_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.residual_mlp.w1.weight.abs( - ).max(dim=1)[0] - - scales[gate_layer_name][ - "x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.residual_mlp.w3.weight.abs( - ).max(dim=1)[0] - - # ================================================================== - layer_name = name + ".residual_mlp.w2" - smoother = smooth_gemm(module.residual_mlp.w2.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.residual_mlp.w2.weight.abs().max( - dim=1)[0] - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - input_ids = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(input_ids) - for h in hooks: - h.remove() - return act_scales - - -def get_weight(named_params, prefix, dtype): - if named_params[prefix + '.weight'].dtype != dtype: - named_params[prefix + - '.weight'].data = named_params[prefix + - '.weight'].to(dtype) - return named_params[prefix + '.weight'].detach() - - -def get_weight_and_scale(named_params, - prefix, - dtype, - mapping=None, - split_scale=False): - if prefix + '.weight_scale' not in named_params: - return get_weight(named_params, prefix, dtype), None - else: - assert named_params[prefix + '.weight'].dtype == torch.float8_e4m3fn - assert named_params[prefix + '.weight_scale'].dtype == torch.float32 - weight_scale = named_params[prefix + '.weight_scale'].detach() - if split_scale: - weight_scale = split(weight_scale, - mapping.tp_size, - mapping.tp_rank, - dim=0) - return named_params[prefix + - '.weight'].detach(), weight_scale.reshape(-1) - - -def get_bias(named_params, prefix, dtype): - if named_params[prefix + '.bias'].dtype != dtype: - named_params[prefix + '.bias'].data = named_params[prefix + - '.bias'].to(dtype) - return named_params[prefix + '.bias'].detach() - - -def get_weight_and_bias(named_params, prefix, dtype): - return get_weight(named_params, prefix, - dtype), get_bias(named_params, prefix, dtype) - - -def fp8_per_channel_quant_weight_gpu(weight, clamp_val, rank=0): - weight = weight.to("cuda:" + str(rank)) - # activation range bound. - x = weight.to(torch.float32).clamp(clamp_val[0], clamp_val[1]) - xmax = x.abs().max(-1, keepdim=True).values - # minimum scaling factor. - torch_weight_scales = (xmax / 448.0).clamp(min=1.0 / (448.0 * 512.0)) - out = x / torch_weight_scales - torch_weight_scales = torch_weight_scales.reshape(-1) - out = torch.clamp(out, -448, 448) - processed_torch_weights = out.to(torch.float8_e4m3fn) - - processed_torch_weights = processed_torch_weights.to( - torch.float8_e4m3fn).cpu() - torch_weight_scales = torch_weight_scales.cpu() - - return processed_torch_weights, torch_weight_scales - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.chunk(q, tp_size, dim=-1) - k_split = torch.chunk(k, tp_size, dim=-1) - v_split = torch.chunk(v, tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = torch.Tensor(vals["weight.int8.col"]) - else: - original_weights = torch.Tensor(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.Tensor([1.0]).to(torch.float32) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_w_quant_orig.col"], - tensor_parallel, - dim=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_w_quant_orig"], - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - else: - if per_channel: - original_weights = torch.Tensor(vals["weight.int8.col"]) - else: - original_weights = torch.Tensor(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_y_accum_quant.col"], - tensor_parallel, - dim=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_y_accum_quant"], - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + - 'per_channel_scale'] = torch.Tensor(cur_per_channel_value).to( - torch.float32).reshape(col_shape).contiguous() - results[prefix + 'act_scale'] = torch.Tensor( - [[vals['scale_y_quant_orig']]]).to(torch.float32).contiguous() - results[last_prefix] = torch.Tensor([vals['scale_x_orig_quant'] - ]).to(torch.float32).contiguous() - - if smoother_value is not None: - cur_smoother_value = torch.chunk(smoother_value, - tensor_parallel, - dim=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def get_prefix_and_param_name_map(architecture, use_safetensors=False): - - key_postfix = "" - if use_safetensors: - key_postfix = ".weight" - - architecture = architecture.lower() - if "exaone" in architecture: - model_prefix = "transformer" - param_name_map = { - "vocab_embedding": "wte" + key_postfix, # vocab_embedding - "lm_head": "lm_head" + key_postfix, # lm_head - "ln_f": "ln_f" + key_postfix, # ln_f - "attention.qkv": "attn.attention", # attention.qkv - "qkv_suffix": "_proj" + key_postfix, # qkv_suffix - "attention.dense": - "attn.attention.out_proj" + key_postfix, # attention.dense - "mlp.gate": "mlp.c_fc_1" + key_postfix, # mlp.gate - "mlp.proj": "mlp.c_proj" + key_postfix, # mlp.proj - "mlp.fc": "mlp.c_fc_0" + key_postfix, # mlp.fc - "input_layernorm": "ln_1" + key_postfix, # input_layernorm - "post_layernorm": "ln_2" + key_postfix, # post_layernorm - } - layer_prefix = 'h' - else: # LLaMA - model_prefix = "model" - param_name_map = { - "vocab_embedding": "embed_tokens" + key_postfix, # vocab_embedding - "lm_head": "lm_head" + key_postfix, # lm_head - "ln_f": "norm" + key_postfix, # ln_f - "attention.qkv": "self_attn", # attention.qkv - "qkv_suffix": "_proj" + key_postfix, # qkv suffix - "attention.dense": - "self_attn.o_proj" + key_postfix, # attention.dense - "mlp.gate": "mlp.up_proj" + key_postfix, # mlp.gate - "mlp.proj": "mlp.down_proj" + key_postfix, # mlp.proj - "mlp.fc": "mlp.gate_proj" + key_postfix, # mlp.fc - "input_layernorm": - "input_layernorm" + key_postfix, # input_layernorm - "post_layernorm": - "post_attention_layernorm" + key_postfix, # post_layernorm - } - layer_prefix = 'layers' - - return model_prefix, layer_prefix, param_name_map - - -def load_hf_llama(model_dir: str, load_model_on_cpu: bool = False): - if "vila" in model_dir: - sys.path.append(model_dir + "/../VILA") - from llava.model import LlavaLlamaConfig, LlavaLlamaModel # noqa - from transformers import AutoModel - model = AutoModel.from_pretrained( - model_dir, - device_map='auto', - trust_remote_code=True, - ) - return model.llm - - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - model_cls = AutoModelForCausalLM - if hf_config.model_type == "llava": - from transformers import LlavaForConditionalGeneration - model_cls = LlavaForConditionalGeneration - if hf_config.model_type == "llava_next": - from transformers import LlavaNextForConditionalGeneration - model_cls = LlavaNextForConditionalGeneration - model = model_cls.from_pretrained( - model_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype='auto', - trust_remote_code=True, - ) - if hf_config.model_type in ["llava", "llava_next"]: - model = model.language_model - return model - - -def load_weights_from_hf_model(hf_model, - config: LLaMAConfig, - act_range: Optional[dict] = None, - qkv_para: Optional[dict] = None, - smoother: Optional[dict] = None): - 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 = (not config.disable_weight_only_quant_plugin) - use_fp8_rowwise = quant_algo in [QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN] - - use_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - fp8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.FP8 - if use_smooth_quant or int8_kv_cache: - assert act_range is not None - assert qkv_para is not None - assert smoother is not None - - weights = {} - tik = time.time() - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - - mapping = config.mapping - moe_config = config.moe - mha_mode = (config.num_key_value_heads == config.num_attention_heads) - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - exclude_layers_id = [0, config.num_hidden_layers - 1] - - model_prefix, layer_prefix, param_name_map = get_prefix_and_param_name_map( - config.architecture) - - def convert_layer(l): - prefix = f'{model_prefix}.{layer_prefix}.{l}.' - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - q_weight = get_weight( - model_params, prefix + f'{param_name_map["attention.qkv"]}.q_proj', - dtype) - k_weight = get_weight( - model_params, prefix + f'{param_name_map["attention.qkv"]}.k_proj', - dtype) - v_weight = get_weight( - model_params, prefix + f'{param_name_map["attention.qkv"]}.v_proj', - dtype) - - # Meta's recipe of not using fp8 rowwise for the first and last layer. - use_fp8_rowwise_in_layer = use_fp8_rowwise and ( - l not in exclude_layers_id) - - if not mha_mode: - if config.num_key_value_heads < mapping.tp_size: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, config.num_key_value_heads, - mapping.tp_size) - v_weight = dup_kv_weight(v_weight, config.num_key_value_heads, - mapping.tp_size) - assert (k_weight.shape[0] % - (mapping.tp_size * config.head_size)) == 0 - assert (v_weight.shape[0] % - (mapping.tp_size * config.head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - split_v = torch.concat((wq, wk, wv)) - - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, config.num_attention_heads, - config.hidden_size, mapping.tp_size, - mapping.tp_rank) - - if prefix + f'{param_name_map["attention.qkv"]}.q_proj.bias' in model_params: - # only used in Internlm 7B models - q_bias = get_bias( - model_params, - prefix + f'{param_name_map["attention.qkv"]}.q_proj', dtype) - k_bias = get_bias( - model_params, - prefix + f'{param_name_map["attention.qkv"]}.k_proj', dtype) - v_bias = get_bias( - model_params, - prefix + f'{param_name_map["attention.qkv"]}.v_proj', dtype) - qkv_bias = torch.cat((q_bias, k_bias, v_bias)) - split_bias_v = split_qkv_bias_tp(qkv_bias, - config.num_attention_heads, - config.hidden_size, - mapping.tp_size, mapping.tp_rank) - else: - split_bias_v = None - - if use_smooth_quant: - qkv_weight = qkv_para[prefix + - f'{param_name_map["attention.qkv"]}.qkv_proj'] - qkv_out_dim = qkv_weight.shape[1] - - if not mha_mode: - local_dim = qkv_weight.shape[0] - kv_hidden_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(local_dim, - local_dim + 2 * kv_hidden_size) - else: - qkv_weight = qkv_weight.reshape(config.hidden_size, 3, - config.hidden_size) - - int8_weights = generate_int8( - qkv_weight, - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.qkv_proj'), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // mapping.tp_size], - mapping.tp_size, - is_qkv=True, - bias=split_bias_v, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'attention.qkv.', - split_bias_v, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise=False)) - - if int8_kv_cache: - qkv_y = torch.cat([ - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.q_proj')["y"], - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.k_proj')["y"], - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.v_proj')["y"] - ], - dim=0) - - int8_kv_scales = qkv_y.max() / 127. - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - weights.update(kv_cache_weights) - elif fp8_kv_cache: - # FIXME: set it to 1.0f for fp8 kv cache. - weights[tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.tensor( - [1.0], dtype=torch.float32) - - attn_dense_weight = get_weight( - model_params, prefix + param_name_map["attention.dense"], dtype) - split_v = split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - if prefix + f'{param_name_map["attention.dense"]}.bias' in model_params: - attn_dense_bias = get_bias( - model_params, prefix + param_name_map["attention.dense"], dtype) - else: - attn_dense_bias = None - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - proj_out_dim = attn_dense_weight.shape[0] - - int8_weights = generate_int8( - attn_dense_weight, - act_range.get(prefix + param_name_map["attention.dense"])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, config.hidden_size], - mapping.tp_size, - is_qkv=False, - bias=attn_dense_bias, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[( - prefix + param_name_map["attention.dense"])], - smoother_shape=[1, proj_out_dim // mapping.tp_size], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'attention.dense.', - attn_dense_bias, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise=False)) - - if moe_config.has_moe(): - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - architecture = config.architecture.lower() - if "granite" not in architecture: - for suffix in ["w1", "w2", "w3"]: - model_params[f'model.layers.{l}.block_sparse_moe.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.block_sparse_moe.experts.{expert}.{suffix}.weight'].detach() - for expert in rank_experts]) - w3 = model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w3.weight'] - w2 = model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w2.weight'] - w1 = model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w1.weight'] - else: - w2 = model_params[ - f'model.layers.{l}.block_sparse_moe.output_linear.weight'] - half_size = model_params[ - f'model.layers.{l}.block_sparse_moe.input_linear.weight'].shape[ - -2] // 2 - w1, w3 = model_params[ - f'model.layers.{l}.block_sparse_moe.input_linear.weight']\ - .split(half_size, dim=-2) - w1 = w1[rank_experts[0]:rank_experts[-1] + 1] - w2 = w2[rank_experts[0]:rank_experts[-1] + 1] - w3 = w3[rank_experts[0]:rank_experts[-1] + 1] - - if mapping.has_moe_tp(): - w3 = split(w3, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - w2 = split(w2, mapping.moe_tp_size, mapping.moe_tp_rank, dim=2) - w1 = split(w1, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - - model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w3w1.weight'] = torch.concat( - [w3, w1], dim=-2) - - model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w2.weight'] = w2 - - ## block_sparse_moe.experts.w2.weight - moe_experts_w2_weights = get_weight( - model_params, prefix + 'block_sparse_moe.experts.w2', dtype) - weights.update( - get_tllm_linear_weight(moe_experts_w2_weights, - tllm_prex + 'mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - ##block_sparse_moe.experts.w3w1.weight - moe_experts_w3w1_weights = get_weight( - model_params, prefix + 'block_sparse_moe.experts.w3w1', dtype) - weights.update( - get_tllm_linear_weight(moe_experts_w3w1_weights, - tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - if config.residual_mlp: - residual_mlp_gate_weights = get_weight( - model_params, prefix + 'residual_mlp.w3', dtype) - if use_smooth_quant: - residual_mlp_gate_weights = residual_mlp_gate_weights.t() - int8_weights = generate_int8( - residual_mlp_gate_weights, - act_range.get(prefix + 'residual_mlp.w3')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'residual_mlp.gate.', - [1, config.hidden_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - split_v = split_matrix_tp(residual_mlp_gate_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'residual_mlp.gate.', - None, use_weight_only, - plugin_weight_only_quant_type, - dtype, use_gemm_woq_plugin)) - - residual_mlp_fc_weight = get_weight(model_params, - prefix + 'residual_mlp.w1', - dtype) - if use_smooth_quant: - residual_mlp_fc_weight = residual_mlp_fc_weight.t( - ) #verified - int8_weights = generate_int8( - residual_mlp_fc_weight, - act_range.get(prefix + 'residual_mlp.w1')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'residual_mlp.fc.', - [1, config.hidden_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - split_v = split_matrix_tp(residual_mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'residual_mlp.fc.', - None, use_weight_only, - plugin_weight_only_quant_type, - dtype, use_gemm_woq_plugin)) - - residual_mlp_proj_weight = get_weight( - model_params, prefix + 'residual_mlp.w2', dtype) - - if use_smooth_quant: - residual_mlp_proj_weight = residual_mlp_proj_weight.t() - int8_weights = generate_int8( - residual_mlp_proj_weight, - act_range.get(prefix + 'residual_mlp.w2')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'residual_mlp.proj.', - [1, config.hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'residual_mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'residual_mlp.w2'], - smoother_shape=[ - 1, config.hidden_size // mapping.tp_size - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - split_v = split_matrix_tp(residual_mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'residual_mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type, - dtype, use_gemm_woq_plugin)) - - architecture = config.architecture.lower() - if "granite" not in architecture: - moe_experts_gate_weights = get_weight( - model_params, prefix + 'block_sparse_moe.gate', - torch.float32) - else: - moe_experts_gate_weights = get_weight( - model_params, prefix + 'block_sparse_moe.router.layer', - torch.float32) - weights.update( - get_tllm_linear_weight( - moe_experts_gate_weights, - tllm_prex + 'mlp.router.', - None, - False, # Router should never be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - else: - mlp_gate_weight, mlp_gate_weight_scale = get_weight_and_scale( - model_params, prefix + param_name_map["mlp.gate"], dtype, - mapping, True) - split_v = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - if use_smooth_quant: - - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8( - # mlp_gate_weight, act_range.get(prefix + 'mlp.up_proj')) - mlp_gate_weight, - act_range.get(prefix + param_name_map["mlp.gate"])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, config.intermediate_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight( - split_v, - tllm_prex + 'mlp.gate.', - None, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise_in_layer, - weight_scale=mlp_gate_weight_scale, - clamp_value=config.quantization.clamp_val)) - - mlp_fc_weight, mlp_fc_weight_scale = get_weight_and_scale( - model_params, prefix + param_name_map["mlp.fc"], dtype, mapping, - True) - split_v = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() #verified - int8_weights = generate_int8( - mlp_fc_weight, - act_range.get(prefix + param_name_map["mlp.fc"])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, config.intermediate_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight( - split_v, - tllm_prex + 'mlp.fc.', - None, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise_in_layer, - weight_scale=mlp_fc_weight_scale, - clamp_value=config.quantization.clamp_val)) - - mlp_proj_weight, mlp_proj_weight_scale = get_weight_and_scale( - model_params, prefix + param_name_map["mlp.proj"], dtype) - split_v = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, - act_range.get(prefix + param_name_map["mlp.proj"])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, config.hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + - param_name_map["mlp.proj"]], - smoother_shape=[ - 1, config.intermediate_size // mapping.tp_size - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight( - split_v, - tllm_prex + 'mlp.proj.', - None, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise_in_layer, - weight_scale=mlp_proj_weight_scale, - clamp_value=config.quantization.clamp_val)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, - prefix + param_name_map["input_layernorm"], - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + param_name_map["post_layernorm"], - dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - if config.residual_mlp: - residual_ln_weight = get_weight(model_params, - prefix + 'residual_layernorm', - dtype) - weights[tllm_prex + - 'residual_layernorm.weight'] = residual_ln_weight - - cur_block_weights = [ - weight_name for weight_name in model_params - if weight_name.find(prefix) != -1 - ] - for weight_name in cur_block_weights: - model_params[weight_name] = None - - for l in layers_range: - convert_layer(l) - release_gc() - - vocab_embedding = get_weight( - model_params, f'{model_prefix}.{param_name_map["vocab_embedding"]}', - dtype) - - if mapping.is_first_pp_rank(): - if config.use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = split_matrix_tp( - vocab_embedding, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - else: - weights['transformer.vocab_embedding.weight'] = vocab_embedding - - if mapping.is_last_pp_rank(): - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - lm_head_weights = vocab_embedding.clone() - else: - lm_head_weights = get_weight(model_params, - param_name_map["lm_head"], dtype) - - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_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', - value=0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, - f'{model_prefix}.{param_name_map["ln_f"]}', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: LLaMAConfig, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - trust_remote_code: bool = True, - calib_batches: int = 512, - calib_max_seq_length: int = 512): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - hf_config = AutoConfig.from_pretrained(hf_model_dir, - trust_remote_code=trust_remote_code) - assert "llava" not in hf_config.model_type, "Smooth quant llava/vila/llava_next is not supported yet" - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map='auto' if device != 'cpu' else 'cpu', - dtype='auto' if not use_smooth_quant else torch.float16, - trust_remote_code=trust_remote_code) - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - use_fast=False, - padding_side='left') - - dataset = load_calib_dataset(calib_dataset) - - if calib_batches == -1: # use the whole dataset if calib_batches is -1 - calib_batches = len(dataset) - - act_range = capture_activation_range(hf_model, - tokenizer, - dataset, - num_samples=calib_batches, - seq_len=calib_max_seq_length) - qkv_para, smoother = {}, {} - if use_smooth_quant: - smooth_llama_model(hf_model, act_range, quant_config.smoothquant_val, - qkv_para, smoother) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model( - hf_model, - config=config, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother, - ) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -class QkvWeightHelper: - """ A helper utility for loading QKV weights from sharded files. """ - - def __init__(self, config: PretrainedConfig): - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.num_kv_heads = config.num_key_value_heads - self.tp_size = config.mapping.tp_size - self.tp_rank = config.mapping.tp_rank - self.is_mha = self.num_heads == self.num_kv_heads - self.head_size = None if not hasattr(config, - "head_size") else config.head_size - self._qkv_weights = {} - self.remove_duplicated_kv_heads = getattr(config, - 'remove_duplicated_kv_heads', - False) - - @staticmethod - def is_qkv_weight(name): - for k in ['q_proj', 'k_proj', 'v_proj']: - if 'self_attn' in name and k in name: - return True - return False - - def add_weight(self, i: int, name: str, weight: torch.Tensor): - if 'q_proj' in name: - tag = 'q' - elif 'k_proj' in name: - tag = 'k' - elif 'v_proj' in name: - tag = 'v' - else: - raise ValueError(f'Got an unexpected parameter of name {name}') - if i not in self._qkv_weights: - self._qkv_weights[i] = {} - self._qkv_weights[i][tag] = weight - - def is_qkv_prepared(self, layer_idx): - if layer_idx not in self._qkv_weights: - return False - weights = self._qkv_weights[layer_idx] - return 'q' in weights and 'k' in weights and 'v' in weights - - def split_qkv_weights(self, layer_idx): - if not self.is_qkv_prepared(layer_idx): - return None - weights = self._qkv_weights.pop(layer_idx) # to prevent memory leak. - q, k, v = (torch.tensor(weights[t]) for t in ['q', 'k', 'v']) - - if self.remove_duplicated_kv_heads: - head_size = self.hidden_size // self.num_heads if self.head_size is None else self.head_size - k = k.reshape( - [k.shape[0] // head_size // 2, 2, head_size, self.hidden_size]) - v = v.reshape( - [v.shape[0] // head_size // 2, 2, head_size, self.hidden_size]) - assert (k[:, 0] == k[:, 1]).all() - assert (v[:, 0] == v[:, 1]).all() - k = k[:, 0].reshape([-1, self.hidden_size]) - v = v[:, 0].reshape([-1, self.hidden_size]) - - if not self.is_mha: - head_size = self.hidden_size // self.num_heads if self.head_size is None else self.head_size - if self.num_kv_heads < self.tp_size: - # duplicate the KV heads up to tensor_parallel - k = dup_kv_weight(k, self.num_kv_heads, self.tp_size) - v = dup_kv_weight(v, self.num_kv_heads, self.tp_size) - assert k.shape[0] % (self.tp_size * head_size) == 0 - assert v.shape[0] % (self.tp_size * head_size) == 0 - wq = split(q, self.tp_size, self.tp_rank) - wk = split(k, self.tp_size, self.tp_rank) - wv = split(v, self.tp_size, self.tp_rank) - fused_qkv = torch.cat((wq, wk, wv), dim=0) - else: - qkv = torch.cat([q, k, v], dim=0) - qkv = qkv.reshape(3, q.shape[0], q.shape[1]) - fused_qkv = split(qkv, self.tp_size, self.tp_rank, dim=1) - fused_qkv = fused_qkv.reshape(3 * (q.shape[0] // self.tp_size), - q.shape[1]) - return fused_qkv - - -def load_weights_from_hf_by_shard(model_dir: str, config: LLaMAConfig): - '''Weights-only quantization is the only supported quantization recipe here.''' - logger.info('Loading weights from HF LLaMA...') - 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 - - weights = {} - tik = time.time() - dtype = getattr(torch, config.dtype) - - mapping = config.mapping - moe_config = config.moe - assert not moe_config.has_moe(), "MoE does not support sharded load" - assert "Exaone" not in config.architecture, "EXAONE model currently not support sharded load" - - from transformers import AutoConfig - hf_config = AutoConfig.from_pretrained(model_dir) - - quant_mode = config.quant_mode - if quant_mode.is_int8_weight_only(): - plugin_weight_only_quant_type = torch.int8 - elif quant_mode.is_int4_weight_only(): - plugin_weight_only_quant_type = torch.quint4x2 - elif config.quant_mode.has_fp8_rowwise(): - plugin_weight_only_quant_type = torch.float8_e4m3fn - else: - plugin_weight_only_quant_type = None - use_weight_only = quant_mode.is_weight_only() - use_fp8_rowwise = quant_mode.has_fp8_rowwise() - # Meta's recipe of not using fp8 rowwise for the first and last layer. - exclude_layers_id = [0, config.num_hidden_layers - 1] - - layers_range = mapping.pp_layers(config.num_hidden_layers) - - qkv_weight_helper = QkvWeightHelper(config) - - def convert_to_dtype(name, param, model_params, dtype): - # fp8 rowwise weights will only load fp8 weights and scales for the mlp layer. - if ('weight_scale' in name or name.replace('weight', 'weight_scale') in model_params) \ - and use_fp8_rowwise: - assert 'mlp' in name, "only MLP layers support fp8 rowwise currently." - return param - else: - return param.to(dtype) - - def fp8_rowwise_quantization(name, - param, - model_params, - clamp_value, - split_scale=False): - # check if weights are already quantized. - loaded_weight_scale = model_params.get( - name.replace('weight', 'weight_scale')) - if loaded_weight_scale is not None: - assert param.dtype == torch.float8_e4m3fn, "weight data type must be torch.float8_e4m3fn" - if split_scale: - assert mapping is not None - loaded_weight_scale = split(loaded_weight_scale, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - return param, loaded_weight_scale.reshape(-1) - else: - return fp8_per_channel_quant_weight_gpu(param, clamp_value) - - for model_file in iterate_shard_files(model_dir, - rank=mapping.tp_rank, - progress_bar=False): - logger.debug(f'Loading file {str(model_file)}...') - model_params = load_state_dict(model_file) - for name, param in model_params.items(): - logger.debug(f'Converting weight {name}...') - layer_idx = retrieved_layer_index_from_name(name) - tllm_prex = f'transformer.layers.{layer_idx}.' - - param = convert_to_dtype(name, param, model_params, dtype) - - if layer_idx is None: - layer = None - else: - if layer_idx not in layers_range: - continue - else: - tllm_prex = f'transformer.layers.{layer_idx - layers_range[0]}.' - - # Meta's recipe of not using fp8 rowwise for the first and last layer. - use_fp8_rowwise_in_layer = use_fp8_rowwise and ( - layer_idx not in exclude_layers_id) - - if 'model.embed_tokens.weight' in name: - if hf_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: - # padding - vocab_size_padded = pad_vocab_size( - config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - param = torch.from_numpy( - np.pad(param.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split( - param, mapping.tp_size, mapping.tp_rank) - if config.use_parallel_embedding: - param = split(param, mapping.tp_size, mapping.tp_rank, - config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = param - elif 'model.norm.weight' in name: - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = param - elif 'lm_head.weight' in name: - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size( - config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - param = torch.from_numpy( - np.pad(param.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split(param, mapping.tp_size, - mapping.tp_rank) - elif 'input_layernorm.weight' in name: - weights[tllm_prex + 'input_layernorm.weight'] = param - elif 'post_attention_layernorm.weight' in name: - weights[tllm_prex + 'post_layernorm.weight'] = param - elif qkv_weight_helper.is_qkv_weight(name): - qkv_weight_helper.add_weight(layer_idx, name, param) - if not qkv_weight_helper.is_qkv_prepared(layer_idx): - continue - split_v = qkv_weight_helper.split_qkv_weights(layer_idx) - if use_weight_only: - param = split_v.transpose() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - param, plugin_weight_only_quant_type) - weights[tllm_prex + - 'attention.qkv.weight'] = processed_torch_weights - weights[ - tllm_prex + - 'attention.qkv.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'attention.qkv.weight'] = split_v - elif 'self_attn.o_proj.weight' in name: - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - weights[tllm_prex + - 'attention.dense.weight'] = processed_torch_weights - weights[ - tllm_prex + - 'attention.dense.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'attention.dense.weight'] = split_v - elif name.endswith('mlp.up_proj.weight'): - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - weights[tllm_prex + - 'mlp.gate.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.gate.per_channel_scale'] = torch_weight_scales - elif use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization( - name, split_v, model_params, - config.quantization.clamp_val, True) - weights[tllm_prex + - 'mlp.gate.weight'] = processed_torch_weights.view( - plugin_weight_only_quant_type) - weights[ - tllm_prex + - 'mlp.gate.per_channel_scale'] = torch_weight_scales.to( - torch.float32) - else: - weights[tllm_prex + 'mlp.gate.weight'] = split_v - elif name.endswith('mlp.down_proj.weight'): - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - weights[tllm_prex + - 'mlp.proj.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.proj.per_channel_scale'] = torch_weight_scales - elif use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization( - name, split_v, model_params, - config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.proj.weight'] = processed_torch_weights.view( - plugin_weight_only_quant_type) - weights[ - tllm_prex + - 'mlp.proj.per_channel_scale'] = torch_weight_scales.to( - torch.float32) - else: - weights[tllm_prex + 'mlp.proj.weight'] = split_v - - elif name.endswith('mlp.gate_proj.weight'): - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - layer.mlp.fc.weight.value = processed_torch_weights - layer.mlp.fc.per_channel_scale.value = torch_weight_scales - weights[tllm_prex + - 'mlp.fc.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.fc.per_channel_scale'] = torch_weight_scales - elif use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization( - name, split_v, model_params, - config.quantization.clamp_val, True) - weights[tllm_prex + - 'mlp.fc.weight'] = processed_torch_weights.view( - plugin_weight_only_quant_type) - weights[ - tllm_prex + - 'mlp.fc.per_channel_scale'] = torch_weight_scales.to( - torch.float32) - else: - weights[tllm_prex + 'mlp.fc.weight'] = split_v - - del model_params - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights - - -def load_weights_from_hf_safetensors(model_dir: str, config: LLaMAConfig): - logger.info('Loading weights from Huggingface {} safetensors...'.format( - config.architecture.split('ForCausalLM')[0])) - tik = time.time() - import json - import os - - import safetensors - weights = {} - - model_dir = model_dir if model_dir.endswith("/") else model_dir + "/" - safetensors_map = {} - has_safetensor_index_json = True - try: - with open(model_dir + "model.safetensors.index.json", 'r') as fr: - sharding_map = json.load(fr) - for k, v in sharding_map['weight_map'].items(): - safetensors_map[k] = int(v[6:11]) - 1 - except FileNotFoundError: - has_safetensor_index_json = False - - shard_files = [] - for name in os.listdir(model_dir): - if name.endswith(".safetensors"): - if has_safetensor_index_json and name not in sharding_map[ - 'weight_map'].values(): - continue - shard_files.append(name) - shard_files.sort() - safetensors_ptrs = [ - safetensors.safe_open(model_dir + shard_file, - framework="pt", - device="cpu") for shard_file in shard_files - ] - - mapping = config.mapping - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - pad_vocab = vocab_size % mapping.tp_size != 0 - vocab_size_padded = pad_vocab_size(config.vocab_size, mapping.tp_size) - dtype = config.dtype - - moe_config = config.moe - - kv_tp_size = None - kv_tp_rank = None - if config.num_key_value_heads < mapping.tp_size: - kv_tp_size = config.num_key_value_heads - kv_tp_rank = mapping.tp_rank * kv_tp_size // mapping.tp_size - - model_prefix, layer_prefix, param_name_map = get_prefix_and_param_name_map( - config.architecture, use_safetensors=True) - - torch_dtype = str_dtype_to_torch(dtype) - - def load(key, - tp_dim=-1, - no_prefix=0, - is_expert_weights=False, - tp_size=None, - tp_rank=None): - if not no_prefix: - key = f'{model_prefix}.' + key - ptr_idx = safetensors_map[key] if key in safetensors_map else 0 - - if key not in safetensors_ptrs[ptr_idx].keys(): - return None - - tensor_slice = safetensors_ptrs[ptr_idx].get_slice(key) - tensor_shape = tensor_slice.get_shape() - if tp_dim == -1: - res = tensor_slice[:] - elif tp_dim >= 0 and tp_dim < len(tensor_shape): - if is_expert_weights: - if tp_size is None: - tp_size = mapping.moe_tp_size - if tp_rank is None: - tp_rank = mapping.moe_tp_rank - else: - if tp_size is None: - tp_size = mapping.tp_size - if tp_rank is None: - tp_rank = mapping.tp_rank - dim_size = tensor_shape[tp_dim] - if dim_size % tp_size != 0: - logger.error( - f"Current weight {key}'s shape {tensor_shape} is invalid at dimension {tp_dim} for TP size {tp_size}" - ) - indices = [slice(None)] * len(tensor_shape) - indices[tp_dim] = slice(dim_size * tp_rank // tp_size, - dim_size * (tp_rank + 1) // tp_size) - res = tensor_slice[indices] - else: - raise ValueError( - f"Invalid TP dim {tp_dim} for weight {key}'s shape {tensor_shape}" - ) - return res.to(torch_dtype).contiguous( - ) if "block_sparse_moe.gate" not in key and "block_sparse_moe.router" not in key else res.to( - torch.float32) - - def load_and_set(target, - key, - tp_dim=-1, - no_prefix=0, - is_expert_weights=False): - res = load(key, tp_dim, no_prefix, is_expert_weights) - weights[target] = res - if "weight" in key: - bias = load(key.replace("weight", "bias"), -1, no_prefix, - is_expert_weights) - if bias is not None: - weights[target.replace("weight", "bias")] = bias - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = load( - param_name_map["vocab_embedding"], config.embedding_sharding_dim - if config.use_parallel_embedding else -1) # vocab_embedding - - if mapping.is_last_pp_rank(): - v = load(param_name_map["lm_head"], -1, 1) if pad_vocab else load( - param_name_map["lm_head"], 0, 1) # lm_head - if v is None: - v = load(param_name_map["vocab_embedding"], - -1 if pad_vocab else 0).clone().detach() - - if pad_vocab: - v = torch.nn.functional.pad( - v, (0, 0, 0, vocab_size_padded - vocab_size), 'constant', 0) - v = split(v, mapping.tp_size, mapping.tp_rank) - weights['lm_head.weight'] = v - weights['transformer.ln_f.weight'] = load( - param_name_map["ln_f"]) # ln_f - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = f'{layer_prefix}.{l}.' - tllm_prex = f'transformer.layers.{layer_idx}' - - # Attention - qkv_list = [] - for comp in ["q", "k", "v"]: - tp_size = kv_tp_size if comp != "q" else None - tp_rank = kv_tp_rank if comp != "q" else None - weight_part = load(prefix + f'{param_name_map["attention.qkv"]}.' + - comp + param_name_map["qkv_suffix"], - 0, - tp_size=tp_size, - tp_rank=tp_rank) - qkv_list.append(weight_part) - bias_part = load( - (prefix + f'{param_name_map["attention.qkv"]}.' + comp + - param_name_map["qkv_suffix"]).replace("weight", "bias"), - 0, - tp_size=tp_size, - tp_rank=tp_rank) - if bias_part is not None: - qkv_list.append(bias_part) - if len(qkv_list) == 3: - # No bias - weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat( - qkv_list, 0) - else: - weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat( - qkv_list[::2], 0) - weights[f'{tllm_prex}.attention.qkv.bias'] = torch.cat( - qkv_list[1::2], 0) - load_and_set(f'{tllm_prex}.attention.dense.weight', - prefix + param_name_map["attention.dense"], - 1) # attention.dense - - # MLP - if not moe_config.has_moe(): - load_and_set(f'{tllm_prex}.mlp.gate.weight', - prefix + param_name_map["mlp.gate"], 0) # mlp.gate - load_and_set(f'{tllm_prex}.mlp.proj.weight', - prefix + param_name_map["mlp.proj"], 1) # mlp.proj - load_and_set(f'{tllm_prex}.mlp.fc.weight', - prefix + param_name_map["mlp.fc"], 0) # mlp.fc - - else: - architecture = config.architecture.lower() - if "granite" not in architecture: - weights[f'{tllm_prex}.mlp.router.weight'] = load( - prefix + 'block_sparse_moe.gate.weight') - else: - weights[f'{tllm_prex}.mlp.router.weight'] = load( - prefix + 'block_sparse_moe.router.layer.weight') - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - - if "granite" not in architecture: - expert_weight_list = [] - for suffix in range(3): - tp_dim = -1 - if mapping.has_moe_tp(): - tp_dim = 1 if suffix == 1 else 0 - expert_weight_list.append( - torch.stack( - list( - load( - prefix + - f'block_sparse_moe.experts.{expert}.w{suffix + 1}.weight', - tp_dim=tp_dim, - is_expert_weights=True) - for expert in rank_experts))) - - w1 = expert_weight_list[0] - w2 = expert_weight_list[1] - w3 = expert_weight_list[2] - else: - w2 = load(prefix + f'block_sparse_moe.output_linear.weight', - is_expert_weights=True) #TODO: correct this - w13 = load(prefix + f'block_sparse_moe.input_linear.weight', - is_expert_weights=True) - - half_size = w13.shape[-2] // 2 - w1, w3 = w13.split(half_size, dim=-2) - w1 = w1[rank_experts[0]:rank_experts[-1] + 1] - w2 = w2[rank_experts[0]:rank_experts[-1] + 1] - w3 = w3[rank_experts[0]:rank_experts[-1] + 1] - - weights[f'{tllm_prex}.mlp.fc.weight'] = \ - torch.concat([w3, w1], dim=-2).contiguous() - weights[f'{tllm_prex}.mlp.proj.weight'] = w2.contiguous() - - load_and_set(f'{tllm_prex}.input_layernorm.weight', prefix + - param_name_map["input_layernorm"]) # input_layernorm - load_and_set(f'{tllm_prex}.post_layernorm.weight', prefix + - param_name_map["post_layernorm"]) # post_layernorm - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - return weights - - -def load_weights_from_gptq(quant_ckpt_path: str, config: LLaMAConfig): - logger.info('Loading weights from groupwise GPTQ LLaMA safetensors...') - weights = {} - tik = time.time() - - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - dtype = config.dtype - mapping = config.mapping - quant_algo = config.quantization.quant_algo - - gptq_llama = safetensors.safe_open(quant_ckpt_path, - framework="pt", - device=0) - gptq_prefix = "model." - gptq_suffix_list = [".qweight", ".qzeros", ".scales"] - gptq_key_list = [ - "embed_tokens.weight", # vocab_embedding - "lm_head.weight", # lm_head - "norm.weight", # ln_f - "self_attn.", # attention.qkv - "_proj", # qkv suffix - "self_attn.o_proj", # attention.dense - "mlp.up_proj", # mlp.gate - "mlp.down_proj", # mlp.proj - "mlp.gate_proj", # mlp.fc - "input_layernorm.weight", # input_layernorm - "post_attention_layernorm.weight", # post_layernorm - ] - split_sym = "." - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = str_dtype_to_torch(dtype) - - def load(key, no_prefix=0): - if no_prefix: - return gptq_llama.get_tensor(key) - else: - return gptq_llama.get_tensor(gptq_prefix + key) - - def torch_split(v, dim): - if v.shape[dim] % mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(mapping.tp_size)) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank] - - def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(v: List[torch.Tensor], - tllm_prex: str, - quant_algo: QuantAlgo, - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in v - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in v - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_UINT8_INPUT = 1 # Set to true if checkpoint store UINT8 weights - USE_GPTQ_FOR_LLAMA = 1 # GPTQ-for-LLaMA added 1 to zeros - - if quant_algo == QuantAlgo.W4A16_GPTQ: - # unpack inputs packed in int32 into int4 and store them in int8 format - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor( - packer(qweight_unpacked_int8), torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - if not USE_UINT4_INPUT: - # Correcting UINT4 values back to INT4 order - mask_negative = qzeros_unpacked_int32[qzeros_unpacked_int32 < 0] - mask_positive = qzeros_unpacked_int32[qzeros_unpacked_int32 >= - 0] - qzeros_unpacked_int32 = qzeros_unpacked_int32 + 16 * mask_negative - 16 * mask_positive - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - else: - # unpack inputs packed in int32 into int8 - qweight_unpacked_int8 = ( - qweight_int32.T.contiguous().view(torch.uint8).T.contiguous() - - 128).to(torch.int8) - qweight_interleaved = preprocessor(qweight_unpacked_int8, - torch.int8, torch.float16).view( - torch.float16) - qzeros_unpacked_int32 = qzeros_int32.view(torch.uint8) - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + - 128 * USE_UINT8_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - results = { - f'{tllm_prex}.weight': qweight_interleaved, - f'{tllm_prex}.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.zero': zeros_x_scales_fp16, - } - return results - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = load(gptq_key_list[0]) - if mapping.is_first_pp_rank(): - # tensorrt_llm_llama.vocab_embedding.weight.value = v.to( - # torch_dtype).cpu().numpy() - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - # 2. lm_head - v = load(gptq_key_list[1], "no_prefix") - if mapping.is_last_pp_rank(): - # tensorrt_llm_llama.lm_head.weight.value = torch_split( - # v, 0).to(torch_dtype).cpu().numpy() - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - v = torch.from_numpy( - np.pad(v.detach().cpu().numpy(), ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = torch_split(v, 0).to(torch_dtype) - - # 3. ln_f - v = load(gptq_key_list[2]) - if mapping.is_last_pp_rank(): - # tensorrt_llm_llama.ln_f.weight.value = v.to(torch_dtype).cpu().numpy() - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - # 4. Weights inside each layer - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = "layers" + split_sym + str(layer_idx) + split_sym - logger.info(f'Process weights in layer: {layer_idx}') - # layer = tensorrt_llm_llama.layers[layer_idx] - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # 4.1 attention.qkv - qkv_weight_list = [] - for suf in gptq_suffix_list: - qkv_list = [] - for comp in ["q", "k", "v"]: - comp_part = load(prefix + gptq_key_list[3] + comp + - gptq_key_list[4] + suf) - comp_part = torch_split(comp_part, 1) - qkv_list.append(comp_part) - qkv_weight_list.append(torch.cat(qkv_list, dim=1)) - - # process_and_assign_weight(layer.attention.qkv, qkv_weight_list) - weights.update( - process_and_assign_weight(qkv_weight_list, - f'{tllm_prex}.attention.qkv', quant_algo)) - # 4.2 attention.dense - v = [load(prefix + gptq_key_list[5] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.attention.dense, v, 0) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.attention.dense', - quant_algo, - tp_dim=0)) - # 4.3 mlp.gate - v = [load(prefix + gptq_key_list[6] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.mlp.gate, v, 1) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.mlp.gate', - quant_algo, - tp_dim=1)) - # 4.4 mlp.proj - v = [load(prefix + gptq_key_list[7] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.mlp.proj, v, 0) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.mlp.proj', - quant_algo, - tp_dim=0)) - # 4.5 mlp.fc - v = [load(prefix + gptq_key_list[8] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.mlp.fc, v, 1) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.mlp.fc', - quant_algo, - tp_dim=1)) - # 4.6 input_layernorm - v = load(prefix + gptq_key_list[9]) - # layer.input_layernorm.weight.value = v.to(torch_dtype).cpu().numpy() - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - - # 4.7 post_layernorm - v = load(prefix + gptq_key_list[10]) - # layer.post_layernorm.weight.value = v.to(torch_dtype).cpu().numpy() - weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - return weights - - -def load_weights_from_deepcompressor(quant_ckpt_path: str, config: LLaMAConfig): - logger.info( - 'Loading weights from DeepCompressor torch checkpoint for QServe W4A8 inference...' - ) - weights = {} - tik = time.time() - - per_group = config.quant_mode.has_per_group_scaling() - group_size = 128 if per_group else -1 - - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - dtype = config.dtype - mapping = config.mapping - torch_dtype = str_dtype_to_torch(dtype) - assert torch_dtype == torch.float16, "Currently QServe only supports float16" - - # weight - fake_quant_weights = torch.load(quant_ckpt_path + '/model.pt', - map_location='cpu', - weights_only=True) - # scale.0, scale.1, zero - quant_params = torch.load(quant_ckpt_path + '/scale.pt', - map_location='cpu', - weights_only=True) - - def load(key): - if 'zero' in key: - v = quant_params[key] - # https://github.com/mit-han-lab/qserve/blob/64ee627dfd747510809998d3592439f05a71ba31/scripts/ckpt_converter/checkpoint_converter.py#L99 - if v.min() < 0: - v = v + 8 - return v - if 'scale' in key: - return quant_params[key] - return fake_quant_weights[key] - - if per_group: - deepcompressor_suffix = [ - 'weight', 'weight.scale.0', 'weight.scale.1', 'weight.scaled_zero' - ] - qserve_suffix = ['weight', 's1_scales', 's2_scales', 's2_szeros'] - else: - deepcompressor_suffix = [ - 'weight', 'weight.scale.0', 'weight.scaled_zero' - ] - qserve_suffix = ['weight', 's1_scales', 's1_szeros'] - - def tp_split_tensor(v: torch.Tensor, dim): - if v.shape[dim] % mapping.tp_size != 0: - logger.error( - f"Current weight shape is invalid for mapping.tp_size={mapping.tp_size}" - ) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank].contiguous() - - def tp_split_weight_and_params(v: List[torch.Tensor], column_linear: bool): - if per_group: - weight, s1_scales, s2_scales, s2_szeros = v - # weight (out_features, in_features) - # weight.scale.0 (out_features, 1, 1, 1) - # weight.scale.1 (out_features, 1, in_features/group_size, 1) - # weight.scaled_zero (out_features, 1, in_features/group_size, 1) - if column_linear: - weight = tp_split_tensor(weight, 0) - s1_scales = tp_split_tensor(s1_scales, 0) - s2_scales = tp_split_tensor(s2_scales, 0) - s2_szeros = tp_split_tensor(s2_szeros, 0) - else: - weight = tp_split_tensor(weight, 1) - s1_scales = s1_scales - s2_scales = tp_split_tensor(s2_scales, 2) - s2_szeros = tp_split_tensor(s2_szeros, 2) - return [weight, s1_scales, s2_scales, s2_szeros] - else: - weight, s1_scales, s1_szeros = v - # weight (out_features, in_features) - # weight.scale.0 (out_features, 1, 1, 1) - # weight.zero (out_features, 1, 1, 1) - if column_linear: - weight = tp_split_tensor(weight, 0) - s1_scales = tp_split_tensor(s1_scales, 0) - s1_szeros = tp_split_tensor(s1_szeros, 0) - else: - weight = tp_split_tensor(weight, 1) - s1_scales = s1_scales - s1_szeros = s1_szeros - return [weight, s1_scales, s1_szeros] - - def process_weight_and_params(v: List[torch.Tensor], tllm_prex: str): - if per_group: - weight, s1_scales, s2_scales, s2_szeros = v - qweight = qserve_quantize_weight_per_group(weight, s1_scales, - s2_scales, s2_szeros, - group_size) - qweight, s1_scales, s2_scales, s2_zeros = qserve_pack_reorder_per_group( - qweight, s1_scales, s2_scales, s2_szeros, group_size) - - return { - # Note: Linear modules in TRTLLM do not use the name 'qweight' - f'{tllm_prex}.{qserve_suffix[0]}': qweight, - f'{tllm_prex}.{qserve_suffix[1]}': s1_scales, - f'{tllm_prex}.{qserve_suffix[2]}': s2_scales, - f'{tllm_prex}.{qserve_suffix[3]}': s2_zeros, - } - else: - weight, s1_scales, s1_szeros = v - qweight = qserve_quantize_weight_per_channel( - weight, s1_scales, s1_szeros) - qweight, s1_scales, s1_szeros = qserve_pack_reorder_per_channel( - qweight, s1_scales, s1_szeros) - - return { - # Note: Linear modules in TRTLLM use the name 'weight' instead of 'qweight' - f'{tllm_prex}.{qserve_suffix[0]}': qweight, - f'{tllm_prex}.{qserve_suffix[1]}': s1_scales, - f'{tllm_prex}.{qserve_suffix[2]}': s1_szeros, - } - - # Load weights - # 1. vocab_embedding - v = load('model.embed_tokens.weight') - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - - # 2. lm_head - v = load('lm_head.weight') - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - v = torch.nn.functional.pad(v, (0, 0, 0, pad_width)) - weights['lm_head.weight'] = tp_split_tensor(v, 0).to(torch_dtype) - - # 3. ln_f - v = load('model.norm.weight') - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - - # 4. Weights inside each layer - layers_range = mapping.pp_layers(num_hidden_layers) - for layer_idx in layers_range: - prefix = f'model.layers.{layer_idx}' - logger.info(f'Processing weights in layer: {layer_idx}') - tllm_prex = f'transformer.layers.{layer_idx - layers_range[0]}' - - # 4.1 attention.qkv - qkv_list = [] - for comp in ["q", "k", "v"]: - v = [ - load(f'{prefix}.self_attn.{comp}_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=True) - qkv_list.append(v) - # Concat qkv - q, k, v = qkv_list - qkv = [ - torch.concat((q[i], k[i], v[i]), dim=0) - for i in range(len(deepcompressor_suffix)) - ] - weights.update( - process_weight_and_params(qkv, f'{tllm_prex}.attention.qkv')) - - # 4.2 attention.dense - v = [ - load(f'{prefix}.self_attn.o_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=False) - weights.update( - process_weight_and_params(v, f'{tllm_prex}.attention.dense')) - - # TODO: The naming here is tricky. - # The implementation of GatedMLP is act(fc(x)) * gate(x). - # However, the common convention is act(gate_proj(x)) * up_proj(x). - - # 4.3 mlp.gate - v = [ - load(f'{prefix}.mlp.up_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=True) - weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.gate')) - - # 4.4 mlp.fc - v = [ - load(f'{prefix}.mlp.gate_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=True) - weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.fc')) - - # 4.5 mlp.proj - v = [ - load(f'{prefix}.mlp.down_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=False) - weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.proj')) - - # 4.6 input_layernorm - v = load(f'{prefix}.input_layernorm.weight') - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - - # 4.7 post_layernorm - v = load(f'{prefix}.post_attention_layernorm.weight') - weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - # TODO: All the RMSNorm weight, including ln_f, input_layernorm, post_layernorm, are actually all 1s - # Could implement a simplified module without weight - return weights - - -def load_torch_meta_ckpt(meta_ckpt_path: Path): - ''' - meta_ckpt_path: The format of meta_ckpt_path is like /consolidated.xx There are two possible cases: - 1. A file like /consolidated.xx.pth, loading it by torch.load directly - 2. A folder like /consolidated.xx/, need to load all weights in the folder. - ''' - file_path = meta_ckpt_path.parent / (meta_ckpt_path.name + ".pth") - if file_path.exists() and file_path.is_file(): - return torch.load(file_path, map_location="cpu") - else: - folder_path = meta_ckpt_path - assert folder_path.exists() and folder_path.is_dir() - - ckpts = list(Path(folder_path).glob("consolidated-*.pth")) - - all_weights = {} - for ckpt in ckpts: - _weight = torch.load(ckpt, map_location="cpu") - all_weights = all_weights | _weight - del _weight - return all_weights - - -def load_weights_from_meta_ckpt(meta_ckpt_dir: str, config: LLaMAConfig): - torch_dtype = str_dtype_to_torch(config.dtype) - mapping = config.mapping - use_fp8_rowwise = config.quant_mode.has_fp8_rowwise() - if config.quant_mode.has_any_quant() and not use_fp8_rowwise: - logger.error( - "Meta ckpts only support fp8_rowwise quantization currently.") - weights = {} - # Meta's recipe of not using fp8 rowwise for the first and last layer. - exclude_layers_id = [0, config.num_hidden_layers - 1] - - def gather_ckpts(ckpts): - gathered = {} - for k in ckpts[0]: - d = 0 - # TODO not sure should we consider tok here. - if any([n in k for n in ["wo", "w2"]]): - d = 1 - if "norm" in k or "rope" in k: # no TP - gathered[k] = ckpts[0][k].clone() - else: - gathered[k] = torch.cat([pt[k] for pt in ckpts], dim=d).clone() - return gathered - - def split_ckpt(ckpt, ranks_per_ckpt, ckpt_rank): - split_ckpt = {} - for k, v in ckpt.items(): - d = 0 - if any(n in k for n in - ["wo", "feed_forward.w2", "tok", "feed_forward.gate"]): - d = 1 - if "norm" in k or "rope" in k: # no TP - split_ckpt[k] = v.clone() - elif config.num_key_value_heads < mapping.tp_size and any( - n in k for n in ["wk", "wv"]): - assert mapping.tp_size % config.num_key_value_heads == 0 - # special case: we need to duplicate KV head - tmp = dup_kv_weight(v, config.num_key_value_heads, - mapping.tp_size) - split_ckpt[k] = torch.split(tmp, - tmp.shape[d] // ranks_per_ckpt, - dim=d)[ckpt_rank].clone() - else: - split_ckpt[k] = torch.split(v, - v.shape[d] // ranks_per_ckpt, - dim=d)[ckpt_rank].clone() - return split_ckpt - - def get_current_weights(num_ckpts): - if num_ckpts > mapping.tp_size: - # combine ckpts - assert (num_ckpts % mapping.tp_size) == 0 - nf = num_ckpts // mapping.tp_size - fs = nf * mapping.tp_rank - file_ids = list(range(fs, fs + nf)) - ckpts = [] - for f in file_ids: - ckpt = load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{f:02d}")) - ckpts.append(ckpt) - return gather_ckpts(ckpts) - elif num_ckpts < mapping.tp_size: - # split ckpt - assert (mapping.tp_size % num_ckpts) == 0 - ranks_per_ckpt = mapping.tp_size // num_ckpts - ckpt_fid = mapping.tp_rank // ranks_per_ckpt - ckpt_rank = mapping.tp_rank % ranks_per_ckpt - nH_per_ckpt = config.num_attention_heads // num_ckpts - assert (nH_per_ckpt % ranks_per_ckpt) == 0 - ckpt = load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{ckpt_fid:02d}")) - return split_ckpt(ckpt, ranks_per_ckpt, ckpt_rank) - - # num_ckpts == tensor_parallel, 1:1 mapping from files to TP - return load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{mapping.tp_rank:02d}")) - - def permute(w, nH, d, dH): - # due to MQA's wk, nH*dH != d could be true - return w.view(nH, dH // 2, 2, d).transpose(1, 2).reshape(nH * dH, d) - - def extract_layer_idx(name): - ss = name.split('.') - for s in ss: - if s.isdigit(): - return s - return None - - if not hasattr(load_weights_from_meta_ckpt, "saved_embed"): - load_weights_from_meta_ckpt.saved_embed = None - - def combine_embeddings(embeds, num_ckpts): - if len(embeds) == 1: - return embeds[0] - assert [ - embeds[i].shape == embeds[i + 1].shape - for i in range(len(embeds) - 1) - ] - if embeds[0].shape[0] == config.vocab_size // num_ckpts: - merge_dim = 0 - elif embeds[0].shape[1] == config.hidden_size // num_ckpts: - merge_dim = 1 - else: - logger.error("Unable to infer embedding split dimension") - assert False, "Unable to infer embedding split dimension" - return torch.cat(embeds, dim=merge_dim) - - def gather_embedding(cur_embed, name: str, num_ckpts): - if mapping.tp_size == 1: - # even if num_ckpts > 1, get_current_weights will already have it gathered - return cur_embed - if load_weights_from_meta_ckpt.saved_embed is None: - embeds = [None] * num_ckpts - for i in range(num_ckpts): - ckpt = load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{i:02d}")) - embeds[i] = ckpt[name] - embed = combine_embeddings(embeds, num_ckpts).to(torch_dtype) - load_weights_from_meta_ckpt.saved_embed = embed - - return load_weights_from_meta_ckpt.saved_embed - - logger.info('Loading weights from Meta LLaMA checkpoints ...') - tik = time.time() - - num_kv_heads = config.num_key_value_heads - mha_mode = (num_kv_heads == config.num_attention_heads) - - ckpts = list(Path(meta_ckpt_dir).glob("consolidated.*")) - num_ckpts = len(ckpts) - # llama/llama2 doesn't have MQA. So, simplifying loader logic by not worrying about it. - assert num_kv_heads > 1 or num_kv_heads >= num_ckpts, \ - f"We don't know how the {num_kv_heads} KV heads are distributed among {num_ckpts} checkpoints." - - tik = time.time() - ckpt = get_current_weights(num_ckpts) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'[{mapping.rank}] get_current_weights. Total time: {t}') - - head_size = config.hidden_size // config.num_attention_heads - layers_range = mapping.pp_layers(config.num_hidden_layers) - - for l in layers_range: - prefix = f'layers.{l}.attention.' - q_weight = permute(ckpt[prefix + 'wq.weight'].clone(), - nH=(config.num_attention_heads // mapping.tp_size), - d=config.hidden_size, - dH=head_size) - if num_kv_heads < mapping.tp_size and num_ckpts >= mapping.tp_size: - assert mapping.tp_size % num_kv_heads == 0 - assert False, "Not supported yet" - k_weight = permute(ckpt[prefix + 'wk.weight'].clone(), - nH=((num_kv_heads + mapping.tp_size - 1) // - mapping.tp_size), - d=config.hidden_size, - dH=head_size) - v_weight = ckpt[prefix + 'wv.weight'].clone() - - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - ckpt[prefix + 'qkv.weight'] = qkv_weight - - for k, v in tqdm(ckpt.items()): - dtype = torch_dtype if 'feed_forward.gate' not in k else torch.float32 - - v = v.to(dtype) - if "tok_embeddings" in k: - if not config.use_parallel_embedding: - v = gather_embedding(v, k, num_ckpts) - elif config.embedding_sharding_dim == 0: - # this needs a gather and then resplit along different dims - v = gather_embedding(v, k, num_ckpts) - v = split(v, mapping.tp_size, mapping.tp_rank, 0) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - elif "output" in k: - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - v = torch.from_numpy( - np.pad(v.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = v.detach().clone() - elif k == "norm.weight": - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v - else: - # layer specific weights - layer_idx = extract_layer_idx(k) - if layer_idx is None or int(layer_idx) not in layers_range: - continue - - # Meta's recipe of not using fp8 rowwise for the first and last layer. - use_fp8_rowwise_in_layer = use_fp8_rowwise and ( - int(layer_idx) not in exclude_layers_id) - idx = int(layer_idx) - layers_range[0] - tllm_prex = f'transformer.layers.{idx}.' - - if 'attention_norm.weight' in k: - weights[tllm_prex + 'input_layernorm.weight'] = v - elif 'ffn_norm.weight' in k: - weights[tllm_prex + 'post_layernorm.weight'] = v - elif 'feed_forward.w3.weight' in k: - if use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu( - v, config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.gate.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.gate.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'mlp.gate.weight'] = v - elif 'feed_forward.w2.weight' in k: - if use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu( - v, config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.proj.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.proj.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'mlp.proj.weight'] = v - elif 'feed_forward.w1.weight' in k: - if use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu( - v, config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.fc.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.fc.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'mlp.fc.weight'] = v - elif 'attention.wo.weight' in k: - weights[tllm_prex + 'attention.dense.weight'] = v - elif 'attention.qkv.weight' in k: - weights[tllm_prex + 'attention.qkv.weight'] = v - elif 'feed_forward.gate' in k: - weights[tllm_prex + 'mlp.router.weight'] = v - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/llama/model.py b/tensorrt_llm/models/llama/model.py deleted file mode 100644 index 2e272772adb0..000000000000 --- a/tensorrt_llm/models/llama/model.py +++ /dev/null @@ -1,614 +0,0 @@ -# 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 os -from typing import Optional, Union - -import transformers - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import (AllReduceFusionOp, AllReduceParams, Tensor, - allgather, concat, constant, div, non_gated_version, - recv, send, unsqueeze) -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, FusedGatedMLP, GatedMLP, - PositionEmbeddingType, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ...quantization.functional import fused_layernorm -from ..convert_utils import has_safetensors -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import LLaMAConfig -from .convert import (load_hf_llama, load_weights_from_deepcompressor, - load_weights_from_gptq, load_weights_from_hf_by_shard, - load_weights_from_hf_model, - load_weights_from_hf_safetensors, - load_weights_from_meta_ckpt) - - -class LLaMADecoderLayer(Module): - - def __init__(self, config: LLaMAConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - layer_idx += config.layer_idx_offset - self.config = config - self.mapping = config.mapping - - if (self.config.use_input_layernorm_in_first_layer - and self.layer_idx == 0) or self.layer_idx > 0: - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - self.is_last_local_layer = layer_idx == layers_range[-1] - self.attention = Attention( - local_layer_idx=self.local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - q_scaling=1.0 / config.attention_multiplier, - quant_mode=config.quant_mode, - cp_group=config.mapping.cp_group, - cp_size=config.mapping.cp_size, - cp_rank=config.mapping.cp_rank) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - ClsMLP = GatedMLP - mlp_kwargs = {} - if config.moe.has_moe(): - ClsMLP = MOE - mlp_kwargs = { - "moe_config": config.moe, - "mapping": config.mapping, - } - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - # Residual MLP that applies on pre-attention input - # TODO: change to self.has_residual_mlp = self.config.residual_mlp after ModelOpt quantize config is updated - self.has_residual_mlp = False - if hasattr(self.config, - "residual_mlp") and self.config.residual_mlp is True: - self.has_residual_mlp = True - - if self.has_residual_mlp: - self.residual_layernorm = RmsNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - ClsMLP = GatedMLP # TODO: may use FusedGatedMLP to further speedup - self.residual_mlp = ClsMLP( - hidden_size=config.hidden_size, - ffn_hidden_size=config. - hidden_size, # residual mlp uses hidden_size - hidden_act=non_gated_version( - config.hidden_act), # back to non-gated - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - next_layer_input_layernorm_args=None): - assert not ( - default_net().plugin_config.reduce_fusion and self.has_residual_mlp - ), "Custom all reduce and residual mlp can't be enabled at the same time." - assert not ( - default_net().plugin_config.reduce_fusion - and default_net().plugin_config.user_buffer - and default_net().plugin_config.pp_reduce_scatter - ), "User buffer reduce fusion enabled with PP reduce scatter is not supported now." - assert not ( - default_net().plugin_config.reduce_fusion - and default_net().plugin_config.norm_quant_fusion - ), "Reduce fusion and quant fusion can't be enabled at the same time." - if default_net( - ).plugin_config.reduce_fusion and self.local_layer_idx > 0: - hidden_states, residual = hidden_states - elif default_net( - ).plugin_config.norm_quant_fusion and self.local_layer_idx > 0: - hidden_states, residual = hidden_states - else: - residual = hidden_states - if (self.config.use_input_layernorm_in_first_layer - and self.layer_idx == 0) or self.layer_idx > 0: - hidden_states = self.input_layernorm(hidden_states) - - reduce_fusion_op = AllReduceFusionOp.NONE - if default_net().plugin_config.reduce_fusion: - if default_net().plugin_config.user_buffer: - if self.config.quant_mode.has_fp8_qdq(): - reduce_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8 - elif self.config.quant_mode.has_nvfp4(): - assert default_net( - ).plugin_config.gemm_plugin == "nvfp4", "UB with nvfp4 model must use nvfp4 gemm plugin" - reduce_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4 - else: - assert False, "UB must enabled with fp8 or nvfp4 model" - else: - reduce_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM - - reduce_fusion_scale = None - if default_net().plugin_config.reduce_fusion and default_net( - ).plugin_config.user_buffer: - if isinstance(self.mlp, FusedGatedMLP): - if self.config.quant_mode.has_fp8_qdq(): - reduce_fusion_scale = constant( - self.mlp.fused_fc.activation_scaling_factor.raw_value. - copy()) - elif self.config.quant_mode.has_nvfp4(): - reduce_fusion_scale = constant( - [1.0] / self.mlp.fused_fc. - activation_global_scaling_factor.raw_value) - else: - if self.config.quant_mode.has_fp8_qdq(): - reduce_fusion_scale = constant( - self.mlp.fc.activation_scaling_factor.raw_value.copy()) - elif self.config.quant_mode.has_nvfp4(): - reduce_fusion_scale = constant( - [1.0] / - self.mlp.fc.activation_global_scaling_factor.raw_value) - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=reduce_fusion_op, - residual=residual, - norm_weight=self.post_layernorm.weight.value, - scale=reduce_fusion_scale, - eps=self.post_layernorm.eps)) - if use_cache: - attention_output, presents = attention_output - - if self.has_residual_mlp: - hidden_states = residual + attention_output - residual_attn = hidden_states - # arctic layer w/ residual mlp - - # residual mlp - hidden_states = self.residual_layernorm(hidden_states) - hidden_states = self.residual_mlp(hidden_states) - residual_mlp = residual_attn + hidden_states - - # parallel moe - # parallel moe layers applies on PRE-ATTENTION input residual, therefore achieving pre-fetching and better parallelism - hidden_states = self.post_layernorm(residual) - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - hidden_states = residual_mlp + hidden_states - else: - if default_net().plugin_config.reduce_fusion: - hidden_states, residual = attention_output - elif default_net().plugin_config.norm_quant_fusion: - hidden_states, residual_attn, act_per_block_scale = fused_layernorm( - input=attention_output, - normalized_shape=self.config.hidden_size, - residual=residual, - weight=self.post_layernorm.weight.value, - scale=div( - 1, self.mlp.fc.activation_global_scaling_factor.value) - if self.mlp.fc.activation_global_scaling_factor.value else - None, - eps=self.post_layernorm.eps, - p_dtype=self.config.dtype) - - hidden_states, residual_attn = ( - hidden_states, act_per_block_scale), residual_attn - assert isinstance(hidden_states, tuple) - else: - hidden_states = residual + attention_output * self.config.residual_multiplier - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - if next_layer_input_layernorm_args is not None: - #this is middle layer - hidden_states = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=reduce_fusion_op, - residual=residual_attn - if default_net().plugin_config.norm_quant_fusion else - residual, - norm_weight=next_layer_input_layernorm_args[0], - scale=next_layer_input_layernorm_args[2], - eps=next_layer_input_layernorm_args[1])) - if default_net().plugin_config.norm_quant_fusion: - hidden_states, residual, act_per_block_scale = fused_layernorm( - input=hidden_states, - normalized_shape=self.config.hidden_size, - residual=residual_attn, - weight=next_layer_input_layernorm_args[0], - scale=div(1, next_layer_input_layernorm_args[2]) - if next_layer_input_layernorm_args[2] else None, - eps=next_layer_input_layernorm_args[1], - p_dtype=self.config.dtype) - hidden_states = (hidden_states, - act_per_block_scale), residual - else: - if default_net( - ).plugin_config.pp_reduce_scatter and self.is_last_local_layer and not self.mapping.is_last_pp_rank( - ): - hidden_states = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - last_local_layer_residual=residual) - else: - if (default_net().plugin_config.reduce_fusion - and default_net().plugin_config.user_buffer): - hidden_states, residual = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.LAST_PROCESS_FOR_UB, - residual=residual)) - else: - hidden_states = self.mlp( - hidden_states, lora_layer_params=lora_layer_params) - hidden_states = residual + hidden_states * self.config.residual_multiplier - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class LLaMAModel(Module): - - def __init__(self, config: LLaMAConfig) -> None: - super().__init__() - - self.mapping = config.mapping - self.vocab_size = config.vocab_size - self.has_partial_lora_mask = config.has_partial_lora_mask - self.hidden_size = config.hidden_size - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.embedding_multiplier = config.embedding_multiplier - - self.layers = DecoderLayerList(LLaMADecoderLayer, config) - - if config.fc_after_embed: - self.fc = ColumnLinear(2 * config.hidden_size, - config.hidden_size, - bias=True, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - if self.mapping.is_last_pp_rank(): - self.ln_f = None - if config.use_last_layernorm: - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - hidden_states_for_embed=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - hidden_states *= self.embedding_multiplier - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - if default_net().plugin_config.pp_reduce_scatter: - hidden_states = allgather(hidden_states, - self.mapping.tp_group, - gather_dim=0) - # reshape to (-1, hidden_size) - hidden_states = hidden_states.view( - concat([-1, self.hidden_size])) - - if hidden_states_for_embed is not None: - hidden_states = concat([hidden_states, hidden_states_for_embed], - dim=-1) - hidden_states = self.fc(hidden_states) - - if lora_params is not None and self.has_partial_lora_mask: - partial_lora_mask = input_ids > (self.vocab_size - 1) - lora_params.partial_lora_mask = unsqueeze(partial_lora_mask, -1) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - if self.ln_f: - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class LLaMAForCausalLM(DecoderModelForCausalLM): - config_class = LLaMAConfig - - def __init__(self, config: LLaMAConfig): - transformer = LLaMAModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(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) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - import transformers - - load_by_shard = kwargs.pop('load_by_shard', False) - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - use_autoawq = kwargs.pop('use_autoawq', None) - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER" - ) is not None and not isinstance( - hf_model_or_dir, transformers.PreTrainedModel): - if "vila" in hf_model_or_dir or "llava" in hf_model_or_dir: - hf_model_or_dir = load_hf_llama(hf_model_or_dir, - load_model_on_cpu) - elif not load_by_shard and not has_safetensors( - hf_model_or_dir) and ( - quant_config is None - or not quant_config.quant_mode.has_any_quant()): - hf_model_or_dir = load_hf_llama(hf_model_or_dir, - load_model_on_cpu) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = LLaMAConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if config.remove_duplicated_kv_heads: - config.num_key_value_heads = config.num_key_value_heads // 2 - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - custom_dict = {} - model_name = hf_model.config.model_type if use_preloading else hf_model_or_dir - if "llava" in model_name: - custom_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head" - } - elif "vila" in model_name: - hf_model_dir += "/llm" - elif "exaone" in model_name.lower(): - custom_dict = { - "transformer": "transformer", - "layers": "h", - "vocab_embedding": "wte", - "lm_head": "lm_head", - "ln_f": "ln_f", - "attention": "attn.attention", - "dense": "out_proj", - "gate": "c_fc_1", - "proj": "c_proj", - "fc": "c_fc_0", - "input_layernorm": "ln_1", - "post_layernorm": "ln_2", - } - elif config.tie_word_embeddings: - custom_dict = {"lm_head": "model.embed_tokens"} - - if quant_ckpt_path is not None: - hf_model_dir = quant_ckpt_path - arg_dict = {"use_autoawq": True} if use_autoawq else {} - - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model, arg_dict) - else: - if use_preloading: - assert not load_by_shard - weights = load_weights_from_hf_model(hf_model, config) - elif load_by_shard: - weights = load_weights_from_hf_by_shard(hf_model_dir, config) - elif has_safetensors( - hf_model_dir) and not config.quant_mode.has_any_quant(): - weights = load_weights_from_hf_safetensors(hf_model_dir, config) - elif quant_ckpt_path is not None: - if quant_config.quant_mode.is_int4_weight_only(): - weights = load_weights_from_gptq(quant_ckpt_path, config) - elif quant_config.quant_mode.is_qserve_w4a8(): - weights = load_weights_from_deepcompressor( - quant_ckpt_path, config) - else: - raise ValueError( - "quant_ckpt_path should be specified only for GPTQ or QServe" - ) - else: - hf_model = load_hf_llama(hf_model_dir, load_model_on_cpu) - weights = load_weights_from_hf_model(hf_model, config) - model = cls(config) - model.load(weights) - return model - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.weight_only_groupwise_quant_matmul_plugin = 'auto' - return plugin_config - - @classmethod - def from_meta_ckpt(cls, - meta_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - config = LLaMAConfig.from_meta_ckpt(meta_ckpt_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - weights = load_weights_from_meta_ckpt(meta_ckpt_dir, config) - - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = LLaMAConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - trust_remote_code = kwargs.pop("trust_remote_code", True) - - convert.quantize(hf_model_dir, - output_dir, - config=config, - device=device, - calib_dataset=calib_dataset, - trust_remote_code=trust_remote_code, - calib_batches=calib_batches, - calib_max_seq_length=calib_max_seq_length) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config) diff --git a/tensorrt_llm/models/mamba/__init__.py b/tensorrt_llm/models/mamba/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/mamba/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/mamba/config.py b/tensorrt_llm/models/mamba/config.py deleted file mode 100644 index 84e2f189c46a..000000000000 --- a/tensorrt_llm/models/mamba/config.py +++ /dev/null @@ -1,313 +0,0 @@ -# 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 json -import os -from enum import Enum -from typing import List, Optional, Union - -import transformers - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class CheckpointType(str, Enum): - mistral_inference = "mistral_inference" - state_spaces = "state_spaces" - hf = "hf" - - -def get_ckpt_type(model_path): - hf_config = transformers.AutoConfig.from_pretrained(model_path, - trust_remote_code=True) - if hasattr(hf_config, "ssm_cfg") and hf_config.ssm_cfg: - return CheckpointType.state_spaces - if os.path.exists(os.path.join(model_path, "params.json")): - return CheckpointType.mistral_inference - return CheckpointType.hf - - -class MambaConfig(PretrainedConfig): - - def __init__(self, - *, - residual_in_fp32: bool = True, - pad_vocab_size_multiple: int = -1, - layer_types: List[str] = ["recurrent"], - **kwargs): - self.residual_in_fp32 = residual_in_fp32 - self.pad_vocab_size_multiple = pad_vocab_size_multiple - self.layer_types = layer_types - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in MambaConfig - - return output - - def update(self, data_dict): - self.__dict__.update(data_dict) - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - ckpt_type = get_ckpt_type(hf_config_or_dir) - - mamba_version = 'Mamba1' - if ckpt_type == CheckpointType.hf: - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, - "pad_vocab_size_multiple", 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - ( - vocab_size % pad_vocab_size_multiple) - return cls(architecture="MambaForCausalLM", - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=mapping.world_size, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - vocab_size=vocab_size, - mamba_version=mamba_version, - hidden_act=hf_config.hidden_act, - rms_norm=hf_config.rms_norm, - residual_in_fp32=hf_config.residual_in_fp32, - pad_vocab_size_multiple=pad_vocab_size_multiple, - rnn_hidden_size=hf_config.intermediate_size, - rnn_conv_dim_size=hf_config.intermediate_size, - state_size=hf_config.state_size, - conv_kernel=hf_config.conv_kernel, - use_bias=hf_config.use_bias, - mapping=mapping, - quantization=quant_config, - **kwargs) - elif ckpt_type == CheckpointType.state_spaces: - - mamba_version = 'Mamba2' - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, - "pad_vocab_size_multiple", 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - ( - vocab_size % pad_vocab_size_multiple) - assert hasattr(hf_config, - 'ssm_cfg') and hf_config.ssm_cfg['layer'] == 'Mamba2' - config = json.load( - open(os.path.join(hf_config_or_dir, 'config.json'))) - ssm_cfg = config.pop('ssm_cfg') - cfg_to_mamba_cfg = { - 'd_model': 'hidden_size', - 'n_layer': 'num_hidden_layers', - 'fused_add_norm': None, - 'tie_embeddings': None, - } - ssm_cfg_to_mamba_cfg = { - 'd_state': 'state_size', - 'd_conv': 'conv_kernel', - 'bias': 'use_bias', - 'headdim': 'head_dim', - 'ngroups': 'n_groups', - 'chunk_size': 'chunk_size', - 'rmsnorm': 'ssm_rmsnorm', - } - for k in cfg_to_mamba_cfg: - if k in config: - v = config.pop(k) - if cfg_to_mamba_cfg[k] is not None: - config[cfg_to_mamba_cfg[k]] = v - for k in ssm_cfg_to_mamba_cfg: - if k in ssm_cfg and ssm_cfg_to_mamba_cfg[k] is not None: - config[ssm_cfg_to_mamba_cfg[k]] = ssm_cfg[k] - - if 'expand' in config: - expand = config['expand'] - hf_config.intermediate_size = expand * config['hidden_size'] - else: - hf_config.intermediate_size = 2 * config['hidden_size'] - mamba2_default_cfg = { - 'n_groups': 1, - 'hidden_size': hf_config.d_model, - 'head_dim': 64, - 'chunk_size': 256, - 'state_size': 128, - } - hf_config.update(mamba2_default_cfg) - - conv_dim = hf_config.intermediate_size + 2 * hf_config.n_groups * hf_config.state_size - ssm_rmsnorm = getattr(hf_config, "ssm_rmsnorm", hf_config.rms_norm) - mamba2_cfg = { - 'rnn_head_size': hf_config.head_dim, - 'rnn_conv_dim_size': conv_dim, - 'ngroups': hf_config.n_groups, - 'chunk_size': hf_config.chunk_size, - 'ssm_rmsnorm': ssm_rmsnorm, - } - hf_config.update(mamba2_cfg) - - return cls(architecture="MambaForCausalLM", - dtype=dtype, - num_hidden_layers=hf_config.n_layer, - num_attention_heads=mapping.world_size - if mapping is not None else 1, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - vocab_size=vocab_size, - mamba_version=mamba_version, - hidden_act=hf_config.hidden_act, - rms_norm=hf_config.rms_norm, - residual_in_fp32=hf_config.residual_in_fp32, - pad_vocab_size_multiple=pad_vocab_size_multiple, - rnn_hidden_size=hf_config.intermediate_size, - rnn_conv_dim_size=hf_config.rnn_conv_dim_size, - state_size=hf_config.state_size, - conv_kernel=hf_config.conv_kernel, - use_bias=hf_config.use_bias, - mapping=mapping, - quantization=quant_config, - rnn_head_size=hf_config.rnn_head_size, - ngroups=hf_config.ngroups, - chunk_size=hf_config.chunk_size, - ssm_rmsnorm=hf_config.ssm_rmsnorm, - **kwargs) - elif ckpt_type == CheckpointType.mistral_inference: - mamba_version = 'Mamba2' - - config = json.load( - open(os.path.join(hf_config_or_dir, 'params.json'))) - cfg_to_mamba_cfg = { - 'dim': 'hidden_size', - 'n_layers': 'num_hidden_layers', - 'n_groups': 'n_groups', - 'fused_add_norm': None, - 'tie_embeddings': None, - 'model_type': None, - } - for k in cfg_to_mamba_cfg: - if k in config: - v = config.pop(k) - if cfg_to_mamba_cfg[k] is not None: - config[cfg_to_mamba_cfg[k]] = v - - config['architecture'] = 'MambaForCuasualLM' - config['dtype'] = dtype - config['num_attention_heads'] = mapping.world_size - - hf_config = MambaConfig(**config) - mamba2_default_cfg = { - 'n_groups': 8, - 'hidden_size': 4096, - 'head_dim': 64, - 'chunk_size': 256, - 'state_size': 128, - 'conv_kernel': 4, - 'use_bias': False - } - - hf_config.update(mamba2_default_cfg) - conv_dim = hf_config.intermediate_size + 2 * hf_config.n_groups * hf_config.state_size - ssm_rmsnorm = getattr(hf_config, "ssm_rmsnorm", hf_config.rms_norm) - mamba2_cfg = { - 'rnn_head_size': hf_config.head_dim, - 'rnn_conv_dim_size': conv_dim, - 'ngroups': hf_config.n_groups, - 'chunk_size': hf_config.chunk_size, - 'ssm_rmsnorm': ssm_rmsnorm, - } - hf_config.update(mamba2_cfg) - - if 'expand' in config: - expand = config['expand'] - hf_config.intermediate_size = expand * hf_config.hidden_size - else: - hf_config.intermediate_size = 2 * hf_config.hidden_size - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, - "pad_vocab_size_multiple", 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - ( - vocab_size % pad_vocab_size_multiple) - - return cls( - architecture="MambaForCausalLM", - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=mapping.world_size, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - # num_key_value_heads=num_key_value_heads, - vocab_size=vocab_size, - mamba_version=mamba_version, - hidden_act=hf_config.hidden_act, - rms_norm=hf_config.rms_norm, - residual_in_fp32=hf_config.residual_in_fp32, - pad_vocab_size_multiple=pad_vocab_size_multiple, - rnn_hidden_size=hf_config.intermediate_size, - rnn_conv_dim_size=hf_config.rnn_conv_dim_size, - state_size=hf_config.state_size, - conv_kernel=hf_config.conv_kernel, - use_bias=hf_config.use_bias, - mapping=mapping, - quantization=quant_config, - rnn_head_size=hf_config.rnn_head_size, - ngroups=hf_config.n_groups, - chunk_size=hf_config.chunk_size, - ssm_rmsnorm=hf_config.ssm_rmsnorm, - **kwargs) - else: - pass - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, "pad_vocab_size_multiple", - 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - (vocab_size % - pad_vocab_size_multiple) diff --git a/tensorrt_llm/models/mamba/convert.py b/tensorrt_llm/models/mamba/convert.py deleted file mode 100644 index f55bda43c7c2..000000000000 --- a/tensorrt_llm/models/mamba/convert.py +++ /dev/null @@ -1,245 +0,0 @@ -# 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 copy -import re -import time -from pathlib import Path -from typing import Union - -import torch - -import tensorrt_llm -from tensorrt_llm.models.convert_utils import (iterate_shard_files, - load_state_dict) - - -def get_weight(config, prefix, dtype): - return config[prefix + '.weight'].to(dtype).detach() - - -def get_bias(config, prefix, dtype): - if (prefix + '.bias') in config: - return config[prefix + '.bias'].to(dtype).detach() - return None - - -def get_weight_and_bias(config, prefix, dtype_w, dtype_b): - return get_weight(config, prefix, - dtype_w), get_bias(config, prefix, dtype_b) - - -def split(v, tp_size, idx, dim=0): - assert v.shape[dim] % tp_size == 0 - split_size = v.shape[dim] // tp_size - if tp_size == 1: - return v - return torch.split(v, split_size, dim=dim)[idx] - - -def rename_hf_to_tllm(name: str): - """ Rename a HF parameter name by the corresponding TRT-LLM style name. """ - # remove model - if 'model.' in name: - name = name.replace('model.', '') - - # change layer name - if 'embeddings.' in name: - name = name.replace('embeddings', 'vocab_embedding') - elif 'embedding.' in name: - name = name.replace('embedding', 'vocab_embedding') - norm_pattern = r'\d\.norm\.' - if 'mixer.' in name: - name = name.replace('mixer.', 'ssm.') - elif re.search(norm_pattern, name): - name = name.replace('norm.', 'input_layernorm.') - elif 'norm_f.' in name: - name = name.replace('norm_f.', 'ln_f.') - - # Parameter names in ssm layers - if 'A_log' in name: - name = name.replace('A_log', 'A') - elif 'dt_proj.bias' in name: - name = name.replace('dt_proj.bias', 'dt_bias') - return name - - -def convert_hf_mamba(hf_mamba, dtype='float32'): - weights = {} - tik = time.time() - - model_params = dict(hf_mamba.named_parameters()) - dtype = getattr(torch, dtype) - - # Parameter names in mamba block - for l in range(hf_mamba.config.num_hidden_layers): - # ssm layer - prefix = f'backbone.layers.{l}.mixer.' - tllm_prex = f'backbone.layers.{l}.ssm.' - for layer in ['conv1d', 'x_proj', 'dt_proj', 'out_proj']: - dtype_b = torch.float32 if layer == 'dt_proj' else dtype - weight, bias = get_weight_and_bias(model_params, prefix + layer, - dtype, dtype_b) - if layer == 'conv1d': - weight = weight.unsqueeze(3) - tllm_weight_name = tllm_prex + layer + '.weight' - tllm_bias_name = tllm_prex + ('dt_bias' if layer == 'dt_proj' else - layer + '.bias') - weights[tllm_weight_name] = weight - if bias is not None: - weights[tllm_bias_name] = bias - # in_proj - weight, bias = get_weight_and_bias(model_params, prefix + 'in_proj', - dtype, dtype) - in_proj_weights = torch.split(weight, weight.size(0) // 2, dim=0) - tllm_weight_name = tllm_prex + 'in_proj.weight' - weights[tllm_weight_name.replace('proj', 'proj_x')] = in_proj_weights[0] - weights[tllm_weight_name.replace('proj', 'proj_z')] = in_proj_weights[1] - if bias is not None: - in_proj_biases = torch.split(bias, bias.size(0) // 2, dim=0) - tllm_bias_name = tllm_prex + 'in_proj.bias' - weights[tllm_bias_name.replace('proj', - 'proj_x')] = in_proj_biases[0] - weights[tllm_bias_name.replace('proj', - 'proj_x')] = in_proj_biases[1] - - # A and D - Aparam = model_params[prefix + 'A_log'].float().detach() - Aparam = Aparam.permute(1, 0).contiguous() - weights[tllm_prex + 'A'] = -torch.exp(Aparam) - weights[tllm_prex + 'D'] = model_params[prefix + 'D'].float().detach() - # norm - prefix = f'backbone.layers.{l}.norm' - tllm_prex = f'backbone.layers.{l}.input_layernorm.' - weight, bias = get_weight_and_bias(model_params, prefix, dtype, dtype) - weights[tllm_prex + 'weight'] = weight - if bias is not None: - weights[tllm_prex + 'bias'] = bias - - # others - for layer in ['backbone.embeddings', 'backbone.norm_f']: - weight, bias = get_weight_and_bias(model_params, layer, dtype, dtype) - layer = layer.replace('embeddings', 'vocab_embedding') - layer = layer.replace('norm_f', 'ln_f') - weights[layer + '.weight'] = weight - if bias is not None: - weights[layer + '.bias'] = bias - weights['lm_head.weight'], _ = get_weight_and_bias(model_params, - 'backbone.embeddings', - dtype, dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def convert_from_hf_checkpoint(mamba_config: dict, model_dir: Union[str, Path]): - - print('Loading weights from HF Mamba...') - tik = time.time() - - tp_rank = mamba_config.mapping.tp_rank - tp_size = mamba_config.mapping.tp_size - d_inner = mamba_config.rnn_hidden_size - d_state = mamba_config.state_size - dtype = mamba_config.dtype - mamba_version = mamba_config.mamba_version - weights = {} - if isinstance(dtype, str): - dtype = tensorrt_llm.str_dtype_to_torch(dtype) - - for model_file in iterate_shard_files(model_dir, 0): - # logger.debug(f'Loading file {str(model_file)}...') - model_params = load_state_dict(model_file, dtype=dtype) - for name, param in model_params.items(): - # logger.debug(f'Converting weight {name}...') - tllm_name = rename_hf_to_tllm(name) - param = param.detach().cpu() - if 'A_log' in name: - param = -torch.exp(param.float()) - if mamba_version == 'Mamba1': - param = param.permute(1, 0).contiguous() - elif 'D' in name: - param = param.float() - elif 'dt_proj.bias' in name: - param = param.float() - elif 'dt_bias' in name: - param = param.float() - elif 'conv1d.weight' in name: - param = param.unsqueeze(3) - - # split in_proj in Mamba1 - if 'in_proj' in name and mamba_version == 'Mamba1': - in_proj_params = torch.split(param, param.size(0) // 2, dim=0) - weights[tllm_name.replace('proj', 'proj_x')] = in_proj_params[0] - weights[tllm_name.replace('proj', 'proj_z')] = in_proj_params[1] - elif 'in_proj' in name and mamba_version == 'Mamba2': - nheads = d_inner // mamba_config.rnn_head_size - ngroups = mamba_config.ngroups - - in_proj_z, in_proj_x, in_proj_b, in_proj_c, in_proj_dt = torch.split( - param, [ - d_inner, d_inner, ngroups * d_state, ngroups * d_state, - nheads - ], - dim=0) - in_proj_z = split(in_proj_z, tp_size, tp_rank, dim=0) - in_proj_x = split(in_proj_x, tp_size, tp_rank, dim=0) - in_proj_b = split(in_proj_b, tp_size, tp_rank, dim=0) - in_proj_c = split(in_proj_c, tp_size, tp_rank, dim=0) - in_proj_dt = split(in_proj_dt, tp_size, tp_rank, dim=0) - in_proj = torch.concat( - [in_proj_z, in_proj_x, in_proj_b, in_proj_c, in_proj_dt]) - weights[tllm_name] = in_proj.contiguous() - elif 'conv1d' in name and mamba_version == 'Mamba2': - ngroups = mamba_config.ngroups - conv_x, conv_b, conv_c = torch.split( - param, [d_inner, ngroups * d_state, ngroups * d_state], - dim=0) - conv_x = split(conv_x, tp_size, tp_rank, dim=0) - conv_b = split(conv_b, tp_size, tp_rank, dim=0) - conv_c = split(conv_c, tp_size, tp_rank, dim=0) - conv = torch.concat([conv_x, conv_b, conv_c]) - weights[tllm_name] = conv.contiguous() - elif any(keyword in name for keyword in ( - 'mixer.norm.weight', - 'A_log', - 'D', - 'dt_proj.bias', - 'dt_bias', - )) and mamba_version == 'Mamba2': - weights[tllm_name] = split(param, tp_size, tp_rank, dim=0) - elif 'out_proj' in name and mamba_version == 'Mamba2': - weights[tllm_name] = split(param, tp_size, tp_rank, - dim=1).contiguous() - else: - weights[tllm_name] = param - del model_params - - # lm_head - emb = weights['backbone.vocab_embedding.weight'] - if 'lm_head.weight' not in weights or weights['lm_head.weight'].data_ptr( - ) == emb.data_ptr(): - weights['lm_head.weight'] = copy.deepcopy(emb) - if mamba_version == 'Mamba2': - weights['lm_head.weight'] = split(weights['lm_head.weight'], - tp_size, - tp_rank, - dim=0) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/mamba/model.py b/tensorrt_llm/models/mamba/model.py deleted file mode 100644 index 0ccd4abd3a38..000000000000 --- a/tensorrt_llm/models/mamba/model.py +++ /dev/null @@ -1,471 +0,0 @@ -# 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 os -from collections import OrderedDict -from typing import List, Optional, Union - -import tensorrt as trt -from transformers import AutoModelForCausalLM - -from ..._common import default_net -from ..._utils import str_dtype_to_trt -from ...functional import (Tensor, arange, cast, concat, expand, - gather_last_token_logits, shape, unsqueeze) -from ...layers import ColumnLinear, Embedding, LayerNorm, Mamba, Mamba2, RmsNorm -from ...mapping import Mapping -from ...module import Module, ModuleList -from ...plugin import current_all_reduce_helper -from ..generation_mixin import GenerationMixin -from ..modeling_utils import PretrainedConfig, PretrainedModel, QuantConfig -from .config import MambaConfig -from .convert import convert_from_hf_checkpoint, convert_hf_mamba - - -class MambaLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.dtype = config.dtype - self.residual_in_fp32 = config.residual_in_fp32 - n_layer = config.num_hidden_layers - self.last_layer = layer_idx == n_layer - 1 - - if config.mamba_version == 'Mamba1': - assert config.mapping.tp_size == 1, "Mamba1 can not support tensor parallelism." - self.ssm = Mamba(config.hidden_size, - config.rnn_hidden_size, - d_state=config.state_size, - d_conv=config.conv_kernel, - bias=config.use_bias, - dtype=config.dtype) - elif config.mamba_version == 'Mamba2': - self.ssm = Mamba2(config.hidden_size, - config.rnn_hidden_size, - d_state=config.state_size, - d_conv=config.conv_kernel, - headdim=config.rnn_head_size, - ngroups=config.ngroups, - chunk_size=config.chunk_size, - bias=config.use_bias, - rmsnorm=config.ssm_rmsnorm, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size) - if config.rms_norm: - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - else: - self.input_layernorm = LayerNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states: Tensor, - residual: Tensor, - conv_state: Tensor, - ssm_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - - hidden_states = self.input_layernorm(hidden_states) - - ssm_out, present_conv, present_ssm = self.ssm( - hidden_states, - conv_state=conv_state, - ssm_state=ssm_state, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - conv_indices=conv_indices) - if self.residual_in_fp32: - residual = residual + cast(ssm_out, 'float32') - hidden_states = cast(residual, self.dtype) - else: - residual = residual + ssm_out - hidden_states = residual - - if self.last_layer: - return hidden_states, None, present_conv, present_ssm - else: - return hidden_states, residual, present_conv, present_ssm - - -class MambaModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.d_conv = config.conv_kernel - self.d_inner = config.rnn_hidden_size // config.mapping.tp_size - n_layer = config.num_hidden_layers - self.residual_in_fp32 = config.residual_in_fp32 - if config.vocab_size % config.pad_vocab_size_multiple != 0: - config.vocab_size += config.pad_vocab_size_multiple - ( - config.vocab_size % config.pad_vocab_size_multiple) - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = ModuleList( - [MambaLayer(config, i) for i in range(n_layer)]) - if config.rms_norm: - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - else: - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - conv_states, - ssm_states, - host_request_types, - last_token_ids, - host_context_lengths, - slot_mapping: Optional[Tensor] = None): - hidden_states = self.vocab_embedding(input_ids) - - # Get conv state indices - indices = None - if not default_net().plugin_config.mamba_conv1d_plugin: - batch_size = shape(input_ids, 0) - indices = expand( - unsqueeze(arange(0, self.d_conv - 1, dtype='int32'), 0), - concat([batch_size, self.d_conv - 1])) - offsets = expand(unsqueeze(last_token_ids, 1), - concat([batch_size, self.d_conv - 1])) - indices = unsqueeze(indices + offsets, 1) - indices = expand( - indices, concat([batch_size, self.d_inner, self.d_conv - 1])) - - residual = cast(hidden_states, - 'float32') if self.residual_in_fp32 else hidden_states - hidden_values = [hidden_states, residual] - present_convs, present_ssms = [], [] - for layer, past_conv, past_ssm in zip(self.layers, conv_states, - ssm_states): - hidden_values = layer(hidden_values[0], hidden_values[1], past_conv, - past_ssm, host_request_types, last_token_ids, - host_context_lengths, slot_mapping, indices) - present_convs.append(hidden_values[2]) - present_ssms.append(hidden_values[3]) - hidden_states = hidden_values[0] - hidden_states = self.ln_f(hidden_states) - return hidden_states, tuple(present_convs), tuple(present_ssms) - - -class MambaForCausalLM(PretrainedModel): - config_class = MambaConfig - - def __init__(self, config: PretrainedConfig): - super().__init__(config) - dtype = config.dtype - logits_dtype = config.logits_dtype - if isinstance(dtype, str): - self.dtype = str_dtype_to_trt(dtype) - else: - assert isinstance(dtype, trt.DataType) - self.dtype = dtype - - self.config = config - self.mamba_version = config.mamba_version - self.d_inner = config.rnn_hidden_size // config.mapping.tp_size - self.d_conv = config.conv_kernel - self.d_state = config.state_size - self.conv_dim = config.rnn_conv_dim_size // config.mapping.tp_size - self.gather_context_logits = False - - if isinstance(logits_dtype, str): - self._logits_dtype = str_dtype_to_trt(logits_dtype) - else: - assert isinstance(logits_dtype, trt.DataType) - self._logits_dtype = logits_dtype - - self.backbone = MambaModel(config) - self.lm_head = ColumnLinear(config.hidden_size, - config.vocab_size, - bias=False, - dtype=dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - def __post_init__(self): - return - - def forward(self, - input_ids, - conv_states, - ssm_states, - host_request_types, - last_token_ids, - last_token_ids_for_logits, - host_context_lengths, - slot_mapping: Optional[Tensor] = None): - hidden_states, present_convs, present_ssms = self.backbone( - input_ids, conv_states, ssm_states, host_request_types, - last_token_ids, host_context_lengths, slot_mapping) - - if not self.gather_context_logits: - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids_for_logits, - default_net().plugin_config.remove_input_padding) - - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output('logits', self._logits_dtype) - if not default_net().plugin_config.paged_state: - for i, present_conv in enumerate(present_convs): - present_conv.mark_output(f'present_conv_state_{i}', self.dtype) - for i, present_ssm in enumerate(present_ssms): - present_ssm.mark_output(f'present_rnn_state_{i}', self.dtype) - - return (lm_logits, present_convs, present_ssms) - - def prepare_inputs( - self, - max_batch_size, - max_input_len, - max_seq_len, - max_num_tokens, - use_cache, - max_beam_width: int = 1, - opt_num_tokens: int = None, - opt_batch_size: int = 0, - prompt_embedding_table_size: int = 0, - max_draft_len: int = 0, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - speculative_decoding_draft_tokens_external: bool = False): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - assert speculative_decoding_draft_tokens_external == False, "Speculative decoding is not supported in Mamba" - assert max_beam_width == 1, "We don't support beam search for the Mamba model." - - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_state = default_net().plugin_config.paged_state - multiple_profiles = default_net().plugin_config.multiple_profiles - use_mamba_conv1d_plugin = default_net( - ).plugin_config.mamba_conv1d_plugin - - self.gather_context_logits = gather_context_logits - mapping = self.config.mapping - - # basic inputs - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gemm_plugin=use_gemm_plugin, - use_mamba_conv1d_plugin=use_mamba_conv1d_plugin, - remove_input_padding=remove_input_padding, - paged_state=paged_state) - - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_num_tokens=max_num_tokens, - max_draft_len=max_draft_len, - opt_batch_size=opt_batch_size, - opt_num_tokens=opt_num_tokens, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles) - - if remove_input_padding: - assert use_mamba_conv1d_plugin, "mamba_conv1d_plugin is needed to support remove_input_padding" - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('num_tokens', ranges['num_tokens_range']), - ])) - else: - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - ranges['bb_range']), - ('input_len', ranges['inlen_range']), - ])) - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor( - mapping, num_profiles) - - # recurrent inputs - conv_states = [] - ssm_states = [] - if use_mamba_conv1d_plugin: - conv_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('kernel_size', [self.d_conv - 1] * num_profiles), - ('dim_size', [self.conv_dim] * num_profiles), - ]) - else: - conv_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('dim_size', [self.conv_dim] * num_profiles), - ('kernel_size', [self.d_conv - 1] * num_profiles), - ]) - - if self.mamba_version == 'Mamba2': - headdim = self.config.rnn_head_size - nheads = self.d_inner // headdim - ssm_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('head_size', [nheads] * num_profiles), - ('state_size', [self.d_state] * num_profiles), - ('headdim_size', [headdim] * num_profiles), - ]) - ssm_state_shape = [-1, nheads, self.d_state, headdim] - else: - ssm_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('state_size', [self.d_state] * num_profiles), - ('dim_size', [self.d_inner] * num_profiles), - ]) - ssm_state_shape = [-1, self.d_state, self.d_inner] - one_dim_range = OrderedDict([ - ('buffer_count', [1] * num_profiles), - ]) - - for i in range(self.config.num_hidden_layers): - if default_net().plugin_config.paged_state: - conv_state = Tensor(name=f'conv_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - - ssm_state = Tensor(name=f'rnn_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - else: - if use_mamba_conv1d_plugin: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, self.d_conv - 1, self.conv_dim], - dim_range=conv_state_dim_range) - else: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, self.conv_dim, self.d_conv - 1], - dim_range=conv_state_dim_range) - - ssm_state = Tensor(name=f'past_rnn_state_{i}', - dtype=self.dtype, - shape=ssm_state_shape, - dim_range=ssm_state_dim_range) - - conv_states.append(conv_state) - ssm_states.append(ssm_state) - - host_request_types = Tensor( - name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', ranges['bb_range'])]), - ) - - if remove_input_padding: - host_context_lengths = Tensor( - name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', ranges['bb_range'])]), - ) - else: - host_context_lengths = None - - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', ranges['bbd_range']), - ]), - ) - last_token_ids_for_logits = None - if not gather_context_logits: - last_token_ids_for_logits = last_token_ids - - return_dict = { - 'input_ids': input_ids, - 'conv_states': conv_states, - 'ssm_states': ssm_states, - 'host_request_types': host_request_types, - 'last_token_ids': last_token_ids, - 'last_token_ids_for_logits': last_token_ids_for_logits, - 'host_context_lengths': host_context_lengths, - } - - if default_net().plugin_config.paged_state: - slot_mapping = Tensor( - name='slot_mapping', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', ranges['bb_range'])]), - ) - return_dict['slot_mapping'] = slot_mapping - - return return_dict - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - config = MambaConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if not os.path.exists(hf_model_dir): - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype="auto", trust_remote_code=True) - - assert isinstance(hf_model, transformers.PreTrainedModel) - weights = convert_hf_mamba(hf_model, dtype) - else: - weights = convert_from_hf_checkpoint(config, hf_model_dir) - - model = cls(config) - model.load(weights) - - return model diff --git a/tensorrt_llm/models/medusa/__init__.py b/tensorrt_llm/models/medusa/__init__.py deleted file mode 100644 index 2a36ca922710..000000000000 --- a/tensorrt_llm/models/medusa/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. diff --git a/tensorrt_llm/models/medusa/config.py b/tensorrt_llm/models/medusa/config.py deleted file mode 100644 index 0da2777244a4..000000000000 --- a/tensorrt_llm/models/medusa/config.py +++ /dev/null @@ -1,114 +0,0 @@ -# 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 json -from typing import Optional, Union - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..llama.config import LLaMAConfig -from ..modeling_utils import PretrainedConfig, QuantConfig -from ..qwen.config import QWenConfig - - -# Medusa-specific config is stored and retrieved from GenericMedusaConfig. -class MedusaConfig(PretrainedConfig): - - def __init__(self, - *, - num_medusa_heads: int = 4, - num_medusa_layers: int = 1, - max_draft_len: int = 63, - **kwargs): - - model_type = str(kwargs.get('model_type', '')).lower() - generic_medusa_config = QWenConfig if 'qwen' in model_type else LLaMAConfig - self.config = generic_medusa_config(**kwargs) - - # Add objects - self.config.num_medusa_heads = num_medusa_heads - self.config.num_medusa_layers = num_medusa_layers - self.config.max_draft_len = max_draft_len - - def to_dict(self): - output = self.config.to_dict() - output['num_medusa_heads'] = self.config.num_medusa_heads - output['num_medusa_layers'] = self.config.num_medusa_layers - output['max_draft_len'] = self.config.max_draft_len - return output - - # Specialization to redirect accesses to self.config - def __getattr__(self, name): - return getattr(self.config, name) - - def __getstate__(self): - return self.__dict__ - - def __setstate__(self, state): - self.__dict__.update(state) - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - trust_remote_code = kwargs.pop('trust_remote_code', True) - speculative_config_or_dir = kwargs.pop('speculative_model_dir', None) - speculative_config = kwargs.pop("speculative_config", None) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - if hasattr(hf_config, "medusa"): - # is modelOpt ckpt - num_medusa_heads = hf_config.medusa["num_medusa_heads"] - num_medusa_layers = hf_config.medusa["num_medusa_layers"] - else: - config_file = speculative_config_or_dir / "config.json" - with open(config_file) as fp: - config = json.load(fp) - - num_medusa_heads = speculative_config.num_medusa_heads if speculative_config is not None else config.get( - 'num_medusa_heads', None) - num_medusa_layers = config.get('medusa_num_layers', None) - - return cls(architecture="MedusaForCausalLM", - 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=hf_config.hidden_act, - norm_epsilon=hf_config.rms_norm_eps, - mapping=mapping, - quantization=quant_config, - num_medusa_heads=num_medusa_heads, - num_medusa_layers=num_medusa_layers, - **kwargs) diff --git a/tensorrt_llm/models/medusa/model.py b/tensorrt_llm/models/medusa/model.py deleted file mode 100644 index 525c4ec391f0..000000000000 --- a/tensorrt_llm/models/medusa/model.py +++ /dev/null @@ -1,267 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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 json -from typing import Optional, Union - -from transformers import AutoModelForCausalLM - -from tensorrt_llm._utils import numpy_to_torch -from tensorrt_llm.models.llama.model import LLaMAForCausalLM -from tensorrt_llm.models.medusa.weight import load_medusa_hf -from tensorrt_llm.models.qwen.model import QWenForCausalLM - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import ACT2FN, stack -from ...layers import ColumnLinear -from ...mapping import Mapping -from ...module import Module, ModuleList -from ..modeling_utils import PretrainedModel, QuantConfig -from .config import MedusaConfig -from .weight import convert_hf_llama - - -class MedusaLayer(Module): - - def __init__( - self, - hidden_size, - hidden_act="silu", - dtype=None, - mapping=Mapping(), - ): - super().__init__() - self.linear = ColumnLinear(hidden_size, - hidden_size, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True) - self.hidden_act = hidden_act - - def forward(self, x): - return x + ACT2FN[self.hidden_act](self.linear(x)) - - -class MedusaHead(Module): - - def __init__( - self, - num_layers, - hidden_size, - vocab_size, - hidden_act="silu", - dtype=None, - mapping=Mapping(), - ): - super().__init__() - self.medusa_layers = ModuleList([ - MedusaLayer(hidden_size=hidden_size, - hidden_act=hidden_act, - dtype=dtype, - mapping=mapping) for _ in range(num_layers) - ]) - self.lm_head = ColumnLinear(hidden_size, - vocab_size, - bias=False, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True) - return - - def forward(self, x): - hidden_states = x - for layer in self.medusa_layers: - hidden_states = layer(hidden_states) - return self.lm_head(hidden_states) - - -# MedusaForCausalLm is a thin wrapper that picks parent class for GenericMedusaForCausalLM. -# All medusa functionality is defined in GenericMedusaForCausalLM. -class MedusaForCausalLm(PretrainedModel): - config_class = MedusaConfig - - def __init__(self, config: MedusaConfig): - super().__init__(config) - - BaseLM = QWenForCausalLM if hasattr( - config, - "model_type") and "qwen" in config.model_type else LLaMAForCausalLM - - class GenericMedusaForCausalLM(BaseLM): - - def __init__(self, config: MedusaConfig): - super().__init__(config) - self.num_medusa_heads = config.num_medusa_heads - self.num_medusa_layers = config.num_medusa_layers - self.hidden_size = config.hidden_size - self.vocab_size = config.vocab_size - vocab_size_padded = pad_vocab_size(self.vocab_size, - config.mapping.tp_size) - self.medusa_heads = ModuleList([ - MedusaHead(num_layers=self.num_medusa_layers, - hidden_size=config.hidden_size, - vocab_size=vocab_size_padded, - hidden_act=config.hidden_act, - dtype=config.dtype, - mapping=config.mapping) - for _ in range(self.num_medusa_heads) - ]) - self.max_medusa_token_len = config.max_draft_len - - def forward(self, *args, **kwargs): - output_original = True - hidden_states = super().forward(*args, **kwargs) - - if kwargs['use_cache']: - if default_net().plugin_config.paged_kv_cache: - lm_logits, hidden_states, _ = hidden_states - else: - lm_logits, presents, hidden_states = hidden_states - - if self.mapping.is_last_pp_rank(): - medusa_logits = [] - for i in range(self.num_medusa_heads): - medusa_logits.append( - self.medusa_heads[i](hidden_states)) - # [num_medusa_heads, batch_size, num_medusa_tokens + 1, padded_vocab_size]. - # Remove padding [num_medusa_heads, batch_size * num_medusa_tokens + 1, padded_vocab_size]. - medusa_logits = stack(medusa_logits, dim=0) - medusa_logits.mark_output('medusa_logits', - self.config.logits_dtype) - else: - hidden_states.mark_output('hidden_states_output', - self.config.dtype) - - if kwargs['use_cache'] and default_net( - ).plugin_config.paged_kv_cache == False: - if self.mapping.is_last_pp_rank(): - if output_original: - return (medusa_logits, lm_logits, presents) - return (medusa_logits, presents) - return (hidden_states, presents) - else: - if self.mapping.is_last_pp_rank(): - if output_original: - return medusa_logits, lm_logits - return medusa_logits - return hidden_states - - def prepare_inputs(self, *args, **kwargs): - kwargs['speculative_decoding_draft_tokens_external'] = False - kwargs['max_draft_len'] = self.max_medusa_token_len - return super().prepare_inputs(*args, **kwargs) - - self.model = GenericMedusaForCausalLM(config) - - # Specialization to redirect accesses to self.model - def __getattribute__(self, name): - if name == 'model' or '__' in name: - return object.__getattribute__(self, name) - else: - model = object.__getattribute__(self, 'model') - return model.__getattribute__(name) - - # Override specialized __setattr__ defined in Module - def __setattr__(self, name, value) -> None: - object.__setattr__(self, name, value) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - speculative_model_dir = kwargs.get('speculative_model_dir', None) - - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = MedusaConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - # ModelOpt ckpt has combined base model and Medusa-head - is_modelopt_ckpt = True if not speculative_model_dir else False - - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if is_modelopt_ckpt: - hf_model = LLaMAForCausalLM.from_hugging_face( - hf_model_dir, - dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - else: - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - dtype="auto", - trust_remote_code=trust_remote_code) - - assert isinstance(hf_model, transformers.PreTrainedModel) - - if is_modelopt_ckpt: - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in hf_model.named_parameters() - } - else: - weights = convert_hf_llama( - hf_model, - config.mapping, - dtype='float16', - use_parallel_embedding=config.use_parallel_embedding) - - model = cls(config) - - if is_modelopt_ckpt: - num_medusa_heads = config.config.num_medusa_heads - num_medusa_layers = config.config.num_medusa_layers - speculative_model_dir = hf_model_or_dir - else: - config_file = speculative_model_dir / "config.json" - with open(config_file) as fp: - model_config = json.load(fp) - - num_medusa_heads = kwargs[ - 'speculative_config'].num_medusa_heads if 'speculative_config' in kwargs else model_config.get( - 'medusa_num_heads', None) - num_medusa_layers = model_config.get('medusa_num_layers', None) - medusa_weights = load_medusa_hf(medusa_path=speculative_model_dir, - num_medusa_heads=num_medusa_heads, - num_medusa_layers=num_medusa_layers, - mapping=mapping, - dtype="float16", - is_modelopt_ckpt=is_modelopt_ckpt) - weights.update(medusa_weights) - model.load(weights) - return model diff --git a/tensorrt_llm/models/medusa/weight.py b/tensorrt_llm/models/medusa/weight.py deleted file mode 100644 index 6964dbdd3ee5..000000000000 --- a/tensorrt_llm/models/medusa/weight.py +++ /dev/null @@ -1,648 +0,0 @@ -import copy -import functools -import time -from collections import defaultdict -from pathlib import Path - -import numpy as np -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers.models.llama.modeling_llama import LlamaDecoderLayer -from transformers.pytorch_utils import Conv1D - -from tensorrt_llm._utils import str_dtype_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (dup_kv_weight, generate_int8, - smooth_gemm, - smooth_gemm_fc1_gate, split, - split_matrix_tp, split_qkv_tp) - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - postfix='weight'): - results = {} - if use_weight_only: - v = weight.t().contiguous().cpu() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + postfix] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + postfix] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def load_medusa_hf(medusa_path: str, - num_medusa_heads: int, - num_medusa_layers: int, - mapping=Mapping(), - dtype='float32', - use_weight_only=False, - plugin_weight_only_quant_type=None, - is_modelopt_ckpt=False): - logger.info("Loading Medusa heads' weights ...") - - if is_modelopt_ckpt: - from safetensors.torch import load_file - state_dict = {} - for filename in sorted(Path(medusa_path).glob("*.safetensors")): - print(f"Loading the weights of Medusa heads from {filename}") - state_dict.update(load_file(filename)) - else: - is_ckpt_safetensors = False - - ckpt_file = Path(medusa_path) / "medusa_lm_head.pt" - if not ckpt_file.exists(): - ckpt_file = Path(medusa_path) / "medusa_lm_head.safetensors" - is_ckpt_safetensors = True - - if is_ckpt_safetensors: - logger.info("Safetensors Found ...") - from safetensors.torch import load_file - state_dict = load_file(ckpt_file) - else: - state_dict = torch.load(ckpt_file, map_location="cpu") - - torch_dtype = str_dtype_to_torch(dtype) - weights = {} - - prefix = "medusa_heads." if is_modelopt_ckpt else "" - for h in range(num_medusa_heads): - for l in range(num_medusa_layers): - w = state_dict[f"{prefix}{h}.{l}.linear.weight"].clone().to( - torch_dtype) - - split_v = split(w, mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight( - split_v, f'medusa_heads.{h}.medusa_layers.{l}.linear.', - None, use_weight_only, plugin_weight_only_quant_type)) - - b = state_dict[f"{prefix}{h}.{l}.linear.bias"].clone().to( - torch_dtype) - - weights['medusa_heads.{}.medusa_layers.{}.linear.bias'.format( - h, l)] = split(b, mapping.tp_size, mapping.tp_rank) - - lm = state_dict[f"{prefix}{h}.{num_medusa_layers}.weight"].clone().to( - torch_dtype) # LM Head - - weights['medusa_heads.{}.lm_head.weight'.format(h)] = split( - lm, mapping.tp_size, mapping.tp_rank) - - # scaling factors - if is_modelopt_ckpt: - scaling_dtype = str_dtype_to_torch("float32") - weights[f'medusa_heads.{h}.medusa_layers.{l}.linear.activation_scaling_factor'] = \ - state_dict[f"{prefix}{h}.{l}.linear.input_scale"].clone().to(scaling_dtype) - - weights[f'medusa_heads.{h}.medusa_layers.{l}.linear.weights_scaling_factor'] = \ - state_dict[f"{prefix}{h}.{l}.linear.weight_scale"].clone().to(scaling_dtype) - - weights['medusa_heads.{}.lm_head.activation_scaling_factor'.format(h)] = \ - state_dict[f"{prefix}{h}.{num_medusa_layers}.input_scale"].clone().to(scaling_dtype) - - weights['medusa_heads.{}.lm_head.weights_scaling_factor'.format(h)] = \ - state_dict[f"{prefix}{h}.{num_medusa_layers}.weight_scale"].clone().to(scaling_dtype) - - return weights - - -@torch.no_grad() -def smooth_llama_model(model, scales, alpha, llama_qkv_para, llama_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, LlamaDecoderLayer): - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - llama_qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - input_ids = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(input_ids) - for h in hooks: - h.remove() - return act_scales - - -def get_weight(config, prefix, dtype): - if config[prefix + '.weight'].dtype != dtype: - config[prefix + '.weight'].data = config[prefix + '.weight'].to(dtype) - return config[prefix + '.weight'] - - -def get_bias(config, prefix, dtype): - if config[prefix + '.bias'].dtype != dtype: - config[prefix + '.bias'].data = config[prefix + '.bias'].to(dtype) - return config[prefix + '.bias'] - - -def get_weight_and_bias(config, prefix, dtype): - return get_weight(config, prefix, dtype), get_bias(config, prefix, dtype) - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = np.split(data, [local_dim, local_dim + head_size], axis=-1) - q_split = np.split(q, tp_size, axis=-1) - k_split = np.split(k, tp_size, axis=-1) - v_split = np.split(v, tp_size, axis=-1) - return [ - np.concatenate((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - original_weights = vals["weight.int8.col"] - - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array(cur_per_channel_value, - dtype=np.float32).reshape(col_shape)).contiguous() - else: - original_weights = np.array(vals["weight.int8"]) - cur_weights = np.split(original_weights, tensor_parallel, - axis=cat_dim)[rank] - - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - cur_per_channel_value = vals["scale_y_accum_quant"] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_llama(hf_model, - mapping, - rank=0, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[], - lora_config=None): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - intermediate_size = hf_model.config.intermediate_size - num_key_value_heads = hf_model.config.num_key_value_heads - mha_mode = (num_key_value_heads == num_attention_heads) - - num_hidden_layers = hf_model.config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = f'model.layers.{l}.' - tllm_prex = f'transformer.layers.{layer_idx}.' - - q_weight = get_weight(model_params, prefix + 'self_attn.q_proj', dtype) - k_weight = get_weight(model_params, prefix + 'self_attn.k_proj', dtype) - v_weight = get_weight(model_params, prefix + 'self_attn.v_proj', dtype) - - if not mha_mode: - head_size = hidden_size // num_attention_heads - if num_key_value_heads < tensor_parallel: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, num_key_value_heads, - tensor_parallel) - v_weight = dup_kv_weight(v_weight, num_key_value_heads, - tensor_parallel) - assert (k_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - assert (v_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - split_v = torch.concat((wq, wk, wv)) - - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - if use_smooth_quant: - qkv_weight = qkv_para[prefix + 'self_attn.qkv_proj'] - - if not mha_mode: - hidden_size = qkv_weight.shape[0] - local_dim = hidden_size - head_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(hidden_size, - local_dim + 2 * head_size) - else: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.qkv_proj'), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.qkv.', [ - 1, 3 * hidden_size // tensor_parallel - if mha_mode else hidden_size // tensor_parallel + - (hidden_size // num_key_value_heads) // - tensor_parallel * 2 - ], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_y = torch.cat([ - act_range.get(prefix + 'self_attn.q_proj')["y"], - act_range.get(prefix + 'self_attn.k_proj')["y"], - act_range.get(prefix + 'self_attn.v_proj')["y"] - ], - dim=0) - - int8_kv_scales = qkv_y.max() / 127. - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'self_attn.o_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'self_attn.o_proj')], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - split_v = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - if use_smooth_quant: - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8(mlp_gate_weight, - act_range.get(prefix + 'mlp.up_proj')) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.gate.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() #verified - int8_weights = generate_int8( - mlp_fc_weight, act_range.get(prefix + 'mlp.gate_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'mlp.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'mlp.down_proj'], - smoother_shape=[1, intermediate_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - v = get_weight(model_params, 'model.embed_tokens', dtype) - - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/mllama/__init__.py b/tensorrt_llm/models/mllama/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/mllama/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/mllama/config.py b/tensorrt_llm/models/mllama/config.py deleted file mode 100644 index cbd7f1b8f387..000000000000 --- a/tensorrt_llm/models/mllama/config.py +++ /dev/null @@ -1,256 +0,0 @@ -# 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 json -from pathlib import Path -from typing import List, Optional, Union - -from ..._utils import get_hf_rope_theta -from ...functional import LayerNormPositionType, LayerNormType, MLPType -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class MLLaMAConfig(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, - cross_attention: bool = True, - cross_attention_layers: List[int] = None, - vision_output_dim: int = 0, - has_position_embedding=False, - type_vocab_size=None, - rescale_before_lm_head=False, - layernorm_type=LayerNormType.RmsNorm, - layernorm_position=LayerNormPositionType.pre_layernorm, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - has_model_final_layernorm=True, - model_type='MLLaMAForCausalLM', - skip_cross_kv=False, - mlp_type=MLPType.GatedMLP, - has_embedding_scale=False, - residual_scaling=1.0, - has_lm_head_bias=False, - num_buckets=None, - max_distance=0, - relative_attention=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 - - assert cross_attention - self.cross_attention = cross_attention - self.cross_attention_layers = cross_attention_layers - assert vision_output_dim != 0 - self.vision_output_dim = vision_output_dim - - self.has_position_embedding = has_position_embedding - self.type_vocab_size = type_vocab_size - self.rescale_before_lm_head = rescale_before_lm_head - self.layernorm_type = layernorm_type - self.layernorm_position = layernorm_position - self.has_attention_qkvo_bias = has_attention_qkvo_bias - self.has_mlp_bias = has_mlp_bias - self.has_model_final_layernorm = has_model_final_layernorm - self.model_type = model_type - self.skip_cross_kv = skip_cross_kv - self.mlp_type = mlp_type - self.has_embedding_scale = has_embedding_scale - self.residual_scaling = residual_scaling - self.has_lm_head_bias = has_lm_head_bias - self.num_buckets = num_buckets - self.max_distance = max_distance - self.relative_attention = relative_attention - self.skip_cross_attn_blocks = True - - kwargs.pop('embed_vocab_size', None) - kwargs.pop('num_kv_heads_per_layer', None) - kwargs.pop('num_kv_heads_per_cross_attn_layer', None) - super().__init__(**kwargs) - - @property - def embed_vocab_size(self): - return self.vocab_size + 8 #FIXME The vocab_size of embedding contains the special tokens for image - - @property - def num_kv_heads_per_layer(self): - num_kv_heads_per_layer = [ - self.num_key_value_heads for _ in range(self.num_hidden_layers) - ] - for layer_idx in self.cross_attention_layers: - num_kv_heads_per_layer[layer_idx] = 0 - return num_kv_heads_per_layer - - @property - def num_kv_heads_per_cross_attn_layer(self): - num_kv_heads_per_cross_attn_layer = [ - 0 for _ in range(self.num_hidden_layers) - ] - for layer_idx in self.cross_attention_layers: - num_kv_heads_per_cross_attn_layer[ - layer_idx] = self.num_key_value_heads - return num_kv_heads_per_cross_attn_layer - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in MLLaMAConfig - 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['cross_attention'] = self.cross_attention - output['cross_attention_layers'] = self.cross_attention_layers - output['vision_output_dim'] = self.vision_output_dim - output['embed_vocab_size'] = self.embed_vocab_size - output['num_kv_heads_per_layer'] = self.num_kv_heads_per_layer - output[ - 'num_kv_heads_per_cross_attn_layer'] = self.num_kv_heads_per_cross_attn_layer - output['skip_cross_attn_blocks'] = self.skip_cross_attn_blocks - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - hf_text_config = hf_config.text_config - hf_vision_config = hf_config.vision_config - num_key_value_heads = getattr(hf_text_config, "num_key_value_heads", - hf_text_config.num_attention_heads) - - hidden_act = hf_text_config.hidden_act if hasattr( - hf_text_config, "hidden_act") else hf_text_config.hidden_activation - norm_epsilon = hf_text_config.rms_norm_eps - - head_dim = getattr( - hf_text_config, "head_dim", - hf_text_config.hidden_size // hf_text_config.num_attention_heads) - head_size = getattr(hf_text_config, "kv_channels", head_dim) - attn_bias = getattr(hf_text_config, 'bias', False) or getattr( - hf_text_config, 'attention_bias', False) - rotary_scaling = getattr(hf_text_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_text_config, 10000.0) - residual_mlp = getattr(hf_text_config, "parallel_attn_mlp_res", False) - disable_weight_only_quant_plugin = kwargs.pop( - 'disable_weight_only_quant_plugin', False) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_text_config.num_hidden_layers, - num_attention_heads=hf_text_config.num_attention_heads, - hidden_size=hf_text_config.hidden_size, - intermediate_size=hf_text_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - head_size=head_size, - vocab_size=hf_text_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_text_config.max_position_embeddings, - hidden_act=hidden_act, - norm_epsilon=norm_epsilon, - attn_bias=attn_bias, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - residual_mlp=residual_mlp, - disable_weight_only_quant_plugin=disable_weight_only_quant_plugin, - mapping=mapping, - quantization=quant_config, - cross_attention_layers=hf_text_config.cross_attention_layers, - vision_output_dim=hf_vision_config.vision_output_dim, - **kwargs) - - @classmethod - def from_meta_ckpt(cls, - meta_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - - with open(Path(meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - - n_embd = meta_config["dim"] - n_head = meta_config["n_heads"] - n_kv_head = meta_config.get("n_kv_heads", n_head) - vocab_size = meta_config.get("vocab_size", 32000) - - # Reset vocab_size to 32000 for LLama v2 checkpoint. - if vocab_size == -1: - vocab_size = 32000 - - if "hidden_dim" in meta_config: - inter_size = meta_config["hidden_dim"] - else: - multiple_of = meta_config.get("multiple_of", 1) - n_embd_ = int(4 * n_embd * 2 / 3) - ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - inter_size = multiple_of * ( - (int(n_embd_ * ffn_dim_multiplier) + multiple_of - 1) // - multiple_of) - - dtype = infer_dtype(dtype, 'bfloat16') - - if meta_config.get('use_scaled_rope'): - rotary_scaling = {"type": "llama3"} - else: - rotary_scaling = meta_config.get("rope_scaling") - - # meta checkpoint don't have vocab_size|hidden_act|rotary_base specified, use same default value as HF - return cls(architecture="MLLaMAForCausalLM", - dtype=dtype, - num_hidden_layers=meta_config["n_layers"], - 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=2048, - hidden_act='silu', - rotary_scaling=rotary_scaling, - rotary_base=meta_config.get('rope_theta', 10000), - norm_epsilon=meta_config["norm_eps"], - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/mllama/model.py b/tensorrt_llm/models/mllama/model.py deleted file mode 100644 index a610a480959d..000000000000 --- a/tensorrt_llm/models/mllama/model.py +++ /dev/null @@ -1,1569 +0,0 @@ -# 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 math -from collections import OrderedDict -from typing import List, Optional, Union - -import tensorrt as trt -import torch - -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import numpy_to_torch, str_dtype_to_torch -from tensorrt_llm.functional import (Conditional, LayerNormPositionType, - LayerNormType, MLPType, - PositionEmbeddingType, Tensor, assertion, - gather_last_token_logits, maximum, minimum, - recv, reduce, send, shape, tanh) -from tensorrt_llm.layers import (MLP, Attention, AttentionMaskParams, - AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, FusedGatedMLP, - GatedMLP, GroupNorm, KeyValueCacheParams, - LayerNorm, LoraParams, RmsNorm) -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.lora_helper import (LoraConfig, - get_default_trtllm_modules_to_hf_modules, - use_lora) -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -from tensorrt_llm.models.modeling_utils import PretrainedModel, QuantConfig -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.parameter import Parameter -from tensorrt_llm.quantization import QuantMode - -from .config import MLLaMAConfig - -layernorm_map = { - LayerNormType.LayerNorm: LayerNorm, - LayerNormType.RmsNorm: RmsNorm, - LayerNormType.GroupNorm: GroupNorm, -} - -mlp_map = { - MLPType.MLP: MLP, - MLPType.GatedMLP: GatedMLP, - MLPType.FusedGatedMLP: FusedGatedMLP, -} - -ADD_DEBUG_TENSOR = False - - -class CrossAttentionTransformerBlock(Module): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.RmsNorm, - layernorm_eps=1e-5, - hidden_act="gated-silu", - mlp_type=MLPType.GatedMLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - skip_cross_kv=False, - use_implicit_relative_attention=False, - rotary_embedding_base=None, - rotary_embedding_scaling=None, - quant_mode=QuantMode(0), - ): - super().__init__() - self.local_layer_idx = local_layer_idx - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - self.layernorm_position = layernorm_position - assert self.layernorm_position == LayerNormPositionType.pre_layernorm - - self.cross_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=True, - relative_attention= - False, # Cross attention has no relative attention bias - max_distance=max_distance, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType. - learned_absolute, # we don't use rope for cross attn - skip_cross_kv=skip_cross_kv, - qk_layernorm=True, - layernorm_type=layernorm_type, - quant_mode=quant_mode, - ) - - self.input_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.gate_attn = Parameter(shape=tuple((1, )), dtype=dtype) - - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - ) - - self.post_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.gate_ffwd = Parameter(shape=tuple((1, )), dtype=dtype) - - self.residual_scaling = residual_scaling - - self.fp16_clamping = fp16_clamping - self.no_ffn = False - - def forward(self, - hidden_states: Tensor, - encoder_output: Optional[Tensor] = None, - attention_mask_params=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - full_text_row_masked_out_mask: Tensor = None, - skip_cross_attn_blocks: Tensor = None): - - assert isinstance(hidden_states, Tensor) - - if encoder_output: - assert isinstance(encoder_output, Tensor) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/1.0: hidden_states', - hidden_states.dtype) - # cross attention - residual = hidden_states * self.residual_scaling - - # skip input_layernorm - if skip_cross_attn_blocks is not None: - input_ln_conditional = Conditional(skip_cross_attn_blocks) - skip_result = input_ln_conditional.add_input(hidden_states) - hidden_states = input_ln_conditional.add_input(hidden_states) - hidden_states = self.input_layernorm(hidden_states) - hidden_states = input_ln_conditional.add_output( - skip_result, hidden_states) - else: - hidden_states = self.input_layernorm(hidden_states) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/2.1: normed_input', - hidden_states.dtype) - # pass full_text_row_masked_out_mask and xattn_mask - attention_output = self.cross_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.cross_attention_mask, - attention_packed_mask=attention_mask_params. - cross_attention_packed_mask, - encoder_output=encoder_output, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - skip_attn=skip_cross_attn_blocks, - ) - - if use_cache: - attention_output, presents_cross = attention_output - if ADD_DEBUG_TENSOR: - attention_output.mark_output( - f'{self.local_layer_idx:2d}/3.1: cross_attention_output', - attention_output.dtype) - - attn_residual_scale = tanh(self.gate_attn.value.cast(trt.float32)).cast( - attention_output.dtype) - - attention_input = hidden_states - hidden_states = residual + attn_residual_scale * attention_output - - # use to skip attention_output with residual - # Since conditional does not work for gpt_attention_plugin, we replace the - # attention_output by hidden_states (input of attention) now. - if skip_cross_attn_blocks is not None: - attn_conditional = Conditional(skip_cross_attn_blocks) - skip_result = attn_conditional.add_input(attention_input) - hidden_states = attn_conditional.add_input(hidden_states) - hidden_states = attn_conditional.add_output(skip_result, - hidden_states) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/3.2: cross_attn_output_with_residual', - hidden_states.dtype) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - # MLP - # skip post_layernorm and mlp - if skip_cross_attn_blocks is not None: - mlp_conditional = Conditional(skip_cross_attn_blocks) - skip_case = mlp_conditional.add_input(hidden_states) - hidden_states = mlp_conditional.add_input(hidden_states) - - attention_output = attention_output * full_text_row_masked_out_mask # TODO should move this mask into attention? - - residual = hidden_states * self.residual_scaling - - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.1: mlp_output', - hidden_states.dtype) - - hidden_states = hidden_states * full_text_row_masked_out_mask - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.2: masked_mlp_output', - hidden_states.dtype) - ffn_residual_scale = tanh(self.gate_ffwd.value.cast(trt.float32)).cast( - hidden_states.dtype) - hidden_states = residual + ffn_residual_scale * hidden_states * float( - not self.no_ffn) - - if skip_cross_attn_blocks is not None: - hidden_states = mlp_conditional.add_output(skip_case, hidden_states) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.4: transformer_out', - hidden_states.dtype) - - if use_cache: - return (hidden_states, presents_cross) - return hidden_states - - -class TransformerBlock(Module): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.RmsNorm, - layernorm_eps=1e-5, - hidden_act="gated-silu", - mlp_type=MLPType.GatedMLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - skip_cross_kv=False, - use_implicit_relative_attention=False, - rotary_embedding_base=None, - rotary_embedding_scaling=None, - quant_mode=QuantMode(0), - ): - super().__init__() - self.local_layer_idx = local_layer_idx - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - self.layernorm_position = layernorm_position - assert self.layernorm_position == LayerNormPositionType.pre_layernorm - - self.self_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=False, - relative_attention=relative_attention, - max_distance=max_distance if use_implicit_relative_attention else 0, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType.relative - if relative_attention else PositionEmbeddingType.rope_gpt_neox, - use_implicit_relative_attention=use_implicit_relative_attention, - rotary_embedding_base=rotary_embedding_base, - rotary_embedding_scaling=rotary_embedding_scaling, - quant_mode=quant_mode, - ) - - self.input_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - ) - - self.post_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.residual_scaling = residual_scaling - - self.fp16_clamping = fp16_clamping - - def forward( - self, - hidden_states: Tensor, - encoder_output: Optional[Tensor] = None, # not used - attention_mask_params=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - full_text_row_masked_out_mask: Tensor = None, # not used - skip_cross_attn_blocks=None, - ): - assert isinstance(hidden_states, Tensor) - - # self-attention - residual = hidden_states * self.residual_scaling - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/1.0: hidden_states', - hidden_states.dtype) - - hidden_states = self.input_layernorm(hidden_states) - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/2.1: normed attn_input', - hidden_states.dtype) - - attention_output = self.self_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.self_attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params) - - if use_cache: - attention_output, presents_self = attention_output - - if ADD_DEBUG_TENSOR: - attention_output.mark_output( - f'{self.local_layer_idx:2d}/3.1: self_attention_output', - attention_output.dtype) - - hidden_states = residual + attention_output - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/3.1: attention_output_with_residual', - hidden_states.dtype) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - # MLP - residual = hidden_states * self.residual_scaling - - hidden_states = self.post_layernorm(hidden_states) - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/3.2: normed_mlp_input', - hidden_states.dtype) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.1: mlp_output', - hidden_states.dtype) - - hidden_states = residual + hidden_states - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.2: mlp_output_residual', - hidden_states.dtype) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if use_cache: - return (hidden_states, presents_self) - return hidden_states - - -class MLLaMAModel(Module): - - def __init__(self, config: MLLaMAConfig) -> None: - super().__init__() - self.config = config - self.position_embedding_type = config.position_embedding_type - - self.mapping = self.config.mapping - - self.layernorm_type = self.config.layernorm_type - ln_type = layernorm_map[self.layernorm_type] - - self.has_attention_qkvo_bias = self.config.has_attention_qkvo_bias - self.has_mlp_bias = self.config.has_mlp_bias - - self.has_model_final_layernorm = self.config.has_model_final_layernorm - self._dtype = self.config.dtype - # no quantization considered for now - self._kv_dtype = self._dtype - self._logits_dtype = self.config.logits_dtype - - self.total_num_layers = self.config.num_hidden_layers - self.num_layers = self.config.num_hidden_layers // self.mapping.pp_size - - self.hidden_size = self.config.hidden_size - self.encoder_hidden_size = self.config.hidden_size - self.num_heads = self.config.num_attention_heads - # num_kv_heads = self.num_heads - num_kv_heads = self.config.num_key_value_heads - if num_kv_heads is None or num_kv_heads <= 0: - num_kv_heads = self.num_heads - self.num_kv_heads = num_kv_heads - self.head_size = self.hidden_size // self.num_heads if self.config.head_size is None else self.config.head_size - - self.fp16_clamping = False - - self.skip_cross_kv = self.config.skip_cross_kv - self.mlp_type = MLPType.MLP if not hasattr( - self.config, "mlp_type") else self.config.mlp_type - self.use_implicit_relative_attention = self.config.use_implicit_relative_attention if hasattr( - self.config, "use_implicit_relative_attention") else False - - self.cross_attention_layers = self.config.cross_attention_layers - - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding( - self.config.embed_vocab_size, - self.config.hidden_size, - dtype=self._dtype, - tp_size=self.mapping.tp_size - if self.config.use_parallel_embedding else 1, - tp_group=self.mapping.tp_group - if self.config.use_parallel_embedding else None, - sharding_dim=self.config.embedding_sharding_dim, - tp_rank=self.mapping.tp_rank) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - _layers = [] - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - args = { - "local_layer_idx": local_layer_idx, - "hidden_size": self.config.hidden_size, - "ffn_hidden_size": self.config.intermediate_size, - "num_attention_heads": self.num_heads, - "num_kv_heads": self.num_kv_heads, - "head_size": self.head_size, - "max_position_embeddings": self.config.max_position_embeddings, - "layernorm_position": self.config.layernorm_position, - "layernorm_eps": self.config.norm_epsilon, - "layernorm_type": self.config.layernorm_type, - "hidden_act": self.config.hidden_act, - "mlp_type": self.mlp_type, - "mapping": self.mapping, - "dtype": self._dtype, - "residual_scaling": self.config.residual_scaling, - "max_distance": self.config.max_distance, - "num_buckets": self.config.num_buckets, - "fp16_clamping": self.fp16_clamping, - "skip_cross_kv": self.skip_cross_kv, - "rotary_embedding_base": self.config.rotary_base, - "rotary_embedding_scaling": self.config.rotary_scaling, - "quant_mode": self.config.quant_mode, - } - if layer_idx in self.cross_attention_layers: - assert layers_range[0] == 0, "not support PP now" - _layers.append(CrossAttentionTransformerBlock(**args)) - else: - _layers.append(TransformerBlock(**args)) - - self.layers = ModuleList(_layers) - - if self.mapping.is_last_pp_rank(): - self.ln_f = None - if self.has_model_final_layernorm: - self.ln_f = ln_type(normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype) - - if self.config.relative_attention and not self.use_implicit_relative_attention: - self.rel_attn_table = Parameter( - shape=(self.config.num_attention_heads // self.mapping.tp_size, - self.config.num_buckets), - dtype=self._dtype) - - def forward( - self, - decoder_input_ids: Tensor, - encoder_output: Tensor, - use_cache=False, - attention_mask_params=None, - last_token_ids=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - lora_params: LoraParams = None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - skip_cross_attn_blocks: Optional[Tensor] = None, - ): - if self.mapping.is_first_pp_rank(): - assert isinstance(decoder_input_ids, Tensor) - else: - assert isinstance(hidden_states, Tensor) - - # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(decoder_input_ids) - self.register_network_output('embedding_layer_output', - hidden_states) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - kv_cache_params.fill_none_tensor_list(len(self.layers)) - - full_text_row_masked_out_mask = reduce( - (attention_mask_params.cross_attention_mask).cast( - hidden_states.dtype), - trt.ReduceOperation.MAX, - dim=-1, - keepdim=True) - if ADD_DEBUG_TENSOR: - full_text_row_masked_out_mask.mark_output( - "full_text_row_masked_out_mask", - full_text_row_masked_out_mask.dtype) - - cross_attention_mask_type = attention_mask_params.cross_attention_mask.dtype - attention_mask_params.cross_attention_mask = ( - attention_mask_params.cross_attention_mask.cast( - full_text_row_masked_out_mask.dtype) * - full_text_row_masked_out_mask).cast(cross_attention_mask_type) - - invert_mask = 1.0 - attention_mask_params.cross_attention_mask.cast( - hidden_states.dtype) - invert_full_text_row_masked_out_mask = 1.0 - full_text_row_masked_out_mask - final_mask = invert_mask - invert_full_text_row_masked_out_mask - attention_mask_params.cross_attention_mask = final_mask.cast( - cross_attention_mask_type) - if ADD_DEBUG_TENSOR: - attention_mask_params.cross_attention_mask.mark_output( - "attention_mask_params.cross_attention_mask", - attention_mask_params.cross_attention_mask.dtype) - - if use_cache: - presents = [] - for i, (decoder_layer, past) in enumerate( - zip(self.layers, kv_cache_params.past_key_value)): - - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(i) - hidden_states = decoder_layer( - hidden_states, - encoder_output=encoder_output, - attention_mask_params=attention_mask_params, - use_cache=use_cache, - kv_cache_params=KeyValueCacheParams( - past_key_value=past, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - cache_indirection=kv_cache_params.cache_indirection, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=kv_cache_params. - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers=kv_cache_params. - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping=kv_cache_params. - host_cross_kv_cache_pool_mapping, - ), - skip_cross_attn_blocks=skip_cross_attn_blocks if isinstance( - decoder_layer, CrossAttentionTransformerBlock) else None, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - full_text_row_masked_out_mask=full_text_row_masked_out_mask, - ) - - if use_cache: - present = hidden_states[1] - presents.append((present)) - hidden_states = hidden_states[0] - - if self.mapping.is_last_pp_rank(): - if self.ln_f: - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - def precompute_relative_attention_bias(self, build_config): - if self.config.relative_attention and not self.use_implicit_relative_attention: - relative_attention_bias_builder = torch.ops.tensorrt_llm.relative_attention_bias - rel_attn_precomputed = torch.zeros( - (self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len + 1, build_config.max_seq_len + 1), - dtype=str_dtype_to_torch(self.config.dtype), - device='cuda') - rel_attn_table = numpy_to_torch( - self.rel_attn_table.raw_value).to('cuda') - relative_attention_bias_builder( - rel_attn_precomputed, - rel_attn_table, - self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len, - self.config.num_buckets, - False, - self.config.max_distance, - ) - for layer_idx in range(self.num_layers): - self.layers[layer_idx].self_attention.set_rel_attn_table( - build_config.max_seq_len, rel_attn_precomputed) - - -# TODO try to inherit the DecoderModelForCausalLM -class MLLaMAForCausalLM(PretrainedModel): - config_class = MLLaMAConfig - - def __init__(self, config: MLLaMAConfig): - super().__init__(config) - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - self.transformer = MLLaMAModel(config) - - self.mapping = self.config.mapping - - self.has_model_final_layernorm = self.config.has_model_final_layernorm - self._dtype = self.config.dtype - self._kv_dtype = self._dtype - self._logits_dtype = self.config.logits_dtype - - if self.mapping.is_last_pp_rank(): - self.lm_head = ColumnLinear( - self.config.hidden_size, - self.config.vocab_size, - bias=False if not hasattr(self.config, "has_lm_head_bias") else - self.config.has_lm_head_bias, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True, - ) - - self.trtllm_modules_to_hf_modules = { - **get_default_trtllm_modules_to_hf_modules(), - "attn_q": "self_attn.q_proj", - "attn_k": "self_attn.k_proj", - "attn_v": "self_attn.v_proj", - "attn_dense": "self_attn.o_proj", - "cross_attn_q": "encoder_attn.q_proj", - "cross_attn_k": "encoder_attn.k_proj", - "cross_attn_v": "encoder_attn.v_proj", - "cross_attn_dense": "encoder_attn.o_proj", - } - - def forward( - self, - decoder_input_ids: Tensor, - encoder_output: Tensor, - use_cache=False, - attention_mask_params=None, - last_token_ids=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - lora_params: LoraParams = None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - skip_cross_attn_blocks: Optional[Tensor] = None, - ): - if self.mapping.is_first_pp_rank(): - assert isinstance(decoder_input_ids, Tensor) - else: - assert isinstance(hidden_states, Tensor) - attention_params = Attention.fill_attention_params( - self, attention_params) - hidden_states = self.transformer( - decoder_input_ids=decoder_input_ids, - encoder_output=encoder_output, - use_cache=use_cache, - attention_mask_params=attention_mask_params, - last_token_ids=last_token_ids, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - hidden_states=hidden_states, - lora_params=lora_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size, - skip_cross_attn_blocks=skip_cross_attn_blocks, - ) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - # [bs, seq, hidden_size] or [num_tokens, hidden_size] -> [bs, hidden_size] - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) - - # [bs, hidden_size] -> [bs, vocab_size] - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output(f'logits', self._logits_dtype) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - hidden_states.mark_output(f'hidden_states_output', self._dtype) - - if use_cache and default_net().plugin_config.paged_kv_cache == False: - for i, present in zip(self.mapping.pp_layers(self.total_num_layers), - presents): - present[0].mark_output(f'present_key_value_{i}', self._kv_dtype) - if default_net().plugin_config.gpt_attention_plugin: - present[1].mark_output(f'cross_present_key_value_{i}', - self._kv_dtype) - if self.mapping.is_last_pp_rank(): - return (lm_logits, tuple(presents)) - return (hidden_states, tuple(presents)) - else: - if self.mapping.is_last_pp_rank(): - return lm_logits - return hidden_states - - def prepare_inputs(self, - max_batch_size, - max_beam_width, - max_decoder_input_len, - max_seq_len, - max_encoder_input_len, - gather_context_logits: bool = False, - gather_generation_logits: bool = False, - lora_target_modules: List[str] = None, - prompt_embedding_table_size: int = 0, - use_cache=True, - *args, - **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - # Prepare inputs - max_output_len = max_decoder_input_len + max_seq_len - - head_size = self.transformer.head_size - num_kv_heads = (self.transformer.num_kv_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - - encoder_head_size = head_size - encoder_num_kv_heads = num_kv_heads - - bb_range = [ - 1, (max_batch_size * max_beam_width + 1) // 2, - max_batch_size * max_beam_width - ] - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - beam_width_range = [1, (max_beam_width + 1) // 2, max_beam_width] - inlen_range = [ - 1, 1, max_decoder_input_len - ] # context phase >= 1 (if forced_input_ids), generation phase = 1 - encoder_inlen_range = [ - 1, (max_encoder_input_len + 1) // 2, max_encoder_input_len - ] - mask_len_range = [1, (max_output_len + 1) // 2 + 1, max_output_len + 1] - max_output_len_range = [0, (max_output_len + 1) // 2, max_output_len] - - encoder_num_tokens_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len * max_batch_size + 1) // 2, - max_encoder_input_len * max_batch_size, - ] - decoder_num_tokens_range = [ - 1, - max_batch_size * max_beam_width, - max(max_decoder_input_len * max_batch_size, - max_beam_width * max_batch_size), - ] - - # No enable_two_optimization_profiles support yet - encoder_input_len_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len + 1) // 2, - max_encoder_input_len - ] - # pack masks into bits (store as int32). - max_cross_packed_mask_dim0 = max_batch_size * ( - (max_decoder_input_len + 128 - 1) // 128) * 128 - max_cross_packed_mask_dim1 = ( - (max_encoder_input_len + 256 - 1) // 256) * 256 // 32 - cross_packed_mask_dim0_range = [ - 1, (max_cross_packed_mask_dim0 + 1) // 2, max_cross_packed_mask_dim0 - ] - cross_packed_mask_dim1_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_cross_packed_mask_dim1 + 1) // 2, - max_cross_packed_mask_dim1 - ] - past_key_value = [] - sequence_length = None - host_past_key_value_lengths = None - attention_mask_params = AttentionMaskParams() - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - remove_input_padding = default_net().plugin_config.remove_input_padding - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - use_lora_plugin = default_net().plugin_config.lora_plugin - kv_cache_type = None - if not use_cache: - kv_cache_type = KVCacheType.DISABLED - else: - if paged_kv_cache: - kv_cache_type = KVCacheType.PAGED - else: - kv_cache_type = KVCacheType.CONTINUOUS - - input_ids, hidden_states = None, None - if remove_input_padding: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, self.hidden_size], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ('hidden_size', [self.hidden_size]), - ])) - else: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('input_len', [inlen_range]), - ])) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, -1, self.hidden_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range - ]), - ('input_len', [inlen_range]), - ('hidden_size', [self.hidden_size]), - ])) - - encoder_input_lengths = Tensor( - name="encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range])]), - ) - encoder_max_input_length = Tensor( - name="encoder_max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("encoder_max_input_length", - [encoder_inlen_range])]), - ) - if remove_input_padding: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, self.config.hidden_size], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("hidden_size", [self.config.hidden_size]), - ]), - ) - else: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, -1, self.config.hidden_size], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("hidden_size", [self.config.hidden_size]), - ]), - ) - - context_lengths = None - host_context_lengths = None - host_request_types = None - host_runtime_perf_knobs = None - host_context_progress = None - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = Tensor(name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - - if use_gpt_attention_plugin: - if kv_cache_type != KVCacheType.DISABLED: - sequence_length = Tensor( - name='sequence_length', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range]) - ]), - ) - host_past_key_value_lengths = Tensor( - name='host_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range]) - ]), - ) - - context_lengths = Tensor(name='context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - host_request_types = Tensor(name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - - host_runtime_perf_knobs = Tensor(name='host_runtime_perf_knobs', - dtype=trt.int64, - shape=[16], - dim_range=OrderedDict([ - ('perf_knob_size', [16]) - ])) - - host_context_progress = Tensor(name='host_context_progress', - dtype=trt.int64, - shape=[1], - dim_range=OrderedDict([ - ('context_progress_size', [1]) - ])) - - last_token_ids = None - if self.mapping.is_last_pp_rank() and not gather_context_logits: - last_token_ids = Tensor( - name="last_token_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_last_token_ids", [bb_range]) - ]), - ) - - attention_mask = None - if not use_gpt_attention_plugin: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('mask_len', [mask_len_range]), - ]), - ) - assert False, "not support non-attention-plugin case now" - - cross_attention_mask = Tensor( - name='cross_attention_mask', - dtype=trt.bool, - shape=[-1, -1], - dim_range=OrderedDict([ - ('decoder_num_tokens_2', - [decoder_num_tokens_range - ]), # TODO should use same name as input_ids - ('encoder_input_len_2', [encoder_input_len_range]), - ]), - ) - - cross_attention_packed_mask = Tensor( - name='cross_attention_packed_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('cross_packed_mask_dim0', [cross_packed_mask_dim0_range]), - ('cross_packed_mask_dim1', [cross_packed_mask_dim1_range]), - ]), - ) - - # create the attention_mask_params. - attention_mask_params = AttentionMaskParams( - attention_mask, None, cross_attention_mask, - cross_attention_packed_mask) - - cache_indirection = Tensor( - name='cache_indirection', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_cache', [bs_range]), - ('beam_width', [beam_width_range]), - ('max_seq_len', [max_output_len_range]), - ]), - ) - - layers_range = self.mapping.pp_layers(self.transformer.total_num_layers) - num_pp_layers = len(layers_range) - - host_max_attention_window_sizes = None - host_sink_token_length = None - if use_gpt_attention_plugin: - host_max_attention_window_sizes = Tensor( - name=f'host_max_attention_window_sizes', - dtype=trt.int32, - shape=[num_pp_layers], - dim_range=OrderedDict([('num_layers', [num_pp_layers])])) - host_sink_token_length = Tensor(name='host_sink_token_length', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('scalar', - [1])])) - # TODO LoRA for mllama is not verified. - lora_weights_pointers = None - lora_ranks = None - lora_params = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - missing_qkv_modules = [] - if any(x in lora_target_modules - for x in ["attn_q", "attn_k", "attn_v"]): - for lora_module in [ - "attn_q", - "attn_k", - "attn_v", - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - if any(x in lora_target_modules - for x in ["cross_attn_q", "cross_attn_k", "cross_attn_v"]): - for lora_module in [ - "cross_attn_q", "cross_attn_k", "cross_attn_v" - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - - # For LoRA - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in (lora_target_modules + missing_qkv_modules): - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([('batch_size_beam_width', - [bb_range]), - ('in_out_scales', [3])])) - lora_weight_pointer_dict.update({ - f'{lora_module}_lora_weights_pointers': - lora_weight_pointer - }) - - lora_rank = Tensor(name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - lora_rank_dict.update( - {f'{lora_module}_lora_ranks': lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - # For cross attention, we need to use encoder_input_lengths (in CPU) to pass - # as the host_context_lengths to the lora_plugin. But for self attention, we - # should keep using the original host_context_lengths. Therefore, we keep both - # of them in the lora_params. - host_encoder_input_lengths = None - if remove_input_padding: - host_encoder_input_lengths = Tensor( - name="host_encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range]) - ]), - ) - - lora_params = LoraParams( - lora_ranks=lora_ranks, - lora_weights_pointers=lora_weights_pointers, - host_context_lengths=host_context_lengths, - max_context_length=max_decoder_input_len, - max_encoder_context_length=max_encoder_input_len, - host_request_types=host_request_types, - host_encoder_input_lengths=host_encoder_input_lengths, - ) - - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - host_kv_cache_pool_pointers = None - host_kv_cache_pool_mapping = None - - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - host_cross_kv_cache_pool_pointers = None - host_cross_kv_cache_pool_mapping = None - - if use_cache: - if not paged_kv_cache: - for i in layers_range: - kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('num_heads', [num_kv_heads]), - ('past_key_len', [max_output_len_range]), - ('head_size', [head_size]), - ]) - kv = Tensor(name=f'past_key_value_{i}', - dtype=self._kv_dtype, - shape=[-1, 2, num_kv_heads, -1, head_size], - dim_range=kv_dim_range) - - past_key_value.append(kv) - - if i in self.transformer.cross_attention_layers: - xa_layer_id = self.transformer.cross_attention_layers.index( - i) + layers_range[-1] - cross_kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('cross_num_heads', [encoder_num_kv_heads]), - ('cross_past_key_len', [encoder_input_len_range]), - ('cross_head_size', [encoder_head_size]), - ]) - cross_kv = Tensor( - name=f'cross_past_key_value_{xa_layer_id}', - dtype=self._kv_dtype, - shape=[ - -1, 2, encoder_num_kv_heads, -1, encoder_head_size - ], - dim_range=cross_kv_dim_range) - past_key_value.append(kv) - - # TODO: Remove this when TRT fix the named dimension - if not remove_input_padding: - assertion( - shape( - input_ids if self.mapping.is_first_pp_rank() else - hidden_states, 0) == shape(kv, 0), 'batch size') - - else: # paged_kv_cache == True - # PagedKV setup for KV cache of self-attention - max_blocks_per_seq_range = [[ - math.ceil(max_output_len_range[0] / tokens_per_block), - math.ceil(max_output_len_range[1] / tokens_per_block), - math.ceil(max_output_len_range[2] / tokens_per_block) - ]] - max_blocks_per_seq_range = [[ - x for x in max_blocks_per_seq_range[0] - ]] - - # PagedKV setup for KV cache of cross-attention - max_cross_blocks_per_seq_range = [[ - math.ceil(encoder_input_len_range[0] / tokens_per_block), - math.ceil(encoder_input_len_range[1] / tokens_per_block), - math.ceil(encoder_input_len_range[2] / tokens_per_block) - ]] - max_cross_blocks_per_seq_range = [[ - x for x in max_cross_blocks_per_seq_range[0] - ]] - - num_kv_cache_pools = 2 - - kv_cache_block_offsets = Tensor( - name=f'kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_block_offsets = Tensor( - name=f'host_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_pool_pointers = Tensor( - name=f'host_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('num_pools', [2]), - ])) - host_kv_cache_pool_mapping = Tensor( - name=f"host_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - # paged blocks for cross kv - cross_kv_cache_block_offsets = Tensor( - name=f'cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_block_offsets = Tensor( - name=f'host_cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_pool_pointers = Tensor( - name=f'host_cross_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('num_pools', [2]), - ])) - host_cross_kv_cache_pool_mapping = Tensor( - name=f"host_cross_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - for i in layers_range: - past_key_value.append(None) - - kv_cache_params = KeyValueCacheParams( - past_key_value=past_key_value, - host_past_key_value_lengths=host_past_key_value_lengths, - host_max_attention_window_sizes=host_max_attention_window_sizes, - host_sink_token_length=host_sink_token_length, - cache_indirection=cache_indirection, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers= - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping= - host_cross_kv_cache_pool_mapping, - ) - - attention_params = AttentionParams( - sequence_length=sequence_length, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_decoder_input_len, - host_request_types=host_request_types, - host_runtime_perf_knobs=host_runtime_perf_knobs, - host_context_progress=host_context_progress, - encoder_input_lengths=encoder_input_lengths, - encoder_max_input_length=encoder_max_input_length, - ) - - cross_kv_cache_gen = Tensor(name='cross_kv_cache_gen', - dtype=trt.bool, - shape=[1], - dim_range=OrderedDict([ - ('boolean', [1]), - ])) - cross_kv_reuse = None - num_heads = (self.transformer.num_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - cross_kv_out_dim = 2 * num_kv_heads * self.transformer.head_size - if self.transformer.skip_cross_kv: - if remove_input_padding: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - else: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, -1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - - skip_cross_attn_blocks = None - if self.config.skip_cross_attn_blocks: - skip_cross_attn_blocks = Tensor(name='skip_cross_attn_blocks', - dtype=trt.bool, - shape=[1], - dim_range=OrderedDict([ - ('boolean', [1]), - ])) - - prompt_embedding_table = None - tasks = None - prompt_vocab_size = None - - if self.mapping.is_first_pp_rank() and prompt_embedding_table_size > 0: - p_embedding_range = [[ - 1, prompt_embedding_table_size // 2, prompt_embedding_table_size - ]] - - prompt_embedding_table = Tensor( - name='prompt_embedding_table', - dtype=self._dtype, - shape=[-1, self.transformer.hidden_size], - dim_range=OrderedDict([ - ('prompt_embedding_table_size', p_embedding_range), - ('hidden_size', [self.transformer.hidden_size]), - ])) - if remove_input_padding: - num_tokens_range = [ - 1, - (max_decoder_input_len * max_batch_size + 1) // 2, - max_decoder_input_len * max_batch_size, - ] - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - else: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([ - ('batch_size', bs_range), - ('broadcast_dim', [1]), - ])) - prompt_vocab_size = Tensor(name='prompt_vocab_size', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('size', [1])])) - - result = { - 'decoder_input_ids': input_ids, - 'encoder_output': encoder_output, - 'use_cache': True, - 'attention_mask_params': attention_mask_params, - 'last_token_ids': last_token_ids, - 'kv_cache_params': kv_cache_params, - 'attention_params': attention_params, - 'hidden_states': hidden_states, - 'lora_params': lora_params, - 'cross_kv_cache_gen': cross_kv_cache_gen, - 'cross_kv_reuse': cross_kv_reuse, - 'prompt_embedding_table': prompt_embedding_table, - 'prompt_tasks': tasks, - 'prompt_vocab_size': prompt_vocab_size, - 'skip_cross_attn_blocks': skip_cross_attn_blocks, - } - - return result - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a MLLaMAForCausalLM object from give parameters - ''' - import transformers - - kwargs.pop('load_by_shard', False) - kwargs.pop('load_model_on_cpu', False) - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = MLLaMAConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - custom_dict = { - "lm_head": "language_model.lm_head", - "transformer.ln_f": "language_model.model.norm", - "transformer": "language_model.model", - "self_attention": "self_attn", - "cross_attention": "cross_attn", - "vocab_embedding": "embed_tokens", - "gate_attn": "cross_attn_attn_gate", - "gate_ffwd": "cross_attn_mlp_gate", - "q_layernorm": "q_norm", - "k_layernorm": "k_norm", - } - - if quant_ckpt_path is not None: - hf_model_dir = quant_ckpt_path - - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model) - - return model diff --git a/tensorrt_llm/models/mmdit_sd3/__init__.py b/tensorrt_llm/models/mmdit_sd3/__init__.py deleted file mode 100644 index 0d106f45ce46..000000000000 --- a/tensorrt_llm/models/mmdit_sd3/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. diff --git a/tensorrt_llm/models/mmdit_sd3/config.py b/tensorrt_llm/models/mmdit_sd3/config.py deleted file mode 100644 index 6d15c4d5c961..000000000000 --- a/tensorrt_llm/models/mmdit_sd3/config.py +++ /dev/null @@ -1,132 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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 Any, Dict, Optional, Sequence, Tuple - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class SD3Transformer2DModelConfig(PretrainedConfig): - - def __init__( - self, - *, - sample_size: int = 128, - patch_size: int = 2, - in_channels: int = 16, - num_layers: int = 24, - attention_head_dim: int = 64, - num_attention_heads: int = 24, - joint_attention_dim: int = 4096, - caption_projection_dim: int = 1536, - pooled_projection_dim: int = 2048, - out_channels: int = 16, - pos_embed_max_size: int = 384, - dual_attention_layers: - Tuple[int] = ( - 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 - ), # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5 - qk_norm: Optional[str] = None, - skip_layers: Optional[Sequence[int]] = None, - use_pretrained_pos_emb: bool = False, - **kwargs): - - kwargs.update({ - 'hidden_size': attention_head_dim * num_attention_heads, - 'num_hidden_layers': num_layers, - 'num_attention_heads': num_attention_heads - }) - super().__init__(**kwargs) - self.sample_size = sample_size - self.patch_size = patch_size - self.in_channels = in_channels - self.num_layers = num_layers - self.attention_head_dim = attention_head_dim - self.num_attention_heads = num_attention_heads - self.joint_attention_dim = joint_attention_dim - self.caption_projection_dim = caption_projection_dim - self.pooled_projection_dim = pooled_projection_dim - self.out_channels = out_channels - self.pos_embed_max_size = pos_embed_max_size - self.dual_attention_layers = dual_attention_layers - self.qk_norm = qk_norm - self.skip_layers = skip_layers - self.use_pretrained_pos_emb = use_pretrained_pos_emb - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in SD3Transformer2DModelConfig - output['sample_size'] = self.sample_size - output['patch_size'] = self.patch_size - output['in_channels'] = self.in_channels - output['num_layers'] = self.num_layers - output['attention_head_dim'] = self.attention_head_dim - output['num_attention_heads'] = self.num_attention_heads - output['joint_attention_dim'] = self.joint_attention_dim - output['caption_projection_dim'] = self.caption_projection_dim - output['pooled_projection_dim'] = self.pooled_projection_dim - output['out_channels'] = self.out_channels - output['pos_embed_max_size'] = self.pos_embed_max_size - output['dual_attention_layers'] = self.dual_attention_layers - output['qk_norm'] = self.qk_norm - output['skip_layers'] = self.skip_layers - output['use_pretrained_pos_emb'] = self.use_pretrained_pos_emb - return output - - @classmethod - def from_hugging_face_config(cls, - hf_config: Dict[str, Any], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - sample_size = hf_config['sample_size'] - patch_size = hf_config['patch_size'] - in_channels = hf_config['in_channels'] - num_layers = hf_config['num_layers'] - attention_head_dim = hf_config['attention_head_dim'] - num_attention_heads = hf_config['num_attention_heads'] - joint_attention_dim = hf_config['joint_attention_dim'] - caption_projection_dim = hf_config['caption_projection_dim'] - pooled_projection_dim = hf_config['pooled_projection_dim'] - out_channels = hf_config['out_channels'] - pos_embed_max_size = hf_config['pos_embed_max_size'] - dual_attention_layers = hf_config['dual_attention_layers'] - qk_norm = hf_config['qk_norm'] - skip_layers = None - use_pretrained_pos_emb = kwargs.get('use_pretrained_pos_emb', False) - dtype = infer_dtype(dtype, hf_config.get('torch_dtype')) - - return cls(architecture='SD3Transformer2DModel', - sample_size=sample_size, - patch_size=patch_size, - in_channels=in_channels, - num_layers=num_layers, - attention_head_dim=attention_head_dim, - num_attention_heads=num_attention_heads, - joint_attention_dim=joint_attention_dim, - caption_projection_dim=caption_projection_dim, - pooled_projection_dim=pooled_projection_dim, - out_channels=out_channels, - pos_embed_max_size=pos_embed_max_size, - dual_attention_layers=dual_attention_layers, - qk_norm=qk_norm, - skip_layers=skip_layers, - use_pretrained_pos_emb=use_pretrained_pos_emb, - dtype=dtype, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/mmdit_sd3/model.py b/tensorrt_llm/models/mmdit_sd3/model.py deleted file mode 100644 index 546abbeade2c..000000000000 --- a/tensorrt_llm/models/mmdit_sd3/model.py +++ /dev/null @@ -1,620 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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 collections import OrderedDict -from typing import Any, Dict, List, Optional - -from ..._utils import str_dtype_to_torch -from ...functional import (Tensor, allgather, chunk, concat, einsum, pad, shape, - unsqueeze) -from ...layers import LayerNorm, Linear -from ...layers.attention import DiffusersAttention -from ...layers.embedding import (CombinedTimestepTextProjEmbeddings, - SD3PatchEmbed) -from ...layers.mlp import (LinearActivation, LinearApproximateGELU, LinearGEGLU, - LinearGELU, LinearSwiGLU) -from ...layers.normalization import (AdaLayerNormContinuous, AdaLayerNormZero, - SD35AdaLayerNormZeroX) -from ...logger import logger -from ...mapping import Mapping -from ...module import Module, ModuleList -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import PretrainedModel -from .config import SD3Transformer2DModelConfig - - -class FeedForward(Module): - - def __init__( - self, - dim: int, - dim_out: Optional[int] = None, - mult: int = 4, - activation_fn: str = "geglu", - inner_dim=None, - bias: bool = True, - mapping=Mapping(), - dtype=None, - ): - super().__init__() - - self.mapping = mapping - self.dtype = dtype - - if inner_dim is None: - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - - if activation_fn == "gelu": - raise NotImplementedError('GELU only support tanh now.') - if activation_fn == "gelu-approximate": - act_fn = LinearGELU(dim, - inner_dim, - approximate="tanh", - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "geglu": - act_fn = LinearGEGLU(dim, - inner_dim, - approximate="tanh", - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "geglu-approximate": - act_fn = LinearApproximateGELU(dim, - inner_dim, - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "swiglu": - act_fn = LinearSwiGLU(dim, - inner_dim, - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "linear-silu": - act_fn = LinearActivation(dim, - inner_dim, - bias=bias, - activation="silu", - mapping=mapping, - dtype=dtype) - - self.net = ModuleList([ - act_fn, - Linear(inner_dim, - dim_out, - bias=bias, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - dtype=self.dtype) - ]) - - def forward(self, hidden_states: Tensor): - for module in self.net: - hidden_states = module(hidden_states) - return hidden_states - - -class JointTransformerBlock(Module): - - def __init__(self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - context_pre_only: bool = False, - qk_norm: Optional[str] = None, - use_dual_attention: bool = False, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.use_dual_attention = use_dual_attention - self.context_pre_only = context_pre_only - context_norm_type = "ada_norm_continous" if context_pre_only else "ada_norm_zero" - - if use_dual_attention: - self.norm1 = SD35AdaLayerNormZeroX(dim, - mapping=mapping, - dtype=dtype) - else: - self.norm1 = AdaLayerNormZero(dim, mapping=mapping, dtype=dtype) - - if context_norm_type == "ada_norm_continous": - self.norm1_context = AdaLayerNormContinuous( - dim, - dim, - elementwise_affine=False, - eps=1e-6, - bias=True, - norm_type="layer_norm", - dtype=dtype) - elif context_norm_type == "ada_norm_zero": - self.norm1_context = AdaLayerNormZero(dim, dtype=dtype) - else: - raise ValueError( - f"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`" - ) - - self.attn = DiffusersAttention( - query_dim=dim, - cross_attention_dim=None, - added_kv_proj_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - context_pre_only=context_pre_only, - bias=True, - qk_norm=qk_norm, - eps=1e-6, - mapping=mapping, - dtype=dtype, - ) - - if use_dual_attention: - self.attn2 = DiffusersAttention( - query_dim=dim, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - bias=True, - qk_norm=qk_norm, - eps=1e-6, - mapping=mapping, - dtype=dtype, - ) - else: - self.attn2 = None - - self.norm2 = LayerNorm(dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - self.ff = FeedForward(dim=dim, - dim_out=dim, - activation_fn="gelu-approximate", - mapping=mapping, - dtype=dtype) - - if not context_pre_only: - self.norm2_context = LayerNorm(dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - self.ff_context = FeedForward(dim=dim, - dim_out=dim, - activation_fn="gelu-approximate", - mapping=mapping, - dtype=dtype) - else: - self.norm2_context = None - self.ff_context = None - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = 0 - - def set_chunk_feed_forward(self, - chunk_size: Optional[int] = None, - dim: int = 0): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - @staticmethod - def _chunked_feed_forward(ff: Module, hidden_states: Tensor, chunk_dim: int, - chunk_size: int): - # "feed_forward_chunk_size" can be used to save memory - if hidden_states.shape[chunk_dim] % chunk_size != 0: - raise ValueError( - f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." - ) - - num_chunks = hidden_states.shape[chunk_dim] // chunk_size - ff_output = concat( - [ - ff(hid_slice) - for hid_slice in chunk(hidden_states, num_chunks, dim=chunk_dim) - ], - dim=chunk_dim, - ) - return ff_output - - def forward(self, - hidden_states: Tensor, - encoder_hidden_states: Tensor, - temb: Tensor, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - *args, - **kwargs): - joint_attention_kwargs = joint_attention_kwargs or {} - if self.use_dual_attention: - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1( - hidden_states, emb=temb) - else: - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( - hidden_states, emb=temb) - - if self.context_pre_only: - norm_encoder_hidden_states = self.norm1_context( - encoder_hidden_states, temb) - else: - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( - encoder_hidden_states, emb=temb) - - # Attention. - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - **joint_attention_kwargs, - ) - - # Process attention outputs for the `hidden_states`. - attn_output = unsqueeze(gate_msa, 1) * attn_output - hidden_states = hidden_states + attn_output - - if self.use_dual_attention: - attn_output2 = self.attn2(hidden_states=norm_hidden_states2, - **joint_attention_kwargs) - attn_output2 = unsqueeze(gate_msa2, 1) * attn_output2 - hidden_states = hidden_states + attn_output2 - - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * ( - 1 + unsqueeze(scale_mlp, 1)) + unsqueeze(shift_mlp, 1) - - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - ff_output = self._chunked_feed_forward(self.ff, norm_hidden_states, - self._chunk_dim, - self._chunk_size) - else: - ff_output = self.ff(norm_hidden_states) - ff_output = unsqueeze(gate_mlp, 1) * ff_output - hidden_states = hidden_states + ff_output - - # Process attention outputs for the `encoder_hidden_states`. - if self.context_pre_only: - encoder_hidden_states = None - else: - context_attn_output = unsqueeze(c_gate_msa, 1) * context_attn_output - encoder_hidden_states = encoder_hidden_states + context_attn_output - - norm_encoder_hidden_states = self.norm2_context( - encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * ( - 1 + unsqueeze(c_scale_mlp, 1)) + unsqueeze(c_shift_mlp, 1) - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - context_ff_output = self._chunked_feed_forward( - self.ff_context, norm_encoder_hidden_states, - self._chunk_dim, self._chunk_size) - else: - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + unsqueeze( - c_gate_mlp, 1) * context_ff_output - - return encoder_hidden_states, hidden_states - - -class SD3Transformer2DModel(PretrainedModel): - config_class = SD3Transformer2DModelConfig - - def __init__(self, config: SD3Transformer2DModelConfig): - super().__init__(config) - self.quant_mode = config.quant_mode - self.mapping = config.mapping - self.dtype = config.dtype - - self.in_channels = config.in_channels - default_out_channels = config.in_channels - self.out_channels = config.out_channels if config.out_channels is not None else default_out_channels - self.inner_dim = config.num_attention_heads * config.attention_head_dim - - self.pos_embed = SD3PatchEmbed( - height=config.sample_size, - width=config.sample_size, - patch_size=config.patch_size, - in_channels=self.in_channels, - embed_dim=self.inner_dim, - pos_embed_max_size=config. - pos_embed_max_size, # hard-code as HF implementation - dtype=self.dtype) - self.time_text_embed = CombinedTimestepTextProjEmbeddings( - embedding_dim=self.inner_dim, - pooled_projection_dim=config.pooled_projection_dim, - mapping=self.mapping, - dtype=self.dtype) - self.context_embedder = Linear(config.joint_attention_dim, - config.caption_projection_dim, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - dtype=self.dtype) - - self.transformer_blocks = ModuleList([ - JointTransformerBlock( - dim=self.inner_dim, - num_attention_heads=config.num_attention_heads, - attention_head_dim=config.attention_head_dim, - context_pre_only=(i == config.num_layers - 1), - qk_norm=config.qk_norm, - use_dual_attention=True - if i in config.dual_attention_layers else False, - mapping=self.mapping, - dtype=self.dtype) for i in range(config.num_layers) - ]) - - self.norm_out = AdaLayerNormContinuous(self.inner_dim, - self.inner_dim, - elementwise_affine=False, - eps=1e-6, - dtype=self.dtype) - self.proj_out = Linear(self.inner_dim, - config.patch_size * config.patch_size * - self.out_channels, - bias=True, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - dtype=self.dtype) - - self.skip_layers = config.skip_layers - self.use_pretrained_pos_emb = config.use_pretrained_pos_emb - self.config = config - - def forward(self, - hidden_states: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - pooled_projections: Optional[Tensor] = None, - timestep: Optional[Tensor] = None, - block_controlnet_hidden_states: List[Tensor] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None): - height, width = hidden_states.shape[-2:] - hidden_states = self.pos_embed( - hidden_states) # takes care of adding positional embeddings too. - temb = self.time_text_embed(timestep, pooled_projections) - encoder_hidden_states = self.context_embedder(encoder_hidden_states) - - if self.mapping.cp_size > 1: - hidden_states = chunk(hidden_states, - chunks=self.mapping.cp_size, - dim=1)[self.mapping.cp_rank] - encoder_redundant = encoder_hidden_states.shape[ - 1] % self.mapping.cp_size - encoder_padding_index = tuple( - [0, 0] * (encoder_hidden_states.ndim() - 2) + - [0, self.mapping.cp_size - encoder_redundant]) - if encoder_redundant != 0: - encoder_hidden_states = pad(encoder_hidden_states, - pad=encoder_padding_index) - encoder_hidden_states = chunk(encoder_hidden_states, - chunks=self.mapping.cp_size, - dim=1)[self.mapping.cp_rank] - for index_block, block in enumerate(self.transformer_blocks): - # Skip specified layers - is_skip = True if self.skip_layers is not None and index_block in self.skip_layers else False - - if not is_skip: - encoder_hidden_states, hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - joint_attention_kwargs=joint_attention_kwargs, - ) - - # controlnet residual - if block_controlnet_hidden_states is not None and block.context_pre_only is False: - interval_control = len(self.transformer_blocks) / len( - block_controlnet_hidden_states) - hidden_states = hidden_states + block_controlnet_hidden_states[ - int(index_block / interval_control)] - - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - if self.mapping.cp_size > 1: - hidden_states = allgather(hidden_states, - group=self.mapping.cp_group, - gather_dim=1) - - # unpatchify - patch_size = self.config.patch_size - height = height // patch_size - width = width // patch_size - - hidden_states = hidden_states.view( - concat([ - shape(hidden_states, 0), height, width, patch_size, patch_size, - self.out_channels - ])) - hidden_states = einsum("nhwpqc->nchpwq", [hidden_states]) - output = hidden_states.view( - concat([ - shape(hidden_states, 0), self.out_channels, height * patch_size, - width * patch_size - ])) - - output.mark_output("output") - return output - - def prepare_inputs(self, max_batch_size, **kwargs): - - def sd3_default_range(max_batch_size): - return [1, max(1, (max_batch_size + 1) // 2), max_batch_size] - - default_range = sd3_default_range - prompt_embeds_len = 256 + 77 # [NOTE] tokenizer_max_length = 77; max_sequence_length = 256 - - hidden_states = Tensor(name='hidden_states', - dtype=self.dtype, - shape=[ - -1, self.in_channels, - self.config.sample_size, - self.config.sample_size - ], - dim_range=OrderedDict([ - ('batch_size', - [default_range(max_batch_size)]), - ('in_channels', [[self.in_channels] * 3]), - ('height', [[self.config.sample_size] * 3]), - ('width', [[self.config.sample_size] * 3]), - ])) - encoder_hidden_states = Tensor( - name='encoder_hidden_states', - dtype=self.dtype, - shape=[-1, prompt_embeds_len, self.config.joint_attention_dim], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('txt_len', [[prompt_embeds_len] * 3]), - ('joint_attention_dim', [[self.config.joint_attention_dim] * 3 - ]), - ])) - pooled_projections = Tensor( - name='pooled_projections', - dtype=self.dtype, - shape=[-1, self.config.pooled_projection_dim], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('pooled_projection_dim', - [[self.config.pooled_projection_dim] * 3]), - ])) - timestep = Tensor(name='timestep', - dtype=self.dtype, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_projections, - "timestep": timestep, - } - - @classmethod - def from_pretrained(cls, - pretrained_model_name_or_path: str, - dtype='float16', - mapping=Mapping(), - **kwargs): - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - - from diffusers import StableDiffusion3Pipeline - - transformer = StableDiffusion3Pipeline.from_pretrained( - pretrained_model_name_or_path, - torch_dtype=str_dtype_to_torch(dtype)).transformer - - config = SD3Transformer2DModelConfig.from_hugging_face_config( - transformer.config, dtype=dtype, mapping=mapping, **kwargs) - - hf_model_dir = transformer.config._name_or_path - custom_dict = {} - if quant_ckpt_path is not None: - hf_model_dir = quant_ckpt_path - - loader = SD3ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model) - return model - - def load(self, weights, from_pruned=False): - required_names = set() - for name, param in self.named_parameters(): - if self.use_pretrained_pos_emb and 'pos_embed' in name: - required_names.add(name) - continue - if param.is_inited(): - continue - if name not in weights: - # Exemption for embedding sharing - if name.endswith('lm_head.weight') and any( - k.endswith('vocab_embedding.weight') - for k in weights.keys()): - continue - if name.endswith('lm_head.per_channel_scale') and any( - k.endswith('vocab_embedding.per_channel_scale') - for k in weights.keys()): - continue - required_names.add(name) - - provided_names = set(weights.keys()) - if not required_names.issubset(provided_names): - raise RuntimeError( - f"Required but not provided tensors:{required_names.difference(provided_names)}" - ) - if not provided_names.issubset(required_names): - logger.warning( - f"Provided but not required tensors: {provided_names.difference(required_names)}" - ) - - for name, param in self.named_parameters(): - if name in provided_names: - if not from_pruned: - try: - param.value = weights[name] - except Exception as e: - raise RuntimeError( - f"Encounter error '{e}' for parameter '{name}'") - else: - param.set_value_or_dummy(weights[name]) - - def enable_forward_chunking(self, - chunk_size: Optional[int] = None, - dim: int = 0): - raise NotImplementedError() - - def disable_forward_chunking(self): - raise NotImplementedError() - - @property - def attn_processors(self): - return None - - def set_attn_processor(self, processor): - raise NotImplementedError() - - def fuse_qkv_projections(self): - raise NotImplementedError() - - def unfuse_qkv_projections(self): - raise NotImplementedError() - - def _set_gradient_checkpointing(self, module, value=False): - raise NotImplementedError() - - -class SD3ModelWeightsLoader(ModelWeightsLoader): - - def translate_to_external_key(self, tllm_key: str, - tllm_to_externel_key_dict: dict): - """Convert and load external checkpoint into a TensorRT LLM model. - """ - trtllm_to_hf_name = { - r"transformer_blocks.(\d+).ff(\w*).net.1.weight": - "transformer_blocks.*.ff*.net.2.weight", - r"transformer_blocks.(\d+).ff(\w*).net.1.bias": - "transformer_blocks.*.ff*.net.2.bias", - } - import re - for k, v in trtllm_to_hf_name.items(): - m = re.match(k, tllm_key) - if m is not None: - matched_pos = m.groups() - placeholders = v.count('*') - assert len(matched_pos) == placeholders - for i in range(len(matched_pos)): - v = v.replace('*', matched_pos[i], 1) - return v - return tllm_key diff --git a/tensorrt_llm/models/model_weights_loader.py b/tensorrt_llm/models/model_weights_loader.py deleted file mode 100644 index a130883e95c7..000000000000 --- a/tensorrt_llm/models/model_weights_loader.py +++ /dev/null @@ -1,402 +0,0 @@ -import glob -import math -import os -import weakref -from enum import Enum -from typing import Callable, List, Optional - -import tensorrt as trt -import torch -from safetensors import safe_open -from tqdm import tqdm -from transformers import PreTrainedModel - -from .._utils import trt_dtype_to_torch -from ..layers.moe import MOEWeightWrapper -from ..logger import logger -from ..quantization.layers import (WeightOnlyGroupwiseQuantColumnLinear, - WeightOnlyGroupwiseQuantRowLinear) - - -class ModelWeightsFormat(Enum): - IN_MEMORY = "in_mem" - SAFETENSORS = "safetensors" - BINARY = "bin" - PYTORCH = "pth" - - -class ModelWeightsLoader: - """Convert and load external checkpoint into a TensorRT LLM model. - - Attributes: - model_dir : Model directory or in-memory torch model. - format : Checkpoint file format. - shards : Shard pointer list (safetensors) or shard dict lists (other types) - shard map : Dict of external checkpoints' keys -> shard index. - tllm_to_externel_key_dict : Dict of TRT-LLM keywords -> External checkpoints' keywords, based on HF LLaMA. - customized_key_dict : Customized dict for updating the default tllm_to_externel_key_dict. - """ - - def __init__(self, model_dir, customized_key_dict: dict = {}) -> None: - - # Checkpoint file format information - self.model_dir = model_dir - self.format = None - self.model = None - self.shards = [] - self.shard_map = {} - - # Key translator vocabulary - self.tllm_to_externel_key_dict = { - "transformer": "model", - "vocab_embedding": "embed_tokens", - "layers": "layers", - "lm_head": "lm_head", - "ln_f": "norm", - "attention": "self_attn", - "qkv": ["q_proj", "k_proj", "v_proj"], - "dense": "o_proj", - "gate": "up_proj", - "proj": "down_proj", - "fc": "gate_proj", - "input_layernorm": "input_layernorm", - "post_layernorm": "post_attention_layernorm", - "kv_cache_scaling_factor": ["k_proj.k_scale", "v_proj.v_scale"], - "kv_cache_rcp_scaling_factor": ["k_proj.k_scale", "v_proj.v_scale"], - } - self.tllm_to_externel_key_dict.update(customized_key_dict) - - self.detect_format() - self.preload() - - def translate_to_external_key( - self, - tllm_key: str, - tllm_to_externel_key_dict: Optional[dict] = None - ) -> str | List[str]: - """Translate TRT-LLM key into HF key or HF key list (e.g. QKV/MoE/GPTQ) - - tllm_key will get translated into HF format section by section. - If one section is responded with multiple hf_keys in a list, \ - the translated keys will also get multiplied accordingly. - tllm_key : "transformer.layers.0.attention. qkv .weight" - | | | | | | - translated: [" model .layers.0.self_attn.q_proj.weight, - " model .layers.0.self_attn.k_proj.weight, - " model .layers.0.self_attn.v_proj.weight] - - Args: - tllm_key (str): Input TRT-LLM key. - tllm_to_externel_key_dict (dict): User specified dict with higher priority. \ - Generated from layer attributes automatically. - - Returns: - hf_keys (str | list[str]) : Translated HF key(s). - """ - tllm_keys = [tllm_key] - d = self.tllm_to_externel_key_dict.copy() - if tllm_to_externel_key_dict is not None: - d.update(tllm_to_externel_key_dict) - for k, v in d.items(): - if k in tllm_key: - # Ensure replacement happen when k covers several full sections in tllm_key - if not any([ - ('.' + k + '.') in tllm_key, - k == tllm_key, - tllm_key.startswith(k) and (k + '.') in tllm_key, - tllm_key.endswith(k) and ('.' + k) in tllm_key, - ]): - continue - if isinstance(v, list): - tllm_keys = [t for t in tllm_keys for _ in range(len(v))] - tllm_keys = [ - s.replace(k, v[idx % len(v)]) - for idx, s in enumerate(tllm_keys) - ] - else: - tllm_keys = [s.replace(k, v) for s in tllm_keys] - - for idx, k in enumerate(tllm_keys): - while ".." in k: - k = k.replace("..", ".") - if k.startswith("."): - k = k[1:] - if k.endswith("."): - k = k[:-1] - tllm_keys[idx] = k - - return tllm_keys[0] if len(tllm_keys) == 1 else tllm_keys - - def detect_format(self): - if isinstance(self.model_dir, dict) or isinstance( - self.model_dir, PreTrainedModel): - self.format = ModelWeightsFormat.IN_MEMORY - elif os.path.isfile(self.model_dir): - if self.model_dir.endswith(".safetensors"): - self.format = ModelWeightsFormat.SAFETENSORS - elif self.model_dir.endswith(".bin"): - self.format = ModelWeightsFormat.BINARY - elif self.model_dir.endswith(".pth"): - self.format = ModelWeightsFormat.PYTORCH - else: - raise NotImplementedError( - "Only safetensors/pickle/binary files are supported.") - elif os.path.isdir(self.model_dir): - file_list = os.listdir(self.model_dir) - if any([f.endswith(".safetensors") for f in file_list]): - self.format = ModelWeightsFormat.SAFETENSORS - elif any([f.endswith(".bin") for f in file_list]): - self.format = ModelWeightsFormat.BINARY - elif any([f.endswith(".pth") for f in file_list]): - self.format = ModelWeightsFormat.PYTORCH - else: - raise NotImplementedError( - "Only safetensors/pickle/binary directories are supported.") - else: - raise NotImplementedError( - "args.model_dir is not a directory, a file or an in-memory module!" - ) - - def preload(self): - # Initialize shards and load_func - if isinstance(self.model_dir, PreTrainedModel): - shard_files = [dict(self.model_dir.named_parameters())] - elif isinstance(self.model_dir, dict): - shard_files = [self.model_dir] - elif os.path.isfile(self.model_dir): - shard_files = [self.model_dir] - elif os.path.isdir(self.model_dir): - shard_files = glob.glob(self.model_dir + "/*." + self.format.value) - else: - raise NotImplementedError( - "args.model_dir is not a directory, a file or an in-memory module!" - ) - shard_files.sort() - if self.format == ModelWeightsFormat.SAFETENSORS: - self.shards = [ - safe_open(f, framework="pt", device="cpu") for f in shard_files - ] - elif self.format == ModelWeightsFormat.BINARY or self.format == ModelWeightsFormat.PYTORCH: - self.shards = [ - torch.load(f, weights_only=True, map_location="cpu", mmap=True) - for f in shard_files - ] - elif self.format == ModelWeightsFormat.IN_MEMORY: - self.shards = [shard_files[0]] - else: - raise NotImplementedError( - "Only *.safetensors/*.pth/*.bin files are supported.") - for idx, shard in enumerate(self.shards): - self.shard_map.update({k: idx for k in shard.keys()}) - - def load_tensor(self, key, tp_size=1, tp_dim=-1, tp_rank=0): - # Retrieve shard index - if key in self.shard_map: - ptr_idx = self.shard_map[key] - else: - if "language_model." + key in self.shard_map: - key = "language_model." + key - ptr_idx = self.shard_map[key] - else: - return None - - if self.format == ModelWeightsFormat.SAFETENSORS: - tensor = self.shards[ptr_idx].get_slice(key) - tensor_shape = tensor.get_shape() - if tensor_shape == []: - tensor = self.shards[ptr_idx].get_tensor(key).unsqueeze(0) - tensor_shape = tensor.shape - else: - tensor = self.shards[ptr_idx][key] - tensor_shape = tensor.shape - - if tp_size <= 1 or tp_dim < 0: - return tensor[:] - else: - if len(tensor_shape) == 1 and (tp_dim > 0 or tensor_shape[0] == 1): - return tensor[:] - else: - width = tensor_shape[tp_dim] - if width == 1: - return tensor[:] - slice_width = math.ceil(width / tp_size) - slice_start = tp_rank * slice_width - slice_end = min((tp_rank + 1) * slice_width, width) - slice_obj = [slice(None)] * len(tensor_shape) - slice_obj[tp_dim] = slice(slice_start, slice_end) - res = tensor[tuple(slice_obj)] - return res - - def load(self, - tllm_key: str, - preprocess: Callable[[int], None] = None, - skip_tp: bool = False, - custom_postprocess_kwargs: dict = {}): - """Load tensor from shards - - This function contains following steps: - 1. Translate tllm_key into external key(s). - 2. Load tensor/tensors partially according to layer attributes. - 3. Call preprocess() if it is not None. - 4. Call layer's post processing function. - 5. Return the dict for updating weight dict. - - Args: - tllm_key (str): TRT-LLM key from model iterators - preprocess (function, Optional): Customized preprocess function for step 3. - skip_tp (bool): Skip TP in case of the derived TP config is inappropriate. - """ - tp_rank = self.model.config.mapping.tp_rank - - sub_module = self.model - for attr in tllm_key.split(".")[:-1]: - sub_module = getattr(sub_module, attr) - param = self.model - for attr in tllm_key.split("."): - param = getattr(param, attr) - if param.is_buffer: - return {} - assert sub_module is not None and param is not None, f"{tllm_key} got Nonetype for parameter or parent module." - - tllm_to_externel_key_dict = getattr(sub_module, - "tllm_to_externel_key_dict", None) - tp_dim = getattr(sub_module, "tp_dim", -1) - require_weight_transpose = ( - isinstance(sub_module, WeightOnlyGroupwiseQuantColumnLinear) - or isinstance(sub_module, WeightOnlyGroupwiseQuantRowLinear)) - if tp_dim >= 0 and require_weight_transpose: - if sub_module.prequant_scaling_factor is not None: - if tllm_key.endswith("prequant_scaling_factor"): - tp_dim = 1 - tp_dim - elif tllm_key.endswith("weights_scaling_factor"): - tp_dim = -1 - elif tllm_key.endswith("weight"): - tp_dim = 1 - tp_dim - tp_size = getattr(sub_module, "tp_size", 1) - # Disable auto TP when num_kv_heads is invalid for split - if getattr(sub_module, "is_qkv", - False) and self.model.config.num_key_value_heads < tp_size: - tp_dim = -1 - tp_size = 1 - if skip_tp: - tp_dim = -1 - tp_size = 1 - if isinstance(sub_module, MOEWeightWrapper): - tp_rank = self.model.config.mapping.moe_tp_rank - external_key = self.translate_to_external_key( - tllm_key, tllm_to_externel_key_dict) - if isinstance(external_key, list): - v = [ - self.load_tensor(k, tp_size, tp_dim, tp_rank) - for k in external_key - ] - else: - v = self.load_tensor(external_key, tp_size, tp_dim, tp_rank) - - if preprocess is not None: - v = preprocess(v) - - if not hasattr(sub_module, "postprocess"): - if isinstance(v, list): - raise ValueError( - f"Param {tllm_key} is translated into {external_key}, post-process function is required." - ) - elif v is None: - weight_dict = {} - else: - weight_dict = {tllm_key: v.to(trt_dtype_to_torch(param.dtype))} - else: - postprocess_kwargs = {"config": self.model.config} - postprocess_kwargs.update(custom_postprocess_kwargs) - v = sub_module.postprocess(tllm_key, v, **postprocess_kwargs) - if isinstance(v, dict): - weight_dict = v - else: - weight_dict = {tllm_key: v} - - for k, v in weight_dict.items(): - if v is not None and not v.is_contiguous(): - weight_dict[k] = v.contiguous() - - return weight_dict - - def update_key_mapping(self, model): - self.model = weakref.ref(model)() - # Auto PP - config = model.config - if config.mapping.has_pp(): - pp_layers = config.mapping.pp_layers(config.num_hidden_layers) - self.tllm_to_externel_key_dict.update({ - f"layers.{tllm_local_layer_idx}": - f"{self.tllm_to_externel_key_dict['layers']}.{hf_global_layer_idx}" - for tllm_local_layer_idx, hf_global_layer_idx in enumerate( - pp_layers) - }) - if self.tllm_to_externel_key_dict['layers'] != 'layers': - del self.tllm_to_externel_key_dict['layers'] - - # Share embedding; only applies to standard structure with lm_head and transformer.vocab_embedding - if hasattr(self.model, 'lm_head') and hasattr( - self.model, 'transformer') and hasattr(self.model.transformer, - 'vocab_embedding'): - lm_head_weights = self.load_tensor( - self.translate_to_external_key('lm_head.weight')) - vocab_embed_weights = self.load_tensor( - self.translate_to_external_key( - 'transformer.vocab_embedding.weight')) - if lm_head_weights is None and vocab_embed_weights is not None: - self.tllm_to_externel_key_dict[ - 'lm_head'] = self.tllm_to_externel_key_dict[ - 'transformer'] + '.' + self.tllm_to_externel_key_dict[ - 'vocab_embedding'] - elif lm_head_weights is not None and vocab_embed_weights is None: - self.tllm_to_externel_key_dict[ - 'vocab_embedding'] = self.tllm_to_externel_key_dict[ - 'lm_head'] - self.model.transformer.vocab_embedding.tllm_to_externel_key_dict = { - 'transformer': '' - } - - def fill(self, weights): - for tllm_key, param in self.model.named_parameters(): - if param.is_buffer: - continue - if tllm_key.endswith('embed_positions_for_gpt_attention'): - continue - w_shape = weights[tllm_key].shape - # WAR for 4bit datatype shape mismatch. - if w_shape != param.shape and param.dtype != trt.fp4: - logger.warning( - f'{tllm_key} has invalid shape {w_shape}. Expected {param.shape}.' - ) - pad = torch.nn.functional.pad - pad_dim = [] - for dim in range(weights[tllm_key].dim()): - current_dim = -1 - dim - pad_dim.append(0) - pad_dim.append( - max(0, param.shape[current_dim] - w_shape[current_dim])) - try: - logger.warning( - f'{tllm_key} is going to be padded by {pad_dim}.') - weights[tllm_key] = pad(weights[tllm_key], - tuple(pad_dim), - value=0) - assert weights[tllm_key].shape == param.shape - except: - raise ValueError( - f'Parameter {tllm_key} has invalid shape {weights[tllm_key].shape} compared with expected shape {param.shape}. Auto padding failed.' - ) - param.value = weights[tllm_key] - - def generate_tllm_weights(self, - model, - custom_postprocess_kwargs: dict = {}): - # For customization, please copy this function and make changes inside the for loop. - self.update_key_mapping(model) - tllm_weights = {} - for tllm_key, _ in tqdm(model.named_parameters()): - tllm_weights.update( - self.load(tllm_key, - custom_postprocess_kwargs=custom_postprocess_kwargs)) - self.fill(tllm_weights) diff --git a/tensorrt_llm/models/modeling_utils.py b/tensorrt_llm/models/modeling_utils.py index b445c128e360..f5a9f27364f7 100644 --- a/tensorrt_llm/models/modeling_utils.py +++ b/tensorrt_llm/models/modeling_utils.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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 copy import dataclasses @@ -7,51 +22,30 @@ import re from enum import IntFlag, auto from functools import cached_property -from pathlib import Path -from typing import (TYPE_CHECKING, Callable, Dict, Generator, List, Optional, - Union) +from typing import TYPE_CHECKING, Dict, Generator, List, Optional, Union -import numpy as np -import safetensors -import torch from pydantic import Field, PrivateAttr -from .._common import default_net -from .._utils import (QuantModeWrapper, get_init_params, numpy_to_torch, - release_gc, str_dtype_to_torch, str_dtype_to_trt, - trt_dtype_to_torch) +from .._utils import QuantModeWrapper from ..bindings.executor import RuntimeDefaults -from ..functional import (PositionEmbeddingType, Tensor, allgather, constant, - cp_split_plugin, gather_last_token_logits, - index_select, tanh, view) -from ..layers import (MLP, AttentionParams, Embedding, FusedGatedMLP, - FusedRgLru, GatedMLP, KeyValueCacheParams, LoraParams, - PromptTuningEmbedding, RgLru) -from ..layers.attention import Attention, BertAttention -from ..layers.linear import ColumnLinear, Linear, RowLinear -from ..layers.lora import Dora, Lora -from ..layers.moe import MOE, MoeOOTB -from ..llmapi.kv_cache_type import KVCacheType +from ..functional import PositionEmbeddingType from ..llmapi.utils import StrictBaseModel from ..logger import logger from ..mapping import Mapping -from ..module import Module, ModuleList -from ..parameter import Parameter -from ..plugin import init_all_reduce_helper -from ..quantization import QuantMode -from ..quantization.functional import preprocess_weights_for_mixed_gemm -from ..quantization.layers import (FP8Linear, Fp8RowwiseFusedGatedMLP, - Fp8RowwiseGatedMLP, - WeightOnlyGroupwiseQuantLinear, - WeightOnlyGroupwiseQuantRowLinear, - WeightOnlyQuantLinear, - WeightOnlyQuantRowLinear) from ..quantization.mode import (KV_CACHE_QUANT_ALGO_LIST, QUANT_ALGO_LIST, - W8A8_SQ_PLUGIN_LIST, QuantAlgo) -from ..quantization.utils import fp4_utils -from ..top_model_mixin import TopModelMixin -from .convert_utils import weight_only_quantize_dict -from .generation_mixin import GenerationMixin + W8A8_SQ_PLUGIN_LIST, QuantAlgo, QuantMode) + +# QuantConfig and LayerQuantConfig live in the (TensorRT-free) +# tensorrt_llm.quantization package; re-exported here for backward +# compatibility with existing import sites. + +__all__ = [ + 'PretrainedConfig', + 'SpeculativeDecodingMode', + 'QuantConfig', + 'LayerQuantConfig', + 'QuantAlgo', +] @dataclasses.dataclass(kw_only=True, frozen=True) @@ -578,1409 +572,3 @@ def for_each_rank(self) -> "Generator[Self, None, None]": config_copy = copy.deepcopy(self) config_copy.set_rank(rank) yield config_copy - - -class DecoderLayerList(ModuleList): - - def __init__(self, cls, config): - self.num_hidden_layers = config.num_hidden_layers - self.layer_list = config.mapping.pp_layers(config.num_hidden_layers) - self.quant_mode = config.quant_mode - super().__init__([cls(config, idx) for idx in self.layer_list]) - - def forward(self, - hidden_states, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - mrope_params=None, - position_ids=None, - lora_params=None, - spec_decoding_params=None, - vision_token_mask=None): - kv_cache_params.fill_none_tensor_list(len(self.layer_list)) - - if use_cache: - presents = [] - - for layer_idx, (layer, past) in enumerate( - zip(self, kv_cache_params.past_key_value)): - - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(layer_idx) - - kwargs = {} - if position_ids is not None: - kwargs['position_ids'] = position_ids - if vision_token_mask is not None: - kwargs['vision_token_mask'] = vision_token_mask - if lora_layer_params is not None: - kwargs['lora_layer_params'] = lora_layer_params - if spec_decoding_params is not None: - kwargs['spec_decoding_params'] = spec_decoding_params - if mrope_params is not None: - kwargs['mrope_params'] = mrope_params - - if default_net().plugin_config.reduce_fusion: - if layer_idx + self.layer_list[0] < self.layer_list[-1]: - qkv_activation_scaling_factor = None - if default_net().plugin_config.user_buffer: - qkv_linear = self[layer_idx + 1].attention.qkv - if self.quant_mode.has_fp8_qdq(): - qkv_activation_scaling_factor = constant( - qkv_linear.activation_scaling_factor.raw_value. - copy()) - elif self.quant_mode.has_nvfp4(): - qkv_activation_scaling_factor = constant( - qkv_linear.activation_global_scaling_factor. - raw_value.copy()) - kwargs['next_layer_input_layernorm_args'] = ( - self[layer_idx + 1].input_layernorm.weight.value, - self[layer_idx + 1].input_layernorm.eps, - qkv_activation_scaling_factor) - else: - kwargs['next_layer_input_layernorm_args'] = None - elif default_net().plugin_config.norm_quant_fusion: - if layer_idx < self.layer_list[-1] - self.layer_list[0]: - try: - activation_scaling_factor = constant( - self[layer_idx + 1].attention.qkv. - activation_global_scaling_factor.raw_value.copy()) - except: - activation_scaling_factor = None - kwargs['next_layer_input_layernorm_args'] = ( - self[layer_idx + 1].input_layernorm.weight.value, - self[layer_idx + 1].input_layernorm.eps, - activation_scaling_factor) - else: - kwargs['next_layer_input_layernorm_args'] = None - - hidden_states = layer( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=KeyValueCacheParams( - past_key_value=[past], - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cache_indirection=kv_cache_params.cache_indirection), - attention_params=attention_params, - **kwargs) - - if use_cache: - presents.append(hidden_states[1]) - hidden_states = hidden_states[0] - - if use_cache: - return hidden_states, presents - return hidden_states - - -class PostInitCaller(type): - - def __call__(cls, *args, **kwargs): - obj = type.__call__(cls, *args, **kwargs) - obj.__post_init__() - return obj - - -class PretrainedModel(Module, - GenerationMixin, - TopModelMixin, - metaclass=PostInitCaller): - - def __init__(self, config: PretrainedConfig): - super().__init__() - init_all_reduce_helper() - self.config = config - - def __post_init__(self): - from ..quantization.quantize import quantize - quantize(self, self.config.quantization) - - # Currently, use_parallel_embedding must be enabled before weight loading; - # otherwise, the model will be inconsistent with the weights loaded from checkpoint. - optimize_model( - self, use_parallel_embedding=self.config.use_parallel_embedding) - - def release(self): - release_gc() - - def __del__(self): - self.release() - - def check_config(self, config): - raise NotImplementedError( - f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." - ) - - @classmethod - def from_config(cls, config: PretrainedConfig): - return cls(config) - - @classmethod - def from_checkpoint( - cls, - ckpt_dir: str, - rank: Optional[int] = None, - config: Optional[PretrainedConfig] = None, - *, - preprocess_weights_hook: Optional[Callable[[Dict[str, Tensor]], - Dict[str, Tensor]]] = None): - if config is None: - config = PretrainedConfig.from_json_file( - os.path.join(ckpt_dir, 'config.json')) - - if rank is not None: - config.set_rank(rank) - - rank = config.mapping.rank - if config.mapping.cp_size > 1: - # cp_tp_pp rank -> tp_pp rank: because different cp ranks share the same ckpt. - cp_size = config.mapping.cp_size - # rank = pp_rank × tp_size × cp_size + tp_rank × cp_size + cp_rank. - # rank // cp_size is equivalent to pp_rank × tp_size + tp_rank. - rank = rank // cp_size - weights_path = os.path.join(ckpt_dir, f'rank{rank}.safetensors') - - assert os.path.isfile(weights_path) - weights = safetensors.torch.load_file(weights_path) - is_checkpoint_pruned = getattr(config, 'is_pruned', False) - - if preprocess_weights_hook is not None: - weights = preprocess_weights_hook(weights) - - weights = preprocess_weights(weights, - config, - from_pruned=is_checkpoint_pruned) - model = cls(config) - model.load(weights, from_pruned=is_checkpoint_pruned) - return model - - def load(self, weights, from_pruned=False): - required_names = set() - for name, param in self.named_parameters(): - if param.is_inited(): - continue - if name not in weights: - # Exemption for embedding sharing - if name.endswith('lm_head.weight') and any( - k.endswith('vocab_embedding.weight') - for k in weights.keys()): - continue - if name.endswith('lm_head.per_channel_scale') and any( - k.endswith('vocab_embedding.per_channel_scale') - for k in weights.keys()): - continue - required_names.add(name) - - provided_names = set(weights.keys()) - - if not required_names.issubset(provided_names): - raise RuntimeError( - f"Required but not provided tensors:{required_names.difference(provided_names)}" - ) - if not provided_names.issubset(required_names): - logger.warning( - f"Provided but not required tensors: {provided_names.difference(required_names)}" - ) - - for name, param in self.named_parameters(): - if name in provided_names: - if not from_pruned: - try: - param.value = weights[name] - except Exception as e: - raise RuntimeError( - f"Encounter error '{e}' for parameter '{name}'") - else: - param.set_value_or_dummy(weights[name]) - - def save_checkpoint(self, output_dir, save_config=True): - # multiple ranks could share same config.json, so adding a save_config parameter to let user avoiding writing config.json in all ranks - rank = self.config.mapping.rank - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in self.named_parameters() - } - # If there are some tensors share memory, this will lead to error when we call "save_file". So, for repeated tensors, we - # clone the tensors to prevent this issue. - data_ptrs = set() - for name, param in weights.items(): - if param.data_ptr() in data_ptrs: - weights[name] = param.clone() - data_ptrs.add(weights[name].data_ptr()) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - if save_config: - self.config.to_json_file(os.path.join(output_dir, 'config.json')) - - def prepare_inputs( - self, - max_batch_size, - max_input_len, - max_seq_len, - max_num_tokens, - use_cache, - max_beam_width: int = 1, - opt_num_tokens: int = None, - prompt_embedding_table_size: int = 0, - position_encoding_2d: bool = False, - max_draft_len: int = 0, - speculative_decoding_draft_tokens_external: bool = False, - spec_decoding_is_generation_length_variable: bool = False, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - opt_batch_size: int = 0, - num_hidden_layers: int = None, - mrope_rotary_cos_sin_size: int = None, - ): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions when using TRT dynamic shapes. - - @return: a list containing values which can be fed into the self.forward() - ''' - - # Prepare inputs - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - use_lora_plugin = default_net().plugin_config.lora_plugin - multiple_profiles = default_net().plugin_config.multiple_profiles - streamingllm = default_net().plugin_config.streamingllm - pp_reduce_scatter = default_net().plugin_config.pp_reduce_scatter - - kv_cache_type = None - if not use_cache: - kv_cache_type = KVCacheType.DISABLED - else: - if paged_kv_cache: - kv_cache_type = KVCacheType.PAGED - else: - kv_cache_type = KVCacheType.CONTINUOUS - - model_inputs = self.prepare_basic_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - hidden_size=self.config.hidden_size, - num_kv_heads=self.config.num_key_value_heads, - head_size=self.config.head_size, - num_layers=num_hidden_layers - if num_hidden_layers is not None else self.config.num_hidden_layers, - kv_dtype=str_dtype_to_trt(self.config.kv_dtype), - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - num_heads=self.config.num_attention_heads, - max_num_tokens=max_num_tokens, - opt_num_tokens=opt_num_tokens, - dtype=str_dtype_to_trt(self.config.dtype), - prompt_embedding_table_size=prompt_embedding_table_size, - position_encoding_2d=position_encoding_2d, - mapping=self.config.mapping, - gather_context_logits=gather_context_logits, - use_lora_plugin=use_lora_plugin, - max_draft_len=max_draft_len, - speculative_decoding_draft_tokens_external= - speculative_decoding_draft_tokens_external, - spec_decoding_is_generation_length_variable= - spec_decoding_is_generation_length_variable, - lora_target_modules=lora_target_modules, - multiple_profiles=multiple_profiles, - streamingllm=streamingllm, - opt_batch_size=opt_batch_size, - pp_reduce_scatter=pp_reduce_scatter, - mrope_rotary_cos_sin_size=mrope_rotary_cos_sin_size) - - result = { - 'input_ids': - model_inputs['input_ids'], - 'position_ids': - model_inputs['position_ids'], - 'use_cache': - kv_cache_type != KVCacheType.DISABLED, - 'last_token_ids': - model_inputs['last_token_ids'], - 'attention_mask': - model_inputs['attention_mask'], - 'kv_cache_params': - KeyValueCacheParams( - past_key_value=model_inputs['past_key_value'], - host_past_key_value_lengths=model_inputs[ - 'host_past_key_value_lengths'], - host_max_attention_window_sizes=model_inputs[ - 'host_max_attention_window_sizes'], - host_sink_token_length=model_inputs['host_sink_token_length'], - kv_cache_block_offsets=model_inputs['kv_cache_block_offsets'], - host_kv_cache_block_offsets=model_inputs[ - 'host_kv_cache_block_offsets'], - host_kv_cache_pool_pointers=model_inputs[ - 'host_kv_cache_pool_pointers'], - host_kv_cache_pool_mapping=model_inputs[ - 'host_kv_cache_pool_mapping'], - cache_indirection=model_inputs['cache_indirection'], - ), - 'attention_params': - AttentionParams( - sequence_length=model_inputs['sequence_length'], - context_lengths=model_inputs['context_lengths'], - host_context_lengths=model_inputs['host_context_lengths'], - max_context_length=max_input_len, - host_request_types=model_inputs['host_request_types'], - host_runtime_perf_knobs=model_inputs['host_runtime_perf_knobs'], - host_context_progress=model_inputs['host_context_progress'], - ) - } - - if prompt_embedding_table_size > 0: - result['prompt_embedding_table'] = model_inputs[ - 'prompt_embedding_table'] - result['prompt_tasks'] = model_inputs['tasks'] - result['prompt_vocab_size'] = model_inputs['prompt_vocab_size'] - if model_inputs['hidden_states_input'] is not None: - result['hidden_states'] = model_inputs['hidden_states_input'] - if use_lora_plugin: - result['lora_params'] = LoraParams( - model_inputs['lora_ranks'], - model_inputs['lora_weights_pointers'], - host_context_lengths=model_inputs['host_context_lengths'], - host_request_types=model_inputs['host_request_types']) - if model_inputs['spec_decoding_params'] is not None: - result['spec_decoding_params'] = model_inputs[ - 'spec_decoding_params'] - if model_inputs['mrope_params'] is not None: - result['mrope_params'] = model_inputs['mrope_params'] - - return result - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - config_cls = getattr(cls, 'config_class', None) - if config_cls is None: - raise NotImplementedError( - f"{cls.__name__} has not implemented corresponding config class, which is needed for correct config parsing." - ) - config: PretrainedConfig = config_cls.from_hugging_face( - hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if config.mapping.moe_ep_size > 1: - raise NotImplementedError( - "Quantization for expert parallelism is not supported") - if not config.quantization._requires_modelopt_quantization: - raise ValueError( - f"The quant_config ({quant_config}) should not call modelopt quantization" - ) - - from ..quantization import quantize_and_export - quantize_and_export( - model_dir=str(hf_model_dir), - device=device, - calib_dataset=calib_dataset, - dtype=config.dtype, - qformat=config.quantization._get_modelopt_qformat(), - kv_cache_dtype=config.quantization._get_modelopt_kv_cache_dtype(), - calib_size=calib_batches, - batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - awq_block_size=config.quantization.group_size, - output_dir=output_dir, - tp_size=config.mapping.tp_size, - pp_size=config.mapping.pp_size, - cp_size=config.mapping.cp_size, - seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length, - ) - - -class DecoderModelForCausalLM(PretrainedModel): - - def __init__(self, config: PretrainedConfig, transformer, lm_head): - super().__init__(config) - self.transformer = transformer - self.lm_head = lm_head - self.mup_width_multiplier = getattr(config, 'mup_width_multiplier', - None) - # Create constant attention parameters to be reused by all layers. - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - last_token_ids=None, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - mrope_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None, - spec_decoding_params=None): - - # fill attention params. - attention_params = Attention.fill_attention_params( - self, attention_params) - - # split the sequence for context parallelism - if self.config.mapping.cp_size > 1: - if len(input_ids.shape) == 1: - # input shape is [-1] - input_ids, cp_join_index = cp_split_plugin( - input_ids, - attention_params.host_request_types, - attention_params.host_context_lengths, - self.config.mapping.cp_size, - self.config.mapping.cp_rank, - ) - else: - assert False, "Context parallelism with non-remove-padding is not supported yet." - - is_gemma_2_cg = self.config.has_config_group(Gemma2ConfigGroup) - is_gemma_3_cg = self.config.has_config_group(Gemma3ConfigGroup) - - kwargs = { - 'input_ids': input_ids, - 'position_ids': position_ids, - 'use_cache': use_cache, - 'attention_mask': attention_mask, - 'kv_cache_params': kv_cache_params, - 'attention_params': attention_params, - } - if lora_params is not None: - kwargs['lora_params'] = lora_params - if hidden_states is not None: - kwargs['hidden_states'] = hidden_states - if prompt_embedding_table is not None: - kwargs['prompt_embedding_table'] = prompt_embedding_table - if prompt_tasks is not None: - kwargs['prompt_tasks'] = prompt_tasks - if prompt_vocab_size is not None: - kwargs['prompt_vocab_size'] = prompt_vocab_size - - if spec_decoding_params is not None: - kwargs['spec_decoding_params'] = spec_decoding_params - if mrope_params is not None: - kwargs['mrope_params'] = mrope_params - - hidden_states = self.transformer.forward(**kwargs) - - if use_cache: - hidden_states, presents = hidden_states - - # All gather and rebuild sequence after transformer layer for context parallelism - if self.config.mapping.cp_size > 1: - if len(hidden_states.shape) == 2: - hidden_states = allgather(hidden_states, - self.config.mapping.cp_group, - gather_dim=0) - hidden_states = view(hidden_states, - [-1, hidden_states.shape[-1]]) - hidden_states = index_select(hidden_states, 0, cp_join_index) - else: - assert False, "Context parallelism with non-remove-padding is not supported yet." - - if self.config.mapping.is_last_pp_rank(): - all_hidden_states = hidden_states - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) - - # [batch_size, hidden_size] -> [batch_size, vocab_size] - lm_logits = self.lm_head(hidden_states) - if hasattr(self.config, 'output_multiplier_scale'): - lm_logits *= getattr(self.config, 'output_multiplier_scale', 1) - if self.mup_width_multiplier is not None: - lm_logits = lm_logits / self.mup_width_multiplier - if is_gemma_2_cg or is_gemma_3_cg: - softcap = self.config.get_config_group( - Gemma2ConfigGroup if not is_gemma_3_cg else - Gemma3ConfigGroup).final_logit_softcapping - if softcap: - lm_logits = lm_logits * float(1 / softcap) - lm_logits = tanh(lm_logits) * float(softcap) - lm_logits.mark_output('logits', self.config.logits_dtype) - else: - hidden_states.mark_output('hidden_states_output', self.config.dtype) - - if use_cache and not default_net().plugin_config.paged_kv_cache: - for i, present in zip( - self.config.mapping.pp_layers( - self.config.num_hidden_layers), presents): - present.mark_output(f'present_key_value_{i}', - self.config.kv_dtype) - if self.config.mapping.is_last_pp_rank(): - return (lm_logits, presents, hidden_states) - return (hidden_states, presents) - else: - if self.config.mapping.is_last_pp_rank(): - return lm_logits, hidden_states, all_hidden_states - return hidden_states - - -def fuse_gate_mlp( - model: PretrainedModel, - gemm_swiglu_plugin_dtype: Optional[str] = None, - low_latency_gemm_swiglu_plugin_dtype: Optional[str] = None, -) -> PretrainedModel: - from ..quantization.quantize import fp8_quantize - - for name, mlp, layer in model.named_modules_with_parent(): - if isinstance(mlp, GatedMLP): - init_params = get_init_params(mlp) - - hidden_act = init_params["hidden_act"] - if hidden_act not in ["silu", "gelu"]: - logger.warning( - f"fuse_gate_mlp cannot be done for {name} due to unsupported activation {hidden_act}. Skipping." - ) - continue - - init_params["inner_layernorm"] = mlp.inner_layernorm is not None - fused_layer = FusedGatedMLP(**init_params) - - fc_name = name + '.fc' - layer_quant_cfg = model.config._get_quant_cfg(fc_name) - layer_quant_algo = layer_quant_cfg.quant_algo - if layer_quant_algo != QuantAlgo.FP8 and layer_quant_algo is not None: - continue - - if isinstance(model.config.quantization.exclude_modules, list) \ - and fc_name in model.config.quantization.exclude_modules: - layer_quant_algo = None - - if layer_quant_algo == QuantAlgo.FP8: - fused_layer = fp8_quantize(fused_layer, layer_quant_cfg) - - if isinstance(mlp.dtype, str): - dtype = str_dtype_to_torch(mlp.dtype) - else: - dtype = trt_dtype_to_torch(mlp.dtype) - - gate_weight = numpy_to_torch(mlp.gate.weight.raw_value) - fc_weight = numpy_to_torch(mlp.fc.weight.raw_value) - assert gate_weight.dtype == fc_weight.dtype - need_qdq = gate_weight.dtype == torch.float8_e4m3fn - - gate_weight = gate_weight.to(dtype) - fc_weight = fc_weight.to(dtype) - # dequantize if needed - if need_qdq: - gate_weight = gate_weight.to(dtype) * numpy_to_torch( - mlp.gate.weights_scaling_factor.raw_value) - fc_weight = fc_weight.to(dtype) * numpy_to_torch( - mlp.fc.weights_scaling_factor.raw_value) - - # concat - fused_weight = torch.cat([gate_weight, fc_weight], dim=0) - - fused_weight_scaling_factor = numpy_to_torch( - max( - mlp.gate.weights_scaling_factor.raw_value, - mlp.fc.weights_scaling_factor.raw_value, - )) - # quantize if needed - if need_qdq: - fused_weight = (fused_weight / - fused_weight_scaling_factor).to( - torch.float8_e4m3fn) - - if gemm_swiglu_plugin_dtype == 'fp8' or low_latency_gemm_swiglu_plugin_dtype == 'fp8': - # gemm_swiglu_plugin needs (k, n) weights - # but weights should still be k-major for fp8 - fused_layer.fused_fc.weight = Parameter( - shape=(fused_layer.fused_fc.in_features, - fused_layer.fused_fc.out_features), - dtype='fp8') - fused_layer.fused_fc.weight.value = fused_weight.view( - fused_layer.fused_fc.in_features, - fused_layer.fused_fc.out_features) - else: - fused_layer.fused_fc.weight.value = fused_weight - fused_layer.fused_fc.weights_scaling_factor.value = fused_weight_scaling_factor - - fused_layer.fused_fc.activation_scaling_factor.value = max( - mlp.gate.activation_scaling_factor.raw_value, - mlp.fc.activation_scaling_factor.raw_value, - ) - - if mlp.bias: - fused_layer.fused_fc.bias.value = np.concatenate( - [mlp.gate.bias.raw_value, mlp.fc.bias.raw_value], - axis=0) - elif layer_quant_algo is None: - fused_layer.fused_fc.weight.value = np.concatenate( - [ - mlp.gate.weight.raw_value, - mlp.fc.weight.raw_value, - ], - axis=0, - ) - if mlp.bias: - fused_layer.fused_fc.bias.value = np.concatenate( - [mlp.gate.bias.raw_value, mlp.fc.bias.raw_value], - axis=0) - else: - raise ValueError(f'Unsupported quant algo: {layer_quant_algo}') - - fused_layer.proj = mlp.proj - fused_layer.inner_layernorm = mlp.inner_layernorm - - _, mlp_name = name.rsplit('.', 1) - setattr(layer, mlp_name, fused_layer) - - elif isinstance(mlp, Fp8RowwiseGatedMLP): - init_params = get_init_params(mlp) - - hidden_act = init_params["hidden_act"] - if hidden_act not in ["silu", "gelu"]: - logger.warning( - f"fuse_gate_mlp cannot be done for {name} due to unsupported activation {hidden_act}. Skipping." - ) - continue - - if mlp.clamp_val is not None: - init_params["clamp_val"] = mlp.clamp_val.raw_value.tolist() - fused_layer = Fp8RowwiseFusedGatedMLP(**init_params) - fused_layer.fused_fc.weight.value = np.concatenate( - [ - mlp.gate.weight.raw_value, - mlp.fc.weight.raw_value, - ], - axis=0, - ) - fused_layer.fused_fc.per_channel_scale.value = np.concatenate( - [ - mlp.gate.per_channel_scale.raw_value, - mlp.fc.per_channel_scale.raw_value, - ], - axis=0, - ) - if mlp.bias: - fused_layer.fused_fc.bias.value = np.concatenate( - [mlp.gate.bias.raw_value, mlp.fc.bias.raw_value], axis=0) - - fused_layer.proj = mlp.proj - _, mlp_name = name.rsplit('.', 1) - setattr(layer, mlp_name, fused_layer) - - return model - - -def unfuse_qkv_gemm(model: PretrainedModel) -> PretrainedModel: - '''Split all the models' Attention layer's QKV GEMM into 3 GEMMs layer.q layer.k, layer.v and return the changed model - ''' - from ..quantization.quantize import quantize - - for name, layer in model.named_modules(): - if isinstance(layer, Attention) and not layer.cross_attention: - assert layer.tp_size == 1, "unfuse_qkv_gemm requires tp_size == 1" - if layer.qkv is None: - continue - qkv_params = get_init_params(layer.qkv, ColumnLinear) - qkv_params["bias"] = qkv_params["bias"] is not None - qkv_params["strict_dtype"] = qkv_params.get( - "strict_dtype") is not None - q = ColumnLinear( - **{ - **qkv_params, - "out_features": - layer.tp_size * layer.num_attention_heads * - layer.attention_head_size, - }) - k = ColumnLinear( - **{ - **qkv_params, - "out_features": - layer.tp_size * layer.num_attention_kv_heads * - layer.attention_head_size, - }) - v = ColumnLinear( - **{ - **qkv_params, - "out_features": - layer.tp_size * layer.num_attention_kv_heads * - layer.attention_head_size, - }) - layer_quant_cfg = model.config._get_quant_cfg(name + '.qkv') - q = quantize(q, layer_quant_cfg) - k = quantize(k, layer_quant_cfg) - v = quantize(v, layer_quant_cfg) - out_features = q.out_features + k.out_features + v.out_features - if isinstance(layer.qkv, ( - WeightOnlyQuantLinear, - WeightOnlyQuantRowLinear, - WeightOnlyGroupwiseQuantLinear, - WeightOnlyGroupwiseQuantRowLinear, - )): - out_dim = 1 - else: - out_dim = 0 - if layer.qkv.weight.is_inited(): - qkv_weight = layer.qkv.weight.raw_value - weights = np.split(qkv_weight, [ - qkv_weight.shape[out_dim] * q.out_features // out_features, - qkv_weight.shape[out_dim] * - (q.out_features + k.out_features) // out_features, - ], - axis=out_dim) - for gemm, weight in zip([q, k, v], weights): - gemm.weight.value = weight - if layer.qkv.bias is not None and layer.qkv.bias.is_inited(): - qkv_bias = layer.qkv.bias.raw_value - biases = np.split(qkv_bias, [ - qkv_bias.shape[out_dim] * q.out_features // out_features, - qkv_bias.shape[out_dim] * - (q.out_features + k.out_features) // out_features, - ], - axis=out_dim) - for gemm, bias in zip([q, k, v], biases): - gemm.bias.value = bias - for name, parameter in layer.qkv._parameters.items(): - if name not in ["weight", "bias"]: - for gemm in [q, k, v]: - setattr(gemm, name, parameter) - layer.q = q - layer.k = k - layer.v = v - layer.qkv = None - return model - - -def fuse_rg_lru(model: PretrainedModel) -> PretrainedModel: - for name, rg_lru, parent in model.named_modules_with_parent(): - if isinstance(rg_lru, RgLru): - fused_layer = FusedRgLru(**get_init_params(rg_lru)) - fused_layer.gate.weight.value = np.concatenate( - [ - rg_lru.input_gate.weight.raw_value, - rg_lru.recurrent_gate.weight.raw_value, - ], - axis=-1, - ) - fused_layer.gate.bias.value = np.concatenate( - [ - rg_lru.input_gate.bias.raw_value, - rg_lru.recurrent_gate.bias.raw_value, - ], - axis=-1, - ) - fused_layer.recurrent_param.value = rg_lru.recurrent_param.raw_value - rg_lru_name = name.rsplit('.', 1)[-1] - setattr(parent, rg_lru_name, fused_layer) - return model - - -def set_prompt_tuning(model: PretrainedModel) -> PretrainedModel: - '''Replace the given models embedding layer with a PromptTuningEmbedding layer in-place, return the changed model - Pre-conditions: vocab_embedding exists - Post-conditions: isinstance(vocab_embedding, PromptTuningEmbedding) - - ''' - for name, embedding, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name == "vocab_embedding" and isinstance(embedding, Embedding): - ptuning_embedding = PromptTuningEmbedding( - **get_init_params(embedding)) - ptuning_embedding.weight.value = embedding.weight.raw_value - parent.vocab_embedding = ptuning_embedding - return model - - -def add_lora(model: PretrainedModel, - max_lora_rank: Optional[int], - with_dora: bool = False) -> PretrainedModel: - ''' Add lora layers to the Attention/BertAttention/Linear/RowLinear/FusedGatedMLP layers to the given model, return the changed model - ''' - for name, layer in model.named_modules(): - max_rank = max_lora_rank - if isinstance(layer, (Attention, BertAttention)): - if max_rank is None: - max_rank = min( - layer.hidden_size, - layer.num_attention_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size) - layer.qkv_lora = Lora( - in_hidden_size=layer.hidden_size, - out_hidden_sizes=[ - layer.num_attention_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size - ], - max_low_rank=max_rank, - ) - - if with_dora: - layer.qkv_dora = Dora(out_hidden_sizes=[ - layer.num_attention_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size - ], ) - - if isinstance(layer, (Linear, RowLinear)): - if max_rank is None: - max_rank = min(layer.in_features, layer.out_features) - layer.lora = Lora( - in_hidden_size=layer.in_features, - out_hidden_sizes=[layer.out_features], - max_low_rank=max_rank, - ) - if with_dora: - layer.dora = Dora(out_hidden_sizes=[layer.out_features]) - - if isinstance(layer, (MLP, FusedGatedMLP)): - if max_rank is None: - max_rank = min(layer.hidden_size, - layer.ffn_hidden_size // layer.tp_size) - layer.lora = Lora( - in_hidden_size=layer.hidden_size, - out_hidden_sizes=[ - layer.ffn_hidden_size // layer.tp_size, - layer.ffn_hidden_size // layer.tp_size - ], - max_low_rank=max_rank, - ) - - if isinstance(layer, FusedGatedMLP): - layer.fused_gate_up_lora = Lora( - in_hidden_size=layer.hidden_size, - out_hidden_sizes=[ - layer.ffn_hidden_size * 2 // layer.tp_size - ], - max_low_rank=max_rank, - ) - - if with_dora: - layer.dora = Dora(out_hidden_sizes=[ - layer.ffn_hidden_size // layer.tp_size, - layer.ffn_hidden_size // layer.tp_size - ], ) - - if isinstance(layer, FusedGatedMLP): - layer.fused_gate_up_dora = Dora(out_hidden_sizes=[ - layer.ffn_hidden_size * 2 // layer.tp_size - ], ) - - if isinstance(layer, MOE): - if max_rank is None: - max_rank = min(layer.hidden_size, - layer.ffn_hidden_size // layer.tp_size) - layer.max_low_rank = max_rank - return model - - -def to_ootb_moe(model: PretrainedModel) -> PretrainedModel: - ''' Use OOTB MoE instead of MoE plugin, return the changed model - ''' - for name, layer, parent in model.named_modules_with_parent(): - if isinstance(layer, MOE): - layer_name = name.rsplit('.', 1)[-1] - ootb_layer = layer.to(MoeOOTB, model.config.quantization) - setattr(parent, layer_name, ootb_layer) - return model - - -def parallelize_embedding(model: PretrainedModel) -> PretrainedModel: - for name, embedding, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if isinstance(embedding, Embedding) and embedding.tp_group is None: - init_params = get_init_params(embedding) - init_params["tp_group"] = model.config.mapping.tp_group - init_params["tp_size"] = model.config.mapping.tp_size - init_params["tp_rank"] = model.config.mapping.tp_rank - init_params["sharding_dim"] = model.config.embedding_sharding_dim - new_embedding = embedding.__class__(**init_params) - setattr(parent, layer_name, new_embedding) - return model - - -def share_embedding(model: PretrainedModel) -> PretrainedModel: - lm_head = None - vocab_embedding = None - for name, layer in model.named_modules(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name == "lm_head": - lm_head = layer - if layer_name == "vocab_embedding": - vocab_embedding = layer - if lm_head is not None and vocab_embedding is not None: - break - - # Cannot find either lm_head or vocab_embedding, e.g., pipeline parallel - if lm_head is None or vocab_embedding is None: - return model - - # lm_head and vocab_embedding have different shapes, e.g., tensor parallel without embedding parallel - if lm_head.weight.shape != vocab_embedding.weight.shape: - return model - - # lm_head can have a different type if quantized - if lm_head.weight.dtype != vocab_embedding.weight.dtype: - return model - - # Don't assume weight can be shared if vocab_embedding is not initialized, e.g., dummy weights - if not vocab_embedding.weight.is_inited(): - return model - - if lm_head.weight.is_inited(): - lm_head_weight = numpy_to_torch(lm_head.weight.raw_value) - vocab_embed_weight = numpy_to_torch(vocab_embedding.weight.raw_value) - # The lm_head and vocab_embedding have different weights - if (lm_head_weight - vocab_embed_weight).abs().max().item() > 1e-6: - return model - - lm_head.weight = vocab_embedding.weight - if getattr(lm_head, 'per_channel_scale', None) and getattr( - vocab_embedding, 'per_channel_scale', None): - lm_head.per_channel_scale = vocab_embedding.per_token_scale - return model - - -def set_fp8_context_fmha(model: PretrainedModel) -> PretrainedModel: - for name, layer in model.named_modules(): - if isinstance(layer, Attention) and hasattr( - layer.dense, 'activation_scaling_factor'): - scale = [1.0] / layer.dense.activation_scaling_factor.raw_value - layer.attention_output_orig_quant_scale = Parameter( - value=scale.astype(np.float32), dtype='float32') - elif isinstance(layer, Attention) and hasattr( - layer.dense, 'activation_global_scaling_factor'): - scale = [1.0 - ] / layer.dense.activation_global_scaling_factor.raw_value - layer.attention_output_orig_quant_scale = Parameter( - value=scale.astype(np.float32), dtype='float32') - - return model - - -def set_fuse_fp4_quant(model: PretrainedModel) -> PretrainedModel: - for name, layer in model.named_modules(): - if isinstance(layer, Attention) and hasattr( - layer.dense, 'activation_global_scaling_factor'): - scale = [1.0 - ] / layer.dense.activation_global_scaling_factor.raw_value - layer.attention_output_sf_scale = Parameter(value=scale.astype( - np.float32), - dtype='float32') - - return model - - -def optimize_model( - model: PretrainedModel, - use_parallel_embedding: bool = False, - share_embedding_table: bool = False, - use_ootb_moe: bool = False, - use_fused_mlp: bool = False, - gemm_swiglu_plugin_dtype: Optional[str] = None, - low_latency_gemm_swiglu_plugin_dtype: Optional[str] = None, - use_fused_rg_lru: bool = False, - use_unfused_qkv_gemm: bool = False, - use_prompt_tuning: bool = False, - use_lora: bool = False, - max_lora_rank: Optional[int] = None, - use_fp8_context_fmha: bool = False, - fuse_fp4_quant: bool = False, - use_optimize_cross_qkv: bool = False, - use_dora: bool = False, -) -> PretrainedModel: - """ - Run optimization passes on model. - There are dependencies between some passes, - so we always run passes in the order of arguments to guarantee the execution order. - """ - # before weight loading - if use_parallel_embedding: - model = parallelize_embedding(model) - - if share_embedding_table: - # if share_embedding_table is enabled, only one copy of the embedding table is stored in converted ckpt - # this pass is required to make lm_head.weight and vocab_embedding.weight point to the same tensor - # however even if share_embedding_table is not enabled, trt would still only keep one copy of the table if the weights are identical - model = share_embedding(model) - - # After weight loading - if use_ootb_moe: - model = to_ootb_moe(model) - if use_fused_mlp: - model = fuse_gate_mlp(model, gemm_swiglu_plugin_dtype, - low_latency_gemm_swiglu_plugin_dtype) - if use_fused_rg_lru: - model = fuse_rg_lru(model) - if use_unfused_qkv_gemm: - model = unfuse_qkv_gemm(model) - if use_prompt_tuning: - model = set_prompt_tuning(model) - if use_lora: - model = add_lora(model, max_lora_rank, with_dora=use_dora) - if use_fp8_context_fmha: - model = set_fp8_context_fmha(model) - if fuse_fp4_quant: - model = set_fuse_fp4_quant(model) - if not use_lora and use_optimize_cross_qkv is True: - # This optimization is not supported when we use lora - model = optimize_cross_qkv(model) - - return model - - -def optimize_cross_qkv(model): - """ - For cross attention layer, we can skip computing the query of encoder_output. - So, add a new attribute 'kv' in the cross_attention layer. This might lead to - additional memory cost on model size, but save the memory usage on runtime. - - Currently, this function only detects ColumnLinear and FP8Linear. It does not support - other quantization now. - """ - for name, attn, layer in model.named_modules_with_parent(): - if isinstance(attn, Attention) and attn.cross_attention and \ - (type(attn.qkv) == ColumnLinear or type(attn.qkv) == FP8Linear): - old_qkv = attn.qkv - linear_class = type(old_qkv) - new_kv = linear_class( - in_features=attn.hidden_size, - out_features=2 * attn.tp_size * attn.num_attention_kv_heads * - attn.attention_head_size, - bias=old_qkv.bias, - dtype=old_qkv.dtype, - tp_group=old_qkv.tp_group, - tp_size=old_qkv.tp_size, - gather_output=old_qkv.gather_output, - prefer_managed_weight=old_qkv.prefer_managed_weight, - is_qkv=old_qkv.is_qkv, - ) - - old_qkv_weight_value = old_qkv.weight.raw_value - if (old_qkv_weight_value.shape == np.asarray([ - (attn.num_attention_heads + 2 * attn.num_attention_kv_heads) * - attn.attention_head_size, attn.hidden_size - ])).all(): - - q_weight, kv_weight = np.array_split( - old_qkv_weight_value.reshape( - attn.num_attention_heads + - 2 * attn.num_attention_kv_heads, - attn.attention_head_size, attn.hidden_size), - [attn.num_attention_heads], - axis=0) - new_kv.weight.value = kv_weight.reshape([ - 2 * attn.num_attention_kv_heads * attn.attention_head_size, - attn.hidden_size - ]) - elif (old_qkv_weight_value.shape == np.asarray([ - attn.hidden_size, - (attn.num_attention_heads + 2 * attn.num_attention_kv_heads) * - attn.attention_head_size - ])).all(): - q_weight, kv_weight = np.array_split( - old_qkv_weight_value.reshape( - attn.hidden_size, attn.num_attention_heads + - 2 * attn.num_attention_kv_heads, - attn.attention_head_size), [attn.num_attention_heads], - axis=1) - new_kv.weight.value = kv_weight.reshape([ - attn.hidden_size, - 2 * attn.num_attention_kv_heads * attn.attention_head_size - ]) - else: - assert False - - if isinstance(attn.qkv, FP8Linear): - new_kv.activation_scaling_factor.value = old_qkv.activation_scaling_factor.raw_value - new_kv.weights_scaling_factor.value = old_qkv.weights_scaling_factor.raw_value - - if old_qkv.bias: - q_bias, kv_bias = np.array_split(old_qkv.bias.raw_value.reshape( - attn.num_attention_heads + 2 * attn.num_attention_kv_heads, - attn.attention_head_size), [attn.num_attention_heads], - axis=0) - new_kv.bias.value = kv_bias.reshape([ - 2 * attn.num_attention_kv_heads * attn.attention_head_size - ]) - setattr(attn, "kv", new_kv) - - return model - - -def preprocess_perlayer_weights(weights, - model_config, - quant_algo, - from_pruned=False): - exclude_modules = model_config.quantization.exclude_modules - - # INT4_AWQ - if quant_algo == QuantAlgo.W4A8_AWQ or quant_algo == QuantAlgo.W4A16_AWQ: - preprocessor = preprocess_weights_for_mixed_gemm - if quant_algo == QuantAlgo.W4A8_AWQ: - activation_type = torch.float8_e4m3fn - elif quant_algo == QuantAlgo.W4A16_AWQ: - activation_type = torch.float16 - for name, param in weights.items(): - if from_pruned and param.numel() == 0: - continue - if name.endswith('weight') and param.dtype == torch.int8: - dtype = torch.float16 - if model_config.dtype == "bfloat16": - dtype = torch.bfloat16 - weights[name] = preprocessor(param.transpose(-1, -2), - torch.quint4x2, - activation_type).view(dtype) - if name.endswith('weights_scaling_factor'): - weights[name] = param.transpose(-1, -2).contiguous().to( - str_dtype_to_torch(model_config.dtype)) - if name.endswith('prequant_scaling_factor'): - if len(weights[name].shape) == 2: - # MoE experts share the same scaling factor. - param = param[0, :] - weights[name] = param.reshape(1, -1) - if model_config.mapping.tp_rank > 0: - if name.endswith('attention.dense.bias') or name.endswith( - 'mlp.proj.bias'): - weights[name] = torch.zeros_like(param) - - if quant_algo == QuantAlgo.W4A8_AWQ: - for name in list(weights): - if name.endswith('weights_scaling_factor'): - activation_scaling_factor = weights.pop( - name.replace('weights_scaling_factor', - 'activation_scaling_factor')) - weights_scaling_factor_2 = weights.pop( - name.replace('weights_scaling_factor', - 'weights_scaling_factor_2')) - weights[name] /= weights_scaling_factor_2 - weights[name] = weights[name].to(torch.float16).view( - str_dtype_to_torch(model_config.dtype)) - weights[name.replace( - 'weights_scaling_factor', - 'prequant_scaling_factor')] /= activation_scaling_factor - weights[name.replace( - 'weights_scaling_factor', 'alpha' - )] = activation_scaling_factor * weights_scaling_factor_2 - weights[name.replace('weights_scaling_factor', - 'activation_scaling_factor' - )] = activation_scaling_factor - - # FP8 - elif quant_algo == QuantAlgo.FP8: - for name, param in weights.items(): - if name.endswith('weight') and param.dtype == torch.int8: - weights[name] = param.view(torch.float8_e4m3fn) - # lm_head is not always quantized to FP8 - if "lm_head.weight" in weights and weights[ - 'lm_head.weight'].dtype is not torch.float8_e4m3fn: - weights.pop('lm_head.weights_scaling_factor', None) - weights.pop('lm_head.activation_scaling_factor', None) - elif quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN: - for name, param in weights.items(): - if name.endswith('weight') and param.dtype == torch.int8: - weights[name] = param.view(torch.float8_e4m3fn) - # lm_head is not quantized to FP8 - if "lm_head.weight" in weights: - assert weights['lm_head.weight'].dtype == str_dtype_to_torch( - model_config.dtype) - weights.pop('lm_head.weights_scaling_factor', None) - weights.pop('lm_head.activation_scaling_factor', None) - # FP4 - elif quant_algo == QuantAlgo.NVFP4: - # Interleave block scale for NVFP4 plugin. - for name in list(weights): - if name.endswith('weights_scaling_factor'): - out_features, in_features = weights[name].shape - nrows = fp4_utils.pad_up(out_features, 128) - ncols = fp4_utils.pad_up(in_features, 4) - new_name = name.replace('weights_scaling_factor', - 'weights_block_scaling_factor') - weights[new_name] = weights[name] - weights[ - new_name + - "_interleaved"] = torch.ops.trtllm.block_scale_interleave( - weights[name].view(fp4_utils.float4_sf_dtype).cpu( - ).contiguous()).reshape(nrows, ncols).view( - fp4_utils.float4_sf_dtype) - weights.pop(name) - if name.endswith('weights_scaling_factor_2'): - new_name = name.replace('weights_scaling_factor_2', - 'weights_global_scaling_factor') - weights[new_name] = weights[name] - weights.pop(name) - if name.endswith('activation_scaling_factor'): - new_name = name.replace('activation_scaling_factor', - 'activation_global_scaling_factor') - weights[new_name] = weights[name] - weights.pop(name) - for name in list(weights): - if name.endswith('weights_global_scaling_factor'): - weight_global_sf = weights[name] - act_global_sf = weights[name.replace( - 'weights_global_scaling_factor', - 'activation_global_scaling_factor')] - weights[name.replace( - 'weights_global_scaling_factor', - 'alpha')] = act_global_sf * weight_global_sf - elif quant_algo in [QuantAlgo.W4A16, QuantAlgo.W8A16]: - weights = weight_only_quantize_dict(weights=weights, - quant_algo=quant_algo, - exclude_modules=exclude_modules, - plugin=True) - - -def preprocess_weights(weights: Dict[str, torch.Tensor], - model_config: PretrainedConfig, - from_pruned=False) -> None: - """This function in-place modifies weights and model_config, making them compatible with each other. - - Note: Typically, it should be called before model creation and weight loading. For example, - preprocess_weights(weights, model_config) - model = XXXForCausalLM(model_config) - model.load(weights) - """ - quant_config = model_config.quantization - quant_algo = quant_config.quant_algo - - pattern_info = ['fc', 'gate', 'proj', 'qkv', 'dense'] - - def process_kv_scaling_factor(weights: Dict[str, torch.Tensor]): - new_entries = {} - names_to_delete = set() - - # If k, v cache scaling factors are stored separately, combine them into kv cache scaling factor. - for name, param in weights.items(): - if name.endswith('.k_cache_scaling_factor'): - v_name = name.replace('k_cache_scaling_factor', - 'v_cache_scaling_factor') - assert v_name in weights, f"{v_name} not found" - kv_name = name.replace('k_cache_scaling_factor', - 'kv_cache_scaling_factor') - new_entries[kv_name] = torch.max(weights[name], weights[v_name]) - names_to_delete.update([name, v_name]) - weights.update(new_entries) - for k in names_to_delete: - del weights[k] - - new_entries = [] - # The unified converter generate_tllm_weights() already generates these rcp weights, but legacy - # converters do not. Handle it here. - for name, param in weights.items(): - if name.endswith('.kv_cache_scaling_factor'): - rcp_name = name.replace('kv_cache_scaling_factor', - 'kv_cache_rcp_scaling_factor') - if rcp_name not in weights: - new_entries.append((rcp_name, torch.reciprocal(param))) - weights.update(new_entries) - - process_kv_scaling_factor(weights) - - per_layer_weights = {} - - for name, param in weights.items(): - in_mode = False - for info in pattern_info: - pattern = rf'(.*?{info}.*?)' - pattern_match = re.match(pattern, name) - if pattern_match: - base_name = pattern_match.group(1) - if base_name not in per_layer_weights.keys(): - per_layer_weights[base_name] = {} - per_layer_weights[base_name][name] = param - in_mode = True - break - if not in_mode: - # [lm_head.weight, ln_f.weight, vocab_embedding.weight] - base_name = name.rsplit('.', 1)[0] - if base_name not in per_layer_weights.keys(): - per_layer_weights[base_name] = {} - per_layer_weights[base_name][name] = param - - new_weights = {} - for base_name, layer_weights in per_layer_weights.items(): - if quant_algo != QuantAlgo.MIXED_PRECISION: - layer_quant_algo = quant_algo - else: - quant_cfg = quant_config._get_quant_cfg(base_name) - if not quant_cfg.quant_algo: - new_weights.update(layer_weights) - continue - - layer_quant_algo = quant_cfg.quant_algo - - preprocess_perlayer_weights(layer_weights, model_config, - layer_quant_algo, from_pruned) - new_weights.update(layer_weights) - - weights = new_weights - for name, param in weights.items(): - if model_config.architecture == 'GPTJForCausalLM': - if model_config.mapping.tp_rank > 0: - if 'attention.dense.bias' in name or 'mlp.proj.bias' in name: - weights[name] = torch.zeros_like(param) - - return weights - - -def get_kv_cache_type_from_legacy(use_cache: bool, - paged_kv_cache: bool) -> KVCacheType: - if use_cache: - if paged_kv_cache: - return KVCacheType.PAGED - else: - return KVCacheType.CONTINUOUS - else: - return KVCacheType.DISABLED - - -def save_config(config: PretrainedConfig, *, output_dir: str, - log: bool) -> None: - config_path = Path(output_dir) / "config.json" - if log: - logger.debug(f"Saving TensorRT LLM configuration to {config_path}") - config_path.parent.mkdir(exist_ok=True, parents=True) - config_path.write_text(json.dumps(config.to_dict(), indent=4)) - - -def save_checkpoint(*, output_dir: str, weights: dict, rank: int) -> None: - """ Checkpoint saver for weight loader.""" - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) diff --git a/tensorrt_llm/models/mpt/__init__.py b/tensorrt_llm/models/mpt/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/mpt/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/mpt/model.py b/tensorrt_llm/models/mpt/model.py deleted file mode 100644 index ecfa343e363f..000000000000 --- a/tensorrt_llm/models/mpt/model.py +++ /dev/null @@ -1,176 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class MPTDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - dtype = config.dtype - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - tp_group = config.mapping.tp_group - layernorm_epsilon = config.norm_epsilon - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - bias=False, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - bias=config.bias, - position_embedding_type=PositionEmbeddingType.alibi, - quant_mode=config.quant_mode, - clip_qkv=config.clip_qkv, - alibi_bias_max=config.alibi_bias_max) - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=config.hidden_act, - dtype=dtype, - bias=config.bias, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - bias=False, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class MPTModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.config = config - - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(MPTDecoderLayer, config) - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - bias=False, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None): - - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_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 MPTForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = MPTModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=config.bias, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('bias', False) - config.set_if_not_exist('clip_qkv', None) - config.set_if_not_exist('alibi_bias_max', 8) diff --git a/tensorrt_llm/models/multimodal_encoders/__init__.py b/tensorrt_llm/models/multimodal_encoders/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/multimodal_encoders/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/multimodal_encoders/config.py b/tensorrt_llm/models/multimodal_encoders/config.py deleted file mode 100644 index 2ae8135caae2..000000000000 --- a/tensorrt_llm/models/multimodal_encoders/config.py +++ /dev/null @@ -1,116 +0,0 @@ -# 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 - -import torch - -from ..._utils import torch_dtype_to_str -from ...logger import logger -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class LlavaNextVisionConfig(PretrainedConfig): - - def __init__(self, - *, - image_size: int, - patch_size: int, - text_hidden_size: int, - projector_hidden_act: str = 'gelu', - num_channels: int = 3, - vision_model_type: str = 'clip_vision_model', - **kwargs): - self.image_size = image_size - self.patch_size = patch_size - self.text_hidden_size = text_hidden_size - self.num_channels = num_channels - self.projector_hidden_act = projector_hidden_act - self.vision_model_type = vision_model_type - - super().__init__(**kwargs) - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - if hf_config.model_type == "llava_next": - from transformers import LlavaNextConfig - hf_config = LlavaNextConfig.from_pretrained(hf_config_dir) - else: - logger.error("Provided model type is not llava_next.") - - text_hidden_size = hf_config.text_config.hidden_size - # Extract only the vision config - llava_next_vision_config = hf_config.vision_config - - # llava-next uses the second last layer as vision output - num_feature_layers = llava_next_vision_config.num_hidden_layers + hf_config.vision_feature_layer + 1 - - vision_model_type = getattr(llava_next_vision_config, - "vision_model_type", "clip_vision_model") - - num_key_value_heads = getattr( - llava_next_vision_config, "num_key_value_heads", - llava_next_vision_config.num_attention_heads) - - # Default configs from HF - hidden_act = 'quick_gelu' - norm_epsilon = 1e-5 - - head_size = llava_next_vision_config.hidden_size // llava_next_vision_config.num_attention_heads - - if dtype == 'auto': - dtype = getattr(hf_config, 'torch_dtype', None) - if dtype is None: - dtype = 'float16' - if isinstance(dtype, torch.dtype): - dtype = torch_dtype_to_str(dtype) - if dtype == 'float32': - dtype = 'float16' - - return cls( - image_size=llava_next_vision_config.image_size, - patch_size=llava_next_vision_config.patch_size, - text_hidden_size=text_hidden_size, - projector_hidden_act=hf_config.projector_hidden_act, - vision_model_type=vision_model_type, - architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=num_feature_layers, - num_attention_heads=llava_next_vision_config.num_attention_heads, - hidden_size=llava_next_vision_config.hidden_size, - intermediate_size=llava_next_vision_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - head_size=head_size, - vocab_size=llava_next_vision_config.vocab_size, - hidden_act=hidden_act, - norm_epsilon=norm_epsilon, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/multimodal_encoders/model.py b/tensorrt_llm/models/multimodal_encoders/model.py deleted file mode 100644 index b60337b0f740..000000000000 --- a/tensorrt_llm/models/multimodal_encoders/model.py +++ /dev/null @@ -1,175 +0,0 @@ -# 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 os -from collections import OrderedDict -from typing import Optional - -import safetensors - -from ..._utils import numpy_to_torch -from ...functional import ACT2FN, Tensor, concat, shape, slice -from ...layers import Linear -from ...logger import logger -from ...mapping import Mapping -from ...models import CLIPVisionTransformer -from ...module import Module -from ...parameter import Parameter -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import PretrainedModel, QuantConfig -from .config import LlavaNextVisionConfig - - -# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/llava_next/modeling_llava_next.py#L149 -class LlavaNextMultiModalProjector(Module): - - def __init__(self, config: LlavaNextVisionConfig): - super().__init__() - - self.linear_1 = Linear(config.hidden_size, - config.text_hidden_size, - dtype=config.dtype) - self.act = ACT2FN[config.projector_hidden_act] - self.linear_2 = Linear(config.text_hidden_size, - config.text_hidden_size, - dtype=config.dtype) - - def forward(self, image_features): - hidden_states = self.linear_1(image_features) - hidden_states = self.act(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states - - -class LlavaNextVisionWrapper(PretrainedModel): - - def __init__(self, config: LlavaNextVisionConfig): - super().__init__(config) - self.vision_tower = None - self.config = config - if config.vision_model_type == "clip_vision_model": - self.vision_tower = CLIPVisionTransformer( - image_size=config.image_size, - num_channels=config.num_channels, - patch_size=config.patch_size, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - max_position_embeddings=config.max_position_embeddings, - norm_epsilon=config.norm_epsilon, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - num_hidden_layers=config.num_hidden_layers, - require_ln_f=False, - mapping=config.mapping, - dtype=config.dtype) - else: - logger.error( - "Currently TRT-LLM only supports CLIP vision transformer.") - - self.multi_modal_projector = LlavaNextMultiModalProjector(config) - self.image_newline = Parameter(shape=(config.text_hidden_size, ), - dtype=config.dtype) - - def forward(self, pixel_values, position_ids=None): - image_features = self.vision_tower(pixel_values) - select_size = concat([ - shape(image_features, 0), image_features.shape[1] - 1, - shape(image_features, 2) - ]) - selected_image_feature = slice(image_features, - starts=[0, 1, 0], - sizes=select_size) # (bs, 576, c) - image_features = self.multi_modal_projector(selected_image_feature) - image_features.mark_output('image_features', self.config.dtype) - return image_features # (bs, 576, c) - - @classmethod - def from_hugging_face(cls, - hf_model_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LlavaNextVisionWrapper object from give parameters - ''' - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is not None: - logger.error( - "Please enable unified converter to convert llava-next checkpoints." - ) - - config = LlavaNextVisionConfig.from_hugging_face( - hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - custom_dict = {} - if "llava" in hf_model_dir: - custom_dict = { - "vision_tower": "vision_tower.vision_model", - "input_layernorm": "layer_norm1", - "post_layernorm": "layer_norm2", - "fc": "fc1", - "proj": "fc2", - "dense": "out_proj", - "pre_layernorm": "pre_layrnorm", - "ln_f": "post_layernorm", - } - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model) - return model - - def save_checkpoint(self, output_dir, save_config=True): - rank = self.config.mapping.rank - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in self.named_parameters() - } - image_newline = { - "image_newline": numpy_to_torch(self.image_newline.raw_value) - } - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - safetensors.torch.save_file( - image_newline, - os.path.join(output_dir, f'image_newlines.safetensors')) - if save_config: - self.config.to_json_file(os.path.join(output_dir, 'config.json')) - - def prepare_inputs(self, max_batch_size, **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - batch_size_range = [ - 1, max(1, (max_batch_size + 1) // 2), max_batch_size - ] - pixel_values = Tensor( - name='pixel_values', - dtype=self.config.dtype, - shape=[ - -1, self.config.num_channels, self.config.image_size, - self.config.image_size - ], - dim_range=OrderedDict([ - ('batch_size', [batch_size_range]), - ('in_channels', [[self.config.num_channels] * 3]), - ('latent_height', [[self.config.image_size] * 3]), - ('latent_width', [[self.config.image_size] * 3]), - ])) - return {'pixel_values': pixel_values} diff --git a/tensorrt_llm/models/nemotron_nas/__init__.py b/tensorrt_llm/models/nemotron_nas/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/nemotron_nas/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/nemotron_nas/config.py b/tensorrt_llm/models/nemotron_nas/config.py deleted file mode 100644 index 11d02df84b07..000000000000 --- a/tensorrt_llm/models/nemotron_nas/config.py +++ /dev/null @@ -1,208 +0,0 @@ -# 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 dataclasses import asdict -from typing import Any, Dict, List, Optional, Union - -from tensorrt_llm._utils import get_hf_rope_theta -from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import infer_dtype -from tensorrt_llm.models.modeling_utils import PretrainedConfig, QuantConfig -from tensorrt_llm.models.nemotron_nas.convert import \ - hf_block_configs_to_layer_configs -from tensorrt_llm.models.nemotron_nas.layer_config import ( - AttentionConfig, AttentionImplementation, DeciLayerConfig, FFNConfig) - - -class DeciConfig(PretrainedConfig): - - def __init__(self, - *, - architecture: str = 'DeciLMForCausalLM', - dtype: str, - hidden_size: int, - num_hidden_layers: int, - num_attention_heads: int, - vocab_size: int, - hidden_act: str = 'gelu', - logits_dtype: str = 'float32', - norm_epsilon: float = 0.00001, - position_embedding_type: Union[ - PositionEmbeddingType, - str] = PositionEmbeddingType.rope_gpt_neox, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - max_position_embeddings: int, - num_key_value_heads: Optional[int] = None, - intermediate_size: Optional[int] = None, - mapping: Optional[Union[Mapping, dict]] = None, - quantization: Optional[Union[QuantConfig, dict]] = None, - use_parallel_embedding: bool = False, - embedding_sharding_dim: int = 0, - head_size: Optional[int] = None, - qk_layernorm: bool = False, - layer_configs: Optional[List[Union[DeciLayerConfig, - Dict[str, - Dict[str, - Any]]]]] = None, - block_configs: Optional[object] = None, - **kwargs): - super().__init__(architecture=architecture, - dtype=dtype, - hidden_size=hidden_size, - num_hidden_layers=num_hidden_layers, - num_attention_heads=num_attention_heads, - vocab_size=vocab_size, - hidden_act=hidden_act, - logits_dtype=logits_dtype, - norm_epsilon=norm_epsilon, - position_embedding_type=position_embedding_type, - max_position_embeddings=max_position_embeddings, - num_key_value_heads=num_key_value_heads, - intermediate_size=intermediate_size, - mapping=mapping, - quantization=quantization, - use_parallel_embedding=use_parallel_embedding, - embedding_sharding_dim=embedding_sharding_dim, - head_size=head_size, - qk_layernorm=qk_layernorm, - **kwargs) - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - - if block_configs is not None: - assert layer_configs is None - self.layer_configs = hf_block_configs_to_layer_configs( - block_configs, - num_attention_heads=num_attention_heads, - hidden_size=hidden_size) - elif layer_configs is not None: - assert len( - layer_configs - ) == num_hidden_layers, f"num_hidden_layers ({num_hidden_layers}) must match len(layer_configs) ({len(layer_configs)})" - - self.layer_configs = self._ensure_layer_configs(layer_configs) - else: - self.layer_configs = None - - # HACK: this is needed for many parts of the code - self.layer_types = [ - AttentionImplementation( - self.get_layer_config(layer_idx).attention.impl).value - for layer_idx in range(self.num_hidden_layers) - ] - - # HACK: this is here since the runtime doesn't parse the layer_configs yet - self.num_kv_heads_per_layer = [] - for layer_idx in range(self.num_hidden_layers): - layer_config = self.get_layer_config(layer_idx) - if layer_config.is_attention_layer: - self.num_kv_heads_per_layer.append( - layer_config.attention.num_key_value_heads) - - def _ensure_layer_configs( - self, layer_configs: List[Union[DeciLayerConfig, Dict[str, Any]]] - ) -> List[DeciLayerConfig]: - return [ - DeciLayerConfig.from_dict(c) if isinstance(c, dict) else c - for c in layer_configs - ] - - def to_dict(self): - output = super().to_dict() - if self.layer_configs is not None: - output["layer_configs"] = [asdict(c) for c in self.layer_configs] - return output - - def get_layer_config(self, layer_idx: int) -> DeciLayerConfig: - if self.layer_configs is not None: - conf = self.layer_configs[layer_idx] - else: - conf = DeciLayerConfig() - - attention_impl = conf.attention.impl - num_key_value_heads = conf.attention.num_key_value_heads or self.num_key_value_heads - ffn_impl = conf.ffn.impl - intermediate_size = conf.ffn.intermediate_size or self.intermediate_size - - return DeciLayerConfig( - attention=AttentionConfig(impl=attention_impl, - num_key_value_heads=num_key_value_heads), - ffn=FFNConfig(impl=ffn_impl, intermediate_size=intermediate_size)) - - def get_layer_num_kv_heads(self, layer_idx) -> int: - layer_config = self.get_layer_config(layer_idx) - assert layer_config.is_attention_layer, f"Layer {layer_idx} is not an attention layer" - return layer_config.attention.num_key_value_heads or self.num_key_value_heads - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - trust_remote_code: bool = True, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_or_dir, trust_remote_code=trust_remote_code) - - assert hf_config.model_type in ( - "deci", - "nemotron-nas"), f"Unsupported model type: {hf_config.model_type}" - - block_configs = getattr(hf_config, "block_configs", None) - if block_configs is not None: - layer_configs = hf_block_configs_to_layer_configs( - block_configs, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size) - else: - # older deci arch - num_key_value_heads_per_layer = getattr( - hf_config, "num_key_value_heads_per_layer", None) - if num_key_value_heads_per_layer is not None: - layer_configs = [ - DeciLayerConfig(attention=AttentionConfig( - num_key_value_heads=num_key_value_heads)) - for num_key_value_heads in num_key_value_heads_per_layer - ] - else: - layer_configs = None - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(dtype=dtype, - hidden_size=hf_config.hidden_size, - hidden_act=hf_config.hidden_act, - intermediate_size=hf_config.intermediate_size, - num_attention_heads=hf_config.num_attention_heads, - num_hidden_layers=hf_config.num_hidden_layers, - num_key_value_heads=hf_config.num_key_value_heads, - norm_epsilon=hf_config.rms_norm_eps, - rotary_scaling=hf_config.rope_scaling, - rotary_base=get_hf_rope_theta(hf_config, 10000.0), - vocab_size=hf_config.vocab_size, - max_position_embeddings=hf_config.max_position_embeddings, - mapping=mapping, - quantization=quant_config, - layer_configs=layer_configs, - **kwargs) diff --git a/tensorrt_llm/models/nemotron_nas/convert.py b/tensorrt_llm/models/nemotron_nas/convert.py deleted file mode 100644 index ba26414ae8d4..000000000000 --- a/tensorrt_llm/models/nemotron_nas/convert.py +++ /dev/null @@ -1,381 +0,0 @@ -# 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 enum -import json -import time -from abc import ABC, abstractmethod -from contextlib import contextmanager -from dataclasses import asdict -from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional, TypedDict, Union - -import safetensors -import torch - -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.logger import logger -from tensorrt_llm.models.convert_utils import dup_kv_weight, split -from tensorrt_llm.models.nemotron_nas.layer_config import ( - AttentionConfig, AttentionImplementation, DeciLayerConfig, FFNConfig, - FFNImplementation) -from tensorrt_llm.quantization.mode import QuantAlgo - - -def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int: - intermediate_size = int(2 * ffn_mult * n_embd / 3) - return _find_multiple(intermediate_size, 256) - - -def _find_multiple(n: int, k: int) -> int: - if n % k == 0: - return n - return n + k - (n % k) - - -# BlockConfig is a custom class defined inside deci huggingface checkpoints, we can't import it -def hf_block_config_to_layer_config(block_config: Union["BlockConfig", dict], - num_attn_heads: int, - hidden_size: int) -> DeciLayerConfig: - """`block_config` (`Union[BlockConfig, dict]`): A `dict` when exported from `ModelOpt`; A `dataclass` at the HF phase - """ - block_config = block_config if isinstance(block_config, - dict) else asdict(block_config) - attn = block_config["attention"] - if attn["no_op"]: - attn_impl = AttentionImplementation.NO_OP - num_key_value_heads = None - elif attn["replace_with_linear"]: - attn_impl = AttentionImplementation.LINEAR - num_key_value_heads = None - elif attn.get("sparsify", None): - raise NotImplementedError("Sparsification is not supported") - else: - attn_impl = AttentionImplementation.ATTENTION - num_key_value_heads = num_attn_heads // attn["n_heads_in_group"] - - ffn = block_config["ffn"] - if ffn["no_op"]: - ffn_impl = FFNImplementation.NO_OP - intermediate_size = None - elif ffn["replace_with_linear"]: - ffn_impl = FFNImplementation.LINEAR - intermediate_size = None - elif ffn.get("sparsify", None): - raise NotImplementedError("Sparsification is not supported") - else: - ffn_impl = FFNImplementation.MLP - intermediate_size = _ffn_mult_to_intermediate_size( - ffn["ffn_mult"], hidden_size) - - return DeciLayerConfig(attention=AttentionConfig( - impl=attn_impl, num_key_value_heads=num_key_value_heads), - ffn=FFNConfig(impl=ffn_impl, - intermediate_size=intermediate_size)) - - -def hf_block_configs_to_layer_configs( - block_configs: Union["BlockConfig", dict], *, num_attention_heads: int, - hidden_size: int) -> List[DeciLayerConfig]: - return [ - hf_block_config_to_layer_config(block_config, num_attention_heads, - hidden_size) - for block_config in block_configs - ] - - -@contextmanager -def timed_loading() -> Iterator[None]: - tik = time.time() - yield - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - -class TpDim(enum.IntEnum): - NO_TP = -1 - COLWISE = 0 - ROWWISE = 1 - - -class SafetensorsIndex(TypedDict): - metadata: Dict[str, Any] - weight_map: Dict[str, str] - - -class WeightsLoader(ABC): - - @abstractmethod - def read_weight(self, name: str) -> torch.Tensor: - ... - - def get_weight(self, - name: str, - tp_dim: TpDim = TpDim.NO_TP, - tp_size: int = 1, - tp_rank: int = 0) -> torch.Tensor: - weight = self.read_weight(name) - if tp_dim != TpDim.NO_TP: - weight = split(weight, tp_size, tp_rank, dim=tp_dim) - return weight - - def get_kv_weight(self, - name: str, - num_heads: int, - tp_size: int = 1, - tp_rank: int = 0) -> torch.Tensor: - weight = self.read_weight(name) - if tp_size > num_heads: - weight = dup_kv_weight(weight, num_heads, tp_size) - if tp_size > 1: - weight = split(weight, tp_size, tp_rank, dim=0) - - return weight - - -class HFModelWeightsLoader(WeightsLoader): - - def __init__(self, *, hf_model: "transformers.PreTrainedModel", - dtype: str) -> None: - self.model_params = dict(hf_model.named_parameters()) - self.dtype = getattr(torch, dtype) - - def read_weight(self, name: str) -> torch.Tensor: - weight = self.model_params[name] - if weight.dtype != self.dtype: - weight = weight.to(self.dtype) - weight = weight.detach() - return weight - - -class SafetensorsWeightsLoader(WeightsLoader): - - def __init__(self, *, model_dir: Path, dtype: str) -> None: - self.model_dir = model_dir - self.dtype = getattr(torch, dtype) - - # the index has a weight map that maps weight names to the files they are found in - safetensor_index_json = self.model_dir / "model.safetensors.index.json" - has_safetensor_index_json = safetensor_index_json.is_file() - if has_safetensor_index_json: - with safetensor_index_json.open("r") as fr: - self.sharding_map: SafetensorsIndex = json.load(fr) - else: - self.sharding_map = SafetensorsIndex(metadata={}, weight_map={}) - - shard_files = {f.name for f in self.model_dir.glob("*.safetensors")} - if has_safetensor_index_json: - # only read the files that have weights according to the index - shard_files &= set(self.sharding_map["weight_map"].values()) - self.shard_files = sorted(list(shard_files)) - - self.safetensors_files = { - shard_file: - safetensors.safe_open(model_dir / shard_file, - framework="pt", - device="cpu") - for shard_file in shard_files - } - - def read_weight(self, name: str) -> torch.Tensor: - shard_filename = self.sharding_map['weight_map'].get( - name, self.shard_files[0]) - res = self.safetensors_files[shard_filename].get_tensor(name) - return res.to(self.dtype).contiguous() - - -def load_model_weights(loader: WeightsLoader, - config: "DeciConfig") -> Dict[str, torch.Tensor]: - mapping = config.mapping - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - pad_vocab = vocab_size % mapping.tp_size != 0 - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - - weights = {} - - def load_weight(name: str, tp_dim: TpDim = TpDim.NO_TP) -> torch.Tensor: - return loader.get_weight(name=name, - tp_dim=tp_dim, - tp_rank=mapping.tp_rank, - tp_size=mapping.tp_size) - - with timed_loading(): - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = load_weight( - "model.embed_tokens.weight", - TpDim(config.embedding_sharding_dim) - if config.use_parallel_embedding else - TpDim.NO_TP) # vocab_embedding - - if mapping.is_last_pp_rank(): - v = load_weight("lm_head.weight", - TpDim.NO_TP) if pad_vocab else load_weight( - "lm_head.weight", TpDim.COLWISE) # lm_head - if pad_vocab: - v = torch.nn.functional.pad( - v, (0, 0, 0, vocab_size_padded - vocab_size), 'constant', 0) - v = split(v, mapping.tp_size, mapping.tp_rank) - weights['lm_head.weight'] = v - weights['transformer.ln_f.weight'] = load_weight( - "model.norm.weight") # ln_f - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_config = config.get_layer_config(l) - layer_idx = l - layers_range[0] - tllm_prex = f'transformer.layers.{layer_idx}' - - # Attention - if layer_config.is_attention_layer: - weights[f'{tllm_prex}.input_layernorm.weight'] = load_weight( - f"model.layers.{l}.input_layernorm.weight" - ) # input_layernorm - - q = load_weight(f"model.layers.{l}.self_attn.q_proj.weight", - TpDim.COLWISE) - k = loader.get_kv_weight( - f"model.layers.{l}.self_attn.k_proj.weight", - num_heads=layer_config.attention.num_key_value_heads, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank) - v = loader.get_kv_weight( - f"model.layers.{l}.self_attn.v_proj.weight", - num_heads=layer_config.attention.num_key_value_heads, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank) - weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat( - [q, k, v], 0) - weights[f'{tllm_prex}.attention.dense.weight'] = load_weight( - f"model.layers.{l}.self_attn.o_proj.weight", - TpDim.ROWWISE) # attention.dense - - elif layer_config.is_linear_attention_layer: - weights[f'{tllm_prex}.input_layernorm.weight'] = load_weight( - f"model.layers.{l}.input_layernorm.weight" - ) # input_layernorm - - weights[f'{tllm_prex}.attention.weight'] = load_weight( - f"model.layers.{l}.self_attn.linear_attn.weight", - TpDim.COLWISE) - - elif not layer_config.is_noop_attention_layer: - raise NotImplementedError( - f"Loading weights for layer with attention of type {layer_config.attention.impl} is not supported" - ) - - # MLP - if layer_config.is_mlp_layer: - weights[f'{tllm_prex}.post_layernorm.weight'] = load_weight( - f"model.layers.{l}.post_attention_layernorm.weight" - ) # post_layernorm - - weights[f'{tllm_prex}.ffn.gate.weight'] = load_weight( - f"model.layers.{l}.mlp.up_proj.weight", - TpDim.COLWISE) # mlp.gate - weights[f'{tllm_prex}.ffn.proj.weight'] = load_weight( - f"model.layers.{l}.mlp.down_proj.weight", - TpDim.ROWWISE) # mlp.proj - weights[f'{tllm_prex}.ffn.fc.weight'] = load_weight( - f"model.layers.{l}.mlp.gate_proj.weight", - TpDim.COLWISE) # mlp.fc - - elif layer_config.is_linear_ffn_layer: - weights[f'{tllm_prex}.post_layernorm.weight'] = load_weight( - f"model.layers.{l}.post_attention_layernorm.weight" - ) # post_layernorm - - weights[f'{tllm_prex}.ffn.weight'] = load_weight( - f"model.layers.{l}.mlp.linear_mlp.weight", TpDim.COLWISE) - - elif not layer_config.is_noop_ffn_layer: - raise NotImplementedError( - f"Loading weights for a layer with FFN of type {layer_config.ffn.impl} is not implemented yet" - ) - - return weights - - -def load_weights_from_hf_model( - hf_model: "transformers.PreTrainedModel", - config: "DeciConfig", - act_range: Optional[dict] = None, - qkv_para: Optional[dict] = None, - smoother: Optional[dict] = None) -> Dict[str, torch.Tensor]: - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - if quant_algo == QuantAlgo.W8A16: - torch.int8 - elif quant_algo == QuantAlgo.W4A16: - torch.quint4x2 - else: - pass - - use_smooth_quant = config.quantization._use_plugin_sq - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - if use_smooth_quant or int8_kv_cache: - assert act_range is not None - assert qkv_para is not None - assert smoother is not None - - # TODO(oargov): add support for these quants - assert not use_weight_only, "WOQ is not supported yet" - assert not use_smooth_quant, "SmoothQuant is not supported yet" - assert not int8_kv_cache, "INT8 KV cache is not supported yet" - - # TODO(oargov): support moe - moe_config = getattr(config, "moe", None) - assert moe_config is None, "MoE is not supported yet" - - # TODO(oargov): implement resisdual mlp - residual_mlp = getattr(config, "residual_mlp", None) - assert not residual_mlp, "Residual MLP is not supported yet" - - loader = HFModelWeightsLoader(hf_model=hf_model, dtype=config.dtype) - logger.info('Converting weights from Huggingface model...') - return load_model_weights(loader=loader, config=config) - - -def load_weights_from_hf_safetensors( - model_dir: Union[str, Path], - config: "DeciConfig") -> Dict[str, torch.Tensor]: - - if isinstance(model_dir, str): - model_dir = Path(model_dir) - - loader = SafetensorsWeightsLoader(model_dir=model_dir, dtype=config.dtype) - logger.info('Loading weights from Huggingface safetensors...') - return load_model_weights(loader=loader, config=config) - - -def update_weights_following_modelopt_optimization( - weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - # Rename MLPs to FFNs to match TRTLLM implementation expectation - weights = {k.replace('.mlp.', '.ffn.'): v for k, v in weights.items()} - - # Move all linear attentions to their expected locations - weights = { - k.replace('.attn_replacing_linear.', '.attention.'): v - for k, v in weights.items() - } - - # Move all linear MLPs to their expected locations - weights = { - k.replace('.mlp_replacing_linear.', '.ffn.'): v - for k, v in weights.items() - } - - return weights diff --git a/tensorrt_llm/models/nemotron_nas/layer_config.py b/tensorrt_llm/models/nemotron_nas/layer_config.py deleted file mode 100644 index 84ed3487da7c..000000000000 --- a/tensorrt_llm/models/nemotron_nas/layer_config.py +++ /dev/null @@ -1,86 +0,0 @@ -# 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 enum -from dataclasses import dataclass, field -from typing import Any, Dict, Optional - - -class AttentionImplementation(str, enum.Enum): - ATTENTION = "attention" - LINEAR = "linear" - NO_OP = "no_op" - - -class FFNImplementation(str, enum.Enum): - MLP = "mlp" - LINEAR = "linear" - NO_OP = "no_op" - - -@dataclass(frozen=True, kw_only=True) -class AttentionConfig: - impl: AttentionImplementation = AttentionImplementation.ATTENTION - num_key_value_heads: Optional[int] = None - - @property - def needs_kv_cache(self) -> bool: - return self.impl == AttentionImplementation.ATTENTION - - -@dataclass(frozen=True, kw_only=True) -class FFNConfig: - impl: FFNImplementation = FFNImplementation.MLP - intermediate_size: Optional[int] = None - - -@dataclass(frozen=True, kw_only=True) -class DeciLayerConfig: - attention: AttentionConfig = field(default_factory=AttentionConfig) - ffn: FFNConfig = field(default_factory=FFNConfig) - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> "DeciLayerConfig": - assert "attention" in d, "Missing attention configuration" - assert "ffn" in d, "Missing mlp configuration" - - return cls( - attention=AttentionConfig(**d["attention"]), - ffn=FFNConfig(**d["ffn"]), - ) - - @property - def is_attention_layer(self) -> bool: - return self.attention.impl == AttentionImplementation.ATTENTION - - @property - def is_mlp_layer(self) -> bool: - return self.ffn.impl == FFNImplementation.MLP - - @property - def is_noop_attention_layer(self) -> bool: - return self.attention.impl == AttentionImplementation.NO_OP - - @property - def is_linear_attention_layer(self) -> bool: - return self.attention.impl == AttentionImplementation.LINEAR - - @property - def is_noop_ffn_layer(self) -> bool: - return self.ffn.impl == FFNImplementation.NO_OP - - @property - def is_linear_ffn_layer(self) -> bool: - return self.ffn.impl == FFNImplementation.LINEAR diff --git a/tensorrt_llm/models/nemotron_nas/model.py b/tensorrt_llm/models/nemotron_nas/model.py deleted file mode 100644 index f6562d3a7649..000000000000 --- a/tensorrt_llm/models/nemotron_nas/model.py +++ /dev/null @@ -1,816 +0,0 @@ -# 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 dataclasses import dataclass -from typing import List, Optional, Tuple, Type, Union - -from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, - AttentionMaskType, PositionEmbeddingType, - Tensor, gather_last_token_logits, recv, - send) -from tensorrt_llm.layers.attention import (Attention, AttentionParams, - KeyValueCacheParams, - SpecDecodingParams) -from tensorrt_llm.layers.embedding import Embedding -from tensorrt_llm.layers.linear import ColumnLinear -from tensorrt_llm.layers.lora import LoraParams -from tensorrt_llm.layers.mlp import GatedMLP -from tensorrt_llm.layers.normalization import RmsNorm -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import has_safetensors -from tensorrt_llm.models.modeling_utils import DecoderModelForCausalLM -from tensorrt_llm.models.nemotron_nas.config import DeciConfig -from tensorrt_llm.models.nemotron_nas.convert import ( - load_weights_from_hf_model, load_weights_from_hf_safetensors, - update_weights_following_modelopt_optimization) -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.plugin.plugin import init_all_reduce_helper - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ..modeling_utils import PretrainedConfig, QuantConfig, preprocess_weights - - -@dataclass -class DeciLMLayerOutput: - hidden_states: Tensor - present_kv: Optional[Tensor] = None - - -@dataclass -class DeciLMLayerListOutput: - hidden_states: Tensor - present_kvs: List[Tensor] - - -class NoOp(Module): - - def forward(self, hidden_states: Tensor, *args, **kwargs) -> int: - return 0 - - -class NoOpAttention(NoOp): - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache: bool = False, - *args, - **kwargs) -> Union[int, Tuple[int, None]]: - out = super().forward(hidden_states=hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - *args, - **kwargs) - if use_cache: - return out, None - return out - - -class LinearAttention(ColumnLinear): - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache: bool = False, - *args, - **kwargs) -> Union[Tensor, Tuple[Tensor, None]]: - out = super().forward(x=hidden_states, - lora_runtime_params=None, - lora_hidden_state=None) - - if use_cache: - return out, None - return out - - -class LinearFFN(ColumnLinear): - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None) -> Tensor: - return super().forward(x=hidden_states, - lora_runtime_params=None, - lora_hidden_state=None) - - -NoOpFFN = NoOp -NoOpLayerNorm = NoOp - - -class DeciLMDecoderLayer(Module): - - def __init__(self, config: DeciConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - - self.layer_config = self.config.get_layer_config(self.layer_idx) - - self._init_attention() - self._init_ffn() - - @property - def input_layernorm_was_fused(self) -> bool: - """ - The previous layer ran our input_layernorm for us if: - 1. The reduce_fusion plugin is enabled and - 2. We are not the first local model layer and - 3. The previous layer is an MLP layer - """ - return default_net( - ).plugin_config.reduce_fusion and self.local_layer_idx > 0 and self.config.get_layer_config( - self.layer_idx - - 1).is_mlp_layer and self.needs_input_layernorm_fusion - - @property - def needs_input_layernorm_fusion(self) -> bool: - """ - This layer needs the previous layer to perform input_layernorm fusion if: - 1. The reduce_fusion plugin is enabled and - 2. This is not a NOOP attention layer (otherwise it has no input_layernorm) - """ - return default_net( - ).plugin_config.reduce_fusion and not self.layer_config.is_noop_attention_layer - - @property - def can_fuse_post_layernorm(self) -> bool: - """ - This layer can fuse attention and post_layernorm if: - 1. The reduce_fusion plugin is enabled and - 2. It is an attention layer and - 3. It is not a NOOP FFN layer (othrewise it has no post_layernorm) - """ - return default_net( - ).plugin_config.reduce_fusion and self.layer_config.is_attention_layer and not self.layer_config.is_noop_ffn_layer - - @property - def can_fuse_input_layernorm(self) -> bool: - """ - This layer can run the next layer's input_layernorm if: - 1. The reduce_fusion plugin is enable and - 2. It is an MLP layer - """ - return default_net( - ).plugin_config.reduce_fusion and self.layer_config.is_mlp_layer - - def _init_attention(self) -> None: - """ - Initialize some attention alternative - """ - # normal attention - if self.layer_config.is_attention_layer: - # according to recurrentgemma, len(layer_types) can be less than num_hidden_layers - # in this case, the list should wrap-around - # for example, if layer_types = ["attention", "recurrent", "recurrent"], and we have 5 layers, we get: - # layer 0 ==> attention - # layer 1 ==> recurrent - # layer 2 ==> recurrent - # layer 3 ==> attention - # layer 4 ==> recurrent - # we check which layers are local to our rank - layers_range = self.config.mapping.pp_layers( - self.config.num_hidden_layers) - # then take the size of layer_types in the config - layer_type_len = len(self.config.layer_types) - # collect the layer types of all the local layers - local_layer_types = [ - self.config.layer_types[layer_id % layer_type_len] - for layer_id in layers_range - ] - # and see how many of them are attention layers to determine our local attention layer idx - local_attn_layer_idx = local_layer_types[:self. - local_layer_idx].count( - "attention") - - # Iterate over all local layer configs, getting num_kv_heads of the attention ones - num_kv_heads_per_local_layer = [ - layer_config.attention.num_key_value_heads for layer_config in - [self.config.layer_configs[idx] for idx in layers_range] - if layer_config.is_attention_layer - ] - - # adjust num heads according to tp size - num_kv_heads_per_local_layer = [ - (nheads + self.config.mapping.tp_size - 1) // - self.config.mapping.tp_size - for nheads in num_kv_heads_per_local_layer - ] - nheads_tp = (self.layer_config.attention.num_key_value_heads + - self.config.mapping.tp_size - - 1) // self.config.mapping.tp_size - - self.input_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.attention = Attention( - local_layer_idx=local_attn_layer_idx, - hidden_size=self.config.hidden_size, - attention_head_size=self.config.head_size, - num_attention_heads=self.config.num_attention_heads, - num_kv_heads=self.layer_config.attention.num_key_value_heads, - max_position_embeddings=self.config.max_position_embeddings, - dtype=self.config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=False, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=self.config.rotary_base, - rotary_embedding_scaling=self.config.rotary_scaling, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - tp_rank=self.config.mapping.tp_rank, - quant_mode=self.config.quant_mode) - - elif self.layer_config.is_noop_attention_layer: - self.input_layernorm = NoOpLayerNorm() - self.attention = NoOpAttention() - - elif self.layer_config.is_linear_attention_layer: - self.input_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.attention = LinearAttention( - in_features=self.config.hidden_size, - out_features=self.config.hidden_size, - bias=False, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True) - - else: - raise NotImplementedError( - f"Attention of type {str(self.layer_config.attention.impl)} is not implemented" - ) - - def _init_ffn(self) -> None: - """ - Initialize some ffn alternative - """ - - if self.layer_config.is_mlp_layer: - intermediate_size = self.layer_config.ffn.intermediate_size or self.config.intermediate_size - mlp_hidden_size = intermediate_size or self.config.hidden_size * 4 - - self.post_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.ffn = GatedMLP( - hidden_size=self.config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=self.config.hidden_act, - bias=False, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - quant_mode=self.config.quant_mode, - ) - - elif self.layer_config.is_noop_ffn_layer: - self.post_layernorm = NoOpLayerNorm() - self.ffn = NoOpFFN() - - elif self.layer_config.is_linear_ffn_layer: - self.post_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.ffn = LinearFFN(in_features=self.config.hidden_size, - out_features=self.config.hidden_size, - bias=False, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True) - - else: - raise NotImplementedError( - f"FFN of type {str(self.layer_config.ffn.impl)} is not implemented" - ) - - def forward(self, - hidden_states: Tensor | Tuple[Tensor, Tensor], - attention_mask: Optional[Tensor] = None, - use_cache: bool = False, - spec_decoding_params=None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - lora_layer_params: Optional[LoraParams] = None, - next_layer_input_layernorm_args: Optional[Tuple[Tensor, - float]] = None): - if self.input_layernorm_was_fused: - # previous layer already performed our layer norm - assert isinstance(hidden_states, tuple) - hidden_states, residual = hidden_states - else: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - if self.can_fuse_post_layernorm: - all_reduce_params = AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=self.post_layernorm.weight.value, - eps=self.post_layernorm.eps) - else: - all_reduce_params = None - - attention_output = self._run_attention( - hidden_states=hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - if use_cache: - attention_output, present_kv = attention_output - else: - present_kv = None - - if self.can_fuse_post_layernorm: - hidden_states, residual = attention_output - else: - hidden_states = residual + attention_output - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - if next_layer_input_layernorm_args is not None: - assert self.can_fuse_input_layernorm - norm_weight, eps = next_layer_input_layernorm_args - all_reduce_params = AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=norm_weight, - eps=eps) - hidden_states = self._run_ffn(hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - else: - hidden_states = self._run_ffn(hidden_states, - lora_layer_params=lora_layer_params) - hidden_states = residual + hidden_states - - return DeciLMLayerOutput(hidden_states=hidden_states, - present_kv=present_kv) - - def _run_attention( - self, - hidden_states: Tensor, - attention_mask: Optional[Tensor] = None, - use_cache: bool = False, - spec_decoding_params=None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - lora_layer_params: Optional[LoraParams] = None, - all_reduce_params: Optional[AllReduceParams] = None - ) -> Union[Tensor, Tuple[Tensor, None]]: - """ - Ideally, this functionality would be encapsulated in a LinearAttention class, but during - FP8 and lower quantization, our linear classes get overrun by ModelOpt, thus we must - control the attention inputs at the DecoderLayer level. - """ - if self.layer_config.is_linear_attention_layer: - out = self.attention(hidden_states) - return out, None if use_cache else out - else: - if not self.layer_config.is_attention_layer: - assert all_reduce_params is None, f"Layer with attention of type {self.layer_config.attention.impl} can't do reduce_fusion" - - return self.attention(hidden_states=hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - def _run_ffn(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None): - """ - Ideally, this functionality would be encapsulated in a LinearMLP class, but during - FP8 and lower quantization, our linear classes get overrun by ModelOpt, thus we must - control the MLP inputs at the DecoderLayer level. - """ - if all_reduce_params is not None: - assert self.layer_config.is_mlp_layer, f"Layer with FFN of type {self.layer_config.ffn.impl} can't do reduce_fusion" - - if self.layer_config.is_linear_ffn_layer: - return self.ffn(hidden_states) - else: - return self.ffn(hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - -class DeciLMDecoderLayerList(ModuleList): - - def __init__(self, cls: Type[DeciLMDecoderLayer], config: DeciConfig): - self.num_hidden_layers = config.num_hidden_layers - # global indices of local layers - self.layer_list = config.mapping.pp_layers(config.num_hidden_layers) - super().__init__([cls(config, idx) for idx in self.layer_list]) - # global indices of local attention layers - self.attention_layer_list = [ - self.layer_list[i] for i, layer in enumerate(self) - if layer.layer_config.is_attention_layer - ] - - def forward( - self, - hidden_states: Tensor, - use_cache: bool, - attention_mask: Optional[Tensor], - kv_cache_params: KeyValueCacheParams, - attention_params: Optional[AttentionParams] = None, - position_ids: Optional[Tensor] = None, - lora_params: Optional[LoraParams] = None, - spec_decoding_params: Optional[SpecDecodingParams] = None, - ) -> DeciLMLayerListOutput: - kv_cache_params.fill_none_tensor_list(len(self.layer_list)) - - presents = [] - - # put None where we don't have attention layers - pkv_iter = iter(kv_cache_params.past_key_value) - - past_key_values = [x for x in pkv_iter] - - for layer_idx, (layer, past) in enumerate(zip(self, past_key_values)): - next_layer_input_layernorm_args = None - if default_net().plugin_config.reduce_fusion: - if layer_idx < self.layer_list[-1]: - # this is not the last layer - next_layer = self[layer_idx + 1] - if layer.can_fuse_input_layernorm and next_layer.needs_input_layernorm_fusion: - # this layer can fuse the next layer's input_layernorm - next_layer_input_layernorm_args = ( - next_layer.input_layernorm.weight.value, - next_layer.input_layernorm.eps) - - layer_out = layer( - hidden_states=hidden_states, - attention_mask=attention_mask, - attention_params=attention_params, - kv_cache_params=KeyValueCacheParams( - past_key_value=[past], - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cache_indirection=kv_cache_params.cache_indirection, - ), - spec_decoding_params=spec_decoding_params, - use_cache=use_cache, - lora_layer_params=lora_params.get_layer_config(layer_idx) - if lora_params is not None - and lora_params.lora_ranks is not None else None, - next_layer_input_layernorm_args=next_layer_input_layernorm_args) - - hidden_states = layer_out.hidden_states - if use_cache and layer_out.present_kv is not None: - presents.append(layer_out.present_kv) - - return DeciLMLayerListOutput(hidden_states=hidden_states, - present_kvs=presents) - - -class DeciLMModel(Module): - - def __init__(self, config: DeciConfig) -> None: - super().__init__() - init_all_reduce_helper() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - # first rank in pipeline-parallel handles token embedding - assert config.vocab_size is not None - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.position_embedding_type = config.position_embedding_type - self.layers = DeciLMDecoderLayerList(DeciLMDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - # last rank in pipeline-parallel handles final norm - self.ln_f = RmsNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype, - ) - - def _vocab_embedding(self, - input_ids: Tensor, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None) -> Tensor: - # prompt tuning - ptuning_args = ([ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else []) - - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - return hidden_states - - def forward( - self, - input_ids, - position_ids=None, - use_cache: bool = False, - attention_mask: Optional[Tensor] = None, - spec_decoding_params=None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - hidden_states: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params: Optional[LoraParams] = None, - ) -> DeciLMLayerListOutput: - - if self.mapping.is_first_pp_rank(): - # first pipeline rank ==> do prompt embedding - hidden_states = self._vocab_embedding( - input_ids=input_ids, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size) - else: - # receive hidden states from prior rank in the pipeline - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - layers_out = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - spec_decoding_params=spec_decoding_params, - ) - - if self.mapping.is_last_pp_rank(): - # last pipeline rank ==> do final norm - hidden_states = self.ln_f(layers_out.hidden_states) - else: - # send hidden states to next rank in the pipeline - hidden_states = send(layers_out.hidden_states, - self.mapping.next_pp_rank()) - - return DeciLMLayerListOutput(hidden_states=hidden_states, - present_kvs=layers_out.present_kvs) - - -class DeciLMForCausalLM(DecoderModelForCausalLM): - config_class = DeciConfig - - def __init__(self, config: DeciConfig): - - transformer = DeciLMModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - if config.mapping.is_last_pp_rank(): - # last pipeline rank needs to do calculate logits - lm_head = ColumnLinear( - 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, - ) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - # Create constant attention parameters to be reused by all layers. - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - - @classmethod - def from_hugging_face(cls, - hf_model_or_dir: Union[ - str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - load_by_shard: bool = False, - load_model_on_cpu: bool = False, - trust_remote_code: bool = True, - **kwargs) -> "DeciLMForCausalLM": - import transformers - - # TODO(oargov): add support for these - assert not load_by_shard, "load_by_shard is not implemented yet" - - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_config_or_dir = hf_model_or_dir.config - else: - hf_config_or_dir = hf_model_or_dir - - config = DeciConfig.from_hugging_face( - hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - trust_remote_code=trust_remote_code, - **kwargs) - - if use_preloading: - assert not load_by_shard - weights = load_weights_from_hf_model(hf_model_or_dir, config) - elif has_safetensors( - hf_model_or_dir) and not config.quant_mode.has_any_quant(): - weights = load_weights_from_hf_safetensors(hf_model_or_dir, config) - else: - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_or_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype=dtype, - trust_remote_code=trust_remote_code, - ) - weights = load_weights_from_hf_model(hf_model, config) - preprocess_weights(weights, config) - - model = DeciLMForCausalLM(config) - model.load(weights) - return model - - @classmethod - def from_checkpoint(cls, - ckpt_dir: str, - rank: Optional[int] = None, - config: Optional["PretrainedConfig"] = None): - return super().from_checkpoint( - ckpt_dir, - rank, - config, - preprocess_weights_hook= - update_weights_following_modelopt_optimization, - ) - - def forward( - self, - input_ids: Tensor, - position_ids: Optional[Tensor] = None, - use_cache: bool = False, - last_token_ids: Optional[Tensor] = None, - attention_mask: Optional[Tensor] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - hidden_states: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params: Optional[LoraParams] = None, - spec_decoding_params=None, - ): - # fill attention params. - attention_params = Attention.fill_attention_params( - self, attention_params) - - model_out = self.transformer.forward( - input_ids=input_ids, - position_ids=position_ids, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - hidden_states=hidden_states, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size, - spec_decoding_params=spec_decoding_params) - hidden_states = model_out.hidden_states - - if self.config.mapping.is_last_pp_rank(): - hidden_states = gather_last_token_logits( - hidden_states, - last_token_ids, - default_net().plugin_config.remove_input_padding, - ) - - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output("logits", self.config.logits_dtype) - else: - hidden_states.mark_output("hidden_states_output", self.config.dtype) - - if use_cache and not default_net().plugin_config.paged_kv_cache: - presents = model_out.present_kvs - for i, present in zip(self.transformer.layers.attention_layer_list, - presents): - present.mark_output(f"present_key_value_{i}", - self.config.kv_dtype) - if self.config.mapping.is_last_pp_rank(): - return (lm_logits, presents, hidden_states) - return (hidden_states, presents) - else: - if self.config.mapping.is_last_pp_rank(): - return lm_logits, hidden_states - return hidden_states - - def prepare_attention_inputs( - self, - *, - max_batch_size: int, - max_beam_width: int, - max_input_len: int, - max_seq_len: int, - num_kv_heads: int, - head_size: int, - num_layers: int, - kv_dtype: str, - kv_cache_type: KVCacheType, - num_profiles: int = 1, - enable_ctx_gen_opt_profiles: bool = False, - remove_input_padding: bool = False, - use_gpt_attention_plugin: bool = False, - paged_kv_cache: bool = False, - tokens_per_block: int = 32, - mapping: Mapping = Mapping(), - use_cache: bool = True, - streamingllm: bool = False, - attn_layer_idx: Optional[List[int]] = None, - opt_batch_size: Optional[int] = None, - num_kv_heads_per_layer: Optional[List[int]] = None): - - if attn_layer_idx is None: - attn_layer_idx, num_kv_heads_per_layer = [], [] - for layer_idx in range(self.config.num_hidden_layers): - layer_config = self.config.get_layer_config(layer_idx) - if layer_config.is_attention_layer: - attn_layer_idx.append(layer_idx) - num_kv_heads_per_layer.append( - layer_config.attention.num_key_value_heads) - - attention_inputs = super().prepare_attention_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - num_kv_heads=num_kv_heads, - head_size=head_size, - num_layers=num_layers, - kv_dtype=kv_dtype, - num_profiles=num_profiles, - kv_cache_type=kv_cache_type, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - tokens_per_block=tokens_per_block, - mapping=mapping, - streamingllm=streamingllm, - attn_layer_idx=attn_layer_idx, - opt_batch_size=opt_batch_size, - num_kv_heads_per_layer=num_kv_heads_per_layer) - - return attention_inputs diff --git a/tensorrt_llm/models/opt/__init__.py b/tensorrt_llm/models/opt/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/opt/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/opt/model.py b/tensorrt_llm/models/opt/model.py deleted file mode 100644 index 2a81a8c29b3e..000000000000 --- a/tensorrt_llm/models/opt/model.py +++ /dev/null @@ -1,181 +0,0 @@ -# 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 ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class OPTDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - self.do_layer_norm_before = self.config.do_layer_norm_before - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - max_position_embeddings=config.max_position_embeddings, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - - attention_input = hidden_states - if self.do_layer_norm_before: - attention_input = self.input_layernorm(hidden_states) - - # At this point the hidden_states object must be a Tensor. - assert isinstance(attention_input, Tensor) - - attention_output = self.attention(attention_input, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - if not self.do_layer_norm_before: - hidden_states = self.input_layernorm(hidden_states) - - residual = hidden_states - if self.do_layer_norm_before: - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if not self.do_layer_norm_before: - hidden_states = self.post_layernorm(hidden_states) - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class OPTModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.do_layer_norm_before = config.do_layer_norm_before - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.position_embedding = Embedding(config.max_position_embeddings, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(OPTDecoderLayer, config) - - if self.do_layer_norm_before: - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - 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, - **kwargs): - - 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) - hidden_states = hidden_states + self.position_embedding(position_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.do_layer_norm_before: - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class OPTForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = OPTModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(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) - - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('do_layer_norm_before', False) diff --git a/tensorrt_llm/models/phi/__init__.py b/tensorrt_llm/models/phi/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/phi/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/phi/config.py b/tensorrt_llm/models/phi/config.py deleted file mode 100644 index 583de15fadf7..000000000000 --- a/tensorrt_llm/models/phi/config.py +++ /dev/null @@ -1,87 +0,0 @@ -# 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 ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class PhiConfig(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in PhiConfig - - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - rotary_scaling = getattr(hf_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - 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=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=hf_config.hidden_act, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - rotary_pct=hf_config.partial_rotary_factor, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/phi/convert.py b/tensorrt_llm/models/phi/convert.py deleted file mode 100644 index 4bf3406c7262..000000000000 --- a/tensorrt_llm/models/phi/convert.py +++ /dev/null @@ -1,145 +0,0 @@ -import torch - -from ..._utils import get_hf_rope_theta, pad_vocab_size, str_dtype_to_torch - - -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() - - -def load_weights_from_hf_model(hf_model, config): - torch_dtype = str_dtype_to_torch(config.dtype) - hf_state_dict = hf_model.state_dict() - weights = {} - is_weight_only = config.quant_mode.is_weight_only() - if config.quant_mode.is_int8_weight_only(): - plugin_weight_only_quant_type = torch.int8 - elif config.quant_mode.is_int4_weight_only(): - plugin_weight_only_quant_type = torch.quint4x2 - # replace key name - for key, value in hf_state_dict.items(): - # Decoder Layers - if "model.layers." in key: - key = key.replace("model.layers.", "transformer.layers.") - key = key.replace("self_attn.", "attention.") - key = key.replace("mlp.fc1.", "mlp.fc.") - key = key.replace("mlp.fc2.", "mlp.proj.") - # Embedding - key = key.replace("model.embed_tokens.weight", - "transformer.vocab_embedding.weight") - # Final Layer norm - key = key.replace("model.final_layernorm.", "transformer.ln_f.") - - weights[key] = value.to(torch_dtype).cpu() - - # merge qkv weights - qkv_keys = ["q_proj", "k_proj", "v_proj"] - - scales = {} - for key in hf_state_dict.keys(): - if 'self_attn.q_proj.weight' in key: - prefix = key.split('self_attn')[0].replace("model.layers.", - "transformer.layers.") - - # [(num_heads x q)|(num_heads x k)|(num_heads x v), hidden_size] - qkv_weights = [] - qkv_bias = [] - - for k in qkv_keys: - split_w = split(weights.pop(f"{prefix}attention.{k}.weight"), - config.mapping.tp_size, config.mapping.tp_rank) - - qkv_weights.append(split_w) - split_b = split(weights.pop(f"{prefix}attention.{k}.bias"), - config.mapping.tp_size, config.mapping.tp_rank) - qkv_bias.append(split_b) - v = torch.cat(qkv_weights, dim=0) - if is_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.t().contiguous().cpu(), plugin_weight_only_quant_type) - weights[ - f"{prefix}attention.qkv.weight"] = processed_torch_weights - scales[ - f"{prefix}attention.qkv.per_channel_scale"] = torch_weight_scales - else: - weights[f"{prefix}attention.qkv.weight"] = v - weights[f"{prefix}attention.qkv.bias"] = torch.cat(qkv_bias, dim=0) - - tp_rank = config.mapping.tp_rank - for weight_name in weights: - loaded_weight = weights[weight_name] - if "attention.dense.weight" in weight_name or "mlp.proj.weight" in weight_name: # RowLinear - v = split(loaded_weight, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=1) - if is_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.t().contiguous().cpu(), plugin_weight_only_quant_type) - weights[weight_name] = processed_torch_weights - scales[weight_name.replace( - '.weight', '.per_channel_scale')] = torch_weight_scales - else: - weights[weight_name] = v - elif "mlp.fc." in weight_name: - - v = split(loaded_weight, config.mapping.tp_size, - config.mapping.tp_rank) - if is_weight_only and "mlp.fc.weight" in weight_name: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.t().contiguous().cpu(), plugin_weight_only_quant_type) - weights[weight_name] = processed_torch_weights - scales[weight_name.replace( - '.weight', '.per_channel_scale')] = torch_weight_scales - else: - weights[weight_name] = v - elif "lm_head." in weight_name: - output_dim = 0 - shard_size = loaded_weight.shape[output_dim] - tp_rank * shard_size - vocab_size = loaded_weight.shape[output_dim] - - if shard_size % config.mapping.tp_size != 0: - pad_width = pad_vocab_size(vocab_size, - config.mapping.tp_size) - vocab_size - loaded_weight = torch.nn.functional.pad(loaded_weight, - (0, 0, 0, pad_width), - 'constant', - value=0) - weights[weight_name] = split(loaded_weight, config.mapping.tp_size, - config.mapping.tp_rank) - - weights.update(scales) - - return weights - - -def convert_hf_config(hf_config, dtype, args): - config = { - 'architecture': hf_config.architectures[0], - 'dtype': dtype, - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_key_value_heads, - 'rotary_pct': hf_config.partial_rotary_factor, - 'rope_theta': get_hf_rope_theta(hf_config, 10000.0), - 'hidden_size': hf_config.hidden_size, - 'intermediate_size': hf_config.intermediate_size, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': hf_config.max_position_embeddings, - 'hidden_act': hf_config.hidden_act, - 'mapping': { - 'world_size': args.tp_size * args.pp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - } - } - return config diff --git a/tensorrt_llm/models/phi/model.py b/tensorrt_llm/models/phi/model.py deleted file mode 100644 index fc5c6ccbc1f9..000000000000 --- a/tensorrt_llm/models/phi/model.py +++ /dev/null @@ -1,217 +0,0 @@ -# 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 AutoModelForCausalLM - -from ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) -from .config import PhiConfig -from .convert import load_weights_from_hf_model - - -class PhiDecoderLayer(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 - - self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - rotary_embedding_percentage=config.rotary_pct, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - 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) - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - ): - residual = hidden_states - - 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=True, - lora_layer_params=lora_layer_params, - ) - - if use_cache: - attention_output, presents = attention_output - - feed_forward_hidden_states = self.mlp(input_layernorm_output, ) - hidden_states = attention_output + feed_forward_hidden_states + residual - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class PhiModel(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(PhiDecoderLayer, config) - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - 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) - - 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 PhiForCausalLM(DecoderModelForCausalLM): - config_class = PhiConfig - - config_class = PhiConfig - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = PhiModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=True, - 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_q": "q_proj", - "attn_k": "k_proj", - "attn_v": "v_proj" - } - - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('partial_rotary_factor', 0.4) - config.set_if_not_exist('rotary_base', 10000.0) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - config = PhiConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype="auto", trust_remote_code=trust_remote_code) - - assert isinstance(hf_model, transformers.PreTrainedModel) - - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/phi3/__init__.py b/tensorrt_llm/models/phi3/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/phi3/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/phi3/config.py b/tensorrt_llm/models/phi3/config.py deleted file mode 100644 index 42d3954092ec..000000000000 --- a/tensorrt_llm/models/phi3/config.py +++ /dev/null @@ -1,143 +0,0 @@ -# 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 ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class Phi3Config(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in PhiConfig - - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - if hasattr(hf_config, "llm_config"): - hf_config = hf_config.llm_config - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - small_variant = hf_config.architectures[0] == "Phi3SmallForCausalLM" - - kwargs['rotary_pct'] = getattr(hf_config, "partial_rotary_factor", 1.0) - kwargs['tie_word_embeddings'] = getattr(hf_config, - "tie_word_embeddings", False) - if small_variant: - kwargs['gegelu_limit'] = getattr(hf_config, "gegelu_limit", None) - kwargs['rotary_base'] = hf_config.rope_embedding_base - kwargs['mup_attn_multiplier'] = getattr(hf_config, - "mup_attn_multiplier", None) - kwargs['mup_embedding_multiplier'] = getattr( - hf_config, "mup_embedding_multiplier", None) - kwargs['mup_use_scaling'] = getattr(hf_config, "mup_use_scaling", - None) - kwargs['mup_width_multiplier'] = getattr(hf_config, - "mup_width_multiplier", - None) - kwargs['blocksparse_block_size'] = getattr( - hf_config, "blocksparse_block_size", None) - kwargs['blocksparse_homo_head_pattern'] = getattr( - hf_config, "blocksparse_homo_head_pattern", None) - kwargs['blocksparse_num_local_blocks'] = getattr( - hf_config, "blocksparse_num_local_blocks", None) - kwargs['blocksparse_vertical_stride'] = getattr( - hf_config, "blocksparse_vert_stride", None) - kwargs['dense_attention_every_n_layers'] = getattr( - hf_config, "dense_attention_every_n_layers", None) - kwargs['norm_epsilon'] = hf_config.layer_norm_epsilon - else: - kwargs['rotary_base'] = get_hf_rope_theta(hf_config, 10000.0) - kwargs['norm_epsilon'] = hf_config.rms_norm_eps - moe_variant = hf_config.architectures[0] == "PhiMoEForCausalLM" - if moe_variant: - kwargs.update({ - 'moe': { - 'num_experts': hf_config.num_local_experts, - 'top_k': hf_config.num_experts_per_tok, - 'normalization_mode': - MoeConfig.ExpertScaleNormalizationMode.SPARSE_MIXER, - 'sparse_mixer_epsilon': hf_config.router_jitter_noise, - }, - 'attention_bias': hf_config.attention_bias - }) - - kwargs['position_embedding_type'] = 'rope_gpt_neox' - if hf_config.max_position_embeddings >= 128000: - kwargs[ - 'original_max_position_embeddings'] = hf_config.original_max_position_embeddings - kwargs['position_embedding_type'] = "long_rope" - kwargs['longrope_scaling_short_factors'] = hf_config.rope_scaling[ - "short_factor"] - kwargs['longrope_scaling_long_factors'] = hf_config.rope_scaling[ - "long_factor"] - if small_variant or moe_variant: - kwargs['longrope_long_mscale'] = hf_config.rope_scaling[ - "long_mscale"] - kwargs['longrope_short_mscale'] = hf_config.rope_scaling[ - "short_mscale"] - - 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=num_key_value_heads, - vocab_size=hf_config.vocab_size, - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act="swiglu" - if hf_config.hidden_act == 'silu' else hf_config.hidden_act, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/phi3/convert.py b/tensorrt_llm/models/phi3/convert.py deleted file mode 100644 index 1f11408b6940..000000000000 --- a/tensorrt_llm/models/phi3/convert.py +++ /dev/null @@ -1,134 +0,0 @@ -import torch - -from ..._utils import str_dtype_to_torch -from .split_weights import shuffle_qkv_weights, split_weights_tp - - -def load_weights_from_hf_model(hf_model, config): - torch_dtype = str_dtype_to_torch(config.dtype) - hf_state_dict = hf_model.state_dict() - weights = {} - - config.quant_mode.is_weight_only() - if config.quant_mode.is_int8_weight_only(): - torch.int8 - elif config.quant_mode.is_int4_weight_only(): - torch.quint4x2 - # replace key name - for key, value in hf_state_dict.items(): - # Decoder Layers - orig_key = key - if "model.layers." in key: - key = key.replace("model.layers.", "transformer.layers.") - #Attention - key = key.replace("self_attn.", "attention.") - key = key.replace("query_key_value.", "qkv.") # small - key = key.replace("Wqkv.weight", "qkv.weight") - key = key.replace("qkv_proj.", "qkv.") #128k - #MLP - key = key.replace("mlp.fc1.", "mlp.fc.") - key = key.replace("mlp.fc2.", "mlp.proj.") - key = key.replace("mlp.gate_up_proj.", "mlp.fc.") - key = key.replace("mlp.up_proj.", "mlp.fc." if config.architecture - == 'Phi3SmallForCausalLM' else "mlp.gate.") #128k - key = key.replace("mlp.down_proj.", "mlp.proj.") #128k - key = key.replace("mlp.gate_proj.", "mlp.fc.") #128k - key = key.replace("o_proj.", "dense.") #128k - # LoRA - key = key.replace("base_layer.", "") - - #MoE - key = key.replace("block_sparse_moe.gate", "mlp.router") - key = key.replace("block_sparse_moe.experts.0.w3", "mlp.fc") - key = key.replace("block_sparse_moe.experts.0.w2", "mlp.proj") - - #Layer norm - key = key.replace("post_attention_layernorm.", - "post_layernorm.") #128k - - # Embedding - key = key.replace("model.embed_tokens.weight", - "transformer.vocab_embedding.weight") - # Final Layer norm - key = key.replace("model.final_layernorm.", "transformer.ln_f.") - key = key.replace("model.norm.", "transformer.ln_f.") #128k - - if "mlp.gate_up_proj." in orig_key: #4k - original_weights = value.contiguous().clone() - half_split = original_weights.shape[0] // 2 - first_half, second_half = original_weights[: - half_split, :], original_weights[ - half_split:, :] - # Swap the halves - value = torch.cat((second_half, first_half), dim=0) - - if config.architecture == "PhiMoEForCausalLM": - num_experts = config.moe["num_experts"] - mlp_hidden_size = config.intermediate_size - num_hidden = config.hidden_size - rank_experts = list(range(num_experts)) - if config.mapping.has_moe_ep(): - rank_experts = config.mapping.ep_experts(num_experts) - - def get_moe_weight(key, suffix): - param = [] - for expert in rank_experts: - name = key.replace(f"0.{suffix}", f"{expert}.{suffix}") - fc_value = hf_state_dict[name] - param.append(fc_value) - w = torch.stack(param) - return w.reshape(-1, mlp_hidden_size, num_hidden) - - if ".0.w3" in orig_key: - w3 = get_moe_weight(orig_key, 'w3') - w1 = get_moe_weight(orig_key.replace("w3", "w1"), 'w1') - value = torch.concat([w3, w1], dim=-2) - elif ".0.w2" in orig_key: - w2 = get_moe_weight(orig_key, 'w2') - value = w2.reshape(-1, num_hidden, mlp_hidden_size) - elif any([k in orig_key for k in ["w1", "w2", "w3"]]): - continue - - if "q_proj" in key: #128k - q_param = value - k_param = hf_state_dict[orig_key.replace("q_proj", "k_proj")] - v_param = hf_state_dict[orig_key.replace("q_proj", "v_proj")] - value = torch.cat([q_param, k_param, v_param], dim=0) - key = key.replace("q_proj", "qkv") - elif "k_proj" in key or "v_proj" in key: - continue - - dtype = torch.float if "router" in key else torch_dtype - weights[key] = value.to(dtype).cpu() - - #This is for InternVL-4B - if config.architecture == 'Phi3ForCausalLM': - keys_to_rename = [ - key for key in weights.keys() if 'language_model.' in key - ] - keys_to_delete = [ - key for key in weights.keys() if 'vision_model.' in key - ] - for key in keys_to_rename: - keys_rename = key.replace('language_model.', '') - weights[keys_rename] = weights[key] - del weights[key] - for key in keys_to_delete: - del weights[key] - - if config.tie_word_embeddings or config.architecture == 'Phi3SmallForCausalLM': - weights['lm_head.weight'] = weights[ - 'transformer.vocab_embedding.weight'].clone() - - if config.architecture == 'Phi3SmallForCausalLM': - # Transform QKV weights from custom Phi3Small format to TRT-LLM format - for key, value in weights.items(): - if "qkv." in key: - weights[key] = shuffle_qkv_weights(weights[key], config) - - if config.architecture in [ - 'Phi3SmallForCausalLM', "PhiMoEForCausalLM", "Phi3ForCausalLM" - ] and config.mapping.has_tp(): - weights = split_weights_tp(config, weights, torch_dtype) - - return weights diff --git a/tensorrt_llm/models/phi3/model.py b/tensorrt_llm/models/phi3/model.py deleted file mode 100644 index b8b1a3283848..000000000000 --- a/tensorrt_llm/models/phi3/model.py +++ /dev/null @@ -1,316 +0,0 @@ -from typing import Optional, Union - -import numpy as np -from transformers import AutoModelForCausalLM - -from ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor -from ...layers import (MLP, MOE, Attention, AttentionMaskType, - BlockSparseAttnParams, ColumnLinear, Embedding, - LayerNorm, MoeConfig, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) -from .config import Phi3Config -from .convert import load_weights_from_hf_model - - -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" - self.moe_variant = config.architecture == "PhiMoEForCausalLM" - 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) - - if self.small_variant or self.moe_variant: - self.input_layernorm = LayerNorm( - normalized_shape=config.hidden_size, - dtype=config.dtype, - eps=config.norm_epsilon) - self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype, - eps=config.norm_epsilon) - 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 or self.moe_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 or self.moe_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, - rotary_embedding_percentage=config.rotary_pct, - block_sparse_params=block_sparse_attn_params) - - ClsMLP = MLP - mlp_kwargs = {} - if hasattr(config, "moe"): - ClsMLP = MOE - moe_config = MoeConfig() - for key, value in config.moe.items(): - setattr(moe_config, key, value) - mlp_kwargs = { - "moe_config": moe_config, - "mapping": config.mapping, - } - - self.mlp = ClsMLP(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, - **mlp_kwargs) - - 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) - if self.small_variant: - feed_forward_hidden_states = self.mlp( - post_attention_output, - gegelu_limit=self.gegelu_limit, - lora_layer_params=lora_layer_params) - else: - feed_forward_hidden_states = self.mlp( - post_attention_output, 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" - self.moe_variant = config.architecture == "PhiMoEForCausalLM" - if self.small_variant or self.moe_variant: - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - if self.small_variant: - 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): - config_class = Phi3Config - - def __init__(self, config: PretrainedConfig): - transformer = Phi3Model(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - self.moe_variant = config.architecture == "PhiMoEForCausalLM" - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=self.moe_variant, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - if self.moe_variant: - self.trtllm_modules_to_hf_modules = { - "attn_q": "q_proj", - "attn_k": "k_proj", - "attn_v": "v_proj", - "attn_dense": "o_proj", - "moe_h_to_4h": "w1", - "moe_4h_to_h": "w2", - "moe_gate": "w3", - "moe_router": "gate", - } - else: - 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 from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - config = Phi3Config.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype="auto", trust_remote_code=trust_remote_code) - - assert isinstance(hf_model, transformers.PreTrainedModel) - - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/phi3/split_weights.py b/tensorrt_llm/models/phi3/split_weights.py deleted file mode 100644 index bca33ba55116..000000000000 --- a/tensorrt_llm/models/phi3/split_weights.py +++ /dev/null @@ -1,218 +0,0 @@ -# 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 torch - -from tensorrt_llm.models.convert_utils import (get_weight, get_weight_and_bias, - split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) - -from ..._utils import pad_vocab_size - - -def shuffle_qkv_weights(weights, config): - # Input weights are organized as - # (q00, q01, ... q0m, k0, v0), (q10, q11, ... q1m, k1, v1), ... (qn0, qn1, ... qnm, kn, vn) - # where n = num_kv_heads, m = num_attention_heads // num_kv_heads (i.e. #q_heads per kv_head) - # - # Output weights will be organized as - # (q00, q01, ..., qnm), (k0, k1, .., kn), (v0, v1, .., vn) - - num_heads = config.num_attention_heads - num_kv_heads = config.num_key_value_heads - num_q_per_kv = num_heads // num_kv_heads - - hidden_size = config.hidden_size - head_dim = hidden_size // num_heads - - input_shape = weights.shape - if weights.dim() < 2: - weights = weights.unsqueeze(1) - - weights = weights.reshape(num_kv_heads, (num_q_per_kv + 2), head_dim, - weights.shape[-1]) - q = weights[:, :-2, :, :] - k = weights[:, -2, :, :] - v = weights[:, -1, :, :] - - # num_heads x head_dim x hidden_size - q = q.reshape(-1, q.shape[2], q.shape[3]) - - # num_heads + (2 * num_kv_heads) x head_dim x hidden_size - weights = torch.cat([q, k, v], dim=0) - weights = weights.reshape(-1, weights.shape[2]) - - weights = weights.squeeze() - assert input_shape == weights.shape - - return weights - - -def split_embedding( - param: torch.Tensor, - tp_size: int, - tp_rank: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_size, tp_rank, dim=sharding_dim) - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + '.weight'] = processed_torch_weights - results[prefix + '.per_channel_scale'] = torch_weight_scales - else: - results[prefix + '.weight'] = weight.contiguous() - - if bias is not None: - results[prefix + '.bias'] = bias - - return results - - -def split_weights_tp(config, weights, dtype): - num_heads = config.num_attention_heads - num_kv_heads = config.num_key_value_heads - hidden_size = config.hidden_size - moe_variant = config.architecture == "PhiMoEForCausalLM" - - mha_mode = num_heads == num_kv_heads - tp_size = config.mapping.tp_size - rank = config.mapping.tp_rank - moe_tp_size = config.mapping.moe_tp_size - moe_tp_rank = config.mapping.moe_tp_rank - use_weight_only = config.quant_mode.is_weight_only() - plugin_weight_only_quant_type = None - if use_weight_only and config.quant_mode.is_int8_weight_only() == 'int8': - plugin_weight_only_quant_type = torch.int8 - elif use_weight_only and config.quant_mode.is_int4_weight_only() == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - - def get_quant_weight(weight, prefix, bias): - return get_tllm_linear_weight(weight, prefix, bias, use_weight_only, - plugin_weight_only_quant_type) - - for layer_id in range(config.num_hidden_layers): - layer_prefix = f"transformer.layers.{layer_id}." - - prefix = layer_prefix + 'attention.qkv' - qkv_weight, qkv_bias = get_weight_and_bias(weights, prefix, dtype) - - if not mha_mode: - num_q_per_kv = num_heads // num_kv_heads - - qkv_weight = qkv_weight.reshape(num_q_per_kv + 2, -1, hidden_size) - q = qkv_weight[:num_q_per_kv, :, :].reshape(-1, hidden_size) - k = qkv_weight[num_q_per_kv:num_q_per_kv + 1, :, :].reshape( - -1, hidden_size) - v = qkv_weight[num_q_per_kv + 1:num_q_per_kv + 2, :, :].reshape( - -1, hidden_size) - split_weight = torch.cat( - [split(x, tp_size, rank) for x in [q, k, v]], dim=0) - if qkv_bias is not None: - qkv_bias = qkv_bias.reshape(num_q_per_kv + 2, -1) - q = qkv_bias[:num_q_per_kv, :].reshape(-1) - k = qkv_bias[num_q_per_kv:num_q_per_kv + 1, :].reshape(-1) - v = qkv_bias[num_q_per_kv + 1:num_q_per_kv + 2, :].reshape(-1) - split_bias = torch.cat( - [split(x, tp_size, rank) for x in [q, k, v]], dim=0) - else: - split_bias = None - else: - split_weight = split_qkv_tp(qkv_weight, num_heads, hidden_size, - tp_size, rank) - if qkv_bias is not None: - split_bias = split_qkv_bias_tp(qkv_bias, num_heads, hidden_size, - tp_size, rank) - else: - split_bias = None - weights.update(get_quant_weight(split_weight, prefix, split_bias)) - - prefix = layer_prefix + 'attention.dense' - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - weights, prefix, dtype) - split_v = split_matrix_tp(attn_dense_weight, tp_size, rank, dim=1) - weights.update(get_quant_weight(split_v, prefix, attn_dense_bias)) - - prefix = layer_prefix + 'mlp.fc' - if not moe_variant: - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - weights, prefix, dtype) - split_v = split_matrix_tp(mlp_fc_weight, tp_size, rank, dim=0) - if mlp_fc_bias is not None: - bias = split_matrix_tp(mlp_fc_bias, tp_size, rank, dim=0) - else: - bias = None - weights.update(get_quant_weight(split_v, prefix, bias)) - else: - mlp_fc_weight = get_weight(weights, prefix, dtype) - w3 = split_matrix_tp(mlp_fc_weight, 2, 0, dim=1) - split_w3 = split_matrix_tp(w3, moe_tp_size, moe_tp_rank, dim=1) - w1 = split_matrix_tp(mlp_fc_weight, 2, 1, dim=1) - split_w1 = split_matrix_tp(w1, moe_tp_size, moe_tp_rank, dim=1) - split_v = torch.concat([split_w3, split_w1], dim=-2) - weights.update(get_quant_weight(split_v, prefix, None)) - - prefix = layer_prefix + 'mlp.proj' - if not moe_variant: - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - weights, prefix, dtype) - split_v = split_matrix_tp(mlp_proj_weight, tp_size, rank, dim=1) - weights.update(get_quant_weight(split_v, prefix, mlp_proj_bias)) - else: - mlp_proj_weight = get_weight(weights, prefix, dtype) - split_v = split_matrix_tp(mlp_proj_weight, - moe_tp_size, - moe_tp_rank, - dim=2) - weights.update(get_quant_weight(split_v, prefix, None)) - - weights['transformer.vocab_embedding.weight'] = split_embedding( - weights['transformer.vocab_embedding.weight'], tp_size, rank) - weights['lm_head.weight'] = split_matrix_tp(weights['lm_head.weight'], - tp_size, - rank, - dim=0) - if moe_variant: - weights['lm_head.bias'] = split_matrix_tp(weights['lm_head.bias'], - tp_size, - rank, - dim=0) - - return weights diff --git a/tensorrt_llm/models/qwen/__init__.py b/tensorrt_llm/models/qwen/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/qwen/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/qwen/config.py b/tensorrt_llm/models/qwen/config.py deleted file mode 100644 index 0f1bd34606be..000000000000 --- a/tensorrt_llm/models/qwen/config.py +++ /dev/null @@ -1,211 +0,0 @@ -# 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 ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class QWenConfig(PretrainedConfig): - - def __init__(self, - *, - mlp_bias: bool = False, - attn_bias: bool = True, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - disable_weight_only_quant_plugin: bool = False, - use_logn_attn: bool = False, - moe: Optional[Union[MoeConfig, dict]] = None, - num_labels: int = 1, - mlp_only_layers: Optional[list] = None, - decoder_sparse_step: int = 1, - **kwargs): - self.mlp_bias = mlp_bias - self.attn_bias = attn_bias - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.disable_weight_only_quant_plugin = disable_weight_only_quant_plugin - self.num_labels = num_labels - self.use_logn_attn = use_logn_attn - self.mlp_only_layers = mlp_only_layers or [] - self.decoder_sparse_step = decoder_sparse_step - if moe is None: - # Legacy MOE config fields - moe = MoeConfig(num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.NONE)) - elif isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in QWenConfig - 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[ - 'disable_weight_only_quant_plugin'] = self.disable_weight_only_quant_plugin - output['use_logn_attn'] = self.use_logn_attn - output['mlp_only_layers'] = self.mlp_only_layers - output['decoder_sparse_step'] = self.decoder_sparse_step - output['moe'] = self.moe.to_dict() - return output - - @classmethod - def from_hugging_face(cls, - hf_config_or_dir: Union[ - str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs) -> "QWenConfig": - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - if hasattr(hf_config, 'llm_config'): - hf_config = hf_config.llm_config - - qwen_type = hf_config.model_type - # lmms llava onevision qwen - if qwen_type == 'llava': - qwen_type = 'qwen2' - if hf_config.architectures and hf_config.architectures[ - 0] == 'LlavaQwenForCausalLM': - hf_config.architectures[0] = 'Qwen2ForCausalLM' - # hf llava onevision qwen - if qwen_type == 'llava_onevision': - hf_config = hf_config.text_config - qwen_type = f'{hf_config.model_type}_llava_onevision' - # Qwen2-Audio - if qwen_type == 'qwen2_audio': - hf_config = hf_config.text_config - hf_config.architectures = ['Qwen2ForCausalLM'] - - valid_types = ('qwen', 'qwen2', 'qwen2_moe', 'qwen2_llava_onevision', - 'qwen2_vl', 'qwen2_audio', 'qwen3', 'qwen3_moe') - assert qwen_type in valid_types, f"Unsupported Qwen type: {qwen_type}, only {valid_types} are acceptable." - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - head_dim = getattr( - hf_config, "head_dim", - hf_config.hidden_size // hf_config.num_attention_heads) - head_size = getattr(hf_config, "kv_channels", head_dim) - hidden_act = getattr(hf_config, "hidden_act", "silu") - if qwen_type in ("qwen2_moe", "qwen3_moe"): - hidden_act = "swiglu" - - # Qwen3 models have no attention bias, while legacy models have bias - if qwen_type in ('qwen3', 'qwen3_moe'): - attn_bias = False # Qwen3 models have no attn bias - else: - attn_bias = True # Legacy Qwen models have attn bias - rotary_scaling = getattr(hf_config, "rope_scaling", None) - seq_length = getattr(hf_config, "seq_length", 8192) - use_logn_attn = getattr(hf_config, "use_logn_attn", False) - disable_weight_only_quant_plugin = kwargs.pop( - 'disable_weight_only_quant_plugin', False) - if qwen_type == "qwen": - rms_norm_eps = hf_config.layer_norm_epsilon - rotary_base = getattr(hf_config, "rotary_emb_base", 10000.0) - else: - rms_norm_eps = hf_config.rms_norm_eps - rotary_base = get_hf_rope_theta(hf_config, 100000.0) - - num_labels = 1 - if hf_config.architectures[0] == "Qwen2ForSequenceClassification": - num_labels = hf_config.num_labels - - moe_num_experts = getattr(hf_config, "num_experts", 0) - moe_top_k = getattr(hf_config, "num_experts_per_tok", 0) - moe_intermediate_size = getattr(hf_config, "moe_intermediate_size", 0) - moe_shared_expert_intermediate_size = getattr( - hf_config, "shared_expert_intermediate_size", 0) - moe_normalization_mode = MoeConfig.ExpertScaleNormalizationMode.NONE - - # Add support for mlp_only_layers and decoder_sparse_step (Qwen3 MoE) - mlp_only_layers = getattr(hf_config, "mlp_only_layers", []) - decoder_sparse_step = getattr(hf_config, "decoder_sparse_step", 1) - - moe_config = MoeConfig(num_experts=moe_num_experts, - top_k=moe_top_k, - normalization_mode=moe_normalization_mode) - moe_config.validate() - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - tie_word_embeddings = getattr(hf_config, 'tie_word_embeddings', False) - - if qwen_type == 'qwen2_vl': - pe_type = 'mrope' - rotary_embedding_percentage = getattr(hf_config, 'rotary_pct', 1.0) - rotary_embedding_dim = getattr( - hf_config, 'rotary_dim', - int(hf_config.hidden_size / hf_config.num_attention_heads * - rotary_embedding_percentage)) - rotary_scaling['type'] = 'mrope' - else: - pe_type = 'rope_gpt_neox' - rotary_embedding_dim = None - - 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=num_key_value_heads, - head_size=head_size, - vocab_size=hf_config.vocab_size, - position_embedding_type=pe_type, - max_position_embeddings=hf_config.max_position_embeddings, - rotary_embedding_dim=rotary_embedding_dim, - hidden_act=hidden_act, - norm_epsilon=rms_norm_eps, - attn_bias=attn_bias, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - disable_weight_only_quant_plugin=disable_weight_only_quant_plugin, - seq_length=seq_length, - use_logn_attn=use_logn_attn, - qwen_type=qwen_type, - moe_intermediate_size=moe_intermediate_size, - moe_shared_expert_intermediate_size= - moe_shared_expert_intermediate_size, - mlp_only_layers=mlp_only_layers, - decoder_sparse_step=decoder_sparse_step, - moe=moe_config, - mapping=mapping, - quantization=quant_config, - num_labels=num_labels, - tie_word_embeddings=tie_word_embeddings, - **kwargs) diff --git a/tensorrt_llm/models/qwen/convert.py b/tensorrt_llm/models/qwen/convert.py deleted file mode 100644 index 23220538aed7..000000000000 --- a/tensorrt_llm/models/qwen/convert.py +++ /dev/null @@ -1,1284 +0,0 @@ -# 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 functools - -import copy -import functools -import json -import os -import time -from collections import defaultdict -from typing import List, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoConfig, AutoTokenizer -from transformers.pytorch_utils import Conv1D - -from ..._utils import pad_vocab_size, str_dtype_to_torch -from ...logger import logger -from ...mapping import Mapping -from ...quantization import QuantAlgo -from ..convert_utils import (dup_kv_bias, dup_kv_weight, generate_int8, - get_weight, get_weight_and_bias, - load_calib_dataset, smooth_gemm, - smooth_gemm_fc1_gate, split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) -from .config import QWenConfig -from .utils import get_qwen_key_list, make_context - - -@torch.no_grad() -def smooth_qwen_model(model, scales, alpha, qwen_qkv_para, qwen_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not module._get_name() == "QWenBlock": - continue - # qkv_proj - layer_name = name + ".attn.c_attn" - smoother = smooth_gemm(module.attn.c_attn.weight, - scales[layer_name]["x"], module.ln_1.weight, - None, alpha) - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.c_attn.weight.abs().max(dim=1)[0] - - # see transpose_weights function - qwen_qkv_para[layer_name] = module.attn.c_attn.weight.transpose( - 0, 1).contiguous() - - # ================================================================= - layer_name = name + ".attn.c_proj" - smoother = smooth_gemm( - module.attn.c_proj.weight, - scales[layer_name]["x"], - None, - None, - alpha=alpha, - ) - qwen_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.c_proj.weight.abs().max(dim=1)[0] - # ================================================================== - fc1_layer_name = name + ".mlp.w1" - gate_layer_name = name + ".mlp.w2" - - smoother = smooth_gemm_fc1_gate(module.mlp.w1.weight, - module.mlp.w2.weight, - scales[fc1_layer_name]["x"], - module.ln_2.weight, None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.w1.weight.abs().max(dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.w2.weight.abs().max(dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.c_proj" - smoother = smooth_gemm(module.mlp.c_proj.weight, - scales[layer_name]["x"], None, None, alpha) - qwen_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.c_proj.weight.abs().max(dim=1)[0] - - -@torch.no_grad() -def smooth_qwen2_model(model, scales, alpha, qwen_qkv_para, qwen_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer - from transformers.models.qwen2_vl.modeling_qwen2_vl import \ - Qwen2VLDecoderLayer - if not isinstance(module, Qwen2DecoderLayer) and not isinstance( - module, Qwen2VLDecoderLayer): - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - qwen_qkv_para[layer_name_qkv] = weight.transpose(0, 1).contiguous() - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - qwen_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - qwen_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - scales_keys_to_rename = [ - key for key in scales.keys() if 'language_model.' in key - ] - - qwen_qkv_para_keys_to_rename = [ - key for key in qwen_qkv_para.keys() if 'language_model.' in key - ] - - qwen_smoother_keys_to_rename = [ - key for key in qwen_smoother.keys() if 'language_model.' in key - ] - - for key in scales_keys_to_rename: - scales[key.replace('language_model.', '')] = scales[key] - del scales[key] - - for key in qwen_qkv_para_keys_to_rename: - qwen_qkv_para[key.replace('language_model.', '')] = qwen_qkv_para[key] - del qwen_qkv_para[key] - - for key in qwen_smoother_keys_to_rename: - qwen_smoother[key.replace('language_model.', '')] = qwen_smoother[key] - del qwen_smoother[key] - - -@torch.no_grad() -def capture_activation_range(model, - qwen_type, - tokenizer, - dataset, - system_prompt, - chat_format, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - if qwen_type == 'qwen': - tokenizer.pad_token_id = tokenizer.im_end_id - else: - tokenizer.pad_token_id = tokenizer.eos_token_id - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - line = dataset[i] - line = line + ' TL;DR: ' - line = line.strip() - line = line.replace(" n't", "n't") - if qwen_type == 'qwen': - _, input_id_list = make_context(tokenizer=tokenizer, - query=line, - history=[], - system=system_prompt, - chat_format=chat_format, - max_input_length=seq_len) - line_encoded = torch.from_numpy( - np.array(input_id_list, - dtype=np.int32)).type(torch.int32).unsqueeze(0) - line_encoded = line_encoded.to(device) - else: - line_encoded = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(line_encoded) - for h in hooks: - h.remove() - return act_scales - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - dtype='float32', - use_gemm_woq_plugin=True, - postfix='weight', - quant_scale_name=None): - results = {} - if use_weight_only: - if weight.dim() > 2: - v = weight.transpose(1, 2).contiguous().clone() - else: - v = weight.t().contiguous().clone() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.cpu(), plugin_weight_only_quant_type) - if not use_gemm_woq_plugin: - results[prefix + postfix] = v.to(dtype) - else: - results[prefix + postfix] = processed_torch_weights - if quant_scale_name is not None: - results[quant_scale_name] = torch_weight_scales - else: - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + postfix] = weight.clone() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.split(q, q.shape[-1] // tp_size, dim=-1) - k_split = torch.split(k, k.shape[-1] // tp_size, dim=-1) - v_split = torch.split(v, v.shape[-1] // tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), dim=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.split( - vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + - 'per_channel_scale'] = cur_per_channel_value.reshape(col_shape) - else: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = torch.split(smoother_value, - smoother_value.shape[-1] // - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def load_hf_qwen(model_dir: str, load_model_on_cpu: bool = False): - config_path = os.path.join(model_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - if config['architectures'] == ['Qwen2ForSequenceClassification']: - from transformers import Qwen2ForSequenceClassification as model_cls - elif config['architectures'] == ['Qwen2VLForConditionalGeneration']: - from transformers import Qwen2VLForConditionalGeneration as model_cls - else: - from transformers import AutoModelForCausalLM as model_cls - - model = model_cls.from_pretrained( - model_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype='auto', - trust_remote_code=True) - return model - - -def convert_hf_qwen(hf_model, - qwen_type, - mapping: Mapping, - vocab_size=32000, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - use_gemm_woq_plugin=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[], - moe_config=None): - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - - dtype = getattr(torch, dtype) - hf_config = hf_model.config - if hasattr(hf_config, 'llm_config'): - hf_config = hf_config.llm_config - - #This is for InternVL2 - 1B - keys_to_rename = [ - key for key in model_params.keys() if 'language_model.' in key - ] - keys_to_delete = [ - key for key in model_params.keys() if 'vision_model.' in key - ] - for key in keys_to_rename: - keys_rename = key.replace('language_model.', '') - model_params[keys_rename] = model_params[key] - del model_params[key] - for key in keys_to_delete: - del model_params[key] - - num_attention_heads = hf_config.num_attention_heads - hidden_size = hf_config.hidden_size - head_size = hidden_size // num_attention_heads - if qwen_type == 'qwen': - intermediate_size = hf_config.intermediate_size // 2 # Qwen version 1 has actual intermediate_size one half of what's in hf_config - else: - intermediate_size = hf_config.intermediate_size - num_key_value_heads = hf_config.num_key_value_heads if hasattr( - hf_config, "num_key_value_heads") else num_attention_heads - mha_mode = (num_key_value_heads == num_attention_heads) - layers_range = mapping.pp_layers(hf_config.num_hidden_layers) - - layer_prefix = "transformer.h." if qwen_type == 'qwen' else "model.layers." - key_list = get_qwen_key_list(qwen_type) - - for l in layers_range: - prefix = layer_prefix + f'{l}.' - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - if qwen_type == 'qwen': - qkv_weight, qkv_bias = get_weight_and_bias(model_params, - prefix + key_list[0], - dtype) - qkv_w = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - qkv_b = split_qkv_bias_tp(qkv_bias, num_attention_heads, - hidden_size, tensor_parallel, - mapping.tp_rank) - else: - q_weight, q_bias = get_weight_and_bias( - model_params, prefix + key_list[0] + 'q_proj', dtype) - k_weight, k_bias = get_weight_and_bias( - model_params, prefix + key_list[0] + 'k_proj', dtype) - v_weight, v_bias = get_weight_and_bias( - model_params, prefix + key_list[0] + 'v_proj', dtype) - if not mha_mode: - if num_key_value_heads < tensor_parallel: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, num_key_value_heads, - tensor_parallel) - v_weight = dup_kv_weight(v_weight, num_key_value_heads, - tensor_parallel) - k_bias = dup_kv_bias(k_bias, num_key_value_heads, - tensor_parallel) - v_bias = dup_kv_bias(v_bias, num_key_value_heads, - tensor_parallel) - assert (k_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - assert (v_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - - if k_bias is not None and v_bias is not None: - assert (k_bias.shape[0] % - (mapping.tp_size * head_size)) == 0 - assert (v_bias.shape[0] % - (mapping.tp_size * head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - qkv_w = torch.concat((wq, wk, wv)) - - if q_bias is not None and k_bias is not None and v_bias is not None: - bq = split(q_bias, mapping.tp_size, mapping.tp_rank) - bk = split(k_bias, mapping.tp_size, mapping.tp_rank) - bv = split(v_bias, mapping.tp_size, mapping.tp_rank) - qkv_b = torch.concat((bq, bk, bv)) - else: - qkv_b = None - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=0) - - qkv_w = split_qkv_tp(qkv_weight, num_attention_heads, - hidden_size, tensor_parallel, - mapping.tp_rank) - qkv_b = split_qkv_bias_tp(qkv_bias, num_attention_heads, - hidden_size, tensor_parallel, - mapping.tp_rank) - - if use_smooth_quant: - qkv_proj_key = key_list[ - 0] if qwen_type == 'qwen' else 'self_attn.qkv_proj' - qkv_weight = qkv_para[prefix + qkv_proj_key] - qkv_out_dim = qkv_weight.shape[1] - - if not mha_mode: - local_dim = qkv_weight.shape[0] - kv_hidden_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(local_dim, - local_dim + 2 * kv_hidden_size) - else: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + qkv_proj_key), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - bias=qkv_b, - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - weights.update( - get_tllm_linear_weight(qkv_w, tllm_prex + 'attention.qkv.', - qkv_b, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - if int8_kv_cache: - if qwen_type == 'qwen': - qkv_y = act_range.get(prefix + key_list[0])["y"] - else: - qkv_y = torch.cat([ - act_range.get(prefix + key_list[0] + 'q_proj')["y"], - act_range.get(prefix + key_list[0] + 'k_proj')["y"], - act_range.get(prefix + key_list[0] + 'v_proj')["y"] - ], - dim=0) - - int8_kv_scales = qkv_y.max() / 127. - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - weights.update(kv_cache_weights) - - attn_dense_weight = get_weight(model_params, prefix + key_list[1], - dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8(attn_dense_weight, - act_range.get(prefix + key_list[1])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + key_list[1])], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - # Qwen3: Add q_norm and k_norm weight conversion - if qwen_type in ('qwen3', 'qwen3_moe'): - # Process q_norm.weight - q_norm_weight = get_weight(model_params, - prefix + key_list[0] + 'q_norm', dtype) - weights.update( - get_tllm_linear_weight( - q_norm_weight, - tllm_prex + 'attention.q_layernorm.', - None, - False, # LayerNorm should not be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - # Process k_norm.weight - k_norm_weight = get_weight(model_params, - prefix + key_list[0] + 'k_norm', dtype) - weights.update( - get_tllm_linear_weight( - k_norm_weight, - tllm_prex + 'attention.k_layernorm.', - None, - False, # LayerNorm should not be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - if moe_config and moe_config.has_moe(): - if qwen_type == "qwen2_moe": - # shared_expert for qwen2_moe - shared_expert_up_proj = model_params[ - f'model.layers.{l}.mlp.shared_expert.up_proj.weight'] - shared_expert_down_proj = model_params[ - f'model.layers.{l}.mlp.shared_expert.down_proj.weight'] - shared_expert_gate = model_params[ - f'model.layers.{l}.mlp.shared_expert.gate_proj.weight'] - shared_expert_up_proj = split(shared_expert_up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_expert_down_proj = split(shared_expert_down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_expert_gate = split(shared_expert_gate, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_expert_gate_up_proj = torch.concat( - [shared_expert_up_proj, shared_expert_gate], - dim=-2).to(dtype) - - ## mlp.shared_expert.gate_up_proj.weight - weights.update( - get_tllm_linear_weight(shared_expert_gate_up_proj, - tllm_prex + 'mlp.shared_expert.fc.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - ## mlp.shared_expert.down_proj.weight - weights.update( - get_tllm_linear_weight( - shared_expert_down_proj.to(dtype), - tllm_prex + 'mlp.shared_expert.proj.', None, - use_weight_only, plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - moe_shared_expert_gate_weights = get_weight( - model_params, prefix + 'mlp.shared_expert_gate', dtype) - weights.update( - get_tllm_linear_weight( - moe_shared_expert_gate_weights, - tllm_prex + 'mlp.shared_expert_gate.', - None, - False, # Router should never be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - ## fine-grained experts - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - for suffix in ["gate_proj", "down_proj", "up_proj"]: - model_params[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.mlp.experts.{expert}.{suffix}.weight'].detach() - for expert in rank_experts]) - w3 = model_params[f'model.layers.{l}.mlp.experts.up_proj.weight'] - w2 = model_params[f'model.layers.{l}.mlp.experts.down_proj.weight'] - w1 = model_params[f'model.layers.{l}.mlp.experts.gate_proj.weight'] - if mapping.has_moe_tp(): - w3 = split(w3, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - w2 = split(w2, mapping.moe_tp_size, mapping.moe_tp_rank, dim=2) - w1 = split(w1, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - - moe_experts_w3w1_weights = torch.concat([w3, w1], dim=-2).to(dtype) - - ## mlp.experts.w2.weight - weights.update( - get_tllm_linear_weight(w2.to(dtype), tllm_prex + 'mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - ## mlp.experts.w3w1.weight - weights.update( - get_tllm_linear_weight(moe_experts_w3w1_weights, - tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - moe_experts_gate_weights = get_weight(model_params, - prefix + 'mlp.gate', - torch.float32) - weights.update( - get_tllm_linear_weight( - moe_experts_gate_weights, - tllm_prex + 'mlp.router.', - None, - False, # Router should never be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - else: - mlp_gate_weight = get_weight(model_params, prefix + key_list[2], - dtype) - split_v = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - if use_smooth_quant: - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8( - mlp_gate_weight, act_range.get(prefix + key_list[2])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.gate.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_fc_weight = get_weight(model_params, prefix + key_list[3], - dtype) - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() #verified - int8_weights = generate_int8( - mlp_fc_weight, act_range.get(prefix + key_list[3])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_proj_weight = get_weight(model_params, prefix + key_list[4], - dtype) - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + key_list[4])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + key_list[4]], - smoother_shape=[ - 1, intermediate_size // tensor_parallel - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + key_list[5], dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, prefix + key_list[6], dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - v = get_weight(model_params, key_list[7], dtype) - - if mapping.is_last_pp_rank(): - if hf_config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - lm_head_weights = v.clone() - else: - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - if mapping.is_last_pp_rank(): - ln_f_w = get_weight(model_params, key_list[8], dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - if hasattr(hf_model, 'score'): - score = get_weight(model_params, 'score', dtype) - weights['lm_head.weight'] = score - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: QWenConfig, - calib_dataset='cnn_dailymail'): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == "INT8" - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - hf_config = AutoConfig.from_pretrained(hf_model_dir, trust_remote_code=True) - if hf_config.architectures == ['Qwen2VLForConditionalGeneration']: - from transformers import Qwen2VLForConditionalGeneration as model_cls - else: - from transformers import AutoModelForCausalLM as model_cls - hf_model = model_cls.from_pretrained( - hf_model_dir, - device_map='auto', - dtype='auto' if not use_smooth_quant else torch.float16, - trust_remote_code=True).half() - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained(hf_model_dir, - trust_remote_code=True, - use_fast=False, - padding_side='left') - dataset = load_calib_dataset(calib_dataset) - - system_prompt = "You are a useful assistant, please directly output the corresponding summary according to the article entered by the user." - gen_config_path = os.path.join(hf_model_dir, 'generation_config.json') - with open(gen_config_path, 'r') as f: - gen_config = json.load(f) - chat_format = getattr(gen_config, 'chat_format', 'chatml') - act_range = capture_activation_range(hf_model, config.qwen_type, tokenizer, - dataset, system_prompt, chat_format) - qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - smoother = {} - if use_smooth_quant: - if config.qwen_type == 'qwen': - smooth_qwen_model(hf_model, act_range, quant_config.smoothquant_val, - qkv_para, smoother) - else: - smooth_qwen2_model(hf_model, act_range, - quant_config.smoothquant_val, qkv_para, smoother) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model(hf_model, - config=config, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -def load_weights_from_hf_model(hf_model, - config: QWenConfig, - act_range: Optional[dict] = None, - qkv_para: Optional[dict] = None, - smoother: Optional[dict] = None): - #TODO: simplify the parameters here - - assert hf_model is not None - plugin_weight_only_quant_type = None # the value does not matter when use_weight_only is False - quant_algo = config.quantization.quant_algo - 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 = (not config.disable_weight_only_quant_plugin) - - mapping = config.mapping - moe_config = config.moe - - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - use_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - qwen_type = config.qwen_type - weights = convert_hf_qwen( - hf_model, - qwen_type, - mapping, - vocab_size=config.vocab_size, - dtype=config.dtype, - use_weight_only=use_weight_only, - use_gemm_woq_plugin=use_gemm_woq_plugin, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=config.use_parallel_embedding, - sharding_dim=config.embedding_sharding_dim, - use_smooth_quant=use_smooth_quant, - per_channel=per_channel, - per_token=per_token, - int8_kv_cache=int8_kv_cache, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother, - moe_config=moe_config) - return weights - - -def load_weights_from_hf_gptq_model(hf_model, config: QWenConfig): - logger.info("loading weights from groupwise GPTQ QWen safetensors...") - weights = {} - tik = time.time() - - qwen_type = config.qwen_type - num_hidden_layers = config.num_hidden_layers - mapping = config.mapping - dtype = config.dtype - - model_params = {k: v for k, v in hf_model.state_dict().items()} - torch.cuda.empty_cache() - valid_types = ('qwen', 'qwen2', 'qwen2_vl', 'qwen3', 'qwen3_moe') - assert qwen_type in valid_types, f"Unsupported Qwen type: {qwen_type}, only {valid_types} are supported for GPTQ." - layer_prefix = "transformer.h." if qwen_type == 'qwen' else "model.layers." - key_list = get_qwen_key_list(qwen_type) - - def torch_split(v, dim): - if v.shape[dim] % mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(mapping.tp_size)) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank] - - def unpack_int32_into_int8(w_packed): - # unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(v: List[torch.Tensor], - tllm_prex: str, - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in v - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in v - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_GPTQ_FOR_QWEN = 1 # GPTQ-for-QWEN added 1 to zeros - - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - if not USE_UINT4_INPUT: - # Correcting UINT4 values back to INT4 order - mask_negative = qzeros_unpacked_int32[qzeros_unpacked_int32 < 0] - mask_positive = qzeros_unpacked_int32[qzeros_unpacked_int32 >= 0] - qzeros_unpacked_int32 = qzeros_unpacked_int32 + 16 * mask_negative - 16 * mask_positive - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_QWEN) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - results = { - f'{tllm_prex}.weight': qweight_interleaved, - f'{tllm_prex}.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.zero': zeros_x_scales_fp16, - } - return results - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = str_dtype_to_torch(dtype) - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = model_params[key_list[7] + '.weight'] - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - - # 2. ln_f - v = model_params[key_list[8] + '.weight'] - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - - # 3. lm_head - v = model_params['lm_head.weight'] - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = torch_split(v, 0).to(torch_dtype) - - # 4. Weights inside each layer - layers_per_pipeline_stage = num_hidden_layers // mapping.pp_size - layers_range = list( - range(mapping.pp_rank * layers_per_pipeline_stage, - (mapping.pp_rank + 1) * layers_per_pipeline_stage, 1)) - suffixs = [".qweight", ".qzeros", ".scales"] - - for l in tqdm(layers_range, desc="loading weight in each layer..."): - layer_idx = l - mapping.pp_rank * layers_per_pipeline_stage - prefix = layer_prefix + str(layer_idx) + "." - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # 4.1 attention.qkv - qkv_weight_list = [] - if qwen_type == 'qwen': - for suf in suffixs: - qkv_part = model_params[prefix + key_list[0] + suf] - q_emb = qkv_part.shape[1] // 3 - model_emb = qkv_part.shape[0] - qkv_part = qkv_part.reshape(model_emb, 3, q_emb) - qkv_part = torch_split(qkv_part, 2) - qkv_part = qkv_part.reshape(model_emb, - 3 * (q_emb // mapping.tp_size)) - qkv_weight_list.append(qkv_part) - else: - for suf in suffixs: - qkv_list = [] - for comp in ["q_proj", "k_proj", "v_proj"]: - comp_part = model_params[prefix + key_list[0] + comp + suf] - comp_part = torch_split(comp_part, 1) - qkv_list.append(comp_part) - qkv_weight_list.append(torch.cat(qkv_list, dim=1)) - weights.update( - process_and_assign_weight(qkv_weight_list, - f'{tllm_prex}.attention.qkv')) - # 4.2 attention.bias - suf = ".bias" - if qwen_type == 'qwen': - qkv_bias = model_params[prefix + key_list[0] + - suf].to(torch_dtype).cpu().contiguous() - q_emb = qkv_bias.shape[0] // 3 - qkv_bias = qkv_bias.reshape(3, q_emb) - split_v = split(qkv_bias, mapping.tp_size, mapping.rank, dim=1) - qkv_bias = split_v.reshape(3 * (q_emb // mapping.tp_size)) - else: - qkv_bias_list = [] - for comp in ["q_proj", "k_proj", "v_proj"]: - comp_part = model_params[prefix + key_list[0] + comp + suf].to( - torch_dtype).cpu().contiguous() - comp_part = torch_split(comp_part, dim=0) - qkv_bias_list.append(comp_part) - qkv_bias = torch.cat(qkv_bias_list, dim=0) - weights[tllm_prex + ".attention.qkv.bias"] = qkv_bias - # 4.3 attention.dense - qkv_dense_list = [] - for suf in suffixs: - qkv_dense_part = model_params[prefix + key_list[1] + suf] - qkv_dense_list.append(qkv_dense_part) - weights.update( - process_and_assign_weight(qkv_dense_list, - f'{tllm_prex}.attention.dense', - tp_dim=0)) - # 4.4 mlp.gate - mlp_gate_list = [] - for suf in suffixs: - mlp_gate_part = model_params[prefix + key_list[2] + suf] - mlp_gate_list.append(mlp_gate_part) - weights.update( - process_and_assign_weight(mlp_gate_list, - f'{tllm_prex}.mlp.gate', - tp_dim=1)) - # 4.5 mlp.fc - mlp_fc_list = [] - for suf in suffixs: - mlp_fc_part = model_params[prefix + key_list[3] + suf] - mlp_fc_list.append(mlp_fc_part) - weights.update( - process_and_assign_weight(mlp_fc_list, - f'{tllm_prex}.mlp.fc', - tp_dim=1)) - # 4.6 mlp.proj - mlp_proj_list = [] - for suf in suffixs: - mlp_proj_part = model_params[prefix + key_list[4] + suf] - mlp_proj_list.append(mlp_proj_part) - weights.update( - process_and_assign_weight(mlp_proj_list, - f'{tllm_prex}.mlp.proj', - tp_dim=0)) - # 4.7 input_layernorm - v = model_params[prefix + key_list[5] + '.weight'] - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - # 4.8 post_layernorm - v = model_params[prefix + key_list[6] + '.weight'] - weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime("%H:%M:%S", time.gmtime(tok - tik)) - logger.info(f"weights loaded. total time: {t}") - - return weights diff --git a/tensorrt_llm/models/qwen/model.py b/tensorrt_llm/models/qwen/model.py deleted file mode 100644 index 028d9e4e3e6f..000000000000 --- a/tensorrt_llm/models/qwen/model.py +++ /dev/null @@ -1,545 +0,0 @@ -# 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 copy -import os -from typing import Optional, Union - -import torch -from tqdm import tqdm - -from ..._utils import pad_vocab_size -from ...functional import LayerNormType, Tensor, recv, send -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, GatedMLP, RmsNorm, SharedMoE) -from ...layers.moe import MOEWeightWrapper -from ...logger import logger -from ...lora_helper import (LoraConfig, - get_default_trtllm_modules_to_hf_modules, use_lora) -from ...mapping import Mapping -from ...module import Module -from ...quantization import QuantAlgo -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import QWenConfig -from .convert import (load_hf_qwen, load_weights_from_hf_gptq_model, - load_weights_from_hf_model) - - -class QWenDecoderLayer(Module): - - def __init__(self, config: QWenConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - dtype = config.dtype - self.tp_group = config.mapping.tp_group - self.tp_size = config.mapping.tp_size - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - # Qwen3: Enable qk_layernorm for Q/K normalization (similar to Gemma3) - qk_layernorm = config.qwen_type in ('qwen3', 'qwen3_moe') - - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_seqlen_for_logn_scaling=config.seq_length, - max_position_embeddings=config.max_position_embeddings, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_rank=config.mapping.tp_rank, - tp_group=self.tp_group, - tp_size=self.tp_size, - cp_rank=config.mapping.cp_rank, - cp_size=config.mapping.cp_size, - cp_group=config.mapping.cp_group, - quant_mode=config.quant_mode, - use_logn_scaling=config.use_logn_attn, - dense_bias=False, - # Qwen3: Add Q/K layer normalization - qk_layernorm=qk_layernorm, - layernorm_type=LayerNormType.RmsNorm - if qk_layernorm else LayerNormType.LayerNorm) - - if config.moe.has_moe(): - mlp_kwargs = {'moe_config': config.moe, 'mapping': config.mapping} - if config.qwen_type == 'qwen2_moe': - # Qwen2 MoE uses SharedMoE with shared expert - ClsMLP = SharedMoE - mlp_kwargs['use_shared_gate'] = True - mlp_kwargs['use_side_stream'] = True - mlp_kwargs['moe_config'].shared_expert_intermediate_size = \ - config.moe_shared_expert_intermediate_size - elif config.qwen_type == 'qwen3_moe': - # Qwen3 MoE uses standard MOE without shared expert - ClsMLP = MOE - else: - ClsMLP = MOE - else: - ClsMLP = GatedMLP - mlp_kwargs = {} - - # Qwen's real inter_size depends on qwen_type - if self.config.qwen_type == 'qwen': - intermediate_size = config.intermediate_size // 2 - elif self.config.qwen_type in ('qwen2_moe', 'qwen3_moe'): - intermediate_size = config.moe_intermediate_size - else: - intermediate_size = config.intermediate_size - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=intermediate_size, - hidden_act=config.hidden_act, - dtype=dtype, - bias=config.mlp_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - mrope_params=None, - ): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - mrope_params=mrope_params, - ) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class QWenModel(Module): - - def __init__(self, config: QWenConfig) -> None: - super().__init__() - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(QWenDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - 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, - spec_decoding_params=None, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - mrope_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - mrope_params=mrope_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class QWenForCausalLM(DecoderModelForCausalLM): - config_class = QWenConfig - - def __init__(self, config: QWenConfig): - transformer = QWenModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - if config.mapping.is_last_pp_rank(): - if config.architecture == 'Qwen2ForSequenceClassification': - lm_head = ColumnLinear(config.hidden_size, - config.num_labels, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = ColumnLinear(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) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - if config.qwen_type == 'qwen': - self.trtllm_modules_to_hf_modules = { - "attn_qkv": "c_attn", - "attn_dense": "attn.c_proj", - "mlp_h_to_4h": "w2", - "mlp_4h_to_h": "mlp.c_proj", - "mlp_gate": "w1", - } - elif config.qwen_type in ('qwen2_moe', 'qwen3_moe'): - self.trtllm_modules_to_hf_modules = copy.copy( - get_default_trtllm_modules_to_hf_modules()) - # Common MoE expert mappings for both Qwen2 and Qwen3 MoE - self.trtllm_modules_to_hf_modules.update({ - "moe_h_to_4h": - "mlp.experts.gate_proj", - "moe_4h_to_h": - "mlp.experts.down_proj", - "moe_gate": - "mlp.experts.up_proj", - }) - # Qwen2 MoE additionally has shared expert - if config.qwen_type == 'qwen2_moe': - self.trtllm_modules_to_hf_modules.update({ - "mlp_h_to_4h": - "mlp.shared_expert.gate_proj", - "mlp_4h_to_h": - "mlp.shared_expert.down_proj", - "mlp_gate": - "mlp.shared_expert.up_proj", - "mlp_router": - "mlp.shared_expert_gate", - }) - else: - self.trtllm_modules_to_hf_modules = None - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a QWenForCausalLM object from give parameters - ''' - import transformers - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - use_autoawq = kwargs.pop('use_autoawq', False) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = QWenConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - arg_dict = {"use_autoawq": True} if use_autoawq else {} - custom_dict = {} - - if config.qwen_type == "qwen": - custom_dict = { - "transformer": "transformer", - "vocab_embedding": "wte", - "ln_f": "ln_f", - "layers": "h", - "attention": "attn", - "qkv": "c_attn", - "dense": "c_proj", - "gate": "w1", - "proj": "c_proj", - "fc": "w2", - "input_layernorm": "ln_1", - "post_layernorm": "ln_2", - } - elif config.qwen_type == "qwen2_moe": - custom_dict = { - "mlp.shared_expert": "mlp.shared_expert", - "mlp.shared_expert_gate": "mlp.shared_expert_gate", - "fc": ["up_proj", "gate_proj"], - } - elif config.qwen_type == "qwen3_moe": - custom_dict = { - "fc": ["up_proj", "gate_proj"], - "q_layernorm": "q_norm", - "k_layernorm": "k_norm", - } - elif config.qwen_type in {"qwen2", "qwen2_vl" - } and config.tie_word_embeddings: - custom_dict = {"lm_head": "model.embed_tokens"} - elif config.architecture == "Qwen2ForSequenceClassification": - custom_dict = { - "lm_head": "score", - } - elif config.qwen_type == "qwen2_llava_onevision": - custom_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head", - } - elif config.qwen_type == "qwen2_audio": - custom_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head", - } - elif config.qwen_type == "qwen3": - custom_dict = { - "q_layernorm": "q_norm", - "k_layernorm": "k_norm", - } - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - if config.qwen_type == "qwen" and model.config.mapping.has_tp(): - - def reshape_qkv(weights): - if weights is None: - return weights - mapping = model.config.mapping - unsqueeze = False - if isinstance(weights, torch.Tensor): - unsqueeze = True - weights = [weights] - - for idx, w in enumerate(weights): - if quant_config.quant_algo == QuantAlgo.W4A16_GPTQ: - w = w.reshape(-1, 3, w.shape[-1] // 3) - w = w.chunk(mapping.tp_size, 2)[mapping.tp_rank] - if w.shape[0] == 1: - weights[idx] = w.reshape(-1) - else: - weights[idx] = w.reshape(w.shape[0], -1) - else: - w = w.reshape(3, w.shape[0] // 3, -1) - w = w.chunk(mapping.tp_size, 1)[mapping.tp_rank] - if w.shape[-1] == 1: - weights[idx] = w.reshape(-1) - else: - weights[idx] = w.reshape(-1, w.shape[-1]) - if unsqueeze: - return weights[0] - else: - return weights - - loader.update_key_mapping(model) - tllm_weights = {} - for tllm_key, _ in tqdm(model.named_parameters()): - if "qkv" in tllm_key: - tllm_weights.update( - loader.load(tllm_key, - reshape_qkv, - skip_tp=True, - custom_postprocess_kwargs=arg_dict)) - else: - tllm_weights.update( - loader.load(tllm_key, - custom_postprocess_kwargs=arg_dict)) - loader.fill(tllm_weights) - elif config.qwen_type in ("qwen2_moe", "qwen3_moe"): - 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"): - tllm_weights.update( - loader.load(tllm_key, - concat_gate_up_proj, - custom_postprocess_kwargs=arg_dict)) - else: - tllm_weights.update( - loader.load(tllm_key, - custom_postprocess_kwargs=arg_dict)) - loader.fill(tllm_weights) - else: - # For Qwen1 w/o TP, Qwen1.5 and Qwen2 w/o MoE - loader.generate_tllm_weights(model, arg_dict) - else: - if not use_preloading: - hf_model = load_hf_qwen(hf_model_dir, load_model_on_cpu) - - logger.debug(f"HuggingFace model: {hf_model}") - - model = QWenForCausalLM(config) - logger.debug(f"TensorRT LLM model: {model}") - - if quant_config.quant_algo == QuantAlgo.W4A16_GPTQ: - weights = load_weights_from_hf_gptq_model(hf_model, config) - else: - weights = load_weights_from_hf_model(hf_model, config) - model.load(weights) - return model - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.weight_only_groupwise_quant_matmul_plugin = 'auto' - return plugin_config - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - calib_dataset='cnn_dailymail', - calib_batches=512, - calib_batch_size=1, - calib_max_seq_length=512, - random_seed=1234, - tokenizer_max_seq_length=2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = QWenConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - convert.quantize(hf_model_dir, - output_dir, - config=config, - calib_dataset=calib_dataset) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/qwen/utils.py b/tensorrt_llm/models/qwen/utils.py deleted file mode 100644 index 1b484898f8d0..000000000000 --- a/tensorrt_llm/models/qwen/utils.py +++ /dev/null @@ -1,116 +0,0 @@ -# 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. - - -def make_context( - tokenizer, - query, - history, - system, - max_input_length, - max_window_size: int = 6144, - chat_format: str = "chatml", -): - if history is None: - history = [] - - if chat_format == "chatml": - im_start, im_end = "<|im_start|>", "<|im_end|>" - im_start_tokens = [tokenizer.im_start_id] - im_end_tokens = [tokenizer.im_end_id] - nl_tokens = tokenizer.encode("\n") - - def _tokenize_str(role, content): - return (f"{role}\n{content}", - tokenizer.encode( - role, - allowed_special=set(), - ) + nl_tokens + tokenizer.encode( - content, - allowed_special=set(), - )) - - system_text, system_tokens_part = _tokenize_str("system", system) - system_tokens = im_start_tokens + system_tokens_part + im_end_tokens - raw_text = "" - context_tokens = [] - - for turn_query, turn_response in reversed(history): - query_text, query_tokens_part = _tokenize_str("user", turn_query) - query_tokens = im_start_tokens + query_tokens_part + im_end_tokens - - response_text, response_tokens_part = _tokenize_str( - "assistant", turn_response) - response_tokens = im_start_tokens + response_tokens_part + im_end_tokens - next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens - prev_chat = ( - f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}" - ) - - current_context_size = (len(system_tokens) + - len(next_context_tokens) + - len(context_tokens)) - if current_context_size < max_window_size: - context_tokens = next_context_tokens + context_tokens - raw_text = prev_chat + raw_text - else: - break - - context_tokens = system_tokens + context_tokens - raw_text = f"{im_start}{system_text}{im_end}" + raw_text - context_tokens += (nl_tokens + im_start_tokens + - _tokenize_str("user", query)[1] + im_end_tokens + - nl_tokens + im_start_tokens + - tokenizer.encode("assistant") + nl_tokens) - raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n" - - elif chat_format == "raw": - raw_text = query - context_tokens = tokenizer.encode(raw_text) - else: - raise NotImplementedError(f"Unknown chat format {chat_format!r}") - # truncate to max_input_length, truncate from the front - return raw_text, context_tokens[-max_input_length:] - - -def get_qwen_key_list(qwen_type): - qwen_key_list = [ - "attn.c_attn", # attention.qkv - "attn.c_proj", # attention.dense - "mlp.w1", # mlp.gate - "mlp.w2", # mlp.fc - "mlp.c_proj", # mlp.proj - "ln_1", # input_layernorm - "ln_2", # post_layernorm - "transformer.wte", # vocabulary embedding - "transformer.ln_f", # final layer norm - ] - qwen2_key_list = [ - "self_attn.", # attention.qkv - "self_attn.o_proj", # attention.dense - "mlp.up_proj", # mlp.gate - "mlp.gate_proj", # mlp.fc - "mlp.down_proj", # mlp.proj - "input_layernorm", # input_layernorm - "post_attention_layernorm", # post_layernorm - "model.embed_tokens", # vocabulary embedding - "model.norm", # final layer norm - ] - key_list = [] - if qwen_type == 'qwen': - key_list.extend(qwen_key_list) - else: - key_list.extend(qwen2_key_list) - return key_list diff --git a/tensorrt_llm/models/recurrentgemma/__init__.py b/tensorrt_llm/models/recurrentgemma/__init__.py deleted file mode 100644 index 2a36ca922710..000000000000 --- a/tensorrt_llm/models/recurrentgemma/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. diff --git a/tensorrt_llm/models/recurrentgemma/model.py b/tensorrt_llm/models/recurrentgemma/model.py deleted file mode 100644 index b5fc842f1359..000000000000 --- a/tensorrt_llm/models/recurrentgemma/model.py +++ /dev/null @@ -1,624 +0,0 @@ -# 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 collections import OrderedDict -from typing import List - -import tensorrt as trt - -from ..._common import default_net -from ..._utils import str_dtype_to_trt -from ...functional import (Tensor, arange, concat, expand, - gather_last_token_logits, shape, tanh, unsqueeze) -from ...layers import (Attention, AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, GatedMLP, KeyValueCacheParams, - PositionEmbeddingType, Recurrent, RmsNorm) -from ...module import Module, ModuleList -from ...plugin import current_all_reduce_helper -from ..generation_mixin import GenerationMixin -from ..modeling_utils import (PretrainedConfig, PretrainedModel, - get_kv_cache_type_from_legacy) - - -class ResidualLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - layer_type_len = len(config.layer_types) - self.temporal_block_type = config.layer_types[layer_idx % - layer_type_len] - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - if self.temporal_block_type == 'recurrent': - self.recurrent = Recurrent(width=config.hidden_size, - lru_width=config.rnn_hidden_size, - d_conv=config.conv_kernel, - num_heads=config.num_attention_heads, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size) - elif self.temporal_block_type == 'attention': - layer_types = config.layer_types * ( - (layer_idx + 1) // layer_type_len) - layer_types = layer_types + config.layer_types[0:( - (layer_idx + 1) % layer_type_len)] - attention_layer_idx = layer_types.count('attention') - 1 - - self.attention = Attention( - local_layer_idx=attention_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_percentage=config.rotary_pct, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode, - bias=False, - dense_bias=True) - else: - raise ValueError( - 'RecurrentGemma only support "recurrent" and "attention" blocks.' - ) - - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - conv_state=None, - lru_state=None, - host_request_types=None, - last_token_ids=None, - host_context_lengths=None, - slot_mapping=None, - conv_indices=None): - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - if self.temporal_block_type == 'recurrent': - temporal_output, present_conv, present_lru = self.recurrent( - hidden_states, - conv_state=conv_state, - lru_state=lru_state, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - conv_indices=conv_indices, - ) - else: - present_conv, present_lru = None, None - - if self.temporal_block_type == 'attention': - temporal_output, present_kv = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - else: - present_kv = None - - hidden_states = residual + temporal_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - return hidden_states, present_kv, present_conv, present_lru - - -class RecurrentGemmaModel(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - self.d_conv = config.conv_kernel - self.lru_width = config.rnn_hidden_size - n_layer = config.num_hidden_layers - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = ModuleList( - [ResidualLayer(config, layer_idx=i) for i in range(n_layer)]) - - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - conv_states=None, - lru_states=None, - host_request_types=None, - last_token_ids=None, - host_context_lengths=None, - slot_mapping=None): - - hidden_states = self.vocab_embedding(input_ids) - - # Get conv state indices - indices = None - if not default_net().plugin_config.mamba_conv1d_plugin: - batch_size = shape(input_ids, 0) - indices = expand( - unsqueeze(arange(0, self.d_conv - 1, dtype='int32'), 0), - concat([batch_size, self.d_conv - 1])) - offsets = expand(unsqueeze(last_token_ids, 1), - concat([batch_size, self.d_conv - 1])) - indices = unsqueeze(indices + offsets, 1) - indices = expand( - indices, concat([batch_size, self.lru_width, self.d_conv - 1])) - - present_kvs, present_convs, present_lrus = [], [], [] - for layer, past_kv, past_conv, past_lru in zip( - self.layers, kv_cache_params.past_key_value, conv_states, - lru_states): - hidden_states, present_kv, present_conv, present_lru = layer( - hidden_states, - use_cache, - attention_mask, - kv_cache_params=KeyValueCacheParams( - past_key_value=[past_kv], - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cache_indirection=kv_cache_params.cache_indirection), - attention_params=attention_params, - conv_state=past_conv, - lru_state=past_lru, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - conv_indices=indices) - present_kvs.append(present_kv) - present_convs.append(present_conv) - present_lrus.append(present_lru) - - hidden_states = self.ln_f(hidden_states) - return hidden_states, tuple(present_kvs), tuple(present_convs), tuple( - present_lrus) - - -class RecurrentGemmaForCausalLM(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - super().__init__(config) - dtype = config.dtype - logits_dtype = config.logits_dtype - if isinstance(dtype, str): - self.dtype = str_dtype_to_trt(dtype) - else: - assert isinstance(dtype, trt.DataType) - self.dtype = dtype - - assert len(config.layer_types) > 0 - layer_types = config.layer_types - layer_types = layer_types * (config.num_hidden_layers // - len(layer_types)) - layer_types = layer_types + layer_types[0:(config.num_hidden_layers % - len(layer_types))] - self.layer_types = layer_types - - self.config = config - self.gather_context_logits = False - self.logits_soft_cap = config.logits_soft_cap - - # Create constant attention parameters to be reused by all layers. - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - - if isinstance(logits_dtype, str): - self._logits_dtype = str_dtype_to_trt(logits_dtype) - else: - assert isinstance(logits_dtype, trt.DataType) - self._logits_dtype = logits_dtype - - self.transformer = RecurrentGemmaModel(config) - self.lm_head = ColumnLinear(config.hidden_size, - config.vocab_size, - bias=False, - dtype=dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - conv_states=None, - rnn_states=None, - host_request_types=None, - last_token_ids=None, - last_token_ids_for_logits=None, - host_context_lengths=None, - slot_mapping=None): - - # fill attention params. - attention_params = Attention.fill_attention_params( - self, attention_params) - - hidden_states, present_kvs, present_convs, present_rnns = self.transformer( - input_ids, use_cache, attention_mask, kv_cache_params, - attention_params, conv_states, rnn_states, host_request_types, - last_token_ids, host_context_lengths, slot_mapping) - - if not self.gather_context_logits: - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids_for_logits, - default_net().plugin_config.remove_input_padding) - - lm_logits = self.lm_head(hidden_states) - lm_logits = tanh( - lm_logits / self.logits_soft_cap) * self.logits_soft_cap - lm_logits.mark_output('logits', self._logits_dtype) - if not default_net().plugin_config.paged_kv_cache: - for i, present_kv in enumerate(present_kvs): - if present_kv is not None: - present_kv.mark_output(f'present_key_value_{i}', self.dtype) - - if not default_net().plugin_config.paged_state: - for i, present_conv in enumerate(present_convs): - if present_conv is not None: - present_conv.mark_output(f'present_conv_state_{i}', - self.dtype) - for i, present_rnn in enumerate(present_rnns): - if present_rnn is not None: - present_rnn.mark_output(f'present_rnn_state_{i}', - str_dtype_to_trt('float32')) - - return (lm_logits, present_kvs, present_convs, present_rnns) - - def prepare_recurrent_inputs(self, max_batch_size, num_profiles, mapping): - use_mamba_conv1d_plugin = default_net( - ).plugin_config.mamba_conv1d_plugin - - default_range = GenerationMixin.default_range - batch_range = [default_range(max_batch_size)] * num_profiles - - conv_states = [] - rnn_states = [] - dim = self.config.rnn_hidden_size // mapping.tp_size - if use_mamba_conv1d_plugin: - conv_state_dim_range = OrderedDict([ - ('batch_size', batch_range), - ('kernel_size', [self.config.conv_kernel - 1] * num_profiles), - ('dim_size', [dim] * num_profiles), - ]) - else: - conv_state_dim_range = OrderedDict([ - ('batch_size', batch_range), - ('dim_size', [dim] * num_profiles), - ('kernel_size', [self.config.conv_kernel - 1] * num_profiles), - ]) - - rnn_state_dim_range = OrderedDict([ - ('batch_size', batch_range), - ('state_size', [1] * num_profiles), - ('dim_size', [dim] * num_profiles), - ]) - one_dim_range = OrderedDict([ - ('buffer_count', [1] * num_profiles), - ]) - - for i in range(self.config.num_hidden_layers): - if self.layer_types[i] == 'recurrent': - if default_net().plugin_config.paged_state: - conv_state = Tensor(name=f'conv_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - - rnn_state = Tensor(name=f'rnn_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - else: - if use_mamba_conv1d_plugin: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, self.config.conv_kernel - 1, dim], - dim_range=conv_state_dim_range) - else: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, dim, self.config.conv_kernel - 1], - dim_range=conv_state_dim_range) - - rnn_state = Tensor(name=f'past_rnn_state_{i}', - dtype=str_dtype_to_trt('float32'), - shape=[-1, 1, dim], - dim_range=rnn_state_dim_range) - else: - conv_state, rnn_state = None, None - conv_states.append(conv_state) - rnn_states.append(rnn_state) - - slot_mapping = None - if default_net().plugin_config.paged_state: - slot_mapping = Tensor( - name='slot_mapping', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', batch_range)]), - ) - - return_dict = { - 'conv_states': conv_states, - 'rnn_states': rnn_states, - 'slot_mapping': slot_mapping, - } - return return_dict - - def prepare_inputs( - self, - max_batch_size, - max_input_len, - max_seq_len, - max_num_tokens, - use_cache, - max_beam_width: int = 1, - opt_num_tokens: int = None, - opt_batch_size: int = 0, - prompt_embedding_table_size: int = 0, - max_draft_len: int = 0, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - speculative_decoding_draft_tokens_external: bool = False): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - assert speculative_decoding_draft_tokens_external == False, \ - "We don't support speculative decoding for the RecurrentGemma model." - assert max_beam_width == 1, "We don't support beam search for the RecurrentGemma model." - - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - multiple_profiles = default_net().plugin_config.multiple_profiles - streamingllm = default_net().plugin_config.streamingllm - use_mamba_conv1d_plugin = default_net( - ).plugin_config.mamba_conv1d_plugin - - self.gather_context_logits = gather_context_logits - mapping = self.config.mapping - kv_cache_type = get_kv_cache_type_from_legacy(use_cache, paged_kv_cache) - - # basic inputs - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type) - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_num_tokens=max_num_tokens, - max_draft_len=max_draft_len, - opt_batch_size=opt_batch_size, - opt_num_tokens=opt_num_tokens, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles, - kv_cache_type=kv_cache_type) - - if remove_input_padding: - assert use_mamba_conv1d_plugin, "mamba_conv1d_plugin is needed to support remove_input_padding" - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('num_tokens', ranges['num_tokens_range']), - ])) - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('position_ids_num_tokens_range', - ranges['num_tokens_range']), - ])) - else: - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - ranges['bb_range']), - ('input_len', ranges['inlen_range']), - ])) - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - ranges['bb_range']), - ('position_ids_inlen_range', - ranges['position_ids_inlen_range']), - ])) - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor( - mapping, num_profiles) - - # attention inputs - attn_layer_idx = [] - for i in range(self.config.num_hidden_layers): - if self.layer_types[i] == 'attention': - attn_layer_idx.append(i) - attention_inputs = self.prepare_attention_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - num_kv_heads=self.config.num_key_value_heads, - head_size=self.config.head_size, - num_layers=self.config.num_hidden_layers, - kv_dtype=str_dtype_to_trt(self.config.kv_dtype), - num_profiles=num_profiles, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - mapping=mapping, - streamingllm=streamingllm, - attn_layer_idx=attn_layer_idx) - - # recurrent inputs - recurrent_inputs = self.prepare_recurrent_inputs( - max_batch_size=max_batch_size, - num_profiles=num_profiles, - mapping=mapping, - ) - - if use_gpt_attention_plugin: - host_request_types = attention_inputs['host_request_types'] - else: - host_request_types = Tensor( - name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', - ranges['bb_range'])]), - ) - - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_last_token_ids', ranges['bbd_range']), - ]), - ) - last_token_ids_for_logits = None - if not gather_context_logits: - last_token_ids_for_logits = last_token_ids - - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = attention_inputs['host_context_lengths'] - elif remove_input_padding: - host_context_lengths = Tensor( - name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', - ranges['bb_range'])]), - ) - else: - host_context_lengths = None - - return_dict = { - 'input_ids': - input_ids, - 'position_ids': - position_ids, - 'use_cache': - True, - 'attention_mask': - attention_inputs['attention_mask'], - 'kv_cache_params': - KeyValueCacheParams( - past_key_value=attention_inputs['past_key_value'], - host_past_key_value_lengths=attention_inputs[ - 'host_past_key_value_lengths'], - host_max_attention_window_sizes=attention_inputs[ - 'host_max_attention_window_sizes'], - host_sink_token_length=attention_inputs[ - 'host_sink_token_length'], - kv_cache_block_offsets=attention_inputs[ - 'kv_cache_block_offsets'], - host_kv_cache_block_offsets=attention_inputs[ - 'host_kv_cache_block_offsets'], - host_kv_cache_pool_pointers=attention_inputs[ - 'host_kv_cache_pool_pointers'], - host_kv_cache_pool_mapping=attention_inputs[ - 'host_kv_cache_pool_mapping'], - cache_indirection=attention_inputs['cache_indirection'], - ), - 'attention_params': - AttentionParams( - sequence_length=attention_inputs['sequence_length'], - context_lengths=attention_inputs['context_lengths'], - host_context_lengths=attention_inputs['host_context_lengths'], - max_context_length=max_input_len, - host_request_types=attention_inputs['host_request_types'], - host_runtime_perf_knobs=attention_inputs[ - 'host_runtime_perf_knobs'], - host_context_progress=attention_inputs['host_context_progress'], - ), - 'conv_states': - recurrent_inputs['conv_states'], - 'rnn_states': - recurrent_inputs['rnn_states'], - 'host_request_types': - host_request_types, - 'last_token_ids': - last_token_ids, - 'last_token_ids_for_logits': - last_token_ids_for_logits, - 'host_context_lengths': - host_context_lengths, - 'slot_mapping': - recurrent_inputs['slot_mapping'], - } - return return_dict diff --git a/tensorrt_llm/models/redrafter/__init__.py b/tensorrt_llm/models/redrafter/__init__.py deleted file mode 100644 index 2a36ca922710..000000000000 --- a/tensorrt_llm/models/redrafter/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. diff --git a/tensorrt_llm/models/redrafter/drafter.py b/tensorrt_llm/models/redrafter/drafter.py deleted file mode 100644 index c85560b3d473..000000000000 --- a/tensorrt_llm/models/redrafter/drafter.py +++ /dev/null @@ -1,117 +0,0 @@ -from typing import Optional - -from tensorrt_llm.functional import Tensor, silu -from tensorrt_llm.layers import ColumnLinear -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.module import Module, ModuleList - -from ..._utils import str_dtype_to_trt - - -class ResBlock(Module): - - def __init__(self, - exit_dim: int, - dtype: Optional[str], - mapping: Mapping = Mapping()): - super().__init__() - self.linear = ColumnLinear( - exit_dim, - exit_dim, - bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - - def forward(self, x: Tensor) -> Tensor: - return x + silu(self.linear(x)) - - -class Drafter(Module): - - def __init__( - self, - num_layers: int, - hidden_size: int, - exit_dim: int, - vocab_size: int, - dtype: Optional[str] = None, - is_rnn: bool = False, - mapping: Mapping = Mapping(), - ): - super().__init__() - self.num_layers = num_layers - self.is_rnn = is_rnn - self.dtype = str_dtype_to_trt(dtype) - - input_dim = 2 * hidden_size - self.input_proj = (None if input_dim == exit_dim else ColumnLinear( - input_dim, - exit_dim, - bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - )) - - self.layers = ModuleList([ - ResBlock(exit_dim, dtype, mapping) for _ in range(self.num_layers) - ]) - self.lm_head = ColumnLinear( - exit_dim, - vocab_size, - bias=False, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - - if is_rnn: - self.rnn_u = ColumnLinear( - hidden_size, - hidden_size, - bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - self.rnn_w = ColumnLinear( - hidden_size, - hidden_size, - bias=False, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - return - - @classmethod - def from_config(cls, config, vocab_size_padded): - kwargs = { - "num_layers": config.redrafter_num_layers, - "hidden_size": config.redrafter_hidden_size, - "exit_dim": config.redrafter_exit_dim, - "vocab_size": vocab_size_padded, - "dtype": config.dtype, - "is_rnn": config.redrafter_is_rnn, - "mapping": config.mapping, - } - return cls(**kwargs) - - def forward(self, x: Tensor) -> Tensor: - hidden_states = self.input_proj(x) if self.input_proj is not None else x - for layer in self.layers: - hidden_states = layer(hidden_states) - - return self.lm_head(hidden_states) - - def rnn_embed(self, x: Tensor, prev: Tensor = None) -> Tensor: - assert self.is_rnn, "This function should not be called when redrafter_is_rnn is false." - w_embd = self.rnn_w(x) - return w_embd if prev is None else w_embd + self.rnn_u(prev) diff --git a/tensorrt_llm/models/redrafter/model.py b/tensorrt_llm/models/redrafter/model.py deleted file mode 100644 index 84c78cc79810..000000000000 --- a/tensorrt_llm/models/redrafter/model.py +++ /dev/null @@ -1,317 +0,0 @@ -# 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 collections import OrderedDict - -import tensorrt as trt - -from tensorrt_llm._common import default_net -from tensorrt_llm.functional import Tensor, cast, categorical_sample -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.models import LLaMAForCausalLM, QWenForCausalLM -from tensorrt_llm.models.generation_mixin import GenerationMixin - -from ..._utils import pad_vocab_size, str_dtype_to_trt -from .drafter import Drafter -from .redrafter_helper import (_beam_search_candidates, _beams2tree, - _process_logits_and_hidden_states) - - -class ReDrafterMixin: - - def __init__(self, config): - - super().__init__(config) - self.dtype = str_dtype_to_trt(config.dtype) - self.vocab_size = config.vocab_size - vocab_size_padded = pad_vocab_size(self.vocab_size, - config.mapping.tp_size) - self.drafter = Drafter.from_config(config, vocab_size_padded) - self.num_beams = config.redrafter_num_beams - self.beam_candidate_length = config.redrafter_draft_len_per_beam - self.beam_length = self.beam_candidate_length + 1 # including true token - self.greedy_search = config.redrafter_greedy_search - self.is_rnn = config.redrafter_is_rnn - assert self.dtype == self.drafter.dtype, f"{self.dtype} != {self.drafter.dtype}" - - def _fwd_helper(self, hidden_states, lm_logits, embedding, drafter, - kwargs: dict): - ''' - Must enable remove_input_padding: - hidden_states [total_tokens, H] - lm_logits [total_tokens, V] - 1. process_logits: context vs gen - a. Context: just return the last hidden states, and logits/probs - b. Gen: - i. verify: use lm_logits, draft_probs, draft_indices, draft_tokens - ii. select hidden state and update probs - 3. Sample token based on probs - 4. Generate candidates using hidden_states, sampled token - 5. Using beams, generate validation buffers, mark them as output - 6. Mark all the outputs - ''' - - num_beams = self.num_beams - beam_length = self.beam_length - - # Get the inputs needed - rand_data_sample = kwargs['rand_data_sample'] - position_ids_base = kwargs['position_ids_base'] - - # Step 1: Process logits and hidden states - # process the base model output (verify for gen-phase) - probs, draft_input, num_accepted_tokens, \ - accepted_beam_index = _process_logits_and_hidden_states( - self, lm_logits, hidden_states, kwargs) - # NOTE: num_accepted_tokens doesn't include true token so add 1 here - num_accepted_tokens = num_accepted_tokens + 1 - - # At this point: - # probs : [bs, V] - # hidden_states : [bs, H] - - # Step 2: Sample token - next_token = categorical_sample(probs, rand_data_sample) - - # Step 3: beam search - new_draft_tokens, new_draft_logits = _beam_search_candidates( - draft_input, next_token, embedding, drafter, self.num_beams, - self.beam_length, self.is_rnn) - - # Step 4: tree processing - active_tokens_flattened, new_draft_token_indices, new_mask, \ - new_position_offsets, packed_position_ids, next_num_gen_tokens, max_gen_token, \ - total_gen_token = _beams2tree(new_draft_tokens, num_beams, beam_length, - position_ids_base + num_accepted_tokens) - - # Step 5: mark all the tensors we need - num_accepted_tokens.mark_output('num_accepted_tokens') - accepted_beam_index.mark_output('accepted_beam_index') - max_gen_token.mark_output('max_gen_token') - total_gen_token.mark_output('total_gen_token') - next_num_gen_tokens.mark_output('next_spec_decoding_generation_lengths') - active_tokens_flattened.mark_output('next_flat_tokens') - new_draft_tokens.mark_output('next_draft_tokens') - new_draft_logits.mark_output('next_draft_probs') - new_draft_token_indices.mark_output('next_draft_indices') - new_mask.mark_output('spec_decoding_mask') - new_position_offsets.mark_output('next_spec_decoding_position_offsets') - packed_position_ids.mark_output('packed_position_ids') - - return next_token, probs, draft_input - - def forward(self, *args, **kwargs): - """ - 0. run base model, get logits, hidden_states - """ - - extra_args = [ - 'draft_tokens', - 'draft_indices', - 'draft_probs', - 'device_request_types', - 'redrafter_inverted_temperature', - 'rand_data_validation', - 'rand_data_sample', - 'position_ids_base', - ] - use_cache = True - base_kwargs = {k: v for k, v in kwargs.items() if k not in extra_args} - if use_cache and default_net().plugin_config.paged_kv_cache is False: - lm_logits, presents, hidden_states = super().forward( - *args, **base_kwargs) - else: - lm_logits, hidden_states, _ = super().forward(*args, **base_kwargs) - - # lm_logits could be in fp32 - lm_logits_cast = cast(lm_logits, self.dtype) # no-op if same type - self.register_network_output("hidden_states", - hidden_states) # debugging - - new_draft_tokens, new_draft_logits, probs = self._fwd_helper( - hidden_states, - lm_logits_cast, - self.transformer.vocab_embedding, - self.drafter, - kwargs=kwargs) - - return new_draft_tokens, new_draft_logits, probs - - def prepare_inputs(self, *args, **kwargs): - """ - Inputs needed: - Assuming, max_gen_tokens = 1 + nb*(bl - 1), counting true token - device_request_types: [bs] - draft_tokens: [bs, nb, bl] - draft_indices: [bs, nb, bl] - draft_probs: [bs, nb, bl-1, V] - spec_decoding_generation_lengths: [bs] - spec_decoding_position_offsets: [bs, max_gen_tokens] - spec_decoding_packed_mask: [bs, max_gen_tokens, packed_length] ** - redrafter_inverted_temperature: [bs] - rand_data_sample: [bs] - rand_data_validation: [bs, nb, bl-1] - - ** The mask is tricky since the boolean mask will need to be - packed in runtime. So, the last dim will be: - packed_length = ceil(max_gen_tokens/32) - """ - default_range = GenerationMixin.default_range - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - max_batch_size = kwargs['max_batch_size'] - assert max_batch_size is not None - bb_range = default_range(max_batch_size) - bb0_range = default_range(max_batch_size, min_range=0, opt_offset=1) - num_beam_tokens = self.num_beams * self.beam_length - max_draft_len = num_beam_tokens - self.num_beams # ignore the true token - max_gen_token_len = 1 + max_draft_len # for the true token - max_gen_token_len_range = default_range(max_gen_token_len) - bb_max_gen_token_len_range = default_range(max_gen_token_len * - max_batch_size, - min_range=0) - - kwargs['speculative_decoding_draft_tokens_external'] = False - kwargs['max_draft_len'] = max_draft_len - kwargs['spec_decoding_is_generation_length_variable'] = True - inputs = super().prepare_inputs(*args, **kwargs) - assert inputs['spec_decoding_params'] is not None - - enable_two_optimization_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=KVCacheType.PAGED - if paged_kv_cache else KVCacheType.CONTINUOUS) - if enable_two_optimization_profiles: - bb_range = [bb_range, bb_range] - bb0_range = [bb0_range, bb0_range] - max_gen_token_len_range = [ - max_gen_token_len_range, max_gen_token_len_range - ] - bb_max_gen_token_len_range = [ - bb_max_gen_token_len_range, bb_max_gen_token_len_range - ] - num_beams_range = [self.num_beams, self.num_beams] - beam_length_range = [self.beam_length, self.beam_length] - candidate_length_range = [ - self.beam_candidate_length, self.beam_candidate_length - ] - vocab_size_range = [self.vocab_size, self.vocab_size] - else: - bb_range = [bb_range] - bb0_range = [bb0_range] - max_gen_token_len_range = [max_gen_token_len_range] - bb_max_gen_token_len_range = [bb_max_gen_token_len_range] - num_beams_range = [self.num_beams] - beam_length_range = [self.beam_length] - candidate_length_range = [self.beam_candidate_length] - vocab_size_range = [self.vocab_size] - - device_request_types = Tensor(name='device_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - draft_tokens = Tensor(name='draft_tokens', - dtype=trt.int32, - shape=[-1, self.num_beams, self.beam_length], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('beam_length', beam_length_range), - ])) - draft_indices = Tensor(name='draft_indices', - dtype=trt.int32, - shape=[-1, self.num_beams, self.beam_length], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('beam_length', beam_length_range), - ])) - draft_probs = Tensor( - name='draft_probs', - dtype=self.dtype, - shape=[-1, self.num_beams, self.beam_length - 1, self.vocab_size], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('candidate_length', candidate_length_range), - ('vocab_size', vocab_size_range), - ])) - redrafter_inverted_temperature = Tensor( - name='redrafter_inverted_temperature', - dtype=self.dtype, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - rand_data_validation = Tensor( - name='rand_data_validation', - dtype=self.dtype, - shape=[-1, self.num_beams, self.beam_length - 1], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('candidate_length', candidate_length_range), - ])) - rand_data_sample = Tensor(name='rand_data_sample', - dtype=self.dtype, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - position_ids_base = Tensor( - name="position_ids_base", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ]), - ) - - inputs[ - 'device_request_types'] = device_request_types # needed by process_logits - inputs['draft_tokens'] = draft_tokens - inputs['draft_indices'] = draft_indices - inputs['draft_probs'] = draft_probs - inputs[ - 'redrafter_inverted_temperature'] = redrafter_inverted_temperature - inputs['rand_data_validation'] = rand_data_validation - inputs['rand_data_sample'] = rand_data_sample - inputs['position_ids_base'] = position_ids_base - return inputs - - -class ReDrafterForQWenLM(ReDrafterMixin, QWenForCausalLM): - """ReDrafter implementation for QWen models. - - Combines: - - Base QWen model functionality from QWenForCausalLM - - Drafting/speculative decoding logic from ReDrafterMixin - """ - - -class ReDrafterForLLaMALM(ReDrafterMixin, LLaMAForCausalLM): - """ReDrafter implementation for LLaMA models. - - Combines: - - Base LLaMA model functionality from LLaMAForCausalLM - - Drafting/speculative decoding logic from ReDrafterMixin - """ diff --git a/tensorrt_llm/models/redrafter/redrafter_helper.py b/tensorrt_llm/models/redrafter/redrafter_helper.py deleted file mode 100644 index 31e930f51a97..000000000000 --- a/tensorrt_llm/models/redrafter/redrafter_helper.py +++ /dev/null @@ -1,759 +0,0 @@ -from typing import Tuple - -import numpy as np - -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import numpy_array - -# isort: off -from tensorrt_llm.functional import ( - Tensor, arange, argmax, cast, concat, constant, constant_to_tensor_, cumsum, - div, eq, exp, expand, expand_dims, floordiv, gather, gather_nd, - index_select, int32_array, log_softmax, lt, max, maximum, masked_select, - minimum, nonzero, not_op, op_and, rand, relu, scatter, select, shape, slice, - silu, softmax, squeeze, stack, sum, topk, transpose, unsqueeze, view, where) -# isort: on -from tensorrt_llm.layers import Embedding -from tensorrt_llm.module import Module - -INT_DTYPE_STR = "int32" -''' -NOTE: - Name differences from Apple's PyTorch Implementation: - `num_candidates` is mapped to `num_beams` and - `candidate_length` is mapped to `beam_length - 1`. - So for each sequence, the paths/beams to verify will be [num_beams, beam_length] tokens where - each beam is a path that includes the true token (1) and the candidate tokens (beam_length - 1). -''' - - -def _unpack_beams(x: Tensor, indices: Tensor, num_beams: int, - beam_length: int) -> Tensor: - """ - x: [bs, S, V] - indices: [bs, nb, bl] - output: - """ - assert x.rank() == 3 - d0 = shape(x, 0, INT_DTYPE_STR) - dl = shape(x, -1, INT_DTYPE_STR) - indices = view(indices, [-1, num_beams * beam_length, 1], False) - res_shape = concat([d0, num_beams, beam_length, dl]) - res = view(gather_nd(x, indices), res_shape, False) # [d0, nb, bl, dl] - return res - - -def _validate_draft_tokens(draft_log_probs: Tensor, - draft_tokens: Tensor, - draft_indices: Tensor, - flattened_logits: Tensor, - num_beams: int, - beam_length: int, - greedy_search: bool, - rand_data: Tensor = None): - ''' - draft_log_probs: [bs, nb, bl-1, V] - draft_tokens: [bs, nb, bl] - draft_indices: [bs, nb, bl] - flattened_logits: [bs, S, V], we need to unflatten it using draft_indices. - The unflattend_logits should be of shape [bs, nb, bl, V] by doing a gather on S. - ''' - batch_size = shape(flattened_logits, 0, INT_DTYPE_STR) - rand_shape = concat([batch_size, num_beams, beam_length - 1]) - if rand_data is None: - rand_data = rand(rand_shape, low=0, high=1, dtype=draft_log_probs.dtype) - - flat_log_probs = log_softmax(flattened_logits, dim=-1) - all_base_log_probs = _unpack_beams(flat_log_probs, draft_indices, num_beams, - beam_length) # [bs, nb, bl, V] - if greedy_search: - all_base_log_probs = _top_1_logits(all_base_log_probs) - - base_log_probs = index_select(all_base_log_probs, - dim=2, - index=constant( - np.arange(beam_length - 1, - dtype=np.int32))) - last_base_log_probs = select(all_base_log_probs, - dim=2, - index=beam_length - 1) - proposed_tokens = unsqueeze(slice(draft_tokens, [0, 0, 1], rand_shape), -1) - - token_base_log_probs = squeeze( - gather(base_log_probs, dim=-1, indices=proposed_tokens), -1) - token_draft_log_probs = squeeze( - gather(draft_log_probs, dim=-1, indices=proposed_tokens), -1) - diff_probs = exp(token_base_log_probs - token_draft_log_probs) - cmp = cast(lt(rand_data, diff_probs), dtype='int32') - ideal_sum = constant(np.arange(1, beam_length, dtype=np.int32)) - cum_sum = cumsum(cmp, dim=-1) - equality = cast((cum_sum == ideal_sum), dtype='int32') - num_accepted = sum(equality, dim=-1) - max_num_accepted_tokens, accepted_beam_index = topk( - num_accepted, k=1, - dim=-1) # need to use topk layer to get both value and index - return squeeze(max_num_accepted_tokens, -1), squeeze(accepted_beam_index, -1),\ - base_log_probs, last_base_log_probs, rand_data - - -def _get_prefix_match_indices(beams, beam_length): - ''' - beams: [bs, nb, bl] - ''' - prefix_target = constant( - np.expand_dims(np.arange(1, beam_length + 1, dtype=np.int32), - [0, 1, 2])) - matches = cast(expand_dims(beams, 1) == expand_dims(beams, 2), beams.dtype) - seq_matches = cast(cumsum(matches, dim=3) == prefix_target, - dtype=beams.dtype) - prefix_match_indices = argmax(seq_matches, dim=2) - return prefix_match_indices - - -def _get_draft_token_indices(prefix_match_indices, num_beams, beam_length): - ''' - prefix_match_indices: [bs, nb, bl] - ''' - pmi_dtype = prefix_match_indices.dtype - segments = cast( - constant(np.expand_dims(np.arange(0, num_beams, dtype=np.int32), - [0, 2])) == prefix_match_indices, pmi_dtype) - segment_lengths = sum(segments, dim=-1) - accum_lengths = cumsum(segment_lengths, dim=-1) - segment_lengths - segment_index = gather(accum_lengths, - dim=1, - indices=view(prefix_match_indices, - shape=[-1, num_beams * beam_length])) - segment_index = view(segment_index, [-1, num_beams, beam_length]) - match = cast( - expand_dims(segment_index, 3) == expand_dims(segment_index, 2), - pmi_dtype) - seq_index = constant(np.arange(beam_length, dtype=np.int32)) - lower_triangle = cast( - expand_dims(seq_index, 1) > expand_dims(seq_index, 0), pmi_dtype) - offset = sum(match * expand_dims(lower_triangle, [0, 1]), dim=-1) - draft_token_indices = segment_index + offset - return draft_token_indices - - -def _get_packed_position_ids( - active_indices: Tensor, - indices: Tensor, - total_lengths: Tensor, - position_ids_base: Tensor, -) -> Tensor: - expand_shape = concat([shape(total_lengths, 0), shape(indices, 0)]) - expanded_indices = expand(unsqueeze(indices, 0), expand_shape) - position_mask = expanded_indices < unsqueeze(total_lengths, 1) - position_ids = active_indices + unsqueeze(position_ids_base, 1) - packed_position_ids = masked_select(position_ids, position_mask) - return packed_position_ids - - -def _get_draft_token_array( - beams: Tensor, - prefix_match_indices: Tensor, - num_beams: int, - beam_length: int, - position_ids_base: Tensor = None, -) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: - ''' - beams: [bs, nb, bl] - prefix_match_indices: [bs, nb, bl] - ''' - prefix_ideal_indices = constant(np.arange(num_beams, dtype=np.int32)) - prefix_ideal_indices = expand_dims(prefix_ideal_indices, [0, 2]) - segments = cast(eq(prefix_match_indices, prefix_ideal_indices), - dtype=beams.dtype) - raw_draft_token_array = view(segments * beams + (segments - 1), - [-1, num_beams * beam_length], False) - raw_active_token_indices = transpose( - nonzero(not_op(raw_draft_token_array == -1)), 0, 1) - active_token_flattened = gather_nd(raw_draft_token_array, - raw_active_token_indices, 0) - - total_lengths = sum(view(segments, [-1, num_beams * beam_length], False), - dim=1) - slice_size = concat([shape(raw_active_token_indices, 0, INT_DTYPE_STR), 1]) - active_token_index_flattened = view( - slice(raw_active_token_indices, starts=[0, 1], sizes=slice_size), [-1], - False) - - max_len = max(total_lengths, dim=0) - total_gen_len = sum(total_lengths, dim=0) - # constant_0 = constant(int32_array(0)) - # offset = arange(constant_0, max_len, dtype='int32') - offset = slice(constant(np.arange(num_beams * beam_length, dtype=np.int32)), - constant_to_tensor_(0), unsqueeze(max_len, 0)) - idx_starts = cumsum(total_lengths, 0) - total_lengths - select_indices = unsqueeze(idx_starts, -1) + unsqueeze(offset, 0) - max_index_allowed = shape(active_token_flattened, 0, INT_DTYPE_STR) - 1 - select_indices = minimum(view(select_indices, [-1], False), - max_index_allowed) - compressed_shape = concat([shape(total_lengths, 0, INT_DTYPE_STR), max_len]) - # draft_token_array = view( - # gather(active_token_flattened, dim=0, indices=select_indices), - # compressed_shape, False) - active_token_indices = view( - gather(active_token_index_flattened, dim=0, indices=select_indices), - compressed_shape, False) - # adding position offsets here - position_offsets = active_token_indices % beam_length - packed_position_ids = constant_to_tensor_(0) # dummy initialization - if position_ids_base is not None: - packed_position_ids = _get_packed_position_ids(position_offsets, offset, - total_lengths, - position_ids_base) - return active_token_flattened, active_token_indices, total_lengths, max_len, total_gen_len, position_offsets, packed_position_ids - - -# FROM APPLE (minor changes by NV) -def _get_mask(draft_token_indices: Tensor, active_token_indices: Tensor, - num_beams: int, beam_length: int) -> Tensor: - """ - Return mask for candidates according to the flattened and compact index. - Args: - draft_token_indices: (batch_size, num_beams, beam_length) - A Mapping of draft candidates index from a stacked representation to a - flattened and compact representation. - active_token_indices: (batch_size, max_len) - A Mapping of draft candidates index from a flattened and compact representation - to a stacked representation. - Returns: - compact_candidate_mask: (batch_size, max_len, max_len) - Output a mask tensor for candidates with a flattened and compact indexing. - """ - - batch_size = shape(draft_token_indices, 0, INT_DTYPE_STR) - max_len = shape(active_token_indices, 1, INT_DTYPE_STR) - all_candidate_len = beam_length * num_beams - - arange_all_candidates = constant( - np.arange(all_candidate_len, dtype=np.int32)) - active_token_beam = div(active_token_indices, beam_length) - beam_blocks = div(arange_all_candidates, beam_length) - - lower_triangle_mask = (unsqueeze(arange_all_candidates, axis=-1) - - unsqueeze(arange_all_candidates, axis=0) >= 0) - block_diagonal_mask = unsqueeze(beam_blocks, axis=-1) - unsqueeze( - beam_blocks, axis=0) == 0 - # `candidates_mask` is the flattened candidates mask - candidates_mask = expand( - expand_dims(op_and(lower_triangle_mask, block_diagonal_mask), [0]), - concat([batch_size, all_candidate_len, all_candidate_len]), - ) - - expanded_active_token_indices = expand( - expand_dims(active_token_indices, [2]), - concat([batch_size, max_len, all_candidate_len])) - raw_token_mask = gather(candidates_mask, - dim=1, - indices=expanded_active_token_indices) - - src_idx = unsqueeze(active_token_beam, axis=-1) * beam_length + expand_dims( - constant(np.arange(beam_length, dtype=np.int32)), [0, 1]) - src_mask = gather(raw_token_mask, dim=2, indices=src_idx) - tgt_idx = gather( - draft_token_indices, - dim=1, - indices=expand(expand_dims(active_token_beam, [2]), - concat([batch_size, max_len, beam_length])), - ) - # `compact_candidate_mask` is the compact and flattened candidates mask - compact_candidate_mask = expand( - expand_dims(cast(constant_to_tensor_(0), dtype="bool"), [0, 1]), - concat([batch_size, max_len, max_len]), - ) - - updated_compact_candidate_mask = scatter( - compact_candidate_mask, - dim=2, - indices=tgt_idx, - updates=src_mask, - ) - - return updated_compact_candidate_mask - - -def _beams2tree( - beams: Tensor, - num_beams: int, - beam_length: int, - position_ids_base: Tensor = None, -) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: - ''' - beams: [bs, nb, bl] - ''' - prefix_match_indices = _get_prefix_match_indices(beams, beam_length) - draft_token_indices = _get_draft_token_indices(prefix_match_indices, - num_beams, beam_length) - active_tokens_flattened, active_token_indices, total_lengths, max_gen_len, \ - total_gen_len, position_offsets, packed_position_ids = _get_draft_token_array( - beams, prefix_match_indices, num_beams, beam_length, position_ids_base) - mask = _get_mask(draft_token_indices, active_token_indices, num_beams, - beam_length) - return active_tokens_flattened, draft_token_indices, mask, position_offsets, packed_position_ids, total_lengths, max_gen_len, total_gen_len - - -def _get_indices_for_gather_beams(batch_size: Tensor, beam_indices: Tensor, - num_beams: int) -> Tensor: - ''' - beam_indices: [bs, nb] - Returns: [bs*nb, 2] - ''' - constant_0 = constant(int32_array(0)) - batch_indices = arange(constant_0, batch_size * num_beams, dtype='int32') - batch_indices = floordiv(batch_indices, num_beams) - - indices = concat([ - view(batch_indices, [-1, 1], False), - view(beam_indices, [-1, 1], False) - ], - dim=1) - return indices - - -def _gather_beams(x: Tensor, indices: Tensor, batch_size: Tensor, - num_beams: int) -> Tensor: - ''' - x: [bs, nb, X] - beam_indices: [bs, nb] - Returns: [bs, nb, X] - ''' - target_shp = [batch_size, constant(int32_array(num_beams))] - for i in range(2, x.ndim()): - target_shp.append(shape(x, i, INT_DTYPE_STR)) - target_shp = concat(target_shp) - return view(gather_nd(x, indices, batch_dims=0), target_shp, False) - - -def _add_decoding_dim(x: Tensor, num_beams: int) -> Tensor: - assert x.ndim() == 1 or x.ndim() == 2 - x = unsqueeze(x, 1) - new_shp = [shape(x, 0, INT_DTYPE_STR), num_beams] if x.ndim() == 2 else [ - shape(x, 0, INT_DTYPE_STR), num_beams, - shape(x, 2, INT_DTYPE_STR) - ] - res = expand(x, concat(new_shp)) - return res - - -def _flatten_decoding_dim(x: Tensor) -> Tensor: - if x.ndim() > 1: - new_shp = [-1 - ] + [shape(x, i, INT_DTYPE_STR) for i in range(2, x.ndim())] - return view(x, concat(new_shp)) - return x - - -def _unflatten_decoding_dim(x: Tensor, num_beams: int) -> Tensor: - ''' - Unflattens the first, flat batch*decoding dimension of a non-scalar array. - x: [bs*num_beams, ...] - ''' - if x.ndim() > 0: - new_shp = [-1, num_beams - ] + [shape(x, i, INT_DTYPE_STR) for i in range(1, x.ndim())] - return view(x, concat(new_shp)) - return x - - -def _beam_search_candidates(prompt_state: Tensor, init_token: Tensor, - embedding: Embedding, drafter: Module, - num_beams: int, beam_length: int, - is_rnn: bool) -> Tuple[Tensor, Tensor]: - """ - This version of beam search matches with ReDrafter GitHub version as of 10/02/2024. - Link: https://github.com/apple/ml-recurrent-drafter/releases/tag/v1.1 - """ - - LOG_0 = -50000.0 - LOG_1 = 0.0 - - def maintain_logits(logits: Tensor) -> Tensor: - max_logits = max(logits, -1, keepdim=True) - max_logits = expand(max_logits, - shape(logits, cast_to_dtype=INT_DTYPE_STR)) - return logits - max_logits - - def warp_logits(logits: Tensor, - top_k: int = 50, - mask_value: float = LOG_0) -> Tensor: - top_k = minimum(top_k, shape(logits, - dim=-1, - cast_to_dtype=INT_DTYPE_STR)) - top_values, _ = topk(logits, k=top_k, dim=-1) # [bs, nb, top_k] - starts = concat([0, 0, top_k - 1]) - sizes = concat([shape(logits, 0), shape(logits, 1), 1]) - lt_mask = logits < slice(top_values, starts=starts, sizes=sizes) - logits = where(lt_mask, - constant_to_tensor_(mask_value, dtype=logits.dtype), - logits) - return logits - - def compute_logits(x: Tensor) -> Tensor: - """ - x: [bs, nb, 2*H] - """ - logits = drafter(x) # [bs, nb, 2*H] => [bs, nb, V] - logits = maintain_logits(logits) # [bs, nb, V] - logits = warp_logits(logits) # [bs, nb, V] - return logits - - assert prompt_state.ndim() == 2 - assert init_token.ndim() == 1 - assert beam_length > 1 - batch_size = shape(prompt_state, 0, INT_DTYPE_STR) - vocab_size = embedding.num_embeddings - dtype = prompt_state.dtype - - log_p_beam = expand( - unsqueeze( - constant( - numpy_array([LOG_1] + [LOG_0] * (num_beams - 1), - trt_dtype=dtype)), 0), # [1, nb] - concat([batch_size, num_beams])) # [bs, nb] - context = _add_decoding_dim(prompt_state, num_beams) # [bs, nb, H] - if init_token.ndim() == 1: - init_token = unsqueeze(init_token, -1) # [bs] => [bs, 1] - beams = _add_decoding_dim(init_token, num_beams) # [bs, nb, 1] - - last_tokens = squeeze(beams, -1) # [bs, nb] - state_shape = shape(context, cast_to_dtype=INT_DTYPE_STR) # [bs, nb, H] - state = expand(expand_dims(constant_to_tensor_(0.0, dtype=dtype), [0, 1]), - state_shape) # [bs, nb, H] - log_p_token_in_beam = None - candidate_length = beam_length - 1 - for _ in range(candidate_length): - state = ( - silu(drafter.rnn_w(embedding(last_tokens)) + - drafter.rnn_u(state)) if is_rnn else embedding(last_tokens) + - state) # [bs, nb, H] - - logits_new_token = compute_logits(concat([context, state], - -1)) # [bs, nb, V] - log_p_new_token = log_softmax(logits_new_token, -1) # [bs, nb, V] - - log_p_beam_new_token = log_p_new_token + unsqueeze(log_p_beam, - 2) # [bs, nb, V] - - tokens_times_beams = view(log_p_beam_new_token, - concat([batch_size, num_beams * vocab_size - ])) # [bs, nb*V] - log_p_beam, topk_indices = topk(tokens_times_beams, k=num_beams, - dim=-1) # [bs, nb] - top_beam_indices = topk_indices // vocab_size # [bs, nb] - # Avoid repeated division for: top_token_ids = topk_indices % vocab_size - top_token_ids = topk_indices - (top_beam_indices * vocab_size - ) # [bs, nb] - - # get the common indices to gather beams - gather_indices = _get_indices_for_gather_beams(batch_size, - top_beam_indices, - num_beams) - - # update running beams, state, logits, and last_tokens - prev_top_beams = _gather_beams(beams, gather_indices, batch_size, - num_beams) # [bs, nb] OR [bs, nb, 1+i] - if prev_top_beams.ndim() == 2: - prev_top_beams = unsqueeze(prev_top_beams, -1) # [bs, nb, 1] - new_tokens = unsqueeze(top_token_ids, -1) # [bs, nb, 1] - beams = concat([prev_top_beams, new_tokens], dim=-1) # [bs, nb, 1+i+1] - - state = _gather_beams(state, gather_indices, batch_size, - num_beams) # [bs, nb, H] - - cur_log_p_token_in_beam = unsqueeze( - _gather_beams(log_p_new_token, gather_indices, batch_size, - num_beams), 2) # [bs, nb, 1, V] - if log_p_token_in_beam is None: # first iteration - log_p_token_in_beam = cur_log_p_token_in_beam - else: - log_p_token_in_beam = concat( - [ - _gather_beams(log_p_token_in_beam, gather_indices, - batch_size, - num_beams), # prev_top_logits [bs, nb, i, V] - cur_log_p_token_in_beam - ], - dim=2) # [bs, nb, i+1, V] - last_tokens = top_token_ids # [bs, nb] - return beams, log_p_token_in_beam - - -def _top_1_logits(logits: Tensor, NINF=-50000.0) -> Tensor: - ''' - logits: [bs, S, V] - ''' - NEG_INF = constant_to_tensor_(NINF, logits.dtype) - # TODO: WAR for bug in max reduction: https://nvbugs/4714485 - # max_values = max(logits, dim=-1, keepdim=True) # [bs, S, 1] - max_values, _ = topk(logits, k=1, dim=-1) # [bs, S, 1] - cmp = not_op(logits == max_values) - res = cast(cmp, dtype=logits.dtype) * NEG_INF - return res - - -def _ctx_logits2probs(logits: Tensor, greedy_search: bool) -> Tensor: - """ - Inputs: - logits: [bs_ctx, V] - Returns: - probs: [bs_ctx, V] - """ - if greedy_search: - logits = _top_1_logits(logits) - probs = softmax(logits, dim=-1) - return probs - - -# Jointly developed with Apple -def _batch_index_select(x: Tensor, batch_index: Tensor) -> Tensor: - """select the tensor by index inside each batch - - Args: - x (Tensor): [batch, ..] - batch_index (Tensor): (batch_size) - - Returns: - Tensor: [batch, ..] Tensors selected by the indices - """ - expanded_shape = concat( - [shape(x, 0, INT_DTYPE_STR), 1] + - [shape(x, i, INT_DTYPE_STR) for i in range(2, x.rank())]) - batch_index = expand( - expand_dims(batch_index, range(1, - x.rank() - batch_index.rank() + 1)), - expanded_shape) - gathered_x = gather(x, dim=1, indices=batch_index) - return squeeze(gathered_x, dim=1) - - -# Jointly developed with Apple -def _prepare_drafter_input( - draft_log_probs: Tensor, - base_log_probs: Tensor, - last_base_log_probs: Tensor, - accepted_beam_index: Tensor, - num_accepted_tokens: Tensor, -) -> Tensor: - """ - Args: - num_accepted_tokens: (batch_size) - Highest count of accepted tokens. - accepted_beam_index: (batch_size) - Beam index with highest count of accepted tokens. - draft_log_probs: (batch_size, num_candidates, candidate_length, vocab_size) - Draft head log probs for draft_tokens. - base_log_probs: (batch_size, num_candidates, candidate_length, vocab_size) - LM log probs for draft_tokens. - last_base_log_probs: (batch_size, num_candidates, vocab_size) - Last token log probs for all candidates to predict the next token beyond each candidate. - Returns: - probs: (batch_size, vocab_size): - Predict next token probability. - - """ - # Select according to the chosen beam index. - candidate_length = shape(draft_log_probs, 2, INT_DTYPE_STR) - selected_draft_log_probs = _batch_index_select(draft_log_probs, - accepted_beam_index) - selected_base_log_probs = _batch_index_select(base_log_probs, - accepted_beam_index) - selected_last_base_log_probs = _batch_index_select(last_base_log_probs, - accepted_beam_index) - - # Check if the entire beam is accepted or not. - entire_beam_accept = unsqueeze(num_accepted_tokens == candidate_length, - axis=-1) - - # If the entire beam is accepted, we use maybe_last_probs to sample next token. - maybe_last_probs = exp(selected_last_base_log_probs) - - # Note the shape of selected_draft_log_probs and selected_base_log_probs is the same - # as [batch_size, candidate_length, vocab_size]. - # Thus, we clamp resample_index to be up to candidate_length - 1. - # Since when num_accepted_tokens == candidate_length, we use maybe_last_probs above. - resample_index = num_accepted_tokens - cast( - eq(num_accepted_tokens, candidate_length), dtype='int32') - sample_draft_log_probs = _batch_index_select(selected_draft_log_probs, - resample_index) - sample_base_log_probs = _batch_index_select(selected_base_log_probs, - resample_index) - # Rejection sampling probs. - probs = relu(exp(sample_base_log_probs) - exp(sample_draft_log_probs)) - probs = where(entire_beam_accept, maybe_last_probs, probs) - - return probs - - -def _process_gen_logits(logits: Tensor, - hidden: Tensor, - draft_probs: Tensor, - draft_tokens: Tensor, - draft_indices: Tensor, - num_beams: int, - beam_length: int, - greedy_search: bool, - rand_data: Tensor = None) -> Tensor: - num_accepted_tokens, accepted_beam_index,\ - base_log_probs, last_base_log_probs, _ = _validate_draft_tokens( - draft_probs, draft_tokens, draft_indices, logits, num_beams, beam_length, - greedy_search, rand_data) - - # need to retrieve flattened index from accepted_beam_index and num_accepted_tokens - indices = stack([accepted_beam_index, num_accepted_tokens], 1) - flat_indices = unsqueeze(gather_nd(draft_indices, indices, batch_dims=1), - -1) - filtered_probs = _prepare_drafter_input(draft_probs, base_log_probs, - last_base_log_probs, - accepted_beam_index, - num_accepted_tokens) - filtered_hidden = gather_nd(hidden, flat_indices, batch_dims=1) - return filtered_probs, filtered_hidden, num_accepted_tokens, accepted_beam_index - - -def _get_gen_token_indices_for_unpack( - num_gen_tokens: Tensor, num_beams: int, beam_length: int, - max_index_allowed: Tensor) -> Tuple[Tensor, Tensor]: - upper_bound = num_beams * beam_length - num_beams + 1 - max_gen_tokens = max(num_gen_tokens, dim=0) - max_gen_tokens = minimum(max_gen_tokens, upper_bound) - max_gen_tokens = maximum(max_gen_tokens, 0) - cum_gen_tokens = cumsum(num_gen_tokens, 0) - gen_token_starts = cum_gen_tokens - num_gen_tokens - gen_unpack_indxs = arange(constant_to_tensor_(0, to_array=False), - max_gen_tokens, - dtype='int32') - gen_unpack_indxs = unsqueeze(gen_unpack_indxs, 0) + unsqueeze( - gen_token_starts, 1) - gen_unpack_indxs = minimum(gen_unpack_indxs, max_index_allowed) - return gen_unpack_indxs, max_gen_tokens - - -def _unpack_gen_data(x: Tensor, num_gen_tokens: Tensor, - gen_unpack_indxs: Tensor, - max_gen_tokens: Tensor) -> Tensor: - """ - x: [sum(num_gen_tokens), V/H] - num_gen_tokens: [gen_bs] - gen_unpack_indxs: [bs, max(num_gen_tokens)] - Returns: - [gen_bs, max_gen_tokens, V/H] where max_gen_tokens = max(num_gen_tokens) - """ - unpacked_x = index_select(x, dim=0, index=view(gen_unpack_indxs, [-1])) - out_shape = concat([ - shape(num_gen_tokens, 0, INT_DTYPE_STR), max_gen_tokens, - shape(x, -1, INT_DTYPE_STR) - ]) - return unpacked_x.view(out_shape, zero_is_placeholder=False) - - -def _process_logits_and_hidden_states( - model: Module, logits: Tensor, hidden_states: Tensor, - kwargs: dict) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - """ - Process the logits and hidden_states correctly. - For logits: - Can be all context, all gen or mixed. - For all context-phase: - the shape is [bs, V], just process to probs - For all gen-phase: - the shape is [sum(num_gen_tokens), V] - gather using num_gen_tokens => [gen_bs, max_gen_tokens, V] - then typical processing as above - For mixed case: - split the logits, do both ctx and gen phase processing - For hidden_states: - context phase: similar processing - gen-phase: filter based on accepted beams and their lengths. - """ - if model is not None: - num_beams = model.num_beams - beam_length = model.beam_length - greedy_search = model.greedy_search - else: - num_beams = kwargs['num_beams'] - beam_length = kwargs['beam_length'] - greedy_search = kwargs.get('greedy_search', False) - device_request_types = kwargs['device_request_types'] - inverted_temperature = kwargs['redrafter_inverted_temperature'] # [bs] - num_gen_tokens = kwargs[ - 'spec_decoding_params'].spec_decoding_generation_lengths - assert default_net( - ).plugin_config.remove_input_padding, "ReDrafter is only supported without input padding." - """ - Split the flattened data: context and generation - Process them separately. - NOTE: Involves processing 0-shaped tensors (if all context or all generation) - """ - # process context - const_0 = constant_to_tensor_(0, to_array=False) - bs = shape(device_request_types, 0, INT_DTYPE_STR) - num_gen = sum(device_request_types, -1) - num_gen = maximum(constant_to_tensor_(0, to_array=False), num_gen) - num_gen = minimum(bs, num_gen) - bs_ctx = bs - num_gen - ctx_idxs = arange(const_0, bs_ctx, dtype='int32') - assert bs_ctx.rank() == 0 - ctx_logits = index_select(logits, dim=0, index=ctx_idxs) - if not greedy_search: - ctx_temperature = index_select(inverted_temperature, - dim=0, - index=ctx_idxs) - ctx_temperature = unsqueeze(ctx_temperature, 1) - ctx_logits = ctx_logits * ctx_temperature - ctx_probs = _ctx_logits2probs(ctx_logits, greedy_search) - ctx_hidden_states = index_select(hidden_states, dim=0, index=ctx_idxs) - # we accept zero draft tokens for ctx-phase - ctx_num_accepted = expand(constant_to_tensor_(0), unsqueeze(bs_ctx, 0)) - ctx_accepted_beam_index = expand(constant_to_tensor_(0), - unsqueeze(bs_ctx, 0)) - - # process generation - # get the logits[bs_ctx:, :] and hidden_states[bs_ctx:, :] - gen_token_idxs = arange(bs_ctx, - shape(logits, 0, INT_DTYPE_STR), - dtype='int32') - gen_logits = index_select(logits, dim=0, index=gen_token_idxs) - gen_hidden = index_select(hidden_states, dim=0, index=gen_token_idxs) - max_index_allowed = shape(gen_logits, 0, INT_DTYPE_STR) - 1 - gen_unpack_idxs, max_gen_tokens = _get_gen_token_indices_for_unpack( - num_gen_tokens, num_beams, beam_length, max_index_allowed) - gen_logits = _unpack_gen_data(gen_logits, num_gen_tokens, gen_unpack_idxs, - max_gen_tokens) - if not greedy_search: - gen_temperature = index_select(inverted_temperature, - dim=0, - index=gen_token_idxs) - gen_temperature = expand_dims(gen_temperature, dim=[1, 2]) - expanded_gen_temperature = expand(gen_temperature, shape(gen_logits)) - gen_logits = gen_logits * expanded_gen_temperature - gen_hidden = _unpack_gen_data(gen_hidden, num_gen_tokens, gen_unpack_idxs, - max_gen_tokens) - - # verify the input draft tokens (from last step) using the gen_logits - gen_probs, gen_hidden_states, gen_num_accepted, gen_accepted_beam_index\ - = _process_gen_logits( - gen_logits, gen_hidden, kwargs['draft_probs'], - kwargs['draft_tokens'], kwargs['draft_indices'], - num_beams, beam_length, greedy_search, - kwargs.get('rand_data_validation', None) - ) - - # combine ctx and gen phase outputs - probs = concat([ctx_probs, gen_probs], dim=0) - drafter_input = concat([ctx_hidden_states, gen_hidden_states], dim=0) - num_accepted_tokens = concat([ctx_num_accepted, gen_num_accepted], dim=0) - accepted_beam_index = concat( - [ctx_accepted_beam_index, gen_accepted_beam_index], dim=0) - - # NOTE: This is needed with shape inference of data-dependent tensors - bs = shape(device_request_types, 0, INT_DTYPE_STR) - const_0 = constant_to_tensor_(0, to_array=False) - bidxs = arange(const_0, bs, dtype='int32') - probs = index_select(probs, dim=0, index=bidxs) - drafter_input = index_select(drafter_input, dim=0, index=bidxs) - num_accepted_tokens = index_select(num_accepted_tokens, dim=0, index=bidxs) - accepted_beam_index = index_select(accepted_beam_index, dim=0, index=bidxs) - return probs, drafter_input, num_accepted_tokens, accepted_beam_index diff --git a/tensorrt_llm/models/stdit/__init__.py b/tensorrt_llm/models/stdit/__init__.py deleted file mode 100644 index 0d106f45ce46..000000000000 --- a/tensorrt_llm/models/stdit/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. diff --git a/tensorrt_llm/models/stdit/config.py b/tensorrt_llm/models/stdit/config.py deleted file mode 100644 index ca16dbdfb490..000000000000 --- a/tensorrt_llm/models/stdit/config.py +++ /dev/null @@ -1,115 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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 Any, Dict, Optional, Sequence - -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class STDiTModelConfig(PretrainedConfig): - - def __init__(self, - architecture: str = 'STDiT3', - checkpoint_path: str = 'pretrained_ckpt/model.safetensors', - vae_type: str = "hpcai-tech/OpenSora-VAE-v1.2", - text_encoder_type: str = "DeepFloyd/t5-v1_1-xxl", - caption_channels: int = 4096, - num_hidden_layers: int = 28, - hidden_size: int = 1152, - width: int = 640, - height: int = 360, - num_frames: int = 102, - latent_size: Sequence[int] = [30, 45, 80], - stdit_patch_size: Sequence[int] = [1, 2, 2], - spatial_patch_size: Sequence[int] = [1, 8, 8], - temporal_patch_size: Sequence[int] = [4, 1, 1], - in_channels: int = 4, - input_sq_size: int = 512, - num_attention_heads: int = 16, - mlp_ratio: float = 4.0, - class_dropout_prob: float = 0.1, - model_max_length: int = 300, - learn_sigma: bool = True, - qk_norm: bool = True, - skip_y_embedder: bool = False, - dtype: Optional[str] = None, - mapping: Mapping = Mapping(), - quant_config: Optional[QuantConfig] = None, - **kwargs): - kwargs.update({ - 'architecture': architecture, - 'num_hidden_layers': num_hidden_layers, - 'num_attention_heads': num_attention_heads, - 'hidden_size': hidden_size, - 'dtype': dtype - }) - - super().__init__(**kwargs) - self.checkpoint_path = checkpoint_path - self.vae_type = vae_type - self.text_encoder_type = text_encoder_type - self.caption_channels = caption_channels - self.width = width - self.height = height - self.num_frames = num_frames - self.latent_size = latent_size - self.stdit_patch_size = stdit_patch_size - self.spatial_patch_size = spatial_patch_size - self.temporal_patch_size = temporal_patch_size - self.in_channels = in_channels - self.input_sq_size = input_sq_size - self.mlp_ratio = mlp_ratio - self.class_dropout_prob = class_dropout_prob - self.model_max_length = model_max_length - self.learn_sigma = learn_sigma - self.qk_norm = qk_norm - self.skip_y_embedder = skip_y_embedder - self.mapping = mapping - self.quant_config = quant_config - - @classmethod - def from_input_config(cls, - input_config: Dict[str, Any], - dtype: str = 'auto', - mapping: Mapping = Mapping(), - quant_config: Optional[QuantConfig] = None, - **kwargs): - return cls(architecture=input_config['architecture'], - checkpoint_path=input_config['checkpoint_path'], - vae_type=input_config['vae_type'], - text_encoder_type=input_config['text_encoder_type'], - caption_channels=input_config['caption_channels'], - num_hidden_layers=input_config['num_hidden_layers'], - width=input_config['width'], - height=input_config['height'], - num_frames=input_config['num_frames'], - latent_size=input_config['latent_size'], - hidden_size=input_config['hidden_size'], - stdit_patch_size=input_config['stdit_patch_size'], - spatial_patch_size=input_config['spatial_patch_size'], - temporal_patch_size=input_config['temporal_patch_size'], - in_channels=input_config['in_channels'], - input_sq_size=input_config['input_sq_size'], - num_attention_heads=input_config['num_attention_heads'], - mlp_ratio=input_config['mlp_ratio'], - class_dropout_prob=input_config['class_dropout_prob'], - model_max_length=input_config['model_max_length'], - learn_sigma=input_config['learn_sigma'], - qk_norm=input_config['qk_norm'], - skip_y_embedder=input_config['skip_y_embedder'], - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/stdit/model.py b/tensorrt_llm/models/stdit/model.py deleted file mode 100644 index 938f80186842..000000000000 --- a/tensorrt_llm/models/stdit/model.py +++ /dev/null @@ -1,1624 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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 builtins -import collections -import functools -import json -import math -import os -import re -from collections import OrderedDict -from typing import Optional - -import numpy as np -import tensorrt as trt -import torch -from tqdm import tqdm - -import tensorrt_llm -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import str_dtype_to_trt, trt_dtype_to_str -from tensorrt_llm.functional import (ACT2FN, AttentionMaskType, LayerNormType, - PositionEmbeddingType, Tensor, - constant_to_tensor_) -from tensorrt_llm.layers import (ColumnLinear, Conv3d, LayerNorm, Linear, - RowLinear) -from tensorrt_llm.layers.attention import (Attention, AttentionParams, - BertAttention, KeyValueCacheParams, - bert_attention, layernorm_map) -from tensorrt_llm.layers.normalization import RmsNorm -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.generation_mixin import GenerationMixin -from tensorrt_llm.models.model_weights_loader import (ModelWeightsFormat, - ModelWeightsLoader) -from tensorrt_llm.models.modeling_utils import PretrainedConfig, PretrainedModel -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.parameter import Parameter -from tensorrt_llm.plugin import current_all_reduce_helper -from tensorrt_llm.quantization import QuantMode - -from ...functional import (allgather, arange, cast, chunk, concat, constant, - cos, div, einsum, exp, expand, expand_dims, - expand_mask, masked_select, matmul, meshgrid2d, pad, - permute, pow, rearrange, repeat, repeat_interleave, - rms_norm, shape, sin, slice, softmax, split, squeeze, - stack, sum, unsqueeze, where) -from .config import STDiTModelConfig - -# [TODO] For now, we only support static shape, which might contains `-1` when inputs are with dynamic shape. -USE_STATIC_SHAPE = True - - -# From PyTorch internals -def _ntuple(n): - - def parse(x): - if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): - return tuple(x) - return tuple([x] * n) - - return parse - - -to_1tuple = _ntuple(1) -to_2tuple = _ntuple(2) -to_3tuple = _ntuple(3) -to_4tuple = _ntuple(4) -to_ntuple = _ntuple - - -# [TODO] make constant `1` compatible with `scale` -def t2i_modulate(x, shift, scale): - return x * (1.0 + scale) + shift - - -class ModuleSequential(ModuleList): - - def __init__(self, modules) -> None: - super(ModuleSequential, self).__init__(modules=modules) - - def forward(self, *args, **kwargs): - module = self.__getitem__(0) - outputs = module(*args, **kwargs) - for idx in range(1, len(self._modules)): - module = self.__getitem__(idx) - outputs = module(outputs) - return outputs - - -class Activation(Module): - - def __init__(self, act_fn='silu'): - super().__init__() - self.act_fn = act_fn - - def forward(self, input: Tensor): - return ACT2FN[self.act_fn](input) - - -class RotaryEmbedder(Module): - - def __init__(self, - dim, - theta=10000, - interpolate_factor=1., - theta_rescale_factor=1., - seq_before_head_dim=False, - cache_if_possible=True, - use_xpos=False, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - theta *= theta_rescale_factor**(dim / (dim - 2)) - freqs = 1. / (theta - **(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) - self.freqs = Parameter(freqs, dtype=dtype) - self.cached_freqs = None - self.seq_before_head_dim = seq_before_head_dim - self.default_seq_dim = -3 if seq_before_head_dim else -2 - self.scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) - assert interpolate_factor >= 1. - self.interpolate_factor = interpolate_factor - - self.cache_if_possible = cache_if_possible - self.use_xpos = use_xpos - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def get_freqs(self, - t: Tensor, - seq_len: Optional[int] = None, - offset: int = 0): - should_cache = self.cache_if_possible and seq_len is not None - if should_cache and isinstance(self.cached_freqs, Tensor): - if (offset + seq_len) <= self.cached_freqs.shape[0]: - return slice(self.cached_freqs, - starts=[offset] + - [0] * len(self.cached_freqs.shape[1:]), - sizes=[seq_len, *self.cached_freqs.shape[1:]]) - freqs = self.freqs.value - freqs = unsqueeze(t, axis=-1) * unsqueeze(freqs, axis=0) - freqs = repeat_interleave(freqs, repeats=2, dim=(freqs.ndim() - 1)) - if should_cache: - self.cached_freqs = freqs - return freqs - - def get_seq_pos(self, seq_len: int, dtype: trt.DataType, offset: int = 0): - return (arange(start=0, end=seq_len, dtype=trt_dtype_to_str(dtype)) + - offset) / self.interpolate_factor - - def rotate_half(self, x: Tensor): - x = x.view([*x.shape[:-1], x.shape[-1] // 2, 2]) - x1, x2 = x.unbind(x.ndim() - 1) - x = stack([-1 * x2, x1], dim=-1) - x = x.view([*x.shape[:-2], x.shape[-2] * x.shape[-1]]) - return x - - def apply_rotary_emb(self, - freqs: Tensor, - t: Tensor, - start_index: int = 0, - scale: int = 1., - seq_dim: int = -2): - if t.ndim() == 3: - seq_len = t.shape[seq_dim] - # freqs = freqs[-seq_len:] - freqs = slice(starts=[freqs.shape[0] - seq_len], sizes=[seq_len]) - rot_dim = freqs.shape[-1] - end_index = start_index + rot_dim - assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size ' + \ - 'to rotate in all the positions {rot_dim}' - t_left = slice(t, - starts=[0] * t.ndim(), - sizes=[*t.shape[:-1], start_index]) - t_right = slice(t, - starts=[0] * (t.ndim() - 1) + [end_index], - sizes=[*t.shape[:-1], t.shape[-1] - end_index]) - - t = (t * cos(freqs) * scale) + (self.rotate_half(t) * sin(freqs) * - scale) - return concat([t_left, t, t_right], dim=-1) - - def rotate_queries_or_keys(self, - t: Tensor, - seq_dim: Optional[int] = None, - offset: int = 0, - freq_seq_len: Optional[int] = None): - seq_dim = self.default_seq_dim if seq_dim is None else seq_dim - assert not self.use_xpos, 'you must use `.rotate_queries_and_keys` method ' + \ - 'instead and pass in both queries and keys, for ' + \ - 'length extrapolatable rotary embeddings' - - seq_len = t.shape[seq_dim] - if freq_seq_len is not None: - assert freq_seq_len >= seq_len - seq_len = freq_seq_len - - freqs = self.get_freqs(self.get_seq_pos(seq_len, - dtype=t.dtype, - offset=offset), - seq_len=seq_len, - offset=offset) - - if seq_dim == -3: - freqs = rearrange(freqs, 'n d -> n 1 d') - rope_output = self.apply_rotary_emb(freqs, t, seq_dim=seq_dim) - return rope_output - - -class STDiTRmsNorm(RmsNorm): - - def __init__(self, - normalized_shape, - num_groups=1, - eps=1e-06, - elementwise_affine=True, - dtype=None): - super().__init__(normalized_shape, num_groups, eps, elementwise_affine, - dtype) - - def forward(self, hidden_states): - weight = None if self.weight is None else self.weight.value - return rms_norm(input=hidden_states, - normalized_shape=self.normalized_shape, - num_groups=self.num_groups, - weight=weight, - eps=self.eps) - - -class STDiTAttention(BertAttention): - - def __init__(self, - hidden_size, - num_attention_heads, - qk_layernorm=True, - layernorm_type=LayerNormType.RmsNorm, - layernorm_eps=1e-06, - bias=True, - rotary_embedding_func=None, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - cp_group=None, - cp_size=1, - quant_mode: QuantMode = QuantMode(0)): - assert hidden_size % num_attention_heads == 0, "hidden_size should be divisible by num_attention_heads" - super().__init__(hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - cp_group=cp_group, - cp_size=cp_size, - quant_mode=quant_mode) - - self.qk_layernorm = qk_layernorm - if self.qk_layernorm: - ln_type = layernorm_map[layernorm_type] - self.q_layernorm = ln_type(self.attention_head_size, - eps=layernorm_eps, - dtype=dtype) - self.k_layernorm = ln_type(self.attention_head_size, - eps=layernorm_eps, - dtype=dtype) - self.rotary_embedding_func = rotary_embedding_func - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - input_lengths=None, - max_input_length: int = None): - - assert isinstance(hidden_states, Tensor) - - B = shape(hidden_states, 0) - N = shape(hidden_states, 1) - C = shape(hidden_states, 2) - input_lengths = expand(unsqueeze(N, 0).cast('int32'), unsqueeze(B, 0)) - - assert (self.qkv is not None) - qkv = self.qkv(hidden_states) - - kv_size = self.attention_head_size * self.num_attention_kv_heads - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - query = query.view( - concat([B, N, self.num_attention_heads, - self.attention_head_size])).permute(dims=[0, 2, 1, 3]) - key = key.view( - concat( - [B, N, self.num_attention_kv_heads, - self.attention_head_size])).permute(dims=[0, 2, 1, 3]) - - if self.qk_layernorm: - query = self.q_layernorm(query) - key = self.k_layernorm(key) - - if self.rotary_embedding_func is not None: - query = self.rotary_embedding_func(query) - key = self.rotary_embedding_func(key) - - # TODO deal with qkv - query = query.permute(dims=[0, 2, 1, 3]).view( - concat([B, N, self.attention_hidden_size])) - key = key.permute(dims=[0, 2, 1, 3]).view(concat([B, N, kv_size])) - qkv = concat([query, key, value], dim=2) - - if default_net().plugin_config.bert_attention_plugin: - # TRT plugin mode - assert input_lengths is not None - assert self.cp_size == 1 - if default_net().plugin_config.remove_input_padding: - qkv = qkv.view( - concat([-1, self.attention_hidden_size + 2 * kv_size])) - max_input_length = constant( - np.zeros([ - max_input_length, - ], dtype=np.int32)) - context = bert_attention(qkv, - input_lengths, - self.num_attention_heads, - self.attention_head_size, - q_scaling=self.q_scaling, - max_distance=self.max_distance, - max_input_length=max_input_length) - else: - # plain TRT mode - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - kv_size = self.attention_head_size * self.num_attention_kv_heads - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - if self.cp_size > 1 and self.cp_group is not None: - key = allgather(key, self.cp_group, gather_dim=1) - value = allgather(value, self.cp_group, gather_dim=1) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - key = key.permute([0, 1, 3, 2]) - attention_scores = matmul(query, key, use_fp32_acc=False) - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - - if attention_mask is not None: - attention_mask = expand_mask(attention_mask, shape(query, 2)) - attention_mask = cast(attention_mask, attention_scores.dtype) - attention_scores = attention_scores + attention_mask - - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([ - shape(context, 0), - shape(context, 1), self.attention_hidden_size - ])) - - context = self.dense(context) - context = context.view(concat([B, N, C])) - return context - - -class STDiTCrossAttention(Attention): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - attention_mask_type=AttentionMaskType.causal, - qkv_bias=True, - dense_bias=True, - position_embedding_type=PositionEmbeddingType.learned_absolute, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - cp_group=[0], - cp_size=1, - cp_rank=0, - quant_mode: QuantMode = QuantMode(0)): - assert hidden_size % num_attention_heads == 0, "hidden_size should be divisible by num_attention_heads" - super().__init__(local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_mask_type=attention_mask_type, - bias=qkv_bias, - dense_bias=dense_bias, - cross_attention=True, - position_embedding_type=position_embedding_type, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - cp_group=cp_group, - cp_size=cp_size, - cp_rank=cp_rank, - quant_mode=quant_mode) - - def forward(self, - hidden_states: Tensor, - encoder_output: Tensor, - use_cache=False, - attention_params: Optional[AttentionParams] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None): - bs = shape(encoder_output, 0) - encoder_input_length = shape(encoder_output, 1) - encoder_hidden_size = shape(encoder_output, 2) - encoder_output = encoder_output.view( - concat([bs * 2, encoder_input_length // 2, encoder_hidden_size])) - - if default_net().plugin_config.remove_input_padding: - B = shape(hidden_states, 0) - N = shape(hidden_states, 1) - C = shape(hidden_states, 2) - hidden_states = hidden_states.view(concat([B * N, C])) - encoder_output = encoder_output.view( - concat([-1, encoder_hidden_size])) - - context = super().forward(hidden_states=hidden_states, - encoder_output=encoder_output, - use_cache=use_cache, - attention_params=attention_params, - kv_cache_params=kv_cache_params) - context = context.view(concat([B, -1, C])) - return context - - -class T2IFinalLayer(Module): - - def __init__(self, - hidden_size, - num_patch, - out_channels, - d_t=None, - d_s=None, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.norm_final = LayerNorm(hidden_size, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - self.linear = Linear(hidden_size, - num_patch * out_channels, - bias=True, - dtype=dtype) - self.scale_shift_table = Parameter(torch.randn(2, hidden_size) / - hidden_size**0.5, - dtype=dtype) - self.out_channels = out_channels - self.d_t = d_t - self.d_s = d_s - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def t_mask_select(self, x_mask, x, masked_x, T: int, S: int): - # x: [B, (T, S), C] - # mased_x: [B, (T, S), C] - # x_mask: [B, T] - x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S) - masked_x = rearrange(masked_x, "B (T S) C -> B T S C", T=T, S=S) - x = where(expand_dims(x_mask, - [x_mask.ndim(), x_mask.ndim() + 1]), x, masked_x) - x = rearrange(x, "B T S C -> B (T S) C") - return x - - def forward(self, - x, - t, - x_mask=None, - t0=None, - T: Optional[int] = None, - S: Optional[int] = None): - if T is None: - T = self.d_t - if S is None: - S = self.d_s - shift, scale = chunk(expand_dims(self.scale_shift_table.value, 0) + - expand_dims(t, 1), - chunks=2, - dim=1) - x = t2i_modulate(self.norm_final(x), shift, scale) - if x_mask is not None: - shift_zero, scale_zero = chunk( - expand_dims(self.scale_shift_table.value, 0) + - expand_dims(t0, 1), - chunks=2, - dim=1) - x_zero = t2i_modulate(self.norm_final(x), shift_zero, scale_zero) - x = self.t_mask_select(x_mask, x, x_zero, T, S) - x = self.linear(x) - self.register_network_output('output', x) - return x - - -class PositionEmbedding2D(Module): - - def __init__(self, - dim: int, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.dim = dim - assert dim % 4 == 0, "dim must be divisible by 4" - half_dim = dim // 2 - self.inv_freq = Parameter( - 1.0 / (10000**(torch.arange(0, half_dim, 2).float() / half_dim)), - is_buffer=True, - dtype=dtype) - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def _get_sin_cos_emb(self, t): - out = einsum("i,d->id", [t, self.inv_freq.value]) - emb_cos = cos(out) - emb_sin = sin(out) - return concat([emb_sin, emb_cos], dim=-1) - - @functools.lru_cache(maxsize=512) - def _get_cached_emb( - self, - dtype, - h: int, - w: int, - scale: Tensor, - base_size: Optional[int] = None, - ): - grid_h = div(arange(0, h, 'float32'), scale.cast('float32')) - grid_w = div(arange(0, w, 'float32'), scale.cast('float32')) - if base_size is not None: - grid_h *= float(base_size) / h - grid_w *= float(base_size) / w - grid_h, grid_w = meshgrid2d(grid_w, grid_h) # here w goes first - grid_h = permute(grid_h, [1, 0]).flatten() - grid_w = permute(grid_w, [1, 0]).flatten() - emb_h = self._get_sin_cos_emb(grid_h) - emb_w = self._get_sin_cos_emb(grid_w) - return unsqueeze(concat([emb_h, emb_w], dim=-1), 0).cast(dtype) - - def forward(self, x, h: int, w: int, scale: Tensor, base_size=None): - pos_embedding = self._get_cached_emb(x.dtype, h, w, scale, base_size) - self.register_network_output('output', pos_embedding) - return pos_embedding - - -class TimestepEmbedder(Module): - - def __init__(self, - hidden_size, - frequency_embedding_size=256, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.mlp = ModuleSequential([ - Linear(frequency_embedding_size, - hidden_size, - bias=True, - dtype=dtype), - Activation('silu'), - Linear(hidden_size, hidden_size, bias=True, dtype=dtype) - ]) - self.frequency_embedding_size = frequency_embedding_size - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - @staticmethod - def timestep_embedding(t, dim, max_period=10000, dtype=None): - half = dim // 2 - freqs = exp( - -math.log(max_period) * - arange(start=0, end=half, dtype=trt_dtype_to_str(trt.float32)) / - constant(np.array([half], dtype=np.float32))) - args = unsqueeze(t, -1).cast(trt.float32) * unsqueeze(freqs, 0) - embedding = concat([cos(args), sin(args)], dim=-1) - if dtype is not None: - embedding = embedding.cast(dtype) - if dim % 2: - embedding = pad(embedding, (0, 0, 0, 1)) - return embedding - - def forward(self, t, dtype): - t_freq = self.timestep_embedding( - t, self.frequency_embedding_size).cast(dtype) - t_emb = self.mlp(t_freq) - return t_emb - - -class SizeEmbedder(TimestepEmbedder): - - def __init__(self, - hidden_size, - frequency_embedding_size=256, - dtype=str_dtype_to_trt("float16"), - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__(hidden_size=hidden_size, - frequency_embedding_size=frequency_embedding_size, - dtype=dtype, - mapping=mapping, - quant_mode=quant_mode) - self.mlp = ModuleSequential([ - Linear(frequency_embedding_size, - hidden_size, - bias=True, - dtype=dtype), - Activation('silu'), - Linear(hidden_size, hidden_size, bias=True, dtype=dtype) - ]) - self.outdim = hidden_size - - def forward(self, s, bs: int): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - if s.ndim() == 1: - s = unsqueeze(s, 1) - assert s.ndim() == 2 - if s.shape[0] != bs: - s = repeat(s, [bs // s.shape[0], 1]) - assert s.shape[0] == bs - b, dims = s.shape[0], s.shape[1] - s = rearrange(s, "b d -> (b d)") - s_freq = self.timestep_embedding(s, self.frequency_embedding_size).cast( - self.dtype) - s_emb = self.mlp(s_freq) - s_emb = rearrange(s_emb, - "(b d) d2 -> b (d d2)", - b=b, - d=dims, - d2=self.outdim) - self.register_network_output('output', s_emb) - return s_emb - - -class CaptionMLP(Module): - - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer="gelu", - bias=True, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - ): - super().__init__() - hidden_act = act_layer - out_features = out_features or in_features - hidden_features = hidden_features or in_features - bias = to_2tuple(bias) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * hidden_features if hidden_act in [ - 'swiglu', 'gegelu' - ] else hidden_features - self.inner_layernorm = LayerNorm(hidden_features, dtype=dtype, - eps=eps) if inner_layernorm else None - - self.fc1 = ColumnLinear(in_features, - fc_output_size, - bias=bias[0], - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=False) - self.fc2 = RowLinear(hidden_features, - out_features, - bias=bias[1], - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - - self.in_features = in_features - self.hidden_features = hidden_features - self.out_features = out_features - self.hidden_act = hidden_act - self.dtype = dtype - self.bias = bias - self.mapping = mapping - self.quant_mode = quant_mode - self.eps = eps - - def forward(self, hidden_states, gegelu_limit=None): - inter = self.fc1(hidden_states) - if self.hidden_act == 'gegelu': - inter = ACT2FN[self.hidden_act](inter, gegelu_limit) - else: - inter = ACT2FN[self.hidden_act](inter) - if self.inner_layernorm is not None: - inter = self.inner_layernorm(inter) - output = self.fc2(inter) - return output - - -class CaptionEmbedder(Module): - - def __init__(self, - in_channels, - hidden_size, - uncond_prob, - act_layer='gelu', - token_num=120, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.y_proj = CaptionMLP( - in_features=in_channels, - hidden_features=hidden_size, - out_features=hidden_size, - act_layer=act_layer, - mapping=mapping, - dtype=dtype, - ) - self.y_embedding = Parameter(torch.randn(token_num, in_channels) / - in_channels**0.5, - dtype=dtype) - self.uncond_prob = uncond_prob - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def token_drop(self, caption, force_drop_ids=None): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - assert (isinstance(force_drop_ids, torch.Tensor) - or isinstance(force_drop_ids, np.array)) - if force_drop_ids is None: - drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob - else: - drop_ids = torch.Tensor(force_drop_ids) == 1 - drop_ids = constant(drop_ids.cpu().numpy()) - caption = where(expand_dims(drop_ids, [1, 2, 3]), - self.y_embedding.value, caption) - return caption - - def forward(self, caption, force_drop_ids=None): - if force_drop_ids is not None: - caption = self.token_drop(caption, force_drop_ids) - caption = self.y_proj(caption) - self.register_network_output('output', caption) - return caption - - -class PatchEmbed3D(Module): - - def __init__(self, - patch_size=(2, 4, 4), - in_chans=3, - embed_dim=96, - norm_layer=None, - flatten=True, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.patch_size = patch_size - self.flatten = flatten - self.in_chans = in_chans - self.embed_dim = embed_dim - - self.proj = Conv3d(in_chans, - embed_dim, - kernel_size=patch_size, - stride=patch_size, - dtype=dtype) - if norm_layer is not None: - self.norm = norm_layer(embed_dim) - else: - self.norm = None - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def forward(self, x): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - _, _, D, H, W = x.shape - if W % self.patch_size[2] != 0: - x = pad(x, (0, self.patch_size[2] - W % self.patch_size[2])) - if H % self.patch_size[1] != 0: - x = pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1])) - if D % self.patch_size[0] != 0: - x = pad( - x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0])) - x = self.proj(x) # (B C T H W) - if self.norm is not None: - D = shape(x, 2) - Wh = shape(x, 3) - Ww = shape(x, 4) - x = x.flatten(2).transpose(1, 2) - x = self.norm(x) - x = x.transpose(1, 2).view([-1, self.embed_dim, D, Wh, Ww]) - if self.flatten: - x = x.flatten(2).transpose(1, 2) # BCTHW -> BNC - self.register_network_output('output', x) - return x - - -class STDiT3Block(Module): - - def __init__(self, - hidden_size, - num_heads, - mlp_ratio=4.0, - rope=None, - qk_norm=False, - temporal=False, - dtype=None, - local_layer_idx=0, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.temporal = temporal - self.hidden_size = hidden_size - - attn_cls = STDiTAttention - mha_cls = STDiTCrossAttention - - self.norm1 = LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=False, - dtype=dtype) - self.attn = attn_cls(hidden_size=hidden_size, - num_attention_heads=num_heads, - qk_layernorm=qk_norm, - bias=True, - rotary_embedding_func=rope, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - cp_group=mapping.cp_group, - cp_size=mapping.cp_size, - quant_mode=quant_mode) - self.cross_attn = mha_cls( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_heads, - attention_mask_type=tensorrt_llm.layers.AttentionMaskType.causal, - qkv_bias=True, - dense_bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - cp_group=mapping.cp_group, - cp_size=mapping.cp_size, - quant_mode=quant_mode) - self.norm2 = LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=False, - dtype=dtype) - self.mlp = CaptionMLP( - in_features=hidden_size, - hidden_features=int(hidden_size * mlp_ratio), - act_layer='gelu', - mapping=mapping, - dtype=dtype, - ) - self.scale_shift_table = Parameter(torch.randn(6, hidden_size) / - hidden_size**0.5, - dtype=dtype) - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def t_mask_select(self, x_mask, x, masked_x, T: int, S: int): - # x: [B, (T, S), C] - # mased_x: [B, (T, S), C] - # x_mask: [B, T] - x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S) - masked_x = rearrange(masked_x, "B (T S) C -> B T S C", T=T, S=S) - x = where(expand_dims(x_mask, [2, 3]), x, masked_x) - x = rearrange(x, "B T S C -> B (T S) C") - return x - - def forward( - self, - x, - y, - t, - x_mask=None, # temporal mask - t0=None, # t with timestamp=0 - T: Optional[int] = None, # number of frames - S: Optional[int] = None, # number of pixel patches - attention_params: Optional[AttentionParams] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - ): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - # prepare modulate parameters - B, N, C = x.shape - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = chunk( - expand_dims(self.scale_shift_table.value, 0) + t.view([B, 6, -1]), - chunks=6, - dim=1) - if x_mask is not None: - shift_msa_zero, scale_msa_zero, gate_msa_zero, shift_mlp_zero, scale_mlp_zero, gate_mlp_zero = chunk( - expand_dims(self.scale_shift_table.value, 0) + - t0.view([B, 6, -1]), - chunks=6, - dim=1) - - # modulate (attention) - x_m = t2i_modulate(self.norm1(x), shift_msa, scale_msa) - if x_mask is not None: - x_m_zero = t2i_modulate(self.norm1(x), shift_msa_zero, - scale_msa_zero) - x_m = self.t_mask_select(x_mask, x_m, x_m_zero, T, S) - - # attention - if self.temporal: - x_m = rearrange(x_m, "B (T S) C -> (B S) T C", T=T, S=S) - x_m = self.attn( - x_m, max_input_length=attention_params.max_context_length) - x_m = rearrange(x_m, "(B S) T C -> B (T S) C", T=T, S=S) - else: - x_m = rearrange(x_m, "B (T S) C -> (B T) S C", T=T, S=S) - x_m = self.attn( - x_m, max_input_length=attention_params.max_context_length) - x_m = rearrange(x_m, "(B T) S C -> B (T S) C", T=T, S=S) - - # modulate (attention) - x_m_s = gate_msa * x_m - if x_mask is not None: - x_m_s_zero = gate_msa_zero * x_m - x_m_s = self.t_mask_select(x_mask, x_m_s, x_m_s_zero, T, S) - - # residual - x = x + x_m_s - - # cross attention - cattn = self.cross_attn( - hidden_states=x, - encoder_output=y, - attention_params=AttentionParams( - sequence_length=attention_params.sequence_length, - context_lengths=attention_params.context_lengths, - host_context_lengths=attention_params.host_context_lengths, - max_context_length=attention_params.max_context_length, - host_request_types=attention_params.host_request_types, - encoder_input_lengths=attention_params.encoder_input_lengths, - encoder_max_input_length=attention_params. - encoder_max_input_length, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - ), - kv_cache_params=KeyValueCacheParams( - past_key_value=kv_cache_params.past_key_value, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - cache_indirection=kv_cache_params.cache_indirection, - kv_cache_block_offsets=None, - host_kv_cache_block_offsets=None, - host_kv_cache_pool_pointers=None, - host_kv_cache_pool_mapping=None, - cross_kv_cache_block_offsets=None, - host_cross_kv_cache_block_offsets=None, - host_cross_kv_cache_pool_pointers=None, - host_cross_kv_cache_pool_mapping=None, - )) - x = x + cattn - - # modulate (MLP) - x_m = t2i_modulate(self.norm2(x), shift_mlp, scale_mlp) - if x_mask is not None: - x_m_zero = t2i_modulate(self.norm2(x), shift_mlp_zero, - scale_mlp_zero) - x_m = self.t_mask_select(x_mask, x_m, x_m_zero, T, S) - - # MLP - x_m = self.mlp(x_m) - - # modulate (MLP) - x_m_s = gate_mlp * x_m - if x_mask is not None: - x_m_s_zero = gate_mlp_zero * x_m - x_m_s = self.t_mask_select(x_mask, x_m_s, x_m_s_zero, T, S) - - # residual - x = x + x_m_s - - return x - - -class STDiT3Model(PretrainedModel): - - def __init__(self, config: STDiTModelConfig): - self.check_config(config) - super().__init__(config) - self.learn_sigma = config.learn_sigma - self.in_channels = config.in_channels - self.out_channels = config.in_channels * 2 if config.learn_sigma else config.in_channels - self.caption_channels = config.caption_channels - self.depth = config.num_hidden_layers - self.mlp_ratio = config.mlp_ratio - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.model_max_length = config.model_max_length - self.latent_size = config.latent_size - self.input_sq_size = config.input_sq_size - self.patch_size = config.stdit_patch_size - self.class_dropout_prob = config.class_dropout_prob - self.qk_norm = config.qk_norm - self.dtype = config.dtype - self.mapping = config.mapping - - self.pos_embed = PositionEmbedding2D(self.hidden_size, dtype=self.dtype) - self.rope = RotaryEmbedder(dim=self.hidden_size // self.num_heads, - dtype=self.dtype) - self.x_embedder = PatchEmbed3D(self.patch_size, - self.in_channels, - self.hidden_size, - dtype=self.dtype) - self.t_embedder = TimestepEmbedder(self.hidden_size, dtype=self.dtype) - self.fps_embedder = SizeEmbedder(self.hidden_size, dtype=self.dtype) - self.t_block = ModuleSequential([ - Activation('silu'), - Linear(self.hidden_size, - 6 * self.hidden_size, - bias=True, - dtype=self.dtype) - ]) - self.y_embedder = CaptionEmbedder(in_channels=self.caption_channels, - hidden_size=self.hidden_size, - uncond_prob=self.class_dropout_prob, - act_layer='gelu', - token_num=self.model_max_length, - dtype=self.dtype) - self.spatial_blocks = ModuleList([ - STDiT3Block(hidden_size=self.hidden_size, - num_heads=self.num_heads, - mlp_ratio=self.mlp_ratio, - qk_norm=self.qk_norm, - dtype=self.dtype, - local_layer_idx=idx, - mapping=self.mapping) for idx in range(self.depth) - ]) - self.temporal_blocks = ModuleList([ - STDiT3Block(hidden_size=self.hidden_size, - num_heads=self.num_heads, - mlp_ratio=self.mlp_ratio, - qk_norm=self.qk_norm, - temporal=True, - rope=self.rope.rotate_queries_or_keys, - dtype=self.dtype, - local_layer_idx=idx, - mapping=self.mapping) for idx in range(self.depth) - ]) - self.final_layer = T2IFinalLayer(self.hidden_size, - np.prod(self.patch_size), - self.out_channels, - dtype=self.dtype, - mapping=self.mapping) - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('caption_channels', 4096) - config.set_if_not_exist('num_hidden_layers', 28) - config.set_if_not_exist('latent_size', [30, 45, 80]) - config.set_if_not_exist('hidden_size', 1152) - config.set_if_not_exist('stdit_patch_size', [1, 2, 2]) - config.set_if_not_exist('in_channels', 4) - config.set_if_not_exist('input_sq_size', 512) - config.set_if_not_exist('num_attention_heads', 16) - config.set_if_not_exist('mlp_ratio', 4.0) - config.set_if_not_exist('class_dropout_prob', 0.1) - config.set_if_not_exist('model_max_length', 300) - config.set_if_not_exist('learn_sigma', True) - config.set_if_not_exist('dtype', None) - config.set_if_not_exist('qk_norm', True) - config.set_if_not_exist('skip_y_embedder', False) - - def __post_init__(self): - return - - def get_dynamic_size(self, x: Tensor): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - _, _, T, H, W = x.shape - if T % self.patch_size[0] != 0: - T += self.patch_size[0] - T % self.patch_size[0] - if H % self.patch_size[1] != 0: - H += self.patch_size[1] - H % self.patch_size[1] - if W % self.patch_size[2] != 0: - W += self.patch_size[2] - W % self.patch_size[2] - T = T // self.patch_size[0] - H = H // self.patch_size[1] - W = W // self.patch_size[2] - return (T, H, W) - - def encode_text(self, y: Tensor, mask: Optional[Tensor] = None): - y = self.y_embedder(y) # [B, 1, N_token, C] - if mask is not None: - if mask.shape[0] != y.shape[0]: - mask = repeat(mask, sizes=(y.shape[0] // mask.shape[0], 1)) - mask = squeeze(squeeze(mask, 1), 1) - y = masked_select( - squeeze(y, 1), - where( - unsqueeze(mask, -1).__eq__( - constant_to_tensor_(0, dtype=mask.dtype)), - constant_to_tensor_(False), - constant_to_tensor_(True))).view((1, -1, self.hidden_size)) - # [TODO] how to convert y_lens to list? - # y_lens = mask.sum(dim=1).tolist() - y_lens = sum(mask, dim=1) - else: - y_lens = constant( - np.array([y.shape[2]] * y.shape[0], dtype=np.int64)) - y = squeeze(y, 1).view((1, -1, self.hidden_size)) - self.register_network_output('encode_text.output.y', y) - self.register_network_output('encode_text.output.y_lens', y_lens) - return y, y_lens - - def unpatchify(self, x: Tensor, N_t: int, N_h: int, N_w: int, R_t: int, - R_h: int, R_w: int): - T_p, H_p, W_p = self.patch_size - x = rearrange( - x, - "B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)", - N_t=N_t, - N_h=N_h, - N_w=N_w, - T_p=T_p, - H_p=H_p, - W_p=W_p, - C_out=self.out_channels, - ) - # unpad - x = slice(x, - starts=[0] * x.ndim(), - sizes=concat([ - shape(x, 0), - shape(x, 1), - constant(np.array([R_t, R_h, R_w]).astype(np.int64)) - ])) - self.register_network_output('unpatchify.output', x) - return x - - def forward(self, - x: Tensor, - timestep: Tensor, - y: Tensor, - fps: Tensor, - height: Tensor, - width: Tensor, - mask: Optional[Tensor] = None, - x_mask: Optional[Tensor] = None, - attention_params: Optional[AttentionParams] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - **kwargs): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - assert tuple(x.shape[2:]) == tuple( - self.latent_size), "For now only static shape is supported." - B = x.shape[0] - x = x.cast(self.dtype) - timestep = timestep.cast(self.dtype) - y = y.cast(self.dtype) - fps = fps.cast(self.dtype) - - # === get pos embed === - _, _, Tx, Hx, Wx = x.shape - T, H, W = self.get_dynamic_size(x) - S = H * W - base_size = round(S**0.5) - resolution_sq = pow( - height.cast('float32') * width.cast('float32'), - constant_to_tensor_(0.5, dtype='float32')) - scale = (resolution_sq / self.input_sq_size).cast(self.dtype) - pos_emb = self.pos_embed(x, h=H, w=W, scale=scale, base_size=base_size) - - # === get timestep embed === - t = self.t_embedder(timestep, dtype=x.dtype) # [B, C] - fps = self.fps_embedder(unsqueeze(fps, 1), bs=B) - t = t + fps - t_mlp = self.t_block(t) - t0 = t0_mlp = None - if x_mask is not None: - t0_timestep = constant( - np.zeros(shape=timestep.shape).astype(np.float32)).cast(x.dtype) - t0 = self.t_embedder(t0_timestep, dtype=x.dtype) - t0 = t0 + fps - t0_mlp = self.t_block(t0) - - # === get y embed === - if self.config.skip_y_embedder: - y_lens = mask - if isinstance(y_lens, Tensor): - y_lens = y_lens.cast('int64') - else: - y, y_lens = self.encode_text(y, mask) - y_lens = None #[11, 11] - - # === get x embed === - x = self.x_embedder(x) # [B, N, C] - x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S) - x = x + pos_emb - x = rearrange(x, "B T S C -> B (T S) C", T=T, S=S) - - # === blocks === - cnt = 0 - for spatial_block, temporal_block in zip(self.spatial_blocks, - self.temporal_blocks): - x = spatial_block( - x, - y, - t_mlp, - x_mask=x_mask, - t0=t0_mlp, - T=T, - S=S, - attention_params=attention_params, - kv_cache_params=kv_cache_params, - ) - x = temporal_block( - x, - y, - t_mlp, - x_mask=x_mask, - t0=t0_mlp, - T=T, - S=S, - attention_params=attention_params, - kv_cache_params=kv_cache_params, - ) - cnt += 1 - - # === final layer === - x = self.final_layer(x, t, x_mask, t0, T, S) - x = self.unpatchify(x, T, H, W, Tx, Hx, Wx) - - # cast to float32 for better accuracy - output = x.cast('float32') - output.mark_output('output', 'float32') - return output - - def prepare_inputs(self, max_batch_size, **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - mapping = self.config.mapping - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor(mapping, 1) - - def stdit_default_batch_range(max_batch_size): - return [max_batch_size, max_batch_size, max_batch_size] - - default_range = stdit_default_batch_range - # [NOTE] For now only static batch size is supported, so we run the model with max_batch_size. - batch_size = max_batch_size - - x = Tensor(name='x', - dtype=self.dtype, - shape=[batch_size, self.in_channels, *self.latent_size], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('in_channels', [[self.in_channels] * 3]), - ('latent_frames', [[self.latent_size[0]] * 3]), - ('latent_height', [[self.latent_size[1]] * 3]), - ('latent_width', [[self.latent_size[2]] * 3]), - ])) - timestep = Tensor(name='timestep', - dtype=self.dtype, - shape=[batch_size], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - y = Tensor( - name='y', - dtype=self.dtype, - shape=[batch_size, 1, self.model_max_length, self.caption_channels], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('mask_batch_size', [[1, 1, 1]]), - ('num_tokens', [[self.model_max_length] * 3]), - ('caption_channels', [[self.caption_channels] * 3]), - ])) - mask = Tensor(name='mask', - dtype=trt.int32, - shape=[1, self.model_max_length], - dim_range=OrderedDict([ - ('mask_batch_size', [[1, 1, 1]]), - ('num_tokens', [[self.model_max_length] * 3]), - ])) - x_mask = Tensor(name='x_mask', - dtype=tensorrt_llm.str_dtype_to_trt('bool'), - shape=[batch_size, self.latent_size[0]], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('latent_frames', [[self.latent_size[0]] * 3]), - ])) - fps = Tensor(name='fps', dtype=self.dtype, shape=[ - 1, - ]) - height = Tensor(name='height', dtype=self.dtype, shape=[ - 1, - ]) - width = Tensor(name='width', dtype=self.dtype, shape=[ - 1, - ]) - - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - remove_input_padding = default_net().plugin_config.remove_input_padding - - cross_attn_batch_size = batch_size - max_cattn_seq_len = int( - np.prod([ - np.ceil(d / p) - for d, p in zip(self.latent_size, self.patch_size) - ])) - max_cattn_enc_len = self.model_max_length - attn_inputs = GenerationMixin().prepare_attention_inputs( - max_batch_size=cross_attn_batch_size, - opt_batch_size=cross_attn_batch_size, - max_beam_width=1, - max_input_len=max_cattn_seq_len, - max_seq_len=max_cattn_seq_len, - num_kv_heads=self.num_heads, - head_size=self.hidden_size // self.num_heads, - num_layers=self.depth, - kv_dtype=self.dtype, - kv_cache_type=KVCacheType.DISABLED, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - enable_ctx_gen_opt_profiles=False, - mapping=self.mapping, - ) - - sequence_length = attn_inputs['sequence_length'] - host_context_lengths = attn_inputs['host_context_lengths'] - host_max_attention_window_sizes = attn_inputs[ - 'host_max_attention_window_sizes'] - host_sink_token_length = attn_inputs['host_sink_token_length'] - context_lengths = attn_inputs['context_lengths'] - host_request_types = attn_inputs['host_request_types'] - - host_past_key_value_lengths = attn_inputs['host_past_key_value_lengths'] - past_key_value = attn_inputs['past_key_value'] - if past_key_value: - past_key_value = past_key_value[0] - cache_indirection = attn_inputs['cache_indirection'] - host_runtime_perf_knobs_tensor = attn_inputs['host_runtime_perf_knobs'] - host_context_progress = attn_inputs['host_context_progress'] - cross_encoder_input_lengths = Tensor(name='encoder_input_lengths', - shape=(cross_attn_batch_size, ), - dtype=str_dtype_to_trt('int32')) - cross_max_encoder_seq_len = Tensor( - name='encoder_max_input_length', - shape=[-1], - dim_range=OrderedDict([ - ("encoder_max_input_length", - [[1, (max_cattn_enc_len + 1) // 2, max_cattn_enc_len]]) - ]), - dtype=str_dtype_to_trt('int32')) - - attention_params = AttentionParams( - sequence_length=sequence_length, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_cattn_seq_len, - host_request_types=host_request_types, - encoder_input_lengths=cross_encoder_input_lengths, - encoder_max_input_length=cross_max_encoder_seq_len, - host_runtime_perf_knobs=host_runtime_perf_knobs_tensor, - host_context_progress=host_context_progress) - - kv_cache_params = KeyValueCacheParams( - past_key_value=past_key_value, - host_past_key_value_lengths=host_past_key_value_lengths, - host_max_attention_window_sizes=host_max_attention_window_sizes, - host_sink_token_length=host_sink_token_length, - cache_indirection=cache_indirection, - kv_cache_block_offsets=None, - host_kv_cache_block_offsets=None, - host_kv_cache_pool_pointers=None, - host_kv_cache_pool_mapping=None, - cross_kv_cache_block_offsets=None, - host_cross_kv_cache_block_offsets=None, - host_cross_kv_cache_pool_pointers=None, - host_cross_kv_cache_pool_mapping=None, - ) - - return { - 'x': x, - 'timestep': timestep, - 'y': y, - 'mask': mask, - 'x_mask': x_mask, - 'fps': fps, - 'height': height, - 'width': width, - 'attention_params': attention_params, - 'kv_cache_params': kv_cache_params, - } - - @classmethod - def from_pretrained(cls, - pretrained_model_dir: str, - dtype='float16', - mapping=Mapping(), - **kwargs): - - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - assert os.path.exists(f"{pretrained_model_dir}/config.json") - with open(f"{pretrained_model_dir}/config.json", 'r') as f: - hf_config = json.load(f) - - hf_tllm_config_remapping = { - 'model_type': 'architecture', - 'depth': 'num_hidden_layers', - 'num_heads': 'num_attention_heads', - 'patch_size': 'stdit_patch_size', - 'pred_sigma': 'learn_sigma', - } - for hf_key, tllm_key in hf_tllm_config_remapping.items(): - hf_config[tllm_key] = hf_config.pop(hf_key) - hf_config.update(kwargs) - - model_config = STDiTModelConfig.from_input_config(hf_config, - dtype=dtype, - mapping=mapping) - - model_dir = pretrained_model_dir - custom_dict = {} - if quant_ckpt_path is not None: - model_dir = quant_ckpt_path - - loader = STDiT3ModelWeightsLoader(model_dir, custom_dict) - model = cls(model_config) - loader.generate_tllm_weights(model) - return model - - -class STDiT3ModelWeightsLoader(ModelWeightsLoader): - - def translate_to_external_key(self, tllm_key: str, - tllm_to_externel_key_dict: dict): - """Convert and load external checkpoint into a TensorRT LLM model. - """ - trtllm_to_hf_name = { - r"spatial_blocks.(\d+).attn.q_layernorm.weight": - "spatial_blocks.*.attn.q_norm.weight", - r"spatial_blocks.(\d+).attn.k_layernorm.weight": - "spatial_blocks.*.attn.k_norm.weight", - r"spatial_blocks.(\d+).attn.dense.weight": - "spatial_blocks.*.attn.proj.weight", - r"spatial_blocks.(\d+).attn.dense.bias": - "spatial_blocks.*.attn.proj.bias", - r"temporal_blocks.(\d+).attn.q_layernorm.weight": - "temporal_blocks.*.attn.q_norm.weight", - r"temporal_blocks.(\d+).attn.k_layernorm.weight": - "temporal_blocks.*.attn.k_norm.weight", - r"temporal_blocks.(\d+).attn.dense.weight": - "temporal_blocks.*.attn.proj.weight", - r"temporal_blocks.(\d+).attn.dense.bias": - "temporal_blocks.*.attn.proj.bias", - r"spatial_blocks.(\d+).cross_attn.dense.weight": - "spatial_blocks.*.cross_attn.proj.weight", - r"spatial_blocks.(\d+).cross_attn.dense.bias": - "spatial_blocks.*.cross_attn.proj.bias", - r"temporal_blocks.(\d+).cross_attn.dense.weight": - "temporal_blocks.*.cross_attn.proj.weight", - r"temporal_blocks.(\d+).cross_attn.dense.bias": - "temporal_blocks.*.cross_attn.proj.bias", - } - - for k, v in trtllm_to_hf_name.items(): - m = re.match(k, tllm_key) - if m is not None: - matched_pos = m.groups() - placeholders = v.count('*') - assert len(matched_pos) == placeholders - for i in range(len(matched_pos)): - v = v.replace('*', matched_pos[i], 1) - return v - return tllm_key - - def load_tensor(self, key, tp_size=1, tp_dim=-1, tp_rank=0): - hidden_size = self.model.config.hidden_size - - if "attn.qkv" in key: - is_cross_attn = "cross_attn.qkv" in key - if is_cross_attn: - # process for cross attention - process_qkv_names = [ - 'q_linear'.join(key.split('qkv')), - 'kv_linear'.join(key.split('qkv')) - ] - else: - process_qkv_names = [key] - qkv_tensors = [] - for qkv_key in process_qkv_names: - # Retrieve shard index - assert qkv_key in self.shard_map - ptr_idx = self.shard_map[qkv_key] - if self.format == ModelWeightsFormat.SAFETENSORS: - # Force to load Pytorch tensor - tensor = self.shards[ptr_idx].get_tensor(qkv_key) - else: - tensor = self.shards[ptr_idx][qkv_key] - qkv_tensors.append(tensor) - if is_cross_attn: - tensor = torch.concat(qkv_tensors, dim=0) - else: - tensor = qkv_tensors[0] - # Post-process weight and bias if tp_size > 1 - if tp_size > 1: - if "weight" in key: - tensor = tensor.reshape(3, hidden_size, hidden_size) - elif "bias" in key: - tensor = tensor.reshape(3, hidden_size) - tp_dim = 1 - tensor_shape = tensor.shape - else: - # Retrieve shard index - if key in self.shard_map: - ptr_idx = self.shard_map[key] - else: - return None - - if self.format == ModelWeightsFormat.SAFETENSORS: - tensor = self.shards[ptr_idx].get_slice(key) - tensor_shape = tensor.get_shape() - if tensor_shape == []: - tensor = self.shards[ptr_idx].get_tensor(key).unsqueeze(0) - tensor_shape = tensor.shape - else: - tensor = self.shards[ptr_idx][key] - tensor_shape = tensor.shape - - if tp_size <= 1 or tp_dim < 0: - return tensor[:] - else: - if len(tensor_shape) == 1 and (tp_dim > 0 or tensor_shape[0] == 1): - return tensor[:] - else: - width = tensor_shape[tp_dim] - if width == 1: - return tensor[:] - slice_width = math.ceil(width / tp_size) - slice_start = tp_rank * slice_width - slice_end = builtins.min((tp_rank + 1) * slice_width, width) - slice_obj = [builtins.slice(None)] * len(tensor_shape) - slice_obj[tp_dim] = builtins.slice(slice_start, slice_end) - res = tensor[tuple(slice_obj)] - if "qkv.weight" in key: - res = res.reshape(3 * (hidden_size // tp_size), hidden_size) - elif "qkv.bias" in key: - res = res.reshape(3 * (hidden_size // tp_size)) - return res - - def generate_tllm_weights(self, - model, - custom_postprocess_kwargs: dict = {}): - self.update_key_mapping(model) - tp_module_patterns = [ - r'.*_blocks.*.attn.qkv.weight$', - r'.*_blocks.*.attn.qkv.bias$', - r'.*_blocks.*.attn.dense.weight$', - r'.*_blocks.*.cross_attn.qkv.weight$', - r'.*_blocks.*.cross_attn.qkv.bias$', - r'.*_blocks.*.cross_attn.dense.weight$', - r'.*_blocks.*.mlp.fc1.weight$', - r'.*_blocks.*.mlp.fc1.bias$', - r'.*_blocks.*.mlp.fc2.weight$', - ] - tllm_weights = {} - for tllm_key, _ in tqdm(model.named_parameters()): - skip_tp = not any([ - re.match(pattern, tllm_key) is not None - for pattern in tp_module_patterns - ]) - tllm_weights.update( - self.load(tllm_key, - custom_postprocess_kwargs=custom_postprocess_kwargs, - skip_tp=skip_tp)) - self.fill(tllm_weights) diff --git a/tensorrt_llm/models/unet/__init__.py b/tensorrt_llm/models/unet/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/unet/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. diff --git a/tensorrt_llm/models/unet/attention.py b/tensorrt_llm/models/unet/attention.py deleted file mode 100644 index c56605106362..000000000000 --- a/tensorrt_llm/models/unet/attention.py +++ /dev/null @@ -1,332 +0,0 @@ -# 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 math -from typing import Optional - -from ..._common import precision -from ...functional import geglu, matmul, softmax, split -from ...layers import Conv2d, GroupNorm, LayerNorm, Linear -from ...module import Module, ModuleList - - -class AttentionBlock(Module): - - def __init__(self, - channels: int, - num_head_channels: Optional[int] = None, - num_groups: int = 32, - rescale_output_factor: float = 1.0, - eps: float = 1e-5): - super().__init__() - self.channels = channels - - self.num_heads = channels // num_head_channels if num_head_channels is not None else 1 - self.num_head_size = num_head_channels - self.group_norm = GroupNorm(num_channels=channels, - num_groups=num_groups, - eps=eps, - affine=True) - - self.qkv = Linear(channels, channels * 3) - self.rescale_output_factor = rescale_output_factor - self.proj_attn = Linear(channels, channels, 1) - - def transpose_for_scores(self, projection): - new_projection_shape = projection.size()[:-1] + (self.num_heads, - self.num_head_size) - # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) -> (B, H, T, D) - new_projection = projection.view(new_projection_shape).permute( - [0, 2, 1, 3]) - return new_projection - - def forward(self, hidden_states): - assert not hidden_states.is_dynamic() - - residual = hidden_states - batch, channel, height, width = hidden_states.size() - - # norm - hidden_states = self.group_norm(hidden_states) - hidden_states = hidden_states.view([batch, channel, - height * width]).transpose(1, 2) - - # proj to q, k, v - qkv_proj = self.qkv(hidden_states) - - query_proj, key_proj, value_proj = split(qkv_proj, channel, dim=2) - - # transpose - query_states = self.transpose_for_scores(query_proj) - key_states = self.transpose_for_scores(key_proj) - value_states = self.transpose_for_scores(value_proj) - - # get scores - with precision('float32'): - attention_scores = matmul(query_states, - (key_states).transpose(-1, -2)) - attention_scores = attention_scores / math.sqrt( - self.channels / self.num_heads) - attention_probs = softmax(attention_scores, dim=-1) - - # compute attention output - hidden_states = matmul(attention_probs, value_states) - hidden_states = hidden_states.permute([0, 2, 1, 3]) - - new_hidden_states_shape = hidden_states.size()[:-2] + (self.channels, ) - hidden_states = hidden_states.view(new_hidden_states_shape) - - # compute next hidden_states - hidden_states = self.proj_attn(hidden_states) - hidden_states = hidden_states.transpose(-1, -2).view( - [batch, channel, height, width]) - - # res connect and rescale - hidden_states = (hidden_states + residual) / self.rescale_output_factor - return hidden_states - - -def _transpose_for_scores(tensor, heads): - batch_size, seq_len, dim = tensor.size() - tensor = tensor.view([batch_size, seq_len, heads, dim // heads]) - tensor = tensor.permute([0, 2, 1, 3]) - return tensor - - -def _attention(query, key, value, scale): - # Multiply scale first to avoid overflow - # Do not use use_fp32_acc or it will be very slow - attention_scores = matmul(query * math.sqrt(scale), - key.transpose(-1, -2) * math.sqrt(scale), - use_fp32_acc=False) - attention_probs = softmax(attention_scores, dim=-1) - hidden_states = matmul(attention_probs, value, use_fp32_acc=False) - hidden_states = hidden_states.permute([0, 2, 1, 3]) - return hidden_states - - -class SelfAttention(Module): - - def __init__(self, - query_dim: int, - heads: int = 8, - dim_head: int = 64, - dtype=None): - super().__init__() - self.inner_dim = dim_head * heads - self.scale = dim_head**-0.5 - self.heads = heads - self._slice_size = None - - self.to_qkv = Linear(query_dim, - 3 * self.inner_dim, - bias=False, - dtype=dtype) - self.to_out = Linear(self.inner_dim, query_dim, dtype=dtype) - - def forward(self, hidden_states, mask=None): - assert not hidden_states.is_dynamic() - - qkv = self.to_qkv(hidden_states) - - query, key, value = split(qkv, self.inner_dim, dim=2) - query = _transpose_for_scores(query, self.heads) - key = _transpose_for_scores(key, self.heads) - value = _transpose_for_scores(value, self.heads) - hidden_states = _attention(query, key, value, self.scale) - - batch_size, seq_len, head_size, head_dim = hidden_states.size() - hidden_states = hidden_states.view( - [batch_size, seq_len, head_size * head_dim]) - return self.to_out(hidden_states) - - -class CrossAttention(Module): - - def __init__(self, - query_dim: int, - context_dim: Optional[int] = None, - heads: int = 8, - dim_head: int = 64, - dtype=None): - super().__init__() - self.inner_dim = dim_head * heads - context_dim = context_dim if context_dim is not None else query_dim - self.scale = dim_head**-0.5 - self.heads = heads - self._slice_size = None - - self.to_q = Linear(query_dim, self.inner_dim, bias=False, dtype=dtype) - self.to_kv = Linear(context_dim, - 2 * self.inner_dim, - bias=False, - dtype=dtype) - self.to_out = Linear(self.inner_dim, query_dim, dtype=dtype) - - def forward(self, hidden_states, context=None, mask=None): - assert not hidden_states.is_dynamic() - query = self.to_q(hidden_states) - is_cross_attn = context is not None - context = context if is_cross_attn else hidden_states - assert not context.is_dynamic() - kv = self.to_kv(context) - - query = _transpose_for_scores(query, self.heads) - key, value = split(kv, self.inner_dim, dim=2) - key = _transpose_for_scores(key, self.heads) - value = _transpose_for_scores(value, self.heads) - hidden_states = _attention(query, key, value, self.scale) - - batch_size, seq_len, head_size, head_dim = hidden_states.size() - hidden_states = hidden_states.view( - [batch_size, seq_len, head_size * head_dim]) - return self.to_out(hidden_states) - - -class FeedForward(Module): - - def __init__(self, - dim: int, - dim_out: Optional[int] = None, - mult: int = 4, - dtype=None): - super().__init__() - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - self.proj_in = Linear(dim, inner_dim * 2, dtype=dtype) - self.proj_out = Linear(inner_dim, dim_out, dtype=dtype) - - def forward(self, hidden_states): - x = self.proj_in(hidden_states) - x = geglu(x) - return self.proj_out(x) - - -class BasicTransformerBlock(Module): - - def __init__( - self, - dim: int, - n_heads: int, - d_head: int, - context_dim: Optional[int] = None, - dtype=None, - ): - super().__init__() - self.attn1 = SelfAttention(query_dim=dim, - heads=n_heads, - dim_head=d_head, - dtype=dtype) # is a self-attention - self.ff = FeedForward(dim, dtype=dtype) - self.attn2 = CrossAttention( - query_dim=dim, - context_dim=context_dim, - heads=n_heads, - dim_head=d_head, - dtype=dtype) # is self-attn if context is none - self.norm1 = LayerNorm(dim, dtype=dtype) - self.norm2 = LayerNorm(dim, dtype=dtype) - self.norm3 = LayerNorm(dim, dtype=dtype) - - def forward(self, hidden_states, context=None): - hidden_states = self.attn1(self.norm1(hidden_states)) + hidden_states - hidden_states = self.attn2(self.norm2(hidden_states), - context=context) + hidden_states - hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states - return hidden_states - - -class Transformer2DModel(Module): - - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - num_layers: int = 1, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - use_linear_projection: bool = False, - dtype=None, - ): - super().__init__() - self.use_linear_projection = use_linear_projection - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - inner_dim = num_attention_heads * attention_head_dim - - self.norm = GroupNorm(num_groups=norm_num_groups, - num_channels=in_channels, - eps=1e-6, - affine=True, - dtype=dtype) - - if use_linear_projection: - self.proj_in = Linear(in_channels, inner_dim, dtype=dtype) - else: - self.proj_in = Conv2d(in_channels, - inner_dim, - kernel_size=(1, 1), - stride=(1, 1), - padding=(0, 0), - dtype=dtype) - - self.transformer_blocks = ModuleList([ - BasicTransformerBlock(inner_dim, - num_attention_heads, - attention_head_dim, - context_dim=cross_attention_dim, - dtype=dtype) for d in range(num_layers) - ]) - - if use_linear_projection: - self.proj_out = Linear(inner_dim, in_channels, dtype=dtype) - else: - self.proj_out = Conv2d(inner_dim, - in_channels, - kernel_size=(1, 1), - stride=(1, 1), - padding=(0, 0), - dtype=dtype) - - def forward(self, hidden_states, context=None): - assert not hidden_states.is_dynamic() - batch, _, height, weight = hidden_states.size() - residual = hidden_states - hidden_states = self.norm(hidden_states) - - if not self.use_linear_projection: - hidden_states = self.proj_in(hidden_states) - inner_dim = hidden_states.size()[1] - hidden_states = hidden_states.permute([0, 2, 3, 1]).view( - [batch, height * weight, inner_dim]) - else: - inner_dim = hidden_states.size()[1] - hidden_states = hidden_states.permute([0, 2, 3, 1]).view( - [batch, height * weight, inner_dim]) - hidden_states = self.proj_in(hidden_states) - - for block in self.transformer_blocks: - hidden_states = block(hidden_states, context=context) - - if not self.use_linear_projection: - hidden_states = hidden_states.view( - [batch, height, weight, inner_dim]).permute([0, 3, 1, 2]) - hidden_states = self.proj_out(hidden_states) - else: - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.view( - [batch, height, weight, inner_dim]).permute([0, 3, 1, 2]) - - return hidden_states + residual diff --git a/tensorrt_llm/models/unet/embeddings.py b/tensorrt_llm/models/unet/embeddings.py deleted file mode 100644 index 8bfe16a408d6..000000000000 --- a/tensorrt_llm/models/unet/embeddings.py +++ /dev/null @@ -1,117 +0,0 @@ -# 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 math - -import tensorrt as trt - -from ..._utils import fp16_array, fp32_array -from ...functional import concat, constant, cos, exp, silu, sin -from ...layers import Linear -from ...module import Module - - -def get_timestep_embedding(timesteps, - embedding_dim, - flip_sin_to_cos=False, - downscale_freq_shift=1.0, - scale=1.0, - max_period=10000, - dtype=None): - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - :param timesteps: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the - embeddings. :return: an [N x dim] Tensor of positional embeddings. - """ - assert timesteps.rank() == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - - exponent = [ - i * -math.log(max_period) / (half_dim - downscale_freq_shift) - for i in range(half_dim) - ] - - if dtype == trt.DataType.HALF: - emb = exp(constant(fp16_array(exponent))) - else: - emb = exp(constant(fp32_array(exponent))) - - ts_shape = list(timesteps.size()) - ts_shape.append(1) - emb_shape = list(emb.size()) - emb_shape.insert(0, 1) - - emb = timesteps.view(ts_shape) * emb.view(emb_shape) - - emb = scale * emb - # concat sine and cosine embeddings - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = concat([cos(emb), sin(emb)], dim=1) - else: - emb = concat([sin(emb), cos(emb)], dim=1) - - #TODO Enable below logic when TensorRT LLM supports pad feature. - # zero pad - # if embedding_dim % 2 == 1: - # emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) - return emb - - -class TimestepEmbedding(Module): - - def __init__(self, channel, time_embed_dim, act_fn="silu", dtype=None): - super().__init__() - - self.linear_1 = Linear(channel, time_embed_dim, dtype=dtype) - self.act = None - if act_fn == "silu": - self.act = silu - self.linear_2 = Linear(time_embed_dim, time_embed_dim, dtype=dtype) - - def forward(self, sample): - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - return sample - - -class Timesteps(Module): - - def __init__(self, - num_channels, - flip_sin_to_cos, - downscale_freq_shift, - dtype=None): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - self.dtype = dtype - - def forward(self, timesteps): - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - dtype=self.dtype) - return t_emb diff --git a/tensorrt_llm/models/unet/pp/attention.py b/tensorrt_llm/models/unet/pp/attention.py deleted file mode 100755 index f64d666b0891..000000000000 --- a/tensorrt_llm/models/unet/pp/attention.py +++ /dev/null @@ -1,107 +0,0 @@ -# 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 ....functional import allgather, split -from ....mapping import Mapping -from ....module import Module -from ..attention import CrossAttention, SelfAttention, _attention - - -class DistriSelfAttentionPP(Module): - - def __init__(self, module: SelfAttention, mapping: Mapping = Mapping()): - super().__init__() - self.mapping = mapping - self.module = module - - def forward(self, hidden_states): - mapping = self.mapping - attn = self.module - - batch_size, sequence_length, _ = hidden_states.shape - - qkv = attn.to_qkv(hidden_states) - - query, kv = split(qkv, [attn.inner_dim, attn.inner_dim * 2], dim=2) - - if mapping.tp_size == 1: - full_kv = kv - else: - full_kv = allgather(kv, group=mapping.tp_group, gather_dim=1) - - key, value = split(full_kv, full_kv.shape[-1] // 2, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - key = key.view([batch_size, -1, attn.heads, head_dim]).transpose(1, 2) - value = value.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - - hidden_states = _attention(query, key, value, attn.scale) - - hidden_states = hidden_states.view( - [batch_size, -1, attn.heads * head_dim]) - - # linear proj - hidden_states = attn.to_out(hidden_states) - - return hidden_states - - -class DistriCrossAttentionPP(Module): - - def __init__(self, module: CrossAttention, mapping: Mapping = Mapping()): - super().__init__() - self.mapping = mapping - self.module = module - self.kv_cache = None - - def forward(self, hidden_states, context): - attn = self.module - recompute_kv = self.kv_cache is None - - if context is None: - context = hidden_states - - batch_size, sequence_length, _ = context.shape - - query = attn.to_q(hidden_states) - - if recompute_kv or self.kv_cache is None: - kv = attn.to_kv(context) - self.kv_cache = kv - else: - kv = self.kv_cache - key, value = split(kv, kv.shape[-1] // 2, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - key = key.view([batch_size, -1, attn.heads, head_dim]).transpose(1, 2) - value = value.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - - hidden_states = _attention(query, key, value, scale=attn.scale) - - hidden_states = hidden_states.view( - [batch_size, -1, attn.heads * head_dim]) - - # linear proj - hidden_states = attn.to_out(hidden_states) - return hidden_states diff --git a/tensorrt_llm/models/unet/pp/conv2d.py b/tensorrt_llm/models/unet/pp/conv2d.py deleted file mode 100755 index a1aa8f7b72c5..000000000000 --- a/tensorrt_llm/models/unet/pp/conv2d.py +++ /dev/null @@ -1,109 +0,0 @@ -# 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 ....functional import allgather, concat, conv2d, slice, stack, unsqueeze -from ....layers import Conv2d -from ....mapping import Mapping -from ....module import Module - - -class DistriConv2dPP(Module): - - def __init__(self, - conv: Conv2d, - mapping: Mapping = Mapping(), - is_first_layer: bool = False): - super().__init__() - self.mapping = mapping - self.conv = conv - self.is_first_layer = is_first_layer - - def sliced_forward(self, x): - mapping = self.mapping - b, c, h, w = x.shape - assert h % mapping.tp_size == 0 - - stride = self.conv.stride[0] - padding = self.conv.padding[0] - - output_h = x.shape[2] // stride // mapping.tp_size - idx = mapping.tp_rank - h_begin = output_h * idx * stride - padding - h_end = output_h * (idx + 1) * stride + padding - pre_padding = [0, padding] - post_padding = [0, padding] - if h_begin < 0: - h_begin = 0 - pre_padding[0] = padding - if h_end > h: - h_end = h - post_padding[0] = padding - sliced_input = slice(x, [0, 0, h_begin, 0], [b, c, h_end - h_begin, w]) - return conv2d(sliced_input, - self.conv.weight.value, - None if self.conv.bias is None else self.conv.bias.value, - stride=self.conv.stride, - padding=(0, 0), - pre_padding=tuple(pre_padding), - post_padding=tuple(post_padding)) - - def forward(self, x, *args, **kwargs): - mapping = self.mapping - if self.is_first_layer: - full_x = x - output = self.sliced_forward(full_x) - else: - boundary_size = self.conv.padding[0] - - def create_padded_x(x, boundaries): - preH = 0 - postH = 0 - if mapping.tp_rank == 0: - b = boundaries.select(0, mapping.tp_rank + 1).select(0, 0) - padded_x = concat([x, b], dim=2) - preH = boundary_size - elif mapping.tp_rank == mapping.tp_size - 1: - b = boundaries.select(0, mapping.tp_rank - 1).select(0, 1) - padded_x = concat([b, x], dim=2) - postH = boundary_size - else: - b0 = boundaries.select(0, mapping.tp_rank - 1).select(0, 1) - b1 = boundaries.select(0, mapping.tp_rank + 1).select(0, 0) - padded_x = concat( - [ - b0, - x, - b1, - ], - dim=2, - ) - return padded_x, preH, postH - - n, c, h, w = x.shape - b0 = slice(x, [0, 0, 0, 0], [n, c, boundary_size, w]) - b1 = slice(x, [0, 0, h - boundary_size, 0], - [n, c, boundary_size, w]) - boundary = stack([b0, b1], dim=0) - - boundaries = allgather(unsqueeze(boundary, 0), - group=mapping.tp_group) - padded_x, preH, postH = create_padded_x(x, boundaries) - output = conv2d(padded_x, - self.conv.weight.value, - self.conv.bias.value, - stride=self.conv.stride, - pre_padding=(preH, self.conv.padding[1]), - post_padding=(postH, self.conv.padding[1])) - return output diff --git a/tensorrt_llm/models/unet/pp/groupnorm.py b/tensorrt_llm/models/unet/pp/groupnorm.py deleted file mode 100755 index 6cbf4d2d0228..000000000000 --- a/tensorrt_llm/models/unet/pp/groupnorm.py +++ /dev/null @@ -1,57 +0,0 @@ -# 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 ....functional import allreduce, pow, select, stack -from ....layers import GroupNorm -from ....mapping import Mapping -from ....module import Module - - -class DistriGroupNorm(Module): - - def __init__(self, - module: GroupNorm, - mapping: Mapping = Mapping(), - is_first_layer: bool = False): - super().__init__() - self.mapping = mapping - self.module = module - - def forward(self, x, *args, **kwargs): - mapping = self.mapping - module = self.module - n, c, h, w = x.shape - num_groups = module.num_groups - group_size = c // num_groups - - x = x.view([n, num_groups, group_size, h, w]) - x_mean = x.mean(dim=4, keepdim=True).mean(dim=(3, 2), keepdim=True) - x2_mean = pow(x, 2.0).mean(dim=4, keepdim=True).mean(dim=(3, 2), - keepdim=True) - mean = stack([x_mean, x2_mean], dim=0) - mean = allreduce(mean, mapping.tp_group) - mean = mean / (mapping.tp_size * 1.0) - x_mean = select(mean, 0, 0) - x2_mean = select(mean, 0, 1) - var = x2_mean - pow(x_mean, 2.0) - num_elements = group_size * h * w - var = var * (num_elements / (num_elements - 1)) - std = (var + module.eps).sqrt() - output = (x - x_mean) / std - output = output.view([n, c, h, w]) - if module.affine: - output = output * module.weight.value.view([1, -1, 1, 1]) - output = output + module.bias.value.view([1, -1, 1, 1]) - - return output diff --git a/tensorrt_llm/models/unet/pp/unet_pp.py b/tensorrt_llm/models/unet/pp/unet_pp.py deleted file mode 100755 index 9b1a7d64a60f..000000000000 --- a/tensorrt_llm/models/unet/pp/unet_pp.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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 ....functional import allgather, concat, slice, stack -from ....layers import Conv2d, GroupNorm -from ....mapping import Mapping -from ....module import Module -from ..attention import CrossAttention, SelfAttention -from ..unet_2d_condition import UNet2DConditionModel -from .attention import DistriCrossAttentionPP, DistriSelfAttentionPP -from .conv2d import DistriConv2dPP -from .groupnorm import DistriGroupNorm - - -class DistriUNetPP(Module): - - def __init__(self, - model: UNet2DConditionModel, - mapping: Mapping = Mapping()): - super().__init__() - self.mapping = mapping - self.model = model - if mapping.tp_size > 1: - for name, module in model.named_modules(): - if isinstance(module, DistriConv2dPP) or isinstance(module, DistriSelfAttentionPP) \ - or isinstance(module, DistriCrossAttentionPP) or isinstance(module, DistriGroupNorm): - continue - for subname, submodule in module.named_children(): - if isinstance(submodule, Conv2d): - kernel_size = submodule.kernel_size - if kernel_size == (1, 1) or kernel_size == 1: - continue - wrapped_submodule = DistriConv2dPP( - submodule, - mapping, - is_first_layer=subname == "conv_in") - setattr(module, subname, wrapped_submodule) - elif isinstance(submodule, SelfAttention): - wrapped_submodule = DistriSelfAttentionPP( - submodule, mapping) - setattr(module, subname, wrapped_submodule) - elif isinstance(submodule, CrossAttention): - wrapped_submodule = DistriCrossAttentionPP( - submodule, mapping) - setattr(module, subname, wrapped_submodule) - elif isinstance(submodule, GroupNorm): - wrapped_submodule = DistriGroupNorm(submodule, mapping) - setattr(module, subname, wrapped_submodule) - - def forward(self, - sample, - timesteps, - encoder_hidden_states, - text_embeds=None, - time_ids=None): - mapping = self.mapping - b, c, h, w = sample.shape - - if mapping.world_size == 1: - output = self.model( - sample, - timesteps, - encoder_hidden_states, - text_embeds=text_embeds, - time_ids=time_ids, - ) - elif mapping.pp_size > 1: - assert b == 2 and mapping.pp_size == 2 - batch_idx = mapping.pp_rank - # sample[batch_idx : batch_idx + 1] - sample = slice(sample, [batch_idx, 0, 0, 0], [1, c, h, w]) - e_shape = encoder_hidden_states.shape - encoder_hidden_states = slice( - encoder_hidden_states, [batch_idx, 0, 0], - [1, e_shape[1], e_shape[2] - ]) # encoder_hidden_states[batch_idx : batch_idx + 1] - if text_embeds: - t_shape = text_embeds.shape - # text_embeds[batch_idx : batch_idx + 1] - text_embeds = slice(text_embeds, [batch_idx, 0], - [1, t_shape[1]]) - if time_ids: - t_shape = time_ids.shape - # time_ids[batch_idx : batch_idx + 1] - time_ids = slice(time_ids, [batch_idx, 0], [1, t_shape[1]]) - output = self.model( - sample, - timesteps, - encoder_hidden_states, - text_embeds=text_embeds, - time_ids=time_ids, - ) - output = allgather( - output, - [i for i in range(mapping.world_size)], - ) - patch_list = [] - for i in range(mapping.tp_size): - patch_list.append(output.select(dim=0, index=i)) - b1 = concat(patch_list, dim=1) - patch_list = [] - for i in range(mapping.tp_size, mapping.world_size): - patch_list.append(output.select(dim=0, index=i)) - b2 = concat(patch_list, dim=1) - output = stack([b1, b2], dim=0) - else: - output = self.model( - sample, - timesteps, - encoder_hidden_states, - text_embeds=text_embeds, - time_ids=time_ids, - ) - output = allgather(output, mapping.tp_group, 2) - - return output - - @property - def add_embedding(self): - return self.model.add_embedding diff --git a/tensorrt_llm/models/unet/resnet.py b/tensorrt_llm/models/unet/resnet.py deleted file mode 100644 index 453e55324f04..000000000000 --- a/tensorrt_llm/models/unet/resnet.py +++ /dev/null @@ -1,254 +0,0 @@ -# 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 functools import partial - -from ...functional import avg_pool2d, interpolate, silu -from ...layers import (AvgPool2d, Conv2d, ConvTranspose2d, GroupNorm, Linear, - Mish) -from ...module import Module - - -class Upsample2D(Module): - - def __init__(self, - channels: int, - use_conv=False, - use_conv_transpose=False, - out_channels=None, - dtype=None) -> None: - super().__init__() - - self.channels = channels - self.out_channels = out_channels - - self.use_conv_transpose = use_conv_transpose - self.use_conv = use_conv - if self.use_conv_transpose: - self.conv = ConvTranspose2d(channels, - self.out_channels, - 4, - 2, - 1, - dtype=dtype) - elif use_conv: - self.conv = Conv2d(self.channels, - self.out_channels, (3, 3), - padding=(1, 1), - dtype=dtype) - else: - self.conv = None - - def forward(self, hidden_states, output_size=None): - assert not hidden_states.is_dynamic() - batch, channels, _, _ = hidden_states.size() - assert channels == self.channels - - if self.use_conv_transpose: - return self.conv(hidden_states) - - if output_size is None: - hidden_states = interpolate(hidden_states, - scale_factor=2.0, - mode="nearest") - else: - hidden_states = interpolate(hidden_states, - size=output_size, - mode="nearest") - - if self.use_conv: - hidden_states = self.conv(hidden_states) - - return hidden_states - - -class Downsample2D(Module): - - def __init__(self, - channels, - use_conv=False, - out_channels=None, - padding=1, - dtype=None) -> None: - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.padding = padding - stride = (2, 2) - - if use_conv: - self.conv = Conv2d(self.channels, - self.out_channels, (3, 3), - stride=stride, - padding=(padding, padding), - dtype=dtype) - else: - assert self.channels == self.out_channels - self.conv = AvgPool2d(kernel_size=stride, stride=stride) - - def forward(self, hidden_states): - assert not hidden_states.is_dynamic() - batch, channels, _, _ = hidden_states.size() - assert channels == self.channels - - #TODO add the missing pad function - hidden_states = self.conv(hidden_states) - - return hidden_states - - -class ResnetBlock2D(Module): - - def __init__( - self, - *, - in_channels, - out_channels=None, - conv_shortcut=False, - dropout=0.0, - temb_channels=512, - groups=32, - groups_out=None, - pre_norm=True, - eps=1e-6, - non_linearity="swish", - time_embedding_norm="default", - kernel=None, - output_scale_factor=1.0, - use_in_shortcut=None, - up=False, - down=False, - dtype=None, - ): - super().__init__() - self.pre_norm = pre_norm - self.pre_norm = True - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.use_conv_shortcut = conv_shortcut - self.time_embedding_norm = time_embedding_norm - self.up = up - self.down = down - self.output_scale_factor = output_scale_factor - - if groups_out is None: - groups_out = groups - - self.norm1 = GroupNorm(num_groups=groups, - num_channels=in_channels, - eps=eps, - affine=True, - dtype=dtype) - self.conv1 = Conv2d(in_channels, - out_channels, - kernel_size=(3, 3), - stride=(1, 1), - padding=(1, 1), - dtype=dtype) - - if temb_channels is not None: - self.time_emb_proj = Linear(temb_channels, - out_channels, - dtype=dtype) - else: - self.time_emb_proj = None - - self.norm2 = GroupNorm(num_groups=groups_out, - num_channels=out_channels, - eps=eps, - affine=True, - dtype=dtype) - self.conv2 = Conv2d(out_channels, - out_channels, - kernel_size=(3, 3), - stride=(1, 1), - padding=(1, 1), - dtype=dtype) - - if non_linearity == "swish": - self.nonlinearity = lambda x: silu(x) - elif non_linearity == "mish": - self.nonlinearity = Mish() - elif non_linearity == "silu": - self.nonlinearity = silu - - self.upsample = self.downsample = None - #TODO add the fir kernel supporting. - if self.up: - if kernel == "sde_vp": - self.upsample = partial(interpolate, - scale_factor=2.0, - mode="nearest") - else: - self.upsample = Upsample2D(in_channels, - use_conv=False, - dtype=dtype) - elif self.down: - - if kernel == "sde_vp": - self.downsample = partial(avg_pool2d, kernel_size=2, stride=2) - else: - self.downsample = Downsample2D(in_channels, - use_conv=False, - padding=1, - name="op", - dtype=dtype) - - self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut - - if self.use_in_shortcut: - self.conv_shortcut = Conv2d(in_channels, - out_channels, - kernel_size=(1, 1), - stride=(1, 1), - padding=(0, 0), - dtype=dtype) - else: - self.conv_shortcut = None - - def forward(self, input_tensor, temb): - hidden_states = input_tensor - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - - if self.upsample is not None: - input_tensor = self.upsample(input_tensor) - hidden_states = self.upsample(hidden_states) - elif self.downsample is not None: - input_tensor = self.downsample(input_tensor) - hidden_states = self.downsample(hidden_states) - - hidden_states = self.conv1(hidden_states) - if self.time_emb_proj is not None: - temb = self.time_emb_proj(self.nonlinearity(temb)) - new_shape = list(temb.size()) - new_shape.extend([1, 1]) #[:,:,None,None] ->view - temb = temb.view(new_shape) - - assert self.time_embedding_norm == "default" - if temb is not None: - hidden_states = hidden_states + temb - - hidden_states = self.norm2(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor) - - output_tensor = (input_tensor + - hidden_states) / self.output_scale_factor - return output_tensor diff --git a/tensorrt_llm/models/unet/unet_2d_blocks.py b/tensorrt_llm/models/unet/unet_2d_blocks.py deleted file mode 100644 index 02479221292e..000000000000 --- a/tensorrt_llm/models/unet/unet_2d_blocks.py +++ /dev/null @@ -1,636 +0,0 @@ -# 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 Tuple, Union - -from ...functional import concat -from ...module import Module, ModuleList -from .attention import AttentionBlock, Transformer2DModel -from .resnet import Downsample2D, ResnetBlock2D, Upsample2D - - -def get_down_block( - down_block_type, - num_layers, - in_channels, - out_channels, - temb_channels, - add_downsample, - resnet_eps, - resnet_act_fn, - attn_num_head_channels, - transformer_layers_per_block=1, - cross_attention_dim=None, - downsample_padding=None, - use_linear_projection=False, - dtype=None, -): - down_block_type = down_block_type[7:] if down_block_type.startswith( - "UNetRes") else down_block_type - if down_block_type == "DownBlock2D": - return DownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - downsample_padding=downsample_padding, - dtype=dtype, - ) - elif down_block_type == "CrossAttnDownBlock2D": - if cross_attention_dim is None: - raise ValueError( - "cross_attention_dim must be specified for CrossAttnUpBlock2D") - return CrossAttnDownBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - downsample_padding=downsample_padding, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attn_num_head_channels, - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - - raise ValueError(f"{down_block_type} does not exist.") - - -def get_up_block( - up_block_type, - num_layers, - in_channels, - out_channels, - prev_output_channel, - temb_channels, - add_upsample, - resnet_eps, - resnet_act_fn, - attn_num_head_channels, - transformer_layers_per_block=1, - cross_attention_dim=None, - use_linear_projection=False, - dtype=None, -): - up_block_type = up_block_type[7:] if up_block_type.startswith( - "UNetRes") else up_block_type - if up_block_type == "UpBlock2D": - return UpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - dtype=dtype, - ) - elif up_block_type == "CrossAttnUpBlock2D": - if cross_attention_dim is None: - raise ValueError( - "cross_attention_dim must be specified for CrossAttnUpBlock2D") - return CrossAttnUpBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attn_num_head_channels, - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - - raise ValueError(f"{up_block_type} does not exist.") - - -class UpBlock2D(Module): - - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor=1.0, - add_upsample=True, - dtype=None, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - self.resnets = ModuleList(resnets) - - if add_upsample: - self.upsamplers = ModuleList([ - Upsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - dtype=dtype) - ]) - else: - self.upsamplers = None - - def forward(self, - hidden_states, - res_hidden_states_tuple, - temb=None, - upsample_size=None): - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = concat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class DownBlock2D(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor=1.0, - add_downsample=True, - downsample_padding=1, - dtype=None, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - self.resnets = ModuleList(resnets) - - if add_downsample: - self.downsamplers = ModuleList([ - Downsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - dtype=dtype) - ]) - else: - self.downsamplers = None - - def forward(self, hidden_states, temb=None): - output_states = () - for resnet in self.resnets: - hidden_states = resnet(hidden_states, temb) - - output_states += (hidden_states, ) - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states, ) - return hidden_states, output_states - - -class UNetMidBlock2D(Module): - - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - attention_type="default", - output_scale_factor=1.0, - dtype=None, - **kwargs, - ): - super().__init__() - - self.attention_type = attention_type - resnet_groups = resnet_groups if resnet_groups is not None else min( - in_channels // 4, 32) - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - ) - ] - attentions = [] - - for _ in range(num_layers): - attentions.append( - AttentionBlock( - in_channels, - num_head_channels=attn_num_head_channels, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - num_groups=resnet_groups, - dtype=dtype, - )) - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - def forward(self, hidden_states, temb=None, encoder_states=None): - hidden_states = self.resnets[0](hidden_states, temb) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - if self.attention_type == "default": - hidden_states = attn(hidden_states) - else: - hidden_states = attn(hidden_states, encoder_states) - hidden_states = resnet(hidden_states, temb) - - return hidden_states - - -class CrossAttnUpBlock2D(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - cross_attention_dim=1280, - attention_type="default", - output_scale_factor=1.0, - add_upsample=True, - use_linear_projection: bool = False, - dtype=None, - ): - super().__init__() - resnets = [] - attentions = [] - - self.attention_type = attention_type - self.attn_num_head_channels = attn_num_head_channels - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * num_layers - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - attentions.append( - Transformer2DModel(in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - num_attention_heads=attn_num_head_channels, - attention_head_dim=out_channels // - attn_num_head_channels, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - cross_attention_dim=cross_attention_dim, - dtype=dtype)) - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - if add_upsample: - self.upsamplers = ModuleList([ - Upsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - dtype=dtype) - ]) - else: - self.upsamplers = None - - def forward( - self, - hidden_states, - res_hidden_states_tuple, - temb=None, - encoder_hidden_states=None, - upsample_size=None, - ): - - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = concat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - hidden_states = attn(hidden_states, context=encoder_hidden_states) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class CrossAttnDownBlock2D(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - cross_attention_dim=1280, - attention_type="default", - output_scale_factor=1.0, - downsample_padding=1, - add_downsample=True, - use_linear_projection: bool = False, - dtype=None, - ): - super().__init__() - resnets = [] - attentions = [] - - self.attention_type = attention_type - self.attn_num_head_channels = attn_num_head_channels - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * num_layers - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - attentions.append( - Transformer2DModel(in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - num_attention_heads=attn_num_head_channels, - attention_head_dim=out_channels // - attn_num_head_channels, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - cross_attention_dim=cross_attention_dim, - dtype=dtype)) - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - if add_downsample: - self.downsamplers = ModuleList([ - Downsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - dtype=dtype) - ]) - else: - self.downsamplers = None - - def forward(self, hidden_states, temb=None, encoder_hidden_states=None): - output_states = () - - for resnet, attn in zip(self.resnets, self.attentions): - - hidden_states = resnet(hidden_states, temb) - - hidden_states = attn(hidden_states, context=encoder_hidden_states) - - output_states += (hidden_states, ) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states, ) - - return hidden_states, output_states - - -class UNetMidBlock2DCrossAttn(Module): - - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - attention_type="default", - output_scale_factor=1.0, - cross_attention_dim=1280, - use_linear_projection: bool = False, - dtype=None, - **kwargs, - ): - super().__init__() - - self.attention_type = attention_type - self.attn_num_head_channels = attn_num_head_channels - resnet_groups = resnet_groups if resnet_groups is not None else min( - in_channels // 4, 32) - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * num_layers - - # there is always at least one resnet - resnets = [ - ResnetBlock2D(in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype) - ] - attentions = [] - - for i in range(num_layers): - attentions.append( - Transformer2DModel(in_channels=in_channels, - num_layers=transformer_layers_per_block[i], - num_attention_heads=attn_num_head_channels, - attention_head_dim=in_channels // - attn_num_head_channels, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - cross_attention_dim=cross_attention_dim, - dtype=dtype)) - resnets.append( - ResnetBlock2D(in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype)) - - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - def forward(self, hidden_states, temb=None, encoder_hidden_states=None): - hidden_states = self.resnets[0](hidden_states, temb) - - for attn, resnet in zip(self.attentions, self.resnets[1:]): - hidden_states = attn(hidden_states, encoder_hidden_states) - hidden_states = resnet(hidden_states, temb) - - return hidden_states diff --git a/tensorrt_llm/models/unet/unet_2d_condition.py b/tensorrt_llm/models/unet/unet_2d_condition.py deleted file mode 100644 index 96531df8db7e..000000000000 --- a/tensorrt_llm/models/unet/unet_2d_condition.py +++ /dev/null @@ -1,249 +0,0 @@ -# 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 - -from ...functional import cast, concat, silu -from ...layers import Conv2d, GroupNorm -from ...module import Module, ModuleList -from .embeddings import TimestepEmbedding, Timesteps -from .unet_2d_blocks import (UNetMidBlock2DCrossAttn, get_down_block, - get_up_block) - - -class UNet2DConditionModel(Module): - - def __init__( - self, - sample_size=None, - in_channels=4, - out_channels=4, - center_input_sample=False, - flip_sin_to_cos=True, - freq_shift=0, - down_block_types=("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", "DownBlock2D"), - up_block_types=("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", - "CrossAttnUpBlock2D"), - block_out_channels=(320, 640, 1280, 1280), - layers_per_block=2, - downsample_padding=1, - mid_block_scale_factor=1.0, - act_fn="silu", - norm_num_groups=32, - norm_eps=1e-5, - cross_attention_dim=1280, - transformer_layers_per_block=1, - attention_head_dim=8, - use_linear_projection=False, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - projection_class_embeddings_input_dim: Optional[int] = None, - dtype=None, - ): - super().__init__() - - self.sample_size = sample_size - self.addition_embed_type = addition_embed_type - time_embed_dim = block_out_channels[0] * 4 - - # input - self.conv_in = Conv2d(in_channels, - block_out_channels[0], - kernel_size=(3, 3), - padding=(1, 1), - dtype=dtype) - # time - self.time_proj = Timesteps(block_out_channels[0], - flip_sin_to_cos, - freq_shift, - dtype=dtype) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding(timestep_input_dim, - time_embed_dim, - dtype=dtype) - - if addition_embed_type == "text_time": - self.add_time_proj = Timesteps(addition_time_embed_dim, - flip_sin_to_cos, - freq_shift, - dtype=dtype) - self.add_embedding = TimestepEmbedding( - projection_class_embeddings_input_dim, - time_embed_dim, - dtype=dtype) - - down_blocks = [] - up_blocks = [] - - if isinstance(attention_head_dim, int): - attention_head_dim = (attention_head_dim, ) * len(down_block_types) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * len(down_block_types) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attention_head_dim[i], - downsample_padding=downsample_padding, - use_linear_projection=use_linear_projection, - dtype=dtype) - down_blocks.append(down_block) - self.down_blocks = ModuleList(down_blocks) - # mid - self.mid_block = UNetMidBlock2DCrossAttn( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - transformer_layers_per_block=transformer_layers_per_block[-1], - resnet_time_scale_shift="default", - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attention_head_dim[-1], - resnet_groups=norm_num_groups, - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_attention_head_dim = list(reversed(attention_head_dim)) - reversed_transformer_layers_per_block = list( - reversed(transformer_layers_per_block)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min( - i + 1, - len(block_out_channels) - 1)] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=layers_per_block + 1, - transformer_layers_per_block= - reversed_transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=time_embed_dim, - add_upsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=reversed_attention_head_dim[i], - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - up_blocks.append(up_block) - prev_output_channel = output_channel - self.up_blocks = ModuleList(up_blocks) - # out - self.conv_norm_out = GroupNorm(num_channels=block_out_channels[0], - num_groups=norm_num_groups, - eps=norm_eps, - dtype=dtype) - self.conv_act = silu - self.conv_out = Conv2d(block_out_channels[0], - out_channels, (3, 3), - padding=(1, 1), - dtype=dtype) - - def forward(self, - sample, - timesteps, - encoder_hidden_states, - text_embeds=None, - time_ids=None): - # time - t_emb = self.time_proj(timesteps) - emb = self.time_embedding(t_emb) - - aug_emb = None - if self.addition_embed_type == "text_time": - assert text_embeds is not None and time_ids is not None - time_embeds = self.add_time_proj(time_ids.view([-1])) - time_embeds = time_embeds.view([text_embeds.shape[0], -1]) - add_embeds = concat([text_embeds, time_embeds], dim=1) - add_embeds = cast(add_embeds, emb.dtype) - aug_emb = self.add_embedding(add_embeds) - - emb = emb + aug_emb if aug_emb is not None else emb - - sample = self.conv_in(sample) - - down_block_res_samples = (sample, ) - for downsample_block in self.down_blocks: - - if hasattr( - downsample_block, - "attentions") and downsample_block.attentions is not None: - - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states) - else: - sample, res_samples = downsample_block(hidden_states=sample, - temb=emb) - down_block_res_samples += res_samples - - sample = self.mid_block(sample, - emb, - encoder_hidden_states=encoder_hidden_states) - - for upsample_block in self.up_blocks: - - res_samples = down_block_res_samples[-len(upsample_block.resnets):] - down_block_res_samples = down_block_res_samples[:-len(upsample_block - .resnets)] - - if hasattr(upsample_block, - "attentions") and upsample_block.attentions is not None: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - ) - else: - sample = upsample_block(hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples) - - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - return sample diff --git a/tensorrt_llm/models/unet/weights.py b/tensorrt_llm/models/unet/weights.py deleted file mode 100644 index 9a8e4c7c63d1..000000000000 --- a/tensorrt_llm/models/unet/weights.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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 time - -import numpy as np - -from ...logger import logger - - -def update_timestep_weight(src, dst): - dst.linear_1.update_parameters(src.linear_1) - dst.linear_2.update_parameters(src.linear_2) - - -def update_crossattn_downblock_2d_weight(src, dst): - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - for index, value in enumerate(src.attentions): - update_transformer_2d_model_weight(dst.attentions[index], value) - - if src.downsamplers: - for index, value in enumerate(src.downsamplers): - dst.downsamplers[index].conv.update_parameters(value.conv) - - -def update_transformer_2d_model_weight(gm, m): - gm.norm.update_parameters(m.norm) - gm.proj_in.update_parameters(m.proj_in) - for i in range(len(gm.transformer_blocks)): - gm.transformer_blocks[i].attn1.to_qkv.weight.value = np.concatenate( - (m.transformer_blocks[i].attn1.to_q.weight.detach().cpu().numpy(), - m.transformer_blocks[i].attn1.to_k.weight.detach().cpu().numpy(), - m.transformer_blocks[i].attn1.to_v.weight.detach().cpu().numpy())) - gm.transformer_blocks[i].attn1.to_out.update_parameters( - m.transformer_blocks[i].attn1.to_out[0]) - - gm.transformer_blocks[i].attn2.to_q.update_parameters( - m.transformer_blocks[i].attn2.to_q) - gm.transformer_blocks[i].attn2.to_kv.weight.value = np.concatenate( - (m.transformer_blocks[i].attn2.to_k.weight.detach().cpu().numpy(), - m.transformer_blocks[i].attn2.to_v.weight.detach().cpu().numpy())) - gm.transformer_blocks[i].attn2.to_out.update_parameters( - m.transformer_blocks[i].attn2.to_out[0]) - - gm.transformer_blocks[i].ff.proj_in.update_parameters( - m.transformer_blocks[i].ff.net[0].proj) - gm.transformer_blocks[i].ff.proj_out.update_parameters( - m.transformer_blocks[i].ff.net[2]) - - gm.transformer_blocks[i].norm1.update_parameters( - m.transformer_blocks[i].norm1) - gm.transformer_blocks[i].norm2.update_parameters( - m.transformer_blocks[i].norm2) - gm.transformer_blocks[i].norm3.update_parameters( - m.transformer_blocks[i].norm3) - - gm.proj_out.update_parameters(m.proj_out) - - -def update_upblock_2d_weight(src, dst): - if src.upsamplers: - for index, value in enumerate(src.upsamplers): - dst.upsamplers[index].conv.update_parameters(value.conv) - - for index, value in enumerate(src.resnets): - dst.resnets[index].norm1.update_parameters(value.norm1) - dst.resnets[index].conv1.update_parameters(value.conv1) - - dst.resnets[index].norm2.update_parameters(value.norm2) - dst.resnets[index].conv2.update_parameters(value.conv2) - - if value.conv_shortcut: - dst.resnets[index].conv_shortcut.update_parameters( - value.conv_shortcut) - - dst.resnets[index].time_emb_proj.update_parameters(value.time_emb_proj) - - -def update_downblock_2d_weight(src, dst): - if src.downsamplers: - for index, value in enumerate(src.downsamplers): - dst.downsamplers[index].conv.update_parameters(value.conv) - - for index, value in enumerate(src.resnets): - dst.resnets[index].norm1.update_parameters(value.norm1) - dst.resnets[index].conv1.update_parameters(value.conv1) - - dst.resnets[index].norm2.update_parameters(value.norm2) - dst.resnets[index].conv2.update_parameters(value.conv2) - - if value.conv_shortcut: - dst.resnets[index].conv_shortcut.update_parameters( - value.conv_shortcut) - - dst.resnets[index].time_emb_proj.update_parameters(value.time_emb_proj) - - -def update_unet_mid_block_2d_weight(src, dst): - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - for index, value in enumerate(src.attentions): - update_transformer_2d_model_weight(dst.attentions[index], value) - - -def update_crossattn_upblock_2d_weight(src, dst): - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - for index, value in enumerate(src.attentions): - update_transformer_2d_model_weight(dst.attentions[index], value) - if src.upsamplers: - for index, value in enumerate(src.upsamplers): - dst.upsamplers[index].conv.update_parameters(value.conv) - - -def update_resnet_block_weight(src, dst): - dst.norm1.update_parameters(src.norm1) - dst.conv1.update_parameters(src.conv1) - - dst.norm2.update_parameters(src.norm2) - dst.conv2.update_parameters(src.conv2) - - dst.time_emb_proj.update_parameters(src.time_emb_proj) - if src.conv_shortcut: - dst.conv_shortcut.update_parameters(src.conv_shortcut) - - -def update_unetmidblock_2d_weight(src, dst): - for index, value in enumerate(src.attentions): - dst.attentions[index].group_norm.update_parameters(value.group_norm) - dst.attentions[index].proj_attn.update_parameters(value.proj_attn) - - dst.attentions[index].qkv.weight.value = np.concatenate( - (value.query.weight.detach().cpu().numpy(), - value.key.weight.detach().cpu().numpy(), - value.value.weight.detach().cpu().numpy())) - dst.attentions[index].qkv.bias.value = np.concatenate( - (value.query.bias.detach().cpu().numpy(), - value.key.bias.detach().cpu().numpy(), - value.value.bias.detach().cpu().numpy())) - - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - -def update_unet_2d_condition_model_weights(src, dst): - dst.conv_in.update_parameters(src.conv_in) - - dst.time_embedding.update_parameters(src.time_embedding) - if src.config.addition_embed_type: - dst.add_embedding.update_parameters(src.add_embedding) - - for index, type in enumerate(src.config.down_block_types): - if type == 'CrossAttnDownBlock2D': - update_crossattn_downblock_2d_weight(src.down_blocks[index], - dst.down_blocks[index]) - elif type == 'DownBlock2D': - update_downblock_2d_weight(src.down_blocks[index], - dst.down_blocks[index]) - - update_unet_mid_block_2d_weight(src.mid_block, dst.mid_block) - - for index, type in enumerate(src.config.up_block_types): - if type == 'CrossAttnUpBlock2D': - update_crossattn_upblock_2d_weight(src.up_blocks[index], - dst.up_blocks[index]) - elif type == 'UpBlock2D': - update_upblock_2d_weight(src.up_blocks[index], dst.up_blocks[index]) - - dst.conv_norm_out.update_parameters(src.conv_norm_out) - - dst.conv_out.update_parameters(src.conv_out) - - -def load_from_hf_unet(src, dst): - logger.info('Loading weights from HF Unet...') - tik = time.time() - - update_unet_2d_condition_model_weights(src, dst) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') diff --git a/tensorrt_llm/module.py b/tensorrt_llm/module.py deleted file mode 100644 index c67674ec15b2..000000000000 --- a/tensorrt_llm/module.py +++ /dev/null @@ -1,288 +0,0 @@ -# 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 operator - -from ._common import default_net -from .logger import logger - - -def _addindent(s_, numSpaces): - s = s_.split("\n") - # don't do anything for single-line stuff - if len(s) == 1: - return s_ - first = s.pop(0) - s = [(numSpaces * " ") + line for line in s] - s = "\n".join(s) - s = first + "\n" + s - return s - - -class Module(object): - def __init__(self) -> None: - self._modules = {} - self._parameters = {} - self._network_outputs = {} - - def forward(self, *args, **kwargs): - raise NotImplementedError - - def __call__(self, *args, **kwargs): - current_net = default_net() - if not current_net._module_call_stack.module_names_set(): - logger.debug("Initializing top level module") - current_net._module_call_stack.set_module_names(self) - unique_name = current_net._module_call_stack.get_mod_name(self) - with current_net._module_call_stack.call_stack_mgr() as stack: - stack.append(unique_name) - start_layer_idx = current_net.trt_network.num_layers - output = self.forward(*args, **kwargs) - end_layer_idx = current_net.trt_network.num_layers - current_net._module_call_stack.set_layer_range( - self, range(start_layer_idx, end_layer_idx) - ) - return output - - def __getattr__(self, name): - parameters = self.__dict__.get("_parameters") - if parameters is not None and name in parameters: - return parameters[name] - - modules = self.__dict__.get("_modules") - if modules is not None and name in modules: - return modules[name] - - raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, name)) - - def __setattr__(self, name, value) -> None: - from .parameter import Parameter - - # Improved module setattr to handle one edge case: - # attribute could be first set to None and later reset to Parameter / Module class - - try: - super().__getattribute__(name) - - except AttributeError: - # if base class doesn't have the attribute, no matter we init or reset: - # - keep Parameter and Module attrs in this Module class - # - leave all other attrs in base class - if isinstance(value, Parameter): - self.__dict__.get("_parameters")[name] = value - elif isinstance(value, Module): - self.__dict__.get("_modules")[name] = value - else: - super().__setattr__(name, value) - - else: - # if base class has the attribute, reset as follows: - # - when reset as Parameter or Module attr, remove from base & add to this Module class - # - other types reset and remain in base class - if isinstance(value, Parameter): - super().__delattr__(name) - self.__dict__.get("_parameters")[name] = value - elif isinstance(value, Module): - super().__delattr__(name) - self.__dict__.get("_modules")[name] = value - else: - super().__setattr__(name, value) - - def named_modules(self, memo=None, prefix="", remove_duplicate=True): - if memo is None: - memo = set() - if self not in memo: - if remove_duplicate: - memo.add(self) - yield prefix, self - for name, module in self._modules.items(): - if module is None: - continue - submodule_prefix = prefix + ("." if prefix else "") + name - for m in module.named_modules(memo, submodule_prefix, remove_duplicate): - yield m - - def named_modules_with_parent(self, memo=None, prefix="", parent=None, remove_duplicate=True): - if memo is None: - memo = set() - if self not in memo: - if remove_duplicate: - memo.add(self) - yield prefix, self, parent - - if parent: - # Use the up-to-date module from the parent, to allow replacing - # layers while iterating this generator. - module_name = prefix.rsplit(".", 1)[-1] - module = getattr(parent, module_name) - if module is None: - return - else: - module = self - - for child_name, child_module in module._modules.items(): - if child_module is None: - continue - submodule_prefix = prefix + ("." if prefix else "") + child_name - for m in child_module.named_modules_with_parent( - memo, submodule_prefix, module, remove_duplicate - ): - yield m - - def named_children(self): - memo = set() - for name, module in self._modules.items(): - if module is not None and module not in memo: - memo.add(module) - yield name, module - - def _named_members(self, get_members_fn, prefix="", recurse=True): - memo = set() - modules = self.named_modules(prefix=prefix) if recurse else [(prefix, self)] - for module_prefix, module in modules: - members = get_members_fn(module) - for k, v in members: - if v is None or v in memo: - continue - memo.add(v) - name = module_prefix + ("." if module_prefix else "") + k - yield name, v - - def parameters(self, recurse=True): - for name, param in self.named_parameters(): - yield param - - def named_parameters(self, prefix="", recurse=True): - gen = self._named_members( - lambda module: module._parameters.items(), prefix=prefix, recurse=recurse - ) - for elem in gen: - yield elem - - def children(self): - for _, module in self.named_children(): - yield module - - def apply(self, fn): - for module in self.children(): - module.apply(fn) - fn(self) - return self - - def _get_name(self): - return self.__class__.__name__ - - def register_parameter(self, name, param): - if param is None: - self._parameters[name] = None - else: - self._parameters[name] = param - - def register_network_output(self, name, value): - self._network_outputs[name] = value - - def named_network_outputs(self): - for name, module in self.named_modules(): - for n, output in module._network_outputs.items(): - yield name + ("." if name else "") + n, output - - def update_parameters(self, torch_module): - m = {k: v for k, v in self.named_parameters()} - tm = {k: v for k, v in torch_module.named_parameters()} - - assert sorted(m.keys()) == sorted(tm.keys()), ( - "The parameter names of the TensorRT LLM module must be the same with the torch module" - ) - - for k, v in self.named_parameters(): - v.value = tm[k].detach().cpu().numpy() - - def _get_name(self): - return self.__class__.__name__ - - def __repr__(self): - # We treat the extra repr like the sub-module, one item per line - child_lines = [] - for key, module in self._modules.items(): - mod_str = repr(module) - mod_str = _addindent(mod_str, 2) - child_lines.append("(" + key + "): " + mod_str) - main_str = self._get_name() + "(" - if child_lines: - # simple one-liner info, which most builtin Modules will use - main_str += "\n " + "\n ".join(child_lines) + "\n" - main_str += ")" - return main_str - - -class ModuleList(Module): - def __init__(self, modules) -> None: - super(ModuleList, self).__init__() - offset = len(self) - for i, module in enumerate(modules): - self._modules[str(offset + i)] = module - - def _get_abs_string_index(self, idx): - """Get the absolute index for the list of modules.""" - idx = operator.index(idx) - if not (-len(self) <= idx < len(self)): - raise IndexError("index {} is out of range".format(idx)) - if idx < 0: - idx += len(self) - return str(idx) - - def __getitem__(self, idx): - if isinstance(idx, slice): - return self.__class__(list(self._modules.values())[idx]) - else: - return self._modules[self._get_abs_string_index(idx)] - - def __setitem__(self, idx, module) -> None: - idx = self._get_abs_string_index(idx) - return setattr(self, str(idx), module) - - def __len__(self): - return len(self._modules) - - def __repr__(self): - """Return a custom repr for ModuleList that compresses repeated module representations.""" - list_of_reprs = [repr(item) for item in self] - if len(list_of_reprs) == 0: - return self._get_name() + "()" - - start_end_indices = [[0, 0]] - repeated_blocks = [list_of_reprs[0]] - for i, r in enumerate(list_of_reprs[1:], 1): - if r == repeated_blocks[-1]: - start_end_indices[-1][1] += 1 - continue - - start_end_indices.append([i, i]) - repeated_blocks.append(r) - - lines = [] - main_str = self._get_name() + "(" - for (start_id, end_id), b in zip(start_end_indices, repeated_blocks): - local_repr = f"({start_id}): {b}" # default repr - - if start_id != end_id: - n = end_id - start_id + 1 - local_repr = f"({start_id}-{end_id}): {n} x {b}" - - local_repr = _addindent(local_repr, 2) - lines.append(local_repr) - - main_str += "\n " + "\n ".join(lines) + "\n" - main_str += ")" - return main_str diff --git a/tensorrt_llm/network.py b/tensorrt_llm/network.py deleted file mode 100644 index e2d3dc453f6b..000000000000 --- a/tensorrt_llm/network.py +++ /dev/null @@ -1,959 +0,0 @@ -# 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 collections -import contextlib -import hashlib -import inspect -import weakref -from collections import OrderedDict, defaultdict -from dataclasses import dataclass, field -from pathlib import Path -from typing import (Any, Dict, Iterable, List, Optional, OrderedDict, Set, - Tuple, Union) - -import numpy as np -import onnx -import onnx_graphsurgeon as gs -import tensorrt as trt - -from tensorrt_llm.module import Module - -from ._common import set_network -from ._utils import get_extra_attr, has_extra_attr, set_extra_attr -from .logger import logger -from .plugin import PluginConfig - - -class _UniqueNameGenerator(object): - - def __init__(self, prefix=''): - self.ids = collections.defaultdict(int) - self.prefix = prefix - - def __call__(self, key, module_name=''): - if module_name != '': - module_name = module_name.replace(".", "/") - key = module_name + '/' + key - tmp = self.ids[key] - self.ids[key] += 1 - return f"{self.prefix}{key}_{tmp}" - - -class PluginInfo: - plugin_creator: trt.IPluginCreator - plugin_name: str - pfc: trt.PluginFieldCollection - - def __init__(self, plugin_creator: trt.IPluginCreator, plugin_name: str, - pfc: trt.PluginFieldCollection): - self.plugin_creator = plugin_creator - self.plugin_name = plugin_name - self.pfc = pfc - self._parse_pfc(pfc) - - def _parse_pfc(self, pfc: trt.PluginFieldCollection): - self.pfc_as_ndarray = {} - self.pfc_as_list = {} - for i in range(len(pfc)): - name, data = pfc[i].name, pfc[i].data - array_data = data - self.pfc_as_ndarray[name] = array_data.copy() - list_data = array_data.tolist() - self.pfc_as_list[name] = list_data - - -def get_plugin_info(trt_network: trt.INetworkDefinition, - layer_name: str) -> PluginInfo: - if not has_extra_attr(trt_network, "plugin_infos"): - return None - plugin_infos = get_extra_attr(trt_network, "plugin_infos") - if layer_name not in plugin_infos: - return None - return plugin_infos[layer_name] - - -def set_plugin_info(trt_network: trt.INetworkDefinition, layer_name: str, - plugin_info: PluginInfo): - if not has_extra_attr(trt_network, "plugin_infos"): - set_extra_attr(trt_network, "plugin_infos", {}) - plugin_infos = get_extra_attr(trt_network, "plugin_infos") - plugin_infos[layer_name] = plugin_info - - -def delete_plugin_info(trt_network: trt.INetworkDefinition, layer_name: str): - if not has_extra_attr(trt_network, "plugin_infos"): - return - plugin_infos = get_extra_attr(trt_network, "plugin_infos") - if layer_name not in plugin_infos: - return - del plugin_infos[layer_name] - - -# TODO: remove this WAR after https://nvbugs/4359151 fixed. -def get_np_weight(trt_network: trt.INetworkDefinition, - layer_name: str) -> np.array: - if not has_extra_attr(trt_network, "np_weights"): - return None - np_weights = get_extra_attr(trt_network, "np_weights") - if layer_name not in np_weights: - return None - return np_weights[layer_name] - - -# TODO: remove this WAR after https://nvbugs/4359151 fixed. -def set_np_weight(trt_network: trt.INetworkDefinition, layer_name: str, - np_weight: np.array): - if not has_extra_attr(trt_network, "np_weights"): - set_extra_attr(trt_network, "np_weights", {}) - np_weights = get_extra_attr(trt_network, "np_weights") - np_weights[layer_name] = np_weight - - -class Network(object): - - def __init__(self, **kwargs): - # intentionally use **kwargs, user should never call this ctor directly - # use Builder.create_network() instead - - # Holds the removed layers and disable them in graph rewriting and other phases. - # This is a hacky way since INetwork python API doesn't provide a way to remove a layer. - # TODO: remove this when TensorRT provides a better way to remove a layer - self._removed_layers: Set[str] = set() - - self.is_graph_altered = False - - from .graph_rewriting import FLayerInfoMemo - self.flayer_memo = FLayerInfoMemo() # holds the functional metadata - self._parameter_tensors = {} # holds the parameter tensors - - def _init(self, trt_network): - self._trt_network = trt_network - self._inputs = {} - self._named_parameters = None - # layer precision of a given scope, this is used together with precision(dtype) context manager - self._dtype = None - self._name_generator = _UniqueNameGenerator() - self._plugin_config = PluginConfig() - self._module_call_stack = _TrtLlmModuleCallStack() - self._registered_ndarrays = [] - self._strongly_typed = trt.INetworkDefinition.get_flag( - self._trt_network, trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) - self._unfilled_weights: Dict[str, Tuple[np.array, np.array]] = {} - - return self - - def _register_unfilled_weights(self, layer_name: str, weights: np.array, - values: np.array): - self._unfilled_weights[layer_name] = (weights, values) - - def _fill_weights(self): - from tensorrt_llm.parameter import Parameter - - for layer_name in list(self._unfilled_weights.keys()): - weights, values = self._unfilled_weights.pop(layer_name) - self.register_ndarray(weights) - if values is not None: - np.copyto(weights, values, casting='no') - else: - Parameter.xavier_init(weights) - - @property - def parameter_tensors(self): - return self._parameter_tensors - - def get_parameter_tensor(self, param): - return self.parameter_tensors.get(param, None) - - def set_parameter_tensor(self, param, tensor): - assert param not in self.parameter_tensors - self.parameter_tensors[param] = tensor - - @property - def dtype(self) -> trt.DataType: - return self._dtype - - @dtype.setter - def dtype(self, dtype: trt.DataType): - assert isinstance(dtype, trt.DataType) or dtype is None - self._dtype = dtype - - @property - def trt_network(self) -> trt.INetworkDefinition: - return self._trt_network - - @property - def plugin_config(self) -> PluginConfig: - return self._plugin_config - - @plugin_config.setter - def plugin_config(self, cfg: PluginConfig): - assert isinstance( - cfg, - PluginConfig), f"Expecting a PluginConfig object, got {type(cfg)}" - self._plugin_config = cfg - - @property - def strongly_typed(self) -> bool: - return self._strongly_typed - - def _add_input(self, - tensor, - name, - dtype, - shape, - dim_range: OrderedDict = None): - assert isinstance(dtype, trt.DataType) - tensor.trt_tensor = self.trt_network.add_input( - name=name, - shape=shape, - dtype=dtype, - ) - assert tensor.trt_tensor is not None, f"Couldn't create TRT tensor for {name} {dtype} {shape}" - if dim_range is not None: - logger.debug( - f'Add input: {name}, shape: {shape}, dtype: {dtype}, dimension names:{list(dim_range.keys())}' - ) - for i, dim_name in enumerate(dim_range.keys()): - tensor.trt_tensor.set_dimension_name(i, str(dim_name)) - else: - logger.debug(f'Add input: {name}, shape: {shape}, dtype: {dtype}') - self._inputs[name] = tensor - - def _mark_output(self, tensor, name, dtype): - from .functional import cast - - # In strongly_typed, if tensor output is not the same, add a cast - if dtype is not None and self.strongly_typed: - tensor = cast(tensor, dtype) - self.trt_network.mark_output(tensor.trt_tensor) - tensor.trt_tensor.name = name - if not self.strongly_typed: - tensor.trt_tensor.dtype = dtype or tensor.trt_tensor.dtype - logger.debug(f'Mark output: {name}, dtype: {dtype}') - - def set_named_parameters(self, named_parameters): - self._named_parameters = named_parameters - - @property - def named_parameters(self): - return self._named_parameters - - def _set_layer_name(self, layer): - original_layer_name = layer.name - layer_name = str(layer.type).split('.')[-1] - current_module = self._module_call_stack.get_current_module() - - func_stack = [] - frame = inspect.currentframe().f_back.f_back - while frame: - func_name = frame.f_code.co_name - line_num = frame.f_lineno - if func_name == "forward": - break - func_stack.insert(0, f"{func_name}_L{line_num}") - if len(func_stack) >= 10: - # NOTE: TRT error messages has a character limit. - # Limiting to only 10 levels helps retain - # the true error message from TRT. - break - frame = frame.f_back - current_module = f"{current_module}.{'.'.join(func_stack)}" - - if layer.type == trt.LayerType.PLUGIN_V2: - layer_name = '_'.join( - [layer_name, - str(layer.plugin.plugin_type).split('.')[-1]]) - elif layer.type in [ - trt.LayerType.UNARY, trt.LayerType.REDUCE, - trt.LayerType.ELEMENTWISE - ]: - layer_name = '_'.join([layer_name, str(layer.op).split('.')[-1]]) - - layer.name = self._name_generator(layer_name, current_module) - for idx in range(layer.num_outputs): - # TRT initializes tensor names from the initial layer's name when the layer is created, - # and does not update tensor names when layer name changed by application, needs to - # change the tensor name to align with the new layer name for better debugging - layer.get_output(idx).name = f"{layer.name}_output_{idx}" - if original_layer_name != layer.name: - if layer.type == trt.LayerType.PLUGIN_V2: - plugin_info = get_plugin_info(self.trt_network, - original_layer_name) - if plugin_info is not None: - set_plugin_info(self.trt_network, layer.name, plugin_info) - delete_plugin_info(self.trt_network, original_layer_name) - - # Set layer metadata to the same as the layer name so that it can show up in NVTX. - layer.metadata = layer.name - - def register_ndarray(self, ndarray: np.ndarray) -> None: - ''' When the functional APIs need to create local numpy array and use as weights for constant or other layers, - they need to register the ndarray objects to the TRT-LLM Network to prolong the lifetime of the ndarray, such that weights are - still valid when functional API returned. - All the weights referenced by the trt Network are weak referenced, it's TRT-LLM's responsibility to keep the weights alive - during the TRT network construction and TRT engine building process. - ''' - self._registered_ndarrays.append(ndarray) - - def _generate_optimization_profiles(self) -> List[trt.IOptimizationProfile]: - input_tensors = self._inputs - if len(input_tensors) == 0: - return [] - num_profiles = len(list(input_tensors.values())[0].profiles) - profiles = [] - for i in range(num_profiles): - logger.debug(f'Adding optimization profile {i+1}/{num_profiles}') - profile = self._trt_network.builder.create_optimization_profile() - for input_name, input_tensor in input_tensors.items(): - shape_profile = input_tensor.profiles[i] - min_shape = list(shape_profile.min) - opt_shape = list(shape_profile.opt) - max_shape = list(shape_profile.max) - if input_tensor.trt_tensor.is_shape_tensor: - profile.set_shape_input(input_name, min_shape, opt_shape, - max_shape) - else: - profile.set_shape(input_name, min_shape, opt_shape, - max_shape) - logger.debug( - f'{input_name}, min: {min_shape}, opt: {opt_shape}, max: {max_shape}' - ) - profiles.append(profile) - return profiles - - def get_inputs(self): - ''' - Get the inputs of the network. - - Returns: - Iterable[Tensor] - ''' - return self._inputs.values() - - def get_outputs(self): - ''' - Get the outputs of the network. - - Returns: - Iterable[Tensor] - ''' - from .functional import Tensor - for i in range(self._trt_network.num_outputs): - tensor = self._trt_network.get_output(i) - yield Tensor(trt_tensor=tensor, - network=self, - is_network_input=False) - - def is_input(self, tensor) -> bool: - ''' - Tell if a tensor is a input of the network. - - Parameters: - tensor: Union[Tensor, str, trt.ITensor] - ''' - from .functional import Tensor - - if isinstance(tensor, str): - tensor_name = tensor - elif isinstance(tensor, (trt.ITensor, Tensor)): - tensor_name = tensor.name - else: - raise ValueError( - f"tensor should be Tensor, str or ITensor, got {tensor}") - - return self._inputs.get(tensor_name, False) - - def is_output(self, tensor) -> bool: - ''' - Tell if a tensor is a output of the network. - - Parameters: - tensor: Tensor - ''' - for i in range(self._trt_network.num_outputs): - if tensor.trt_tensor is self._trt_network.get_output(i): - return True - return False - - def get_layers(self) -> Iterable["Layer"]: - ''' - Get all the layers of network. - - Returns: - Iterable[Layer] - ''' - from .graph_rewriting import Layer - for i in range(self._trt_network.num_layers): - layer = Layer(network=self, - trt_layer=self._trt_network.get_layer(i)) - yield layer - - def get_layer_by_name(self, name: str) -> Optional["Layer"]: - state = self._get_graph() - return state.name_to_layer.get(name, None) - - def get_tensor_users(self, tensor) -> Iterable["Layer"]: - ''' - Get the layers those consumes this tensor. - ''' - state = self._get_graph() - for layer in state.tensor_to_consumers[tensor]: - yield layer - - def get_tensor_parent(self, tensor) -> Optional["Layer"]: - ''' - Get the layer that produces this tensor. - ''' - state = self._get_graph() - return state.tensor_to_producer.get(tensor, None) - - def mark_removed_layer(self, layer: "Layer"): - from .graph_rewriting import FLayerInfoMemo - self._removed_layers.add(layer.name) - - # Try to delete the layer if it is a Plugin - FLayerInfoMemo.instance().remove(layer.name) - - def is_removed_layer(self, layer: "Layer") -> bool: - return layer.name in self._removed_layers - - @property - def removed_layers(self) -> Iterable["Layer"]: - for layer_name in self._removed_layers: - layer = self.get_layer_by_name(layer_name) - assert layer, "Invalid layer name" - yield layer - - def to_dot(self, path: Union[str, Path] = None) -> Optional[str]: - ''' - Get a graphviz representation of the network. - - NOTE, the graph might be redundancy since TRT's INetwork won't clean the unused inputs and layers - automatically. - TODO: add an flag to hide all the removed layers and their output tensors - TODO: replace this when TensorRT provides a better way to get the graph of INetworkDefinition - TODO: a little feature, add blocks in the figure to highlight the subgraphes of Modules - - Parameters: - path: the path to save the graphviz file, if not provided, will return the graphviz source code - ''' - format = 'text' if not path else path.split('.')[-1] - - try: - import graphviz - except ImportError: - logger.error( - "Failed to import graphviz, please install graphviz to enable Network.to_dot()" - ) - return - - dot = graphviz.Digraph( - comment= - f'TensorRT Graph of {self._get_network_hash(lightweight=False)}', - format=format if format != 'text' else None) - - inputs_names = set([x.name for x in self.get_inputs()]) - output_names = set([x.name for x in self.get_outputs()]) - - node_style = dict( - shape='box', - style='rounded,filled,bold', - fontname='Arial', - fillcolor='#ffffff', - color='#303A3A', - width='1.3', - height='0.84', - ) - - hl_node_style = dict( - shape='box', - style='rounded,filled,bold', - fontname='Arial', - fillcolor='lightblue', - color='#303A3A', - width='1.3', - height='0.84', - ) - - state = self._get_graph() - nodes = set() - tensor_to_alias = {} - tensor_id = [0] - - def get_alias(tensor, tensor_id): - if tensor not in tensor_to_alias: - if (tensor not in inputs_names) and (tensor - not in output_names): - tensor_to_alias[tensor] = f"t{tensor_id[0]}" - tensor_id[0] += 1 - else: - tensor_to_alias[tensor] = tensor - - return tensor_to_alias[tensor] - - def create_tensor_node(tensor: str, dtype=None, shape=None): - tensor_alias = get_alias(tensor, tensor_id) - if tensor_alias not in nodes: - dot.node(tensor_alias, - str(dtype) + "\n" + tensor_alias + "\n" + str(shape), - **node_style) - nodes.add(tensor_alias) - return tensor_alias - - def create_layer_node(layer: str): - if layer not in nodes: - dot.node(layer, layer, **hl_node_style) - nodes.add(layer) - - for tensor, layer in state.tensor_to_producer.items(): - tensor_alias = create_tensor_node(tensor.name, tensor.dtype, - tensor.shape) - create_layer_node(layer.name) - dot.edge(layer.name, tensor_alias) - for tensor, layers in state.tensor_to_consumers.items(): - tensor_alias = create_tensor_node(tensor.name, tensor.dtype, - tensor.shape) - for layer in layers: - create_layer_node(layer.name) - dot.edge(tensor_alias, layer.name) - - if format == "text": - return dot.source - dot.save(path) - - def to_onnx(self, path: str = None) -> None: - ''' - Export the network into a "ONNX-like" file for visualization. - - Parameters: - path: the path to the output file - ''' - if path is None: - return - - network = self.trt_network - b_onnx_type = True # Use ONNX style operator names - - def layer_type_to_class(layer: trt.ILayer = None) -> trt.ILayer: - ''' - Convert trt.ILayer into a exact TensorRT layer - ''' - layer_type_name = str(layer.type)[10:] - # Some special cases - if layer_type_name == "ELEMENTWISE": - return trt.IElementWiseLayer - if layer_type_name == "LRN": - return trt.ILRNLayer - if layer_type_name == "NMS": - return trt.INMSLayer - if layer_type_name == "PARAMETRIC_RELU": - return trt.IParametricReLULayer - if layer_type_name == "PLUGIN": - return None # IPluginLayer is not supported any more - if layer_type_name == "RAGGED_SOFTMAX": - return trt.IRaggedSoftMaxLayer - if layer_type_name == "SOFTMAX": - return trt.ISoftMaxLayer - if layer_type_name == "TOPK": - return trt.ITopKLayer - - # general cases, e.g. MATRIX_MULTIPLY -> MatrixMultiply - name = "".join(name[0] + name[1:].lower() - for name in layer_type_name.split("_")) - return trt.__builtins__["getattr"](trt, f"I{name}Layer") - - def convert_type_to_onnx(node_type: str = "", - attribution: OrderedDict = OrderedDict()): - ''' - Convert layer names of TensorRT style into ONNX style - ''' - if node_type == "ACTIVATION": - convert_list = { - "RELU": "Relu", - "SIGMOID": "Sigmoid", - "TANH": "Tanh", - "LEAKY_RELU": "LeakyRelu", - "ELU": "Elu", - "SELU": "Selu", - "SOFTSIGN": "Softsign", - "SOFTPLUS": "Softplus", - "CLIP": "Clip", - "HARD_SIGMOID": "HardSigmoid", - "SCALED_TANH": "ScaledTanh", - "THRESHOLDED_RELU": "ThresholdedRelu", - } - # No corresponding operator for GELU_ERF, GELU_TANH - if "algo-type" in attribution.keys() and attribution[ - "algo-type"].split(".")[-1] in convert_list.keys(): - return convert_list[attribution["algo-type"].split(".")[-1]] - return node_type - if node_type == "CAST": - return "Cast" - if node_type == "CONCATENATION": - return "Concat" - if node_type == "CONSTANT": - return "Constant" - if node_type == "CONVOLUTION": - return "Conv" - if node_type == "DECONVOLUTION": - return "Deconv" - if node_type == "ELEMENTWISE": - convert_list = { - "SUM": "Add", - "PROD": "Mul", - "MAX": "Max", - "MIN": "Min", - "SUB": "Sub", - "DIV": "Div", - "POW": "Pow", - "AND": "And", - "OR": "Or", - "XOR": "Xor", - "EQUAL": "Equal", - "GREATER": "Greater", - "LESS": "Less", - } - if "op" in attribution.keys() and attribution["op"].split( - ".")[-1] in convert_list.keys(): - return convert_list[attribution["op"].split(".")[-1]] - return node_type - if node_type == "GATHER": - return "Gather" - if node_type == "LOOP": - return "Loop" - if node_type == "MATRIX_MULTIPLY": - return "Gemm" - if node_type == "POOLING": - convert_list = {"MAX": "MaxPool", "AVERAGE": "AveragePool"} - if "algo-type" in attribution.keys() and attribution[ - "algo-type"].split(".")[-1] in convert_list.keys(): - return convert_list[attribution["algo-type"].split(".")[-1]] - return node_type - if node_type == "REDUCE": - convert_list = { - "SUM": "ReduceSum", - "PROD": "ReduceProd", - "MAX": "ReduceMax", - "MIN": "ReduceMin", - "AVG": "ReduceMean", - } - if "op" in attribution.keys() and attribution["op"].split( - ".")[-1] in convert_list.keys(): - return convert_list[attribution["op"].split(".")[-1]] - return node_type - if node_type == "SELECT": - return "Where" - if node_type == "SHUFFLE": - return "Reshape" - if node_type == "SHAPE": - return "Shape" - if node_type == "SLICE": - return "Slice" - if node_type == "SOFTMAX": - return "Softmax" - if node_type == "TOPK": - return "TopK" - if node_type == "UNARY": - convert_list = {"SQRT": "Sqrt", "NOT": "Not"} - if "op" in attribution.keys() and attribution["op"].split( - ".")[-1] in convert_list.keys(): - return convert_list[attribution["op"].split(".")[-1]] - return node_type - return node_type - - def add_node_for_trt_network( - graph: gs.Graph = None, - node_name: str = "", - node_type: str = "", - input_list: List[gs.Variable] = [], - attribution: OrderedDict = OrderedDict(), - name_list: Union[str, List[str]] = "", - datatype_list: Union[np.dtype, List[np.dtype]] = [], - shape_list: Union[list, List[list]] = [], - number: int = 0, - b_onnx_type: bool = False, - ) -> Tuple[gs.Variable, int]: - ''' - Simplify version of function `add_node`, and we do some beautify to it. - ''' - if isinstance(name_list, list) or isinstance( - datatype_list, list) or isinstance( - shape_list, list): # Case of multi-output - assert len(name_list) == len(datatype_list) - assert len(name_list) == len(shape_list) - else: # Case of single-output - name_list = [name_list] - datatype_list = [datatype_list] - shape_list = [shape_list] - - n_output = len(name_list) - output_list = [] - for i in range(n_output): - tensor = gs.Variable(name_list[i], datatype_list[i], - shape_list[i]) - output_list.append(tensor) - - if b_onnx_type: - node_type = convert_type_to_onnx(node_type, attribution) - - node = gs.Node(node_type, - node_name, - inputs=input_list, - outputs=output_list, - attrs=attribution) - graph.nodes.append(node) # Update graph inside `add_node` - - if len(output_list) == 1: # Case of single-output - output_list = output_list[0] - return output_list, number + 1 - - graph = gs.Graph(nodes=[], inputs=[], outputs=[]) - graph.name = "" if network.name == "Unnamed Network 0" else network.name - n = 0 - - # mapping from TRT tensor (trt.ITensor) to GS tensor (gs.Variable) - global_tensor_map = {} - - for i in range(network.num_inputs): - trt_tensor = network.get_input(i) - gs_tensor = gs.Variable(trt_tensor.name, - trt.nptype(trt_tensor.dtype), - trt_tensor.shape) - global_tensor_map[trt_tensor] = gs_tensor - if gs_tensor not in graph.inputs: - graph.inputs.append(gs_tensor) - - for i in range(network.num_layers): - layer = network.get_layer(i) - - input_tensor_list = [] - for j in range(layer.num_inputs): - trt_tensor = layer.get_input(j) - if trt_tensor is None: # Useful for constant layer - gs_tensor = None - elif trt_tensor in global_tensor_map.keys( - ): # already in the map - gs_tensor = global_tensor_map[trt_tensor] - else: - logger.debug( - f"[ExportONNX]Layer input tensor not in global_tensor_map: {trt_tensor.name}" - ) # ■ - gs_tensor = gs.Variable(trt_tensor.name, - trt.nptype(trt_tensor.dtype), - trt_tensor.shape) - global_tensor_map[trt_tensor] = gs_tensor - input_tensor_list.append(gs_tensor) - - output_name_list = [] - output_datatype_list = [] - output_shape_list = [] - for i in range(layer.num_outputs): - trt_tensor = layer.get_output(i) - # Don't do this check because we need this trt_tensor to overwrite the placeholder tensor in ■ - # if trt_tensor in global_tensor_map.keys(): - # gs_tensor = global_tensor_map[trt_tensor] - output_name_list.append(trt_tensor.name) - output_datatype_list.append(trt.nptype(trt_tensor.dtype)) - output_shape_list.append(trt_tensor.shape) - - attr = OrderedDict() - # Set attribution of ILayer - for key in dir(layer): - if not (key.startswith("_") - or callable(layer.__getattribute__(key))): - attr[key] = str(layer.__getattribute__(key)) - # Set attribution of exact layer type - layer.__class__ = layer_type_to_class(layer) - for key in dir(layer): - if key in dir(trt.ILayer) and key != "type": - continue - if key == "type" and not isinstance(layer.type, trt.LayerType): - attr["algo-type"] = str(layer.type) - continue - value = layer.__getattribute__(key) - if isinstance(value, np.ndarray): - # Convert all attributions into string besides weights - value = value.astype(np.float32) # In case of overflow - ss = f"shape={value.shape}, SumAbs={np.sum(abs(value)):.5e}, Var={np.var(value):.5f}, " - ss += f"Max={np.max(value):.5f}, Min={np.min(value):.5f}, SAD={np.sum(np.abs(np.diff(value.reshape(-1)))):.5f}, " - ss += f"[:5]={value.reshape(-1)[:5]}, [-5:]={value.reshape(-1)[-5:]}" - attr[key] = ss - else: - attr[key] = str(value) - - output_tensor_list, n = add_node_for_trt_network(graph, layer.name, attr["type"][10:], input_tensor_list, attr, \ - output_name_list, output_datatype_list, output_shape_list, n, b_onnx_type) - - if layer.num_outputs == 1: - global_tensor_map[layer.get_output(0)] = output_tensor_list - else: - for i in range(layer.num_outputs): - global_tensor_map[layer.get_output( - i)] = output_tensor_list[i] - - for i in range(network.num_outputs): - gs_tensor = global_tensor_map[network.get_output(i)] - if gs_tensor not in graph.outputs: - graph.outputs.append(gs_tensor) - - onnx.save(gs.export_onnx(graph), - path + "/network.onnx", - save_as_external_data=False) - return - - def _get_graph(self) -> "Network._GraphState": - ''' - Get the graph of the network. - - Returns: - Network._GraphState - ''' - return self._get_graph_impl(self._get_network_hash()) - - #TODO: using one LRU cache here can cause the Network object to be leaked, need a way to speed this function w/o using global lru cache. - def _get_graph_impl(self, network_hash: bytes) -> "Network._GraphState": - graph = Network._GraphState() - graph.build(self) - return graph - - @dataclass - class _GraphState: - # Tensor to Layers - tensor_to_consumers: Dict[Any, List["Layer"]] = field( - default_factory=lambda: defaultdict(list)) - # Tensor to Layer - tensor_to_producer: Dict[Any, "Layer"] = field(default_factory=dict) - inputs: Dict[str, Any] = field(default_factory=OrderedDict) - outputs: Dict[str, Any] = field(default_factory=OrderedDict) - name_to_layer: Dict[str, "Layer"] = field(default_factory=dict) - - def build(self, network: "Network") -> None: - from .graph_rewriting import Layer - self.inputs = network.get_inputs() - self.outputs = network.get_outputs() - - for layer in network.get_layers(): - self.name_to_layer[layer.name] = Layer( - network=network, trt_layer=layer.trt_layer) - for i in range(layer.num_inputs): - input_tensor = layer.get_inputs(i)[0] - if input_tensor.is_trt_wrapper(): - self.tensor_to_consumers[input_tensor].append(layer) - for i in range(layer.num_outputs): - output_tensor = layer.get_outputs(i)[0] - if output_tensor.is_trt_wrapper(): - self.tensor_to_producer[output_tensor] = layer - - def _get_network_hash(self, lightweight=True) -> bytes: - # TODO: Ask TensorRT team to add a hash function for INetworkDefinition instead of using this hacky way - num_layers = self.trt_network.num_layers - - # Some special layers, such as slice, may be associated with tensors that do not have the `trt_tensor` member. - get_tensor_tag = lambda tensor: tensor.trt_tensor.name if tensor.is_trt_wrapper( - ) else 'None' - - if lightweight and not self.is_graph_altered: - return num_layers - self.is_graph_altered = False - - data = hashlib.sha256() - # network layer count - data.update(str(num_layers).encode()) - # network inputs - data.update(','.join( - [get_tensor_tag(tensor) for tensor in self.get_inputs()]).encode()) - # network outputs - data.update(','.join( - [get_tensor_tag(tensor) for tensor in self.get_outputs()]).encode()) - # layer names - data.update(','.join( - [layer.trt_layer.name for layer in self.get_layers()]).encode()) - - # layer -> output - data.update(','.join([ - f'{layer.trt_layer.name}->{get_tensor_tag(tensor)}' - for layer in self.get_layers() for tensor in layer.get_outputs() - ]).encode()) - - # input -> layer - data.update(','.join([ - f'{get_tensor_tag(tensor)}->{layer.trt_layer.name}' - for layer in self.get_layers() for tensor in layer.get_inputs() - ]).encode()) - - return data.hexdigest() - - -@contextlib.contextmanager -def net_guard(network): - from ._common import net - assert isinstance( - network, Network - ), f"Invalid network, can only guard Network instance, got: {network}" - - old_net = net - set_network(network) - yield - set_network(old_net) - - -class _TrtLlmModuleCallStack(object): - - def __init__(self): - super().__init__() - self.call_stack = [] - self.module_name_map = weakref.WeakKeyDictionary() - self.module_to_layer_range_map: Dict[str, range] = {} - self.mod_names_set = False - - def module_names_set(self): - return self.mod_names_set - - def set_module_names(self, top_level_module): - assert top_level_module, "Expected a top level module" - for name, mod in top_level_module.named_modules( - prefix=top_level_module._get_name()): - if mod not in self.module_name_map: - self.module_name_map[mod] = name - self.mod_names_set = True - return - - def get_current_module(self): - mod_name = '' - if len(self.call_stack): - mod_name = self.call_stack[-1] - return mod_name - - def get_mod_name(self, mod_obj): - name = '' - if mod_obj in self.module_name_map: - name = self.module_name_map[mod_obj] - return name - - def set_layer_range(self, mod_obj: Module, layer_range: range): - if mod_obj in self.module_name_map: - name = self.module_name_map[mod_obj] - self.module_to_layer_range_map[name] = layer_range - - def get_stack(self): - return self.call_stack - - @contextlib.contextmanager - def call_stack_mgr(self): - call_stack = self.get_stack() - try: - yield call_stack - finally: - call_stack.pop() diff --git a/tensorrt_llm/parameter.py b/tensorrt_llm/parameter.py deleted file mode 100644 index d740cbb0bdb9..000000000000 --- a/tensorrt_llm/parameter.py +++ /dev/null @@ -1,281 +0,0 @@ -# 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 math -import weakref -from typing import Optional, Sequence, Union - -import numpy as np - -# isort: off -import torch -import tensorrt as trt -# isort: on - -from ._common import default_net -from ._utils import (copy_torch_to_numpy, np_dtype_to_trt, str_dtype_to_trt, - torch_to_numpy, trt_dtype_to_np, trt_dtype_to_torch) -from .functional import Tensor, constant -from .logger import logger -from .network import Network - - -class Parameter: - _DEFAULT_DTYPE = trt.DataType.FLOAT - - def __init__(self, - value: Optional[Union[np.ndarray, torch.Tensor]] = None, - shape: Sequence[int] = None, - dtype: Union[str, trt.DataType] = None, - is_buffer: bool = False, - prefer_managed=False): - if dtype is None: - logger.warning( - f'Parameter dtype is None, using default dtype: {self._DEFAULT_DTYPE}, it is recommended to always specify dtype explicitly' - ) - dtype = self._DEFAULT_DTYPE if dtype is None else dtype - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - self._dtype: trt.DataType = dtype - if value is None: - assert isinstance(shape, ( - list, - tuple)), f"shape must be list or tuple, receive {(type(shape))}" - self._shape = tuple(shape) - self._value = None - else: - self._shape = value.shape - self._value = self._regularize_value(value) - self.is_buffer = is_buffer - self._prefer_managed = prefer_managed - self._tensor: Tensor = None - self._network: weakref.ref = None - self._name = None - self.need_transpose = False - - @property - def shape(self): - return self._shape - - @property - def dtype(self): - return self._dtype - - @property - def name(self): - return self._name - - def _create_managed_tensor(self, network) -> Tensor: - num = len(network._inputs) - self._name = f"managed_constant_{num}" - - if self._value is None or (isinstance(self._value, np.ndarray) - and not self._value.flags['C_CONTIGUOUS']): - value_old = self._value - self._value = np.empty(self._shape, trt_dtype_to_np(self._dtype)) - network._register_unfilled_weights( - # use updated self._shape here - self._name, - self._value, - value_old) - return Tensor(name=self._name, dtype=self._dtype, shape=self._shape) - - def get_managed_tensor(self, network: Network) -> Tensor: - if self._network is None or self._network() != network: - self._network = weakref.ref(network) - self._tensor = network.get_parameter_tensor(self) - if self._tensor is None: - self._tensor = self._create_managed_tensor(network) - network.set_parameter_tensor(self, self._tensor) - return self._tensor - - def _create_constant_tensor(self) -> Tensor: - if (self._value is not None and isinstance(self._value, np.ndarray) - and self._value.flags['C_CONTIGUOUS']): - lower_type = None - lower_shape = None - # workaround for reinterpreted data type - dtype = self._value.dtype - if (self.dtype == trt.fp4 or self.dtype - == trt.fp8) and (dtype == np.uint8 or dtype == np.int8 - or dtype == np.int32 or dtype == np.int64): - lower_type = self.dtype - lower_shape = self.shape - - self._value = constant(self._value, lower_type, lower_shape) - return self._value - elif self._value is None or isinstance(self._value, np.ndarray): - if self._dtype == trt.fp4: - shape = list(self._shape) - assert shape[ - -1] % 16 == 0, "For FP4, the last dimension of the shape should be multiple of 16" - shape[-1] = shape[-1] // 16 - dtype = np.int64 - else: - shape = self._shape - dtype = trt_dtype_to_np(self._dtype) - ndarray = np.empty(shape, dtype) - tensor = constant(ndarray, self._dtype, self._shape) - default_net()._register_unfilled_weights(tensor.producer.name, - ndarray, self._value) - return tensor - - def get_constant_tensor(self, network: Network) -> Tensor: - if self._network is None or self._network() != network: - self._network = weakref.ref(network) - self._tensor = network.get_parameter_tensor(self) - if self._tensor is None: - self._tensor = self._create_constant_tensor() - self._name = self._tensor.producer.name - network.set_parameter_tensor(self, self._tensor) - return self._tensor - - def get_tensor(self, network) -> Tensor: - if self.is_managed(network): - return self.get_managed_tensor(network) - else: - return self.get_constant_tensor(network) - - def is_managed(self, network): - if network is None: - network = default_net() - return self._prefer_managed and network.plugin_config.manage_weights - - @property - def value(self) -> Tensor: - return self.get_tensor(default_net()) - - @classmethod - def xavier_init(cls, weights: np.ndarray): - shape = weights.shape - dtype = np_dtype_to_trt(weights.dtype) - if len(shape) == 2: - # Xavier initialization see https://paperswithcode.com/method/xavier-initialization - v_range = math.sqrt(6) / math.sqrt(shape[0] + shape[1]) - else: - v_range = 0.1 - - if dtype == trt.DataType.INT8 or dtype == trt.DataType.INT32 or dtype == trt.DataType.INT64: - range_map = { - trt.DataType.INT8: 128, - trt.DataType.INT32: 2**31, - trt.DataType.INT64: 2**63 - } - upper = math.ceil(range_map[dtype] * v_range) - value = torch.randint(-upper, - upper, (shape), - dtype=trt_dtype_to_torch(dtype), - device='cuda') - # value ~ U[int(-128 * v_range), int(128 * v_range)] - elif dtype == trt.DataType.FP8: - value = torch.rand((shape), device='cuda') * 2 - 1 - # value ~ U[-v_range, v_range] - value = value * v_range - value = value.to(trt_dtype_to_torch(dtype)) - else: - value = torch.rand( - (shape), dtype=trt_dtype_to_torch(dtype), device='cuda') * 2 - 1 - # value ~ U[-v_range, v_range] - value = value * v_range - - copy_torch_to_numpy(value, weights) - - def is_inited(self) -> bool: - return self._value is not None - - @property - def raw_value(self) -> np.ndarray: - if self._value is None: - dtype = trt_dtype_to_np(self.dtype) - self._value = np.empty(self.shape, dtype) - Parameter.xavier_init(self._value) - assert isinstance( - self._value, np.ndarray - ), "Must be np.ndarray. Proper usage: get parameter.raw_value before getting parameter.value" - return self._value - - @value.setter - def value(self, v: Union[np.ndarray, torch.Tensor]): - v = self._regularize_value(v) - - if v.shape != self.shape and v.ndim == 0 and max(self.shape) == 1: - # convert the scalar into a tensor which each dim is 1. - v = v.reshape(self.shape) - - if self.dtype == trt.fp4: - assert v.shape[:-1] == self.shape[:-1] and v.shape[-1] == self.shape[-1] // 2 // v.dtype.itemsize, \ - f'For FP4, the shape of the value should be the same as the original shape, ' \ - f'except the last dimension should be half of the original shape. ' \ - f'Updated: {v.shape}, original: {self.shape}' - else: - assert v.shape == self.shape, \ - f'The value updated is not the same shape as the original. ' \ - f'Updated: {v.shape}, original: {self.shape}' - if (self.dtype == trt.fp4 or self.dtype - == trt.fp8) and (v.dtype == np.int8 or v.dtype == np.uint8 - or v.dtype == np.int32 or v.dtype == np.int64): - pass - else: - dtype = np_dtype_to_trt(v.dtype) - if self.dtype != dtype: - logger.warning( - f"Parameter was initialized as {self.dtype} but set to {dtype}" - ) - self._dtype = dtype - self._value = v - - def set_value_or_dummy(self, v: Union[np.ndarray, torch.Tensor]): - v = self._regularize_value(v) - if v.shape != self._shape: - self.value = np.empty(self._shape, trt_dtype_to_np(self._dtype)) - return - - self.value = v - - def set_name(self, name: str, network: Network): - self._name = name - if self.is_managed(network): - self._get_weights(network).name = name - return True - else: - weights = self._get_weights(network) - # TensorRT bindings may return numpy array instead of trt.Weights - if isinstance(weights, np.ndarray): - trt_dtype = np_dtype_to_trt( - weights.dtype - ) if weights.dtype != np.object_ else self._dtype - trt_count = int(np.prod(weights.shape)) - weights = trt.Weights(trt_dtype, weights.ctypes.data, trt_count) - return network.trt_network.set_weights_name(weights, name) - - def _get_weights(self, network: Network) -> trt.Weights | Tensor | None: - tensor = network.get_parameter_tensor(self) - if self.is_managed(network): - return tensor - elif tensor is not None: - tensor.producer.__class__ = trt.IConstantLayer - return tensor.producer.weights - else: - return None - - def _regularize_value(self, value): - if isinstance(value, np.ndarray): - return value - - elif isinstance(value, torch.distributed.tensor.DTensor): - return value.to_local().cpu().numpy() - elif isinstance(value, torch.Tensor): - return torch_to_numpy(value) - raise TypeError( - f'Expected numpy.ndarray or torch.Tensor, got {type(value)}') diff --git a/tensorrt_llm/plugin/__init__.py b/tensorrt_llm/plugin/__init__.py deleted file mode 100644 index 55bffdee6b85..000000000000 --- a/tensorrt_llm/plugin/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# 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 .plugin import (TRT_LLM_PLUGIN_NAMESPACE, PluginConfig, _load_plugin_lib, - add_plugin_argument, current_all_reduce_helper, - init_all_reduce_helper, plugin_lib_path) - -__all__ = [ - 'TRT_LLM_PLUGIN_NAMESPACE', '_load_plugin_lib', 'PluginConfig', - 'add_plugin_argument', 'plugin_lib_path', "current_all_reduce_helper", - "init_all_reduce_helper" -] diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py deleted file mode 100644 index f8434cf11bec..000000000000 --- a/tensorrt_llm/plugin/plugin.py +++ /dev/null @@ -1,780 +0,0 @@ -# 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 ctypes -import os -import platform -from collections import OrderedDict -from enum import IntEnum -from pathlib import Path -from textwrap import dedent -from typing import (Any, List, Literal, Optional, Tuple, Union, get_args, - get_origin) - -import tensorrt as trt -from pydantic import (ConfigDict, Field, PrivateAttr, ValidationInfo, - field_validator, model_validator) - -from .._ipc_utils import IpcMemory, can_access_peer -from .._utils import get_sm_version -from ..bindings.internal.runtime import (lamport_initialize, - lamport_initialize_all, - max_workspace_size_lowprecision) -from ..llmapi.utils import StrictBaseModel -from ..logger import logger -from ..mapping import Mapping - -TRT_LLM_PLUGIN_NAMESPACE = 'tensorrt_llm' - - -def plugin_lib_path() -> str: - project_dir = Path(__file__).parent.parent.absolute() - dyn_lib = "libnvinfer_plugin_tensorrt_llm.so" if platform.system( - ) != "Windows" else "nvinfer_plugin_tensorrt_llm.dll" - return str(project_dir.joinpath("libs", dyn_lib)) - - -def _load_plugin_lib(): - on_windows = platform.system() == "Windows" - winmode = 0 if on_windows else None - handle = ctypes.CDLL(plugin_lib_path(), - mode=ctypes.RTLD_GLOBAL, - winmode=winmode) - try: - handle.initTrtLlmPlugins.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - handle.initTrtLlmPlugins.restype = ctypes.c_bool - except AttributeError as err: - raise ImportError('TensorRT LLM Plugin is unavailable') from err - - try: - assert handle.initTrtLlmPlugins( - None, TRT_LLM_PLUGIN_NAMESPACE.encode('utf-8')) - except OSError as e: - windows_err = """ - The error above may be caused by an outdated Microsoft Visual C++ Redistributable Version. - Please install the latest MSVC from the link below and re-launch. - - https://learn.microsoft.com/en-us/cpp/windows/latest-supported-vc-redist?view=msvc-170#latest-microsoft-visual-c-redistributable-version - """ - err_msg = dedent(windows_err if on_windows else "Unknown error") - raise RuntimeError(err_msg) from e - except Exception as e: - raise e - - -class ContextFMHAType(IntEnum): - disabled = 0 - # FP16 I/O, FP16 Accumulation - enabled = 1 - # FP16 I/O, FP32 Accumulation - enabled_with_fp32_acc = 2 - - -DefaultPluginDtype = Literal["auto", "float16", "float32", "bfloat16", "int32", - None] - - -class PluginConfig(StrictBaseModel): - """The config that manages plugin-related options. - - There are two option categories: - * Plugin options (typically with xxx_plugin naming). These options can be assigned with: - * "float16"/"bfloat16"/"float32"/"int32", which means the plugin is enabled with the specified precision; (Some plugins only support limited dtype, i.e., gemm_swiglu_plugin and low_latency_gemm_swiglu_plugin only supports fp8 now) - * "auto", which means the plugin is enabled with the precision of `dtype` field (the `dtype` field must be same to model dtype, i.e., the one in PretrainedConfig); - * None, which means the plugin is disabled. - * Other features. These options can be assigned with boolean: - * True, which means the plugin is enabled; - * False, which means the plugin is disabled. - """ - model_config = ConfigDict(validate_assignment=True, extra="forbid") - - dtype: str = Field(default="float16", - description="Base dtype for the model and plugins") - - # Plugins - bert_attention_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "The plugin that uses efficient kernels and enables an in-place update of the KV cache for attention layer of BERT-like encoder models." - ) - gpt_attention_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "The plugin that uses efficient kernels and enables an in-place update of the KV cache for attention layer of GPT-like decoder models." - ) - gemm_plugin: Optional[Literal[ - "auto", "float16", "float32", "bfloat16", "int32", "fp8", "nvfp4", - None]] = Field( - default=None, - description= - "The GEMM plugin that utilizes NVIDIA cuBLASLt to perform GEMM operations. " - "Note: it's only affective for non-quantized gemm operations (except FP8)." - "Note: For FP8, it also requires same calibration in checkpoint.") - _explicitly_disable_gemm_plugin: bool = PrivateAttr(default=False) - gemm_swiglu_plugin: Optional[Literal["fp8", None]] = Field( - default=None, - description= - "The GEMM + SwiGLU fusion in Gated-MLP combines two Matmul operations and " - "one SwiGLU operation into a single kernel. Currently this is only supported for FP8 precision on Hopper." - ) - fp8_rowwise_gemm_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "The quantized GEMM for fp8, which uses per token dynamic scales for " - "activation and per channel static scales for weights." - "Note: It also requires same calibration in checkpoint.") - qserve_gemm_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "The quantized GEMM from [QServe](https://arxiv.org/abs/2405.04532), " - "which employs 4-bit quantization for weights and 8-bit quantization for activations." - ) - identity_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "The identity plugin simply copies inputs to outputs, it's used mostly for debugging purpose." - ) - nccl_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "The NCCL plugin wraps NCCL operators to support multi-GPU and even multi-nodes." - ) - lora_plugin: Optional[DefaultPluginDtype] = Field( - default=None, description="Enable LoRA.") - dora_plugin: bool = Field(default=False, description="Enable DoRA.") - weight_only_groupwise_quant_matmul_plugin: Optional[ - DefaultPluginDtype] = Field( - default=None, - description= - "Enable weight-only groupwise quantization matmul operators.") - weight_only_quant_matmul_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description="Enable weight-only quantization matmul operators.") - smooth_quant_plugins: bool = Field( - default=True, - description="Enable a group of plugins to support smooth quantization.") - smooth_quant_gemm_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "Enable plugin that supports smooth quantization gemm kernels.") - layernorm_quantization_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description="Enable plugin that supports layernorm quantization kernels." - ) - rmsnorm_quantization_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description="Enable plugin that supports rmsnorm quantization kernels.") - quantize_per_token_plugin: bool = Field( - default=False, - description="Enable plugin that supports per-token quantization.") - quantize_tensor_plugin: bool = Field( - default=False, - description="Enable plugin that supports per-tensor quantization.") - moe_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "Enable some customized kernels to speed up the MoE layer of MoE models." - ) - mamba_conv1d_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "Enable customized kernels to speed up conv1d operator for Mamba.") - low_latency_gemm_plugin: Optional[Literal["fp8", None]] = Field( - default=None, - description= - "The GEMM plugin that optimized specially for low latency scenarios.") - low_latency_gemm_swiglu_plugin: Optional[Literal["fp8", None]] = Field( - default=None, - description= - "The GEMM + SwiGLU fusion plugin that optimized specially for low latency scenarios." - ) - gemm_allreduce_plugin: Optional[Literal[ - "float16", "bfloat16", - None]] = Field(default=None, - description="The GEMM + AllReduce kernel fusion plugin.") - - # Features - context_fmha: bool = Field( - default=True, - description= - "Enable the fused multi-head attention during the context phase, " - "will trigger a kernel that performs the MHA/MQA/GQA block using a single kernel." - ) - bert_context_fmha_fp32_acc: bool = Field( - default=False, - description= - "Enable the FP32 accumulator for context FMHA in the bert_attention_plugin. " - "If disabled, FP16 is used, better performance but potentially worse accuracy is expected." - ) - paged_kv_cache: Optional[bool] = Field( - default=None, - description= - "Enable paged KV cache, which helps manage memory for the KV cache more efficiently, " - "and usually leads to an increase in the batch size and an improved efficiency." - ) - remove_input_padding: bool = Field( - default=True, - description= - "Pack different tokens together, which reduces both the amount of computations and memory consumption." - ) - norm_quant_fusion: bool = Field( - default=False, - description= - "Fuse the LayerNorm and quantization kernels into a single kernel, " - "resulting in improved end-to-end performance.") - reduce_fusion: bool = Field( - default=False, - description= - "Fuse the ResidualAdd and LayerNorm kernels after AllReduce into a single kernel, " - "resulting in improved end-to-end performance.") - user_buffer: bool = Field( - default=False, - description= - "Eliminate extra copies from the local buffer to the shared buffer " - "in the communication kernel, leading to improved end-to-end performance. " - "This feature must be enabled with `--reduce_fusion enable` and " - "is currently only supported for the FP8 LLAMA model.") - tokens_per_block: int = Field( - default=32, - description= - "Define how many tokens are contained in each paged kv cache block.") - use_paged_context_fmha: bool = Field( - default=True, - description= - "Allow advanced features like KV cache reuse and chunked context.") - use_fp8_context_fmha: bool = Field( - default=True, - description= - "When FP8 quantization is activated, the attention can be further accelerated by enabling FP8 Context FMHA" - ) - fuse_fp4_quant: bool = Field( - default=False, - description="Whether to fuse FP4 quantization into attention kernel.") - multiple_profiles: bool = Field( - default=False, - description= - "Enables multiple TensorRT optimization profiles in the built engines, " - "will benefits the performance especially when GEMM plugin is disabled, " - "because more optimization profiles help TensorRT have more chances to select better kernels. " - "Note: This feature increases engine build time but no other adverse effects are expected." - ) - paged_state: bool = Field( - default=True, - description= - "Enable paged state, which helps manage memory for the RNN state more efficiently." - ) - streamingllm: bool = Field( - default=False, - description= - "Enable [StreamingLLM](https://arxiv.org/abs/2309.17453), which uses a window attention to perform efficient and stable LLM on long texts." - ) - manage_weights: bool = Field( - default=False, - description= - "Enable TensorRT LLM managed weights to speed up engine building process." - ) - use_fused_mlp: bool = Field( - default=True, - description= - "Enable horizontal fusion in Gated-MLP that combines two Matmul " - "operations into a single one followed by a separate SwiGLU kernel.") - pp_reduce_scatter: bool = Field( - default=False, - description="Enable a pipeline parallelism optimization with " - "ReduceScatter + AllGather targeting large MoE models.") - - def __getattribute__(self, name: str) -> Any: - """Override to resolve 'auto' values to dtype field. - - When a plugin field has value 'auto', return the value of dtype instead. - """ - # Use object.__getattribute__ to avoid infinite recursion - value = object.__getattribute__(self, name) - - if name != "dtype" and value == "auto": - return self.dtype - - return value - - @field_validator("dtype") - @classmethod - def validate_dtype_not_auto(cls, v: str) -> str: - if v == "auto": - raise ValueError("Plugin dtype cannot be 'auto'") - return v - - @field_validator("*", mode="before") - @classmethod - def convert_enable_disable(cls, value, info: ValidationInfo): - """Allow passing enable/disable strings which map to boolean/None values.""" - if value == "enable": - return True - elif value == "disable": - annotation = cls.model_fields[info.field_name].annotation - if annotation is bool or (get_origin(annotation) is Union - and bool in get_args(annotation)): - return False - return None - return value - - @field_validator("*", mode="after") - @classmethod - def log_field_changes(cls, v: Any, info: ValidationInfo) -> Any: - """Log all field changes for debugging.""" - logger.info(f"Set {cls.__name__}.{info.field_name} to {v}.") - return v - - @classmethod - def from_arguments(cls, args: argparse.Namespace): - """Create a PluginConfig from argparse arguments.""" - args = vars(args) - # Filter to only include fields that are part of PluginConfig - valid_fields = set(cls.model_fields.keys()) - filtered_args = {k: v for k, v in args.items() if k in valid_fields} - obj = cls(**filtered_args) - - # We want to know if the user explicitly disabled the gemm_plugin - # because nvfp4 gemm uses plugin by default currently - if 'gemm_plugin' in args and args['gemm_plugin'] == 'disable': - obj._explicitly_disable_gemm_plugin = True - - return obj - - def to_legacy_setting(self): - """Legacy setting means that all of the plugins and features are - disabled, this is needed for the legacy `build.py` script, which will be - migrated to the centralized building script `tensorrt_llm/commands/build.py`. - - After the migration is done, this function may or may not be deleted. - """ - for field_name, field_value in self: - if field_name == "dtype": - continue - elif isinstance(field_value, str): - setattr(self, field_name, None) - elif isinstance(field_value, - bool) or field_name == "paged_kv_cache": - setattr(self, field_name, False) - - @model_validator(mode="after") - def _validate_sm_compatibility(self) -> "PluginConfig": - unsupported_plugins = { - # bert_attention_plugin is handled within BertAttention - 100: [ - 'gemm_swiglu_plugin', 'fp8_rowwise_gemm_plugin', - 'low_latency_gemm_plugin', 'low_latency_gemm_swiglu_plugin', - 'bert_context_fmha_fp32_acc' - ] - } - sm = get_sm_version() - if sm in unsupported_plugins: - for plugin in unsupported_plugins[sm]: - val = getattr(self, plugin, None) - if val is not None and val != False: - raise ValueError( - f"{plugin}={val} is not supported on SM {sm}.") - return self - - @property - def context_fmha_type(self): - if self.bert_context_fmha_fp32_acc: - return ContextFMHAType.enabled_with_fp32_acc - elif self.context_fmha: - return ContextFMHAType.enabled - else: - return ContextFMHAType.disabled - - def is_context_fmha_enabled(self): - return self.context_fmha_type != ContextFMHAType.disabled - - @context_fmha_type.setter - def context_fmha_type(self, value): - if value == ContextFMHAType.disabled: - self.context_fmha = False - self.bert_context_fmha_fp32_acc = False - else: - self.context_fmha = True - if value == ContextFMHAType.enabled: - self.bert_context_fmha_fp32_acc = False - elif value == ContextFMHAType.enabled_with_fp32_acc: - self.bert_context_fmha_fp32_acc = True - - def set_smooth_quant_plugins(self, dtype: str = "auto"): - self.smooth_quant_gemm_plugin = dtype - self.rmsnorm_quantization_plugin = dtype - self.layernorm_quantization_plugin = dtype - self.quantize_per_token_plugin = True - self.quantize_tensor_plugin = True - return self - - def set_qserve_plugins(self, dtype: str = "auto"): - self.qserve_gemm_plugin = dtype - self.rmsnorm_quantization_plugin = dtype - self.quantize_per_token_plugin = True - return self - - def set_fp8_rowwise_quant_plugins(self, dtype: str = "auto"): - self.fp8_rowwise_gemm_plugin = dtype - self.rmsnorm_quantization_plugin = dtype - self.layernorm_quantization_plugin = dtype - self.quantize_per_token_plugin = True - self.quantize_tensor_plugin = True - return self - - def set_context_fmha(self, context_fmha_type=ContextFMHAType.enabled): - assert type(context_fmha_type) == ContextFMHAType - self.context_fmha_type = context_fmha_type - return self - - def enable_paged_kv_cache(self, tokens_per_block: int = 32): - self.paged_kv_cache = True - self.tokens_per_block = tokens_per_block - return self - - def set_nccl_plugin(self, dtype: str = "auto"): - self.nccl_plugin = dtype - init_all_reduce_helper() - return self - - def set_lora_plugin(self, dtype: str = None): - self.lora_plugin = dtype - return self - - def set_dora_plugin(self, enable: bool = False): - self.dora_plugin = enable - return self - - -# Only plugin configs in this list will be exposed as `trtllm-build` arguments, -# others are automatically enabled when needed, no need for users to control. -cli_plugin_args = [ - # Plugins - "bert_attention_plugin", - "gpt_attention_plugin", - "gemm_plugin", - "gemm_swiglu_plugin", - "fp8_rowwise_gemm_plugin", - "lora_plugin", - "dora_plugin", - "moe_plugin", - "mamba_conv1d_plugin", - "nccl_plugin", - "low_latency_gemm_plugin", - "low_latency_gemm_swiglu_plugin", - "gemm_allreduce_plugin", - - # Features - "context_fmha", - "bert_context_fmha_fp32_acc", - "remove_input_padding", - "tokens_per_block", - "use_paged_context_fmha", - "use_fp8_context_fmha", - "fuse_fp4_quant", - "multiple_profiles", - "paged_state", - "streamingllm", - "norm_quant_fusion", - "reduce_fusion", - "user_buffer", - "use_fused_mlp", - "pp_reduce_scatter", -] - - -def add_plugin_argument(parser: argparse.ArgumentParser): - for field_name, field_info in PluginConfig.model_fields.items(): - if field_name not in cli_plugin_args: - continue - help_message = field_info.description - if not help_message: - raise AttributeError(f"Please add help message for {field_name}.") - annotation = field_info.annotation - - # Extract choices from the Optional[Literal[...]] type - plugin_dtype_options = None - if get_origin(annotation) is Union: - args = get_args(annotation) - for arg in args: - if get_origin(arg) is Literal: - plugin_dtype_options = list(get_args(arg)) - if type(None) in args: - plugin_dtype_options.append(None) - break - - if plugin_dtype_options is not None: - if field_name == "gemm_plugin": - default = field_info.default - else: - default = field_info.default if field_info.default else "disable" - parser.add_argument( - "--" + field_name, - type=str, - default=default, - choices=[x if x else "disable" for x in plugin_dtype_options], - help=help_message) - elif annotation is bool: - parser.add_argument( - "--" + field_name, - type=str, - default="enable" if field_info.default else "disable", - choices=["enable", "disable"], - help=help_message) - else: - parser.add_argument("--" + field_name, - type=annotation, - default=field_info.default, - help=help_message) - return parser - - -def force_all_reduce_deterministic(): - return os.getenv("FORCE_DETERMINISTIC", "0") == "1" or os.getenv( - "FORCE_ALL_REDUCE_DETERMINISTIC", "0") == "1" - - -class CustomAllReduceHelper: - """ - Globally visible class to help usage of custom_all_reduce plugin. - Provides the following utilities: - - workspace: Tensor - When using CUSTOM or AUTO mode, a tensor containing pointers to memory - visible to all GPUs. It should be 3 pointers per TP rank - - ptr to data buffer, ptr to barriers in, ptr to barriers out. - It must be initialized using IpcMemory class. - - Usage: - - Set custom_all_reduce_helper.workspace with the required tensor. - Then, each instance of allreduce will reference that tensor automatically. - """ - POINTERS_PER_RANK = 7 - POINTERS_OF_COUNTER = 3 - - def __init__(self) -> None: - self.workspace: Optional[Tensor] = None - - def set_workspace_tensor(self, - mapping: Mapping, - num_profiles: Optional[int] = None): - from ..functional import Tensor - workspace_size = self.POINTERS_PER_RANK * mapping.tp_size + self.POINTERS_OF_COUNTER - - dim_range = None - if num_profiles is not None: - dim_range = OrderedDict([('all_reduce_size', - [workspace_size] * num_profiles)]) - - self.workspace = Tensor( - name='all_reduce_workspace', - dtype=trt.int64, - shape=[workspace_size], - dim_range=dim_range, - ) - - @staticmethod - def max_workspace_size_auto(tp_size: int, - support_deterministic=True) -> int: - """Calculate workspace size for allreduce fusion kernel. - - The workspace is used for lamport buffers in the fusion kernel. - Required size calculation: - - Each GPU needs 3 sub-buffers (for triple buffering) - - Each sub-buffer stores: max_num_tokens * hidden_size * dtype_size (bf16=2) - - The lamport allocation multiplies by tp_size, so: - lamport_size = 3 * size * tp_size (per GPU) - - Example: Llama 8B (hidden=4096), max_tokens=8192, bf16, TP=4 - - Data per sub-buffer: 8192 * 4096 * 2 = 64 MiB - - Total lamport: 3 * 64MB * 4 = 768 MiB per GPU - - Required 'size' parameter: 64 MiB (gets multiplied by tp_size in allocation) - - Default (67,108,864 = 64 MiB) supports: - - Models up to hidden_size=4096 with max_num_tokens=8192 - - Or hidden_size=8192 with max_num_tokens=4096 - - Override with TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE env var if needed for larger models. - """ - if force_all_reduce_deterministic() and support_deterministic: - workspace_size = os.getenv("FORCE_ALLREDUCE_KERNEL_WORKSPACE_SIZE", - "1000000000") - return int(workspace_size) - - # Allow override via environment variable for edge cases - workspace_size_env = os.getenv("TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE") - if workspace_size_env: - size = int(workspace_size_env) - logger.info( - f"Using custom allreduce fusion workspace size: {size} bytes ({size / (1024**2):.1f} MiB)" - ) - return size - - # Default: 64 MiB - supports most common model configurations - # Increase via env var if you see CUDA illegal memory access errors with large models - default_size = 67_108_864 # Exactly 64 MiB - return default_size - - @staticmethod - def max_workspace_size_lowprecision(tp_size: int) -> int: - return max_workspace_size_lowprecision(tp_size) - - @staticmethod - def initialize_lowprecision_buffers(workspace: "torch.tensor", - tp_size: int) -> None: - import torch - return torch.ops.trtllm.initialize_static_lowprecision_buffers( - workspace, tp_size) - - @staticmethod - def allocate_workspace(mapping: Mapping, - size: int) -> Tuple[List[IpcMemory], "torch.tensor"]: - import torch - - # Force pull mode and disable lamport when force deterministic is enabled, for reducing device memory usage. - force_deterministic = force_all_reduce_deterministic() - is_p2p_supported = can_access_peer(mapping) - ipc_buffers_size = size if force_deterministic else size * mapping.tp_size - ipc_buffers_ping = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_buffers_pong = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_barriers_in = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2 * - mapping.tp_size, is_p2p_supported) - ipc_barriers_out = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2 * - mapping.tp_size, is_p2p_supported) - lamport_buffers_size = 1 if force_deterministic else size * mapping.tp_size - lamport_buffers_0 = IpcMemory(mapping, lamport_buffers_size, - is_p2p_supported) - lamport_buffers_1 = IpcMemory(mapping, lamport_buffers_size, - is_p2p_supported) - lamport_buffers_2 = IpcMemory(mapping, lamport_buffers_size, - is_p2p_supported) - # TODO: it seems we may need to initialize lamport buffers for all tp groups - # just like its cpp counterpart (AllReduceBuffers::AllReduceBuffers()) does. - if is_p2p_supported: - lamport_initialize_all( - lamport_buffers_0.local_ptr, - lamport_buffers_1.local_ptr, - lamport_buffers_2.local_ptr, - lamport_buffers_size, - ) - buffers = [ - ipc_buffers_ping, - ipc_buffers_pong, - ipc_barriers_in, - ipc_barriers_out, - lamport_buffers_0, - lamport_buffers_1, - lamport_buffers_2, - # Start from 1 since 0 represents released state for barrier at the beginning of the all_reduce. - # The last element is the barrier flag counter. - torch.tensor([1, 1, 0], dtype=torch.int64, device="cuda") - ] - - return buffers, torch.tensor( - ipc_buffers_ping.serialize() + ipc_buffers_pong.serialize() + - ipc_barriers_in.serialize() + ipc_barriers_out.serialize() + - lamport_buffers_0.serialize() + lamport_buffers_1.serialize() + - lamport_buffers_2.serialize() + [buffers[-1].data_ptr()] + - [buffers[-1][1:].data_ptr()] + [buffers[-1][2:].data_ptr()], - dtype=torch.int64, - device="cpu") - - @staticmethod - def allocate_lowprecision_workspace( - mapping: Mapping, - size: int) -> Tuple[List[IpcMemory], "torch.tensor"]: - import torch - - # Force pull mode and disable lamport when force deterministic is enabled, for reducing device memory usage. - is_p2p_supported = can_access_peer(mapping) - ipc_buffers_size = size - ipc_buffers_ping = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_buffers_pong = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_barriers_in = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, - is_p2p_supported) - ipc_barriers_out = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, - is_p2p_supported) - buffers = [ - ipc_buffers_ping, ipc_buffers_pong, ipc_barriers_in, - ipc_barriers_out - ] - - return buffers, torch.tensor( - ipc_buffers_ping.serialize() + ipc_buffers_pong.serialize() + - ipc_barriers_in.serialize() + ipc_barriers_out.serialize() + [0] + - [0], - dtype=torch.int64, - device="cpu") - - @staticmethod - def allocate_allreduce_fusion_workspace( - mapping: Mapping, - size: int) -> Tuple[List[IpcMemory], "torch.tensor"]: - import torch - is_p2p_supported = can_access_peer(mapping) - ipc_buffers_size = size * mapping.tp_size - ipc_buffers = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) - ipc_barriers = IpcMemory(mapping, 256 * mapping.tp_size, - is_p2p_supported) - lamport_buffers_size = size * mapping.tp_size - lamport_buffers = IpcMemory(mapping, 3 * lamport_buffers_size, - is_p2p_supported) - if is_p2p_supported: - lamport_initialize( - lamport_buffers.local_ptr, - 3 * lamport_buffers_size, - ) - # flag_buffer[0], atomic flag read counter - # flag_buffer[1], non-lamport flag - # flag_buffer[2], lamport flag - flag_buffer = torch.tensor([0, 0, 0], dtype=torch.int, device="cuda") - # layout_buffer[0], clear size for next lamport kernel - # layout_buffer[1], triple buffer offset for lamport kernel - layout_buffer = torch.tensor([0, lamport_buffers_size], - dtype=torch.int64, - device="cuda") - - buffers = [ - ipc_buffers, ipc_barriers, lamport_buffers, flag_buffer, - layout_buffer - ] - - return buffers, torch.tensor( - ipc_buffers.serialize() + ipc_barriers.serialize() + - lamport_buffers.serialize() + [flag_buffer.data_ptr()] + - [layout_buffer.data_ptr()], - dtype=torch.int64, - device="cuda") - - -custom_all_reduce_helper = None - - -def init_all_reduce_helper(): - global custom_all_reduce_helper - custom_all_reduce_helper = CustomAllReduceHelper() - - -def current_all_reduce_helper(): - global custom_all_reduce_helper - assert custom_all_reduce_helper is not None, "You must call `init_all_reduce_helper` first" - return custom_all_reduce_helper diff --git a/tensorrt_llm/profiler.py b/tensorrt_llm/profiler.py index 2ef29c845f2f..10432bfe27cf 100644 --- a/tensorrt_llm/profiler.py +++ b/tensorrt_llm/profiler.py @@ -16,10 +16,7 @@ from functools import partial from typing import Literal, Optional, Tuple, Union -# isort: off import torch -import tensorrt as trt -# isort: on try: import psutil @@ -29,12 +26,9 @@ import pynvml except ImportError: pynvml = None -import traceback from tensorrt_llm.logger import logger -from ._common import _is_building - if psutil is None: logger.warning( "A required package 'psutil' is not installed. Will not " @@ -203,76 +197,3 @@ def print_memory_usage( f"Device {_format(alloc_device_mem, unit)}" ) _print_mem_message(msg, tag) - - -@_is_building -def check_gpt_mem_usage( - engine, - kv_dtype, - use_gpt_attention_plugin, - paged_kv_cache, - max_batch_size, - max_beam_width, - max_seq_len, - local_num_kv_heads, - head_size, - num_layers, -) -> int: - # Get the amount of memory - runtime = trt.Runtime(logger.trt_logger) - # 1. TensorRT engine activation memory - activation_size = 0 - try: - cuda_engine = runtime.deserialize_cuda_engine(engine) - assert cuda_engine is not None - activation_size = cuda_engine.device_memory_size_v2 / 1024 / 1024 - del cuda_engine - except Exception: - logger.warning(f"Exception when deserializing engine: {traceback.format_exc()}") - logger.warning("Activation memory size will be regarded as 0.") - logger.info(f"Activation memory size: {activation_size:.2f} MiB") - - # 2. Weights - weights_size = bytes_to_target_unit(engine.nbytes, "MiB") - logger.info(f"Weights memory size: {weights_size:.2f} MiB") - - # 3. Estimated max KV Cache size - kv_cache_size = ( - max_batch_size - * max_beam_width - * 2 - * local_num_kv_heads - * max_seq_len - * head_size - * num_layers - * kv_dtype.itemsize - ) - # without plugin, we need two set of kv cache buffers, - # one for inputs, and the other for outputs. - if not use_gpt_attention_plugin: - kv_cache_size *= 2 - kv_cache_size = bytes_to_target_unit(kv_cache_size, "MiB") - logger.info(f"Max KV Cache memory size: {kv_cache_size:.2f} MiB") - - # Estimated total amount of memory - est_memory_size = activation_size + weights_size + kv_cache_size - logger.info(f"Estimated max memory usage on runtime: {est_memory_size:.2f} MiB") - _, _, total_mem = device_memory_info(torch.cuda.current_device()) - total_mem = bytes_to_target_unit(total_mem, "MiB") - if est_memory_size > total_mem: - logger.warning( - f"Engine is successfully built, but GPU Memory ({total_mem:.2f} MB)" - " may not be enough when running inference on max shape." - ) - if paged_kv_cache: - logger.warning( - "Since paged_kv_cache is enabled, the max KV Cache " - "memory size is a estimate for very extreme cases, " - "it's possible that most cases won't meet OOM." - ) - else: - logger.warning( - "Enabling `--paged_kv_cache` could help reduce the GPU memory usage on runtime." - ) - - return est_memory_size diff --git a/tensorrt_llm/python_plugin.py b/tensorrt_llm/python_plugin.py deleted file mode 100644 index 3c6e4bc62d03..000000000000 --- a/tensorrt_llm/python_plugin.py +++ /dev/null @@ -1,578 +0,0 @@ -# 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 inspect -import pickle # nosec B403 -import typing -from copy import deepcopy -from dataclasses import dataclass -from typing import Sequence, Type, Union - -import numpy as np -import tensorrt as trt -import torch - -from ._common import default_trtnet -from ._utils import ( - TensorWrapper, - np_dtype_to_trt, - str_dtype_to_trt, - torch_dtype_to_trt, - trt_dtype_to_torch, -) -from .functional import Tensor, _create_tensor -from .plugin.plugin import TRT_LLM_PLUGIN_NAMESPACE - -_plugin_registered = dict() - - -@dataclass(slots=True, frozen=True) -class PluginInfo: - trt_plugin_version: int - plugin_namespace: str - plugin_name: str - plugin_version: str - plugin_num_outputs: int - - def __hash__(self): - return hash((self.plugin_name, self.plugin_namespace, self.plugin_version)) - - def __eq__(self, obj): - if not isinstance(obj, PluginInfo): - return False - return ( - self.plugin_name == obj.plugin_name - and self.plugin_namespace == obj.plugin_namespace - and self.plugin_version == obj.plugin_version - ) - - -def make_expr( - exprBuilder: Union[trt.IExprBuilder, Type[None]], - dim: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]], -) -> Union[trt.IDimensionExpr, Type[None]]: - """Make a dimension expression. - - Parameters: - exprBuilder: The trt.exprBuilder object. Using it to check whether dim has the same exprBuilder - or to create trt.IDimensionExpr if necessary. - dim: The input dim. - - Returns: - A trt.IDimensionExpr object. - """ - if isinstance(dim, DimensionExpr): - assert exprBuilder == dim.exprBuilder - return dim.expr - elif isinstance(dim, int): - return exprBuilder.constant(dim) - elif dim is None: - return None - elif isinstance(dim, trt.IDimensionExpr): - return dim - else: - raise Exception - - -def expr_operation( - a: Union[trt.IDimensionExpr, Type[None]], - b: Union[trt.IDimensionExpr, Type[None]], - operation: trt.DimensionOperation, - exprBuilder: trt.IExprBuilder, -): - """The function to do expr operation with None support.""" - if exprBuilder is None or a is None or b is None: - expr = None - else: - expr = exprBuilder.operation(operation, a, b) - return DimensionExpr(expr, exprBuilder) - - -class DimensionExpr: - """The class to wrap `trt.IDimensionExpr` to support more pythonic methods.""" - - def __init__( - self, - expr: Union[trt.IDimensionExpr, int, Type[None]], - exprBuilder: Union[trt.IExprBuilder, Type[None]], - ): - self.exprBuilder = exprBuilder - self.expr = expr - - @property - def expr(self): - return self._expr - - @expr.setter - def expr(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - self._expr = make_expr(self.exprBuilder, expr) - - def __add__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.SUM, self.exprBuilder) - - def __radd__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - return self.__add__(expr) - - def __mul__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.PROD, self.exprBuilder) - - def __rmul__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - return self.__mul__(expr) - - def __sub__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.SUB, self.exprBuilder) - - def __rsub__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(expr, self.expr, trt.DimensionOperation.SUB, self.exprBuilder) - - def __eq__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.EQUAL, self.exprBuilder) - - def __lt__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.LESS, self.exprBuilder) - - def __floordiv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.FLOOR_DIV, self.exprBuilder) - - def __rfloordiv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(expr, self.expr, trt.DimensionOperation.FLOOR_DIV, self.exprBuilder) - - def __truediv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.CEIL_DIV, self.exprBuilder) - - def __rtruediv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(expr, self.expr, trt.DimensionOperation.CEIL_DIV, self.exprBuilder) - - def max(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.MAX, self.exprBuilder) - - def min(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.MIN, self.exprBuilder) - - -class ShapeExpr: - """The class to Wrap `trt.DimsExprs` to support more pythonic methods.""" - - def __init__( - self, - dims: Union[Sequence[trt.IDimensionExpr], Sequence[int], Sequence[type[None]]], - exprBuilder: Union[trt.IExprBuilder, type[None]], - ): - self.exprBuilder = exprBuilder - self.dims = dims - - @property - def dims(self): - return self._dims - - @dims.setter - def dims( - self, - dims: Sequence[Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]], - ): - if dims is not None: - self._dims = [ - DimensionExpr(make_expr(self.exprBuilder, i), self.exprBuilder) for i in dims - ] - else: - self._dims = None - - def __getitem__(self, index: int): - if self._dims is not None: - return self._dims[index] - else: - return DimensionExpr(None, self.exprBuilder) - - def __setitem__( - self, - index: int, - value: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]], - ): - if self._dims is None: - return - assert index < len(self._dims) - value = DimensionExpr(make_expr(self.exprBuilder, value), self.exprBuilder) - self._dims[index] = value - - def __len__(self): - if self._dims is None: - return 0 - else: - return len(self._dims) - - def to_trt(self) -> trt.DimsExprs: - return trt.DimsExprs([i.expr for i in self.dims]) - - -class SymTensor: - """The class to represent symbolic tensors. - - Only contains dtype and shape information for users to write their own shape/dtype inference function. - """ - - def __init__( - self, - dtype: Union[torch.dtype, np.dtype, str, trt.DataType, Type[None]], - shape: Union[ShapeExpr, Sequence[int]], - ): - self.dtype = dtype - self.shape = shape - - @property - def shape(self) -> Union[ShapeExpr, Sequence[int]]: - return self._shape - - @shape.setter - def shape(self, shape: Union[ShapeExpr, Sequence[int]]): - assert isinstance(shape, (ShapeExpr, list, tuple)) - if isinstance(shape, (list, tuple)): - for i in shape: - assert isinstance(i, int) - self._shape = shape - - @property - def dtype(self) -> Union[trt.DataType, Type[None]]: - return self._dtype - - @dtype.setter - def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType, Type[None]]): - if isinstance(dtype, torch.dtype): - self._dtype = torch_dtype_to_trt(dtype) - elif isinstance(dtype, str): - self._dtype = str_dtype_to_trt(dtype) - elif isinstance(dtype, np.dtype): - self._dtype = np_dtype_to_trt(dtype) - elif isinstance(dtype, trt.DataType): - self._dtype = dtype - elif dtype is None: - self._dtype = None - else: - raise TypeError(f"Unsupported dtype: {dtype}") - - -def _convert_return_value_to_list(ret): - if not isinstance(ret, (list, tuple)): - return [ret] - assert isinstance(ret, (list, tuple)) - return ret - - -class PluginBase( - trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, trt.IPluginV3OneRuntime -): - """The base class of TRT-LLM plugin. - - All TRT-LLM plugin should inherit this class and at least rewrite `forward` and `shape_dtype_inference` - function. `forward` defines the plugin's compute flow while `shape_dtype_inference` defines how would - the output tensor's shape and dtype be inferenced from the input tensor. - """ - - _plugin_creator = None - _no_serialize_attr = {"_current_stream", "_workspace"} - - def __init__(self): - cls = type(self) - # Runtime check for plugin decorator - assert cls._plugin_creator is not None, ( - "Please make sure the plugin is registered through `@trtllm_plugin`" - ) - assert cls != PluginBase - - trt.IPluginV3.__init__(self) - trt.IPluginV3OneCore.__init__(self) - trt.IPluginV3OneBuild.__init__(self) - trt.IPluginV3OneRuntime.__init__(self) - - self.plugin_phase = trt.TensorRTPhase.BUILD - self.num_outputs = self._num_outputs - self.plugin_namespace = self._plugin_namespace - self.plugin_name = self._plugin_name - self.plugin_version = self._plugin_version - self.current_stream = -1 - self.workspace = 0 # nullptr - - @property - def current_stream(self): - if self._current_stream == -1: - return torch.cuda.current_stream().cuda_stream - else: - return self._current_stream - - @current_stream.setter - def current_stream(self, stream: int): - assert isinstance(stream, int) - self._current_stream = stream - - @property - def workspace(self) -> int: - buffer = self._workspace - return buffer if isinstance(buffer, int) else buffer.data_ptr() - - @workspace.setter - def workspace(self, workspace: Union[int, torch.Tensor]): - assert isinstance(workspace, (int, torch.Tensor)) - self._workspace = workspace - - def clone(self): - cls = type(self) - cloned_plugin = cls.__new__(cls) - super(cls, cloned_plugin).__init__() - cloned_plugin.__dict__.update(self._get_dict_to_serialize()) - return cloned_plugin - - def get_capability_interface(self, type): - return self - - def configure_plugin(self, input_desc, output_desc): - pass - - def get_output_data_types(self, input_types): - ret = self.shape_dtype_inference([SymTensor(i, ShapeExpr(None, None)) for i in input_types]) - - ret = _convert_return_value_to_list(ret) - assert len(ret) == self.num_outputs - for i in ret: - assert isinstance(i, SymTensor) - - return [i.dtype for i in ret] - - def get_output_shapes(self, inputs, shape_inputs, exprBuilder): - assert len(shape_inputs) == 0, "Currently we do not support shape inputs" - - ret = self.shape_dtype_inference( - [SymTensor(None, ShapeExpr(i, exprBuilder)) for i in inputs] - ) - - ret = _convert_return_value_to_list(ret) - assert len(ret) == self.num_outputs - for i in ret: - assert isinstance(i, SymTensor) - - return [i.shape.to_trt() for i in ret] - - def supports_format_combination(self, pos, in_out, num_inputs): - """By default, TRT-LLM plugin supports all dtype and linear format. - - It is the users responsibility to check the dtype the plugin supported in `forward` function. - """ - assert pos < len(in_out) - - desc = in_out[pos].desc - if desc.format != trt.TensorFormat.LINEAR: - return False - - return True - - def attach_to_context(self, context): - return self.clone() - - def get_fields_to_serialize(self): - buffer = pickle.dumps(self._get_dict_to_serialize()) - return trt.PluginFieldCollection( - [trt.PluginField("__plugin_pickle_obj__", buffer, trt.PluginFieldType.UNKNOWN)] - ) - - def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream): - torch_stream = torch.cuda.ExternalStream(stream_ptr=stream) - self.workspace = workspace - self.current_stream = stream - - with torch.cuda.stream(torch_stream): - self.forward( - tuple( - TensorWrapper.from_trt_desc(input_desc[i], inputs[i]) - for i in range(len(input_desc)) - ), - tuple( - TensorWrapper.from_trt_desc(output_desc[i], outputs[i]) - for i in range(len(output_desc)) - ), - ) - - self.current_stream = -1 - - def __call__(self, *args: Union[Sequence[TensorWrapper], Sequence[torch.Tensor]]): - is_trtllm = True - for i in args: - is_trtllm &= isinstance(i, Tensor) - - if not is_trtllm: - for i in args: - assert isinstance(i, torch.Tensor), ( - "Plugin inputs must be `tensorrt_llm.Tensor`s or `torch.Tensor`s" - ) - sym_tensors = self.shape_dtype_inference( - [SymTensor(i.dtype, [j for j in i.shape]) for i in args] - ) - sym_tensors = _convert_return_value_to_list(sym_tensors) - ret = [ - torch.empty(sym_tensor.shape, dtype=trt_dtype_to_torch(sym_tensor.dtype)) - for sym_tensor in sym_tensors - ] - self.current_stream = torch.cuda.current_stream().cuda_stream - self.workspace = torch.empty(self.workspace) - self.forward(args, ret) - else: - args = [i.trt_tensor for i in args] - layer_plugin = default_trtnet().add_plugin_v3(args, [], self) - ret = [ - _create_tensor(layer_plugin.get_output(i), layer_plugin) - for i in range(self.num_outputs) - ] - - if len(ret) == 1: - return ret[0] - - return ret - - def on_shape_change(self, input_desc, output_desc): - pass - - def get_valid_tactics(self): - return [] - - def set_tactic(self, index): - if index != 0: - raise RuntimeError( - "By default TRT should not set tactics since PluginBase do not provide custom tactic." - ) - - def forward(self, inputs: Sequence[TensorWrapper], outputs: Sequence[TensorWrapper]): - """Expect users to rewrite this function to define the compute flow. - - There are a few special attributes for users to get access to some resources. - - `self.workspace`: The workspace address of TRT managed workspace. - `self.current_stream`: The CUDA stream this plugin is expected to execute on. By default - `PluginBase` set the torch.cuda.current_stream() to this stream. This attribute is for the - toolkit that doesn't work with torch's stream. - """ - raise NotImplementedError - - def shape_dtype_inference(self, inputs: Sequence[SymTensor]): - """Expect users to rewrite this function to define the shape dtype inference for output tensors.""" - raise NotImplementedError - - def _get_dict_to_serialize(self): - ret = {} - for k, v in self.__dict__.items(): - if k not in self._no_serialize_attr: - ret[k] = deepcopy(v) if self.deepcopy_clone else v - return ret - - -class PluginCreatorBase(trt.IPluginCreatorV3One): - def __init__(self): - super().__init__() - - def create_plugin(self, name, fc, phase): - if len(fc) == 1 and fc[0].name == "__plugin_pickle_obj__": - data = fc[0].data - plugin_dict = pickle.loads(data) # nosec B301 - plugin = self.plugin_cls.__new__(self.plugin_cls) - super(self.plugin_cls, plugin).__init__() - plugin.__dict__.update(plugin_dict) - else: - raise RuntimeError("Expect to be called by TRT") - plugin.plugin_phase = phase - return plugin - - -def trtllm_plugin( - plugin_name: str, - *, - plugin_version: str = "1", - plugin_namespace: str = TRT_LLM_PLUGIN_NAMESPACE, - plugin_num_outputs: Union[int, Type[None]] = None, - deepcopy_clone: bool = True, - no_serialize_attr: Sequence[str] = set(), -): - def plugin_registration(plugin_cls): - assert issubclass(plugin_cls, PluginBase) - assert hasattr(plugin_cls, "__dict__"), ( - "Plugin wrapper uses `__dict__` to track plugin states" - ) - nonlocal plugin_num_outputs - - annotation = inspect.signature(plugin_cls.shape_dtype_inference).return_annotation - origin_annotation = typing.get_origin(annotation) - if origin_annotation is tuple or annotation is SymTensor: - if origin_annotation is tuple: - element_types = typing.get_args(annotation) - for ty in element_types: - assert ty == SymTensor, ( - f"Plugin {plugin_name}'s `shape_dtype_inference` return annotation must be SymTensor " - "or a tuple of SymTensor" - ) - infered_num_outputs = len(element_types) - else: - infered_num_outputs = 1 - if plugin_num_outputs is not None: - assert plugin_num_outputs == infered_num_outputs, ( - f"Plugin {plugin_name}'s `_num_outputs` and return annotation mismatch, " - f"{plugin_cls._num_outputs} != {infered_num_outputs}" - ) - plugin_num_outputs = infered_num_outputs - else: - assert plugin_num_outputs is not None, ( - "Must specify `num_outputs` or valid `shape_dtype_inference` return annotation for " - f"{plugin_name}. The valid types are SymTensor or a tuple of SymTensor, got {annotation}." - ) - - plugin_info = PluginInfo( - 3, plugin_namespace, plugin_name, plugin_version, plugin_num_outputs - ) - assert plugin_info not in _plugin_registered, ( - f"Redefine plugin with info: {plugin_info} which is previously defined as " - f"{_plugin_registered[plugin_info]}" - ) - - _plugin_registered[plugin_info] = plugin_info - plugin_cls._plugin_name = plugin_name - plugin_cls._plugin_version = plugin_version - plugin_cls._plugin_namespace = plugin_namespace - plugin_cls._num_outputs = plugin_num_outputs - plugin_cls.deepcopy_clone = deepcopy_clone - plugin_cls._no_serialize_attr.update(no_serialize_attr) - - plugin_registry = trt.get_plugin_registry() - - plugin_creator = PluginCreatorBase() - plugin_creator.name = plugin_cls._plugin_name - plugin_creator.plugin_namespace = plugin_cls._plugin_namespace - plugin_creator.plugin_version = plugin_cls._plugin_version - plugin_creator.field_names = trt.PluginFieldCollection([]) - plugin_creator.plugin_cls = plugin_cls - - plugin_cls._plugin_creator = plugin_creator - ret = plugin_registry.register_creator(plugin_creator, plugin_cls._plugin_namespace) - - assert ret, f"Plugin: {plugin_cls} register failed, please check the error log." - - return plugin_cls - - return plugin_registration diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py index ddaa7394eb5b..2eb1eab32624 100644 --- a/tensorrt_llm/quantization/functional.py +++ b/tensorrt_llm/quantization/functional.py @@ -12,950 +12,19 @@ # 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, Tuple, Union -import numpy as np -import tensorrt as trt import torch -import torch.nn.functional as F -from .._common import default_net, default_trtnet -from .._utils import (get_sm_version, str_dtype_to_np, str_dtype_to_trt, - trt_dtype_to_np) -from ..functional import (Tensor, _add_plugin_info, _create_tensor, cast, clip, - constant, flatten, layer_norm, matmul, - repeat_interleave, rms_norm, round, sum, view) -from ..layers.linear import ColumnLinear -from ..parameter import Parameter -from ..plugin import TRT_LLM_PLUGIN_NAMESPACE -from .mode import QuantMode - - -def smooth_quant_gemm(input: Tensor, weights: Tensor, scales_a: Tensor, - scales_b: Tensor, per_token_scaling: bool, - per_channel_scaling: bool, dtype: str) -> Tensor: - if not default_net().plugin_config.smooth_quant_gemm_plugin: - if per_token_scaling and input.size(0) == -1: - # WAR for DQ per-token scaling doesn't support dynamic shapes - - scale_one = constant(np.array(1.0, dtype=np.float32)) - input = dequantize(input, scale_one, 0, 'float32') - weights = dequantize(weights, scale_one, 0, 'float32') - result = matmul(input, weights, False, True, False) - scales = matmul(scales_a, scales_b, False, False, False) - result = result * scales - result = cast(result, dtype) - return result - else: - if not per_token_scaling: - scales_a = view(scales_a, []) - else: - scales_a = flatten(scales_a) - if not per_channel_scaling: - scales_b = view(scales_b, []) - else: - scales_b = flatten(scales_b) - input = dequantize(input, scales_a, 0, dtype) - weights = dequantize(weights, scales_b, 0, dtype) - result = matmul(input, weights, False, True, False) - return result - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'SmoothQuantGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - per_channel_scaling = 1 if per_channel_scaling else 0 - per_channel_scaling = trt.PluginField( - "has_per_channel_scaling", - np.array(per_channel_scaling, dtype=np.int32), - trt.PluginFieldType.INT32) - - per_token_scaling = 1 if per_token_scaling else 0 - per_token_scaling = trt.PluginField( - "has_per_token_scaling", np.array(per_token_scaling, - dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.smooth_quant_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [per_channel_scaling, per_token_scaling, pf_type]) - gemm_plug = plg_creator.create_plugin("sq_gemm", pfc) - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, scales_a.trt_tensor, - scales_b.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "sq_gemm", pfc) - if not default_net().strongly_typed: - layer.get_input(0).set_dynamic_range(-127, 127) - layer.get_input(1).set_dynamic_range(-127, 127) - return _create_tensor(layer.get_output(0), layer) - - -def qserve_gemm_per_group(input: Tensor, - act_scales: Tensor, - weights: Tensor, - s1_scales: Tensor, - s2_scales: Tensor, - s2_zeros: Tensor, - group_size: int = 128) -> Tensor: - if not default_net().plugin_config.qserve_gemm_plugin: - raise TypeError("QServe Quant GEMM is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QServeGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.qserve_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_group_size = trt.PluginField("group_size", - np.array([group_size], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, pf_group_size]) - gemm_plug = plg_creator.create_plugin("qserve_gemm", pfc) - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, s2_zeros.trt_tensor, - s2_scales.trt_tensor, s1_scales.trt_tensor, act_scales.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "qserve_gemm", pfc) - if not default_net().strongly_typed: - # Useless. But must be kept otherwise leads to the following TRT API Usage error: - # input/output with DataType Int8 in network without Q/DQ layers must have dynamic range set when no calibrator is used - layer.get_input(0).set_dynamic_range(-128, 127) - layer.get_input(1).set_dynamic_range(-128, 127) - layer.get_input(2).set_dynamic_range(-128, 127) - layer.get_input(3).set_dynamic_range(-128, 127) - return _create_tensor(layer.get_output(0), layer) - - -def qserve_gemm_per_channel(input: Tensor, act_scales: Tensor, act_sums: Tensor, - weights: Tensor, s1_scales: Tensor, - s1_szeros: Tensor) -> Tensor: - if not default_net().plugin_config.qserve_gemm_plugin: - raise TypeError("QServe Quant GEMM is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QServeGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.qserve_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_group_size = trt.PluginField("group_size", np.array([-1], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, pf_group_size]) - gemm_plug = plg_creator.create_plugin("qserve_gemm", pfc) - - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, s1_scales.trt_tensor, - s1_szeros.trt_tensor, act_sums.trt_tensor, act_scales.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "qserve_gemm", pfc) - - if not default_net().strongly_typed: - # Useless. But must be kept otherwise leads to the following TRT API Usage error: - # input/output with DataType Int8 in network without Q/DQ layers must have dynamic range set when no calibrator is used - layer.get_input(0).set_dynamic_range(-128, 127) - layer.get_input(1).set_dynamic_range(-128, 127) - - return _create_tensor(layer.get_output(0), layer) - - -def fp8_rowwise_gemm(input: Tensor, weights: Tensor, scales_a: Tensor, - scales_b: Tensor, per_token_scaling: bool, - per_channel_scaling: bool) -> Tensor: - if not default_net().plugin_config.fp8_rowwise_gemm_plugin: - raise TypeError("Fp8 Rowwise GEMM is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Fp8RowwiseGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - per_channel_scaling = 1 if per_channel_scaling else 0 - per_channel_scaling = trt.PluginField( - "has_per_channel_scaling", - np.array(per_channel_scaling, dtype=np.int32), - trt.PluginFieldType.INT32) - - per_token_scaling = 1 if per_token_scaling else 0 - per_token_scaling = trt.PluginField( - "has_per_token_scaling", np.array(per_token_scaling, - dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.fp8_rowwise_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [per_channel_scaling, per_token_scaling, pf_type]) - gemm_plug = plg_creator.create_plugin("fp8_rowwise_gemm", pfc) - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, scales_a.trt_tensor, - scales_b.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "fp8_rowwise_gemm", pfc) - if not default_net().strongly_typed: - layer.get_input(0).set_dynamic_range(-448, 448) - layer.get_input(1).set_dynamic_range(-448, 448) - return _create_tensor(layer.get_output(0), layer) - - -def weight_only_quant_matmul(input: Tensor, - weights: Tensor, - scales: Tensor, - weightTypeId: int, - dtype: str = 'float16', - transa: bool = False, - transb: bool = False) -> Tensor: - if not default_net( - ).plugin_config.weight_only_quant_matmul_plugin or transa or transb: - scale_axis = 0 if transb else 1 - if weights.dtype != trt.int8: - # Q->DQ - weights = quantize(weights, scales, dtype='int8', axis=1) - weights = dequantize(weights, scales, scale_axis, input.dtype) - else: - weights = dequantize(weights, scales, scale_axis, input.dtype) - - res = matmul(input, weights, transa=transa, transb=transb) - return cast(res, dtype) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'WeightOnlyQuantMatmul', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - weight_type_id = trt.PluginField("weight_type_id", - np.array(weightTypeId, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.weight_only_quant_matmul_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, weight_type_id]) - matmul_plug = plg_creator.create_plugin("woq_matmul", pfc) - plug_inputs = [input.trt_tensor, weights.trt_tensor, scales.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, matmul_plug) - _add_plugin_info(layer, plg_creator, "woq_matmul", pfc) - if not default_net().strongly_typed: - layer.get_input(1).set_dynamic_range(-127, 127) - return _create_tensor(layer.get_output(0), layer) - - -def weight_only_groupwise_quant_matmul(input: Tensor, - pre_quant_scale: Tensor, - weights: Tensor, - scales: Tensor, - zeros: Tensor, - biases: Tensor, - alpha: Parameter, - quant_algo: int, - group_size: int, - dtype: str = 'float16') -> Tensor: - if not default_net( - ).plugin_config.weight_only_groupwise_quant_matmul_plugin: - scales = repeat_interleave(scales, group_size, 0) - weights = quantize(weights, scales, dtype='int8', axis=1) - weights = dequantize(weights, scales, 1, input.dtype) - - if quant_algo & 8: - # fp8_alpha - input = input * alpha.value - if quant_algo & 4: - # pre quant - input = input * pre_quant_scale - elif quant_algo & 2: - # zero - zeros = repeat_interleave(zeros, group_size, 0) - weights += zeros - res = matmul(input, weights) - if quant_algo & 1: - # bias - res += biases - - return cast(res, dtype) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'WeightOnlyGroupwiseQuantMatmul', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - quant_algo_ = trt.PluginField("quant_algo", - np.array(quant_algo, dtype=np.int32), - trt.PluginFieldType.INT32) - group_size_ = trt.PluginField("group_size", - np.array(group_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - if alpha: - alpha.is_buffer = True - alpha_value = alpha.raw_value[0] - else: - alpha_value = 1.0 - - alpha_ = trt.PluginField("alpha", np.array(alpha_value, - dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - p_dtype = default_net( - ).plugin_config.weight_only_groupwise_quant_matmul_plugin - pf_type_ = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [pf_type_, quant_algo_, group_size_, alpha_]) - - matmul_plug = plg_creator.create_plugin("woq_groupwise_matmul", pfc) - - # quant_algo = use_int8_weight * 16 + fp8_alpha * 8 + pre_quant_scale * 4 + zero * 2 + bias - plug_inputs = [input.trt_tensor] - - # Flags for indicating whether the corresponding inputs are applied in quant_algo - # quant_algo = use_int8_weight * INT8_WEIGHT + fp8_alpha * FP8_ALPHA + pre_quant_scale * PRE_QUANT_SCALE + zero * ZERO + bias * BIAS - # Here use_int8_weight, pre_quant_scale, zero and bias are boolean type - BIAS = 1 - ZERO = 2 - PRE_QUANT_SCALE = 4 - - if quant_algo & PRE_QUANT_SCALE: - plug_inputs += [pre_quant_scale.trt_tensor] - - plug_inputs += [weights.trt_tensor, scales.trt_tensor] - - if quant_algo & ZERO: - plug_inputs += [zeros.trt_tensor] - if quant_algo & BIAS: - plug_inputs += [biases.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, matmul_plug) - _add_plugin_info(layer, plg_creator, "woq_groupwise_matmul", pfc) - - return _create_tensor(layer.get_output(0), layer) - - -# TODO: Should be renamed to layer_norm_quantize. -def smooth_quant_layer_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - eps: float = 1e-05, - use_diff_of_squares: bool = True, - dynamic_act_scaling: bool = False) -> Tensor: - if not default_net().plugin_config.layernorm_quantization_plugin: - dtype = trt_dtype_to_np(input.dtype) - if weight is None: - weight = constant(np.ones(normalized_shape, dtype=dtype)) - if bias is None: - bias = constant(np.zeros(normalized_shape, dtype=dtype)) - result = layer_norm(input, normalized_shape, weight, bias, eps, - use_diff_of_squares) - if not dynamic_act_scaling: - return quantize_tensor(result, scale) - else: - return quantize_per_token(result) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LayernormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.int8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.use_smooth_quant(per_token=True))], - np.int32), trt.PluginFieldType.INT32) - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - use_diff_of_squares = trt.PluginField( - "use_diff_of_squares", - np.array([int(use_diff_of_squares)], dtype=np.int32), - trt.PluginFieldType.INT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.layernorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([ - eps, use_diff_of_squares, dyn_act_scaling, pf_type, output_type, - quant_mode - ]) - layernorm_plug = plg_creator.create_plugin("layernorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - - # LayerNorm plugin only supports float32 scale - scale = cast(scale, "float32") - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, layernorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "layernorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - - -# TODO: Should be renamed to rms_norm_quantize. This is also used by QServe. -def smooth_quant_rms_norm( - input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - clamp_val: Optional[Tensor] = None, - eps: float = 1e-05, - dynamic_act_scaling: bool = False, - scale_dtype='float32', - sum_per_token: bool = False, - sum_dtype='float32' -) -> Tensor | tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: - if sum_per_token and not dynamic_act_scaling: - raise ValueError( - "sum_per_token is only allowed if dynamic_act_scaling is enabled!") - - if not default_net().plugin_config.rmsnorm_quantization_plugin: - result = rms_norm(input, normalized_shape, 1, weight, eps) - if bias is not None: - result += bias - if not dynamic_act_scaling: - return quantize_tensor(result, scale) - else: - return quantize_per_token(result, clamp_val, scale_dtype, - sum_per_token, sum_dtype) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'RmsnormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.int8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.use_smooth_quant(per_token=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int32), - trt.PluginFieldType.INT32) - - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - - sum_per_token_pf = trt.PluginField( - "sum_per_token", np.array([int(sum_per_token)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.rmsnorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([ - eps, dyn_act_scaling, sum_per_token_pf, clamp_enabled, quant_mode, - pf_type, output_type - ]) - rmsnorm_plug = plg_creator.create_plugin("rmsnorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - - # TODO: Why not fuse scale (which seems to be a per-tensor scaling factor of the original values) into weight? - if scale is None: - scale = constant(np.ones(1, dtype=str_dtype_to_np(p_dtype))) - - # RMS Norm Plugin only supports float32 scale - scale = cast(scale, "float32") - - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, rmsnorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "rmsnorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - output_quantized = _create_tensor(layer.get_output(0), layer) - output_scales = _create_tensor(layer.get_output(1), layer) - - # TODO: The plugin should be able to directly output float16 scales - if str_dtype_to_trt(scale_dtype) != output_scales.dtype: - output_scales = cast(output_scales, scale_dtype) - - if not sum_per_token: - return output_quantized, output_scales - - output_sums = _create_tensor(layer.get_output(2), layer) - # TODO: The plugin should be able to directly output float16 sums - if str_dtype_to_trt(sum_dtype) != output_sums.dtype: - output_sums = cast(output_sums, sum_dtype) - - return output_quantized, output_scales, output_sums - - -def fp8_rowwise_rms_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - clamp_val: Optional[Tensor] = None, - eps: float = 1e-05, - dynamic_act_scaling: bool = True) -> Tensor: - if not default_net().plugin_config.rmsnorm_quantization_plugin: - raise TypeError("Fp8 Rowwise Rms Norm is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'RmsnormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.fp8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.from_description(use_fp8_rowwise=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int32), - trt.PluginFieldType.INT32) - - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(False)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.rmsnorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - eps, dyn_act_scaling, sum_per_token_pf, clamp_enabled, quant_mode, - pf_type, output_type - ]) - - rmsnorm_plug = plg_creator.create_plugin("rmsnorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if scale is None: - scale = constant(np.ones((1, ), dtype=str_dtype_to_np(p_dtype))) - - # RMS Norm Plugin only supports float32 scale - scale = cast(scale, "float32") - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, rmsnorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-448, 448) - _add_plugin_info(layer, plg_creator, "rmsnorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - - -def fp8_rowwise_layer_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - clamp_val: Optional[Tensor] = None, - eps: float = 1e-05, - dynamic_act_scaling: bool = True) -> Tensor: - if not default_net().plugin_config.layernorm_quantization_plugin: - raise TypeError("Fp8 Rowwise Layer Norm is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LayernormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.fp8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.from_description(use_fp8_rowwise=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int32), - trt.PluginFieldType.INT32) - - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(False)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.layernorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - eps, dyn_act_scaling, sum_per_token_pf, clamp_enabled, quant_mode, - pf_type, output_type - ]) - - layernorm_plug = plg_creator.create_plugin("layernorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if scale is None: - scale = constant(np.ones((1, ), dtype=str_dtype_to_np(p_dtype))) - - # Layer Norm Plugin only supports float32 scale - scale = cast(scale, "float32") - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, layernorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-448, 448) - _add_plugin_info(layer, plg_creator, "layernorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - - -def fused_layernorm( - input: Tensor, - normalized_shape: Union[int, Tuple[int]], - residual: Optional[Tensor] = None, - weight: Optional[Tensor] = None, - # beta: Optional[Tensor] = None, - # bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - eps: float = 1e-05, - p_dtype: str = 'float16', - need_fp32_output: bool = False) -> Tensor: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'FusedLayernorm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - need_fp32_output_value = need_fp32_output - need_fp32_output = trt.PluginField( - "need_fp32_output", np.array([int(need_fp32_output_value)], np.int32), - trt.PluginFieldType.INT32) - need_quantize_value = scale is not None - need_quantize = trt.PluginField( - "need_quantize", np.array([int(need_quantize_value)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection( - [eps, need_fp32_output, need_quantize, pf_type]) - fused_layernorm_plug = plg_creator.create_plugin("fused_layernorm", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - # if beta is None: - # beta = constant( - # np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - # if bias is None: - # bias = constant( - # np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if need_quantize_value: - plug_inputs = [ - input.trt_tensor, residual.trt_tensor, weight.trt_tensor, - scale.trt_tensor - ] - else: - plug_inputs = [ - input.trt_tensor, - residual.trt_tensor, - weight.trt_tensor, - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, fused_layernorm_plug) - _add_plugin_info(layer, plg_creator, "fused_layernorm", pfc) - if not need_quantize_value: - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - return _create_tensor(layer.get_output(0), layer), _create_tensor( - layer.get_output(1), layer), _create_tensor(layer.get_output(2), layer) - - -def quantize(input: Tensor, - scale_factor: Tensor, - dtype: str, - axis: int = -1) -> Tensor: - layer = default_trtnet().add_quantize(input.trt_tensor, - scale_factor.trt_tensor, - str_dtype_to_trt(dtype)) - layer.axis = axis - - output = _create_tensor(layer.get_output(0), layer) - - return output - - -def dequantize(input: Tensor, - scale_factor: Tensor, - axis: int = -1, - output_type: Union[str, trt.DataType] = 'float16') -> Tensor: - - if isinstance(output_type, str): - output_type = str_dtype_to_trt(output_type) - - layer = default_trtnet().add_dequantize(input.trt_tensor, - scale_factor.trt_tensor, - output_type) - layer.axis = axis - - if not default_net().strongly_typed: - layer.precision = input.dtype - - output = _create_tensor(layer.get_output(0), layer) - - return output - - -def quantize_per_token( - x: Tensor, - clamp_val: Optional[Tensor] = None, - scale_dtype='float32', - sum_per_token: bool = False, - sum_dtype='float32', -) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: - if not default_net().plugin_config.quantize_per_token_plugin: - x = cast(x, 'float32') - xmax = x.abs().max(-1, keepdim=True) - scales = xmax / 127.0 - out = x * 127.0 / xmax - out = round(out) - out = clip(out, -128, 127) - quantized = cast(out, 'int8') - if not sum_per_token: - return quantized, scales - sums = sum(x, -1, keepdim=True) - if sum_dtype is not None and str_dtype_to_trt(sum_dtype) != sums.dtype: - sums = cast(sums, sum_dtype) - return quantized, scales, sums - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizePerToken', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("type_id", np.array([int(trt.int8)], - np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.use_smooth_quant(per_token=True))], np.int32), - trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField("clamp_enabled", - np.array([clamp_val is not None], np.int8), - trt.PluginFieldType.INT8) - - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(sum_per_token)], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [output_type, quant_mode, clamp_enabled, sum_per_token_pf]) - quantize_plug = plg_creator.create_plugin("quantize_per_token_plugin", pfc) - - plug_inputs = [x.trt_tensor] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "quantize_per_token_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - scales = _create_tensor(layer.get_output(1), layer) - - # TODO: The plugin should be able to directly output float16 scales to avoid a cast - if scale_dtype is not None and str_dtype_to_trt( - scale_dtype) != scales.dtype: - scales = cast(scales, scale_dtype) - if not sum_per_token: - return quantized, scales - - sums = _create_tensor(layer.get_output(2), layer) - # TODO: The plugin should be able to directly output float16 sums to avoid a cast - if sum_dtype is not None and str_dtype_to_trt(sum_dtype) != sums.dtype: - sums = cast(sums, sum_dtype) - - return quantized, scales, sums - - -def quantize_fp8_per_token(x: Tensor, - clamp_val: Optional[Tensor] = None) -> Tuple[Tensor]: - if not default_net().plugin_config.quantize_per_token_plugin: - x = cast(x, 'float32') - xmax = x.abs().max(-1, keepdim=True) - scale = xmax / 448.0 - out = x * 448.0 / xmax - out = round(out) - out = clip(out, -448, 448) - quantized_out = cast(out, 'fp8') - return quantized_out, scale - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizePerToken', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("type_id", - np.array([int(trt.fp8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.from_description(use_fp8_rowwise=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int8), - trt.PluginFieldType.INT8) - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(False)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection( - [output_type, quant_mode, clamp_enabled, sum_per_token_pf]) - quantize_plug = plg_creator.create_plugin("quantize_per_token_plugin", - pfc) - - plug_inputs = [x.trt_tensor] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-448, 448) - _add_plugin_info(layer, plg_creator, "quantize_per_token_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - scales = _create_tensor(layer.get_output(1), layer) - - return quantized, scales - - -def quantize_tensor(x, scale): - if not default_net().plugin_config.quantize_tensor_plugin: - if scale.dtype == str_dtype_to_trt('float32'): - x = cast(x, 'float32') - scaled = x * scale - rounded = round(scaled) - clipped = clip(rounded, -128, 127) - quantized = cast(clipped, 'int8') - else: - scale = cast(scale, 'float32') - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizeTensor', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pfc = trt.PluginFieldCollection([]) - quantize_plug = plg_creator.create_plugin("quantize_tensor_plugin", pfc) - - plug_inputs = [x.trt_tensor, scale.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "quantize_tensor_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - return quantized - - -def symmetric_quantize_last_axis_of_batched_matrix(weight, quant_mode): - amax = weight.abs().max(dim=0)[0].to(weight.dtype) - if quant_mode == torch.int8: - scale = amax / 128. - qweight = torch.clamp((weight / scale).round(), -128, 127).char() - qweight = qweight.T.reshape(weight.shape) - else: - scale = amax / 8. - qweight = torch.clamp((weight / scale).round(), -8, 7).char() - qweight[qweight < 0] += 16 - qweight = qweight.T.view(torch.uint8) - qweight = (qweight[:, 1::2] * 16 + qweight[:, ::2]).view(torch.int8) - qweight = qweight.reshape(weight.shape[0], weight.shape[1] // 2) - return qweight, scale +from .._utils import get_sm_version def preprocess_weights_for_mixed_gemm( - tensor: torch.Tensor, - quant_mode: torch.dtype, - act_dtype: torch.dtype, - sm_: int = -1, - do_weight_interleave: bool = True) -> torch.Tensor: + tensor: torch.Tensor, + quant_mode: torch.dtype, + act_dtype: torch.dtype, + sm_: int = -1, + do_weight_interleave: bool = True, +) -> torch.Tensor: sm_ = sm_ if sm_ > 0 else get_sm_version() # 3-D inputs (MoE) on Hopper+ and any input on SM120/SM121 reuse the SM80 # interleaved layout. Check the original rank before unsqueeze. @@ -969,13 +38,73 @@ def preprocess_weights_for_mixed_gemm( permutation_map = { "16_8": [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15], "16_4": [ - 0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, - 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31 + 0, + 1, + 8, + 9, + 16, + 17, + 24, + 25, + 2, + 3, + 10, + 11, + 18, + 19, + 26, + 27, + 4, + 5, + 12, + 13, + 20, + 21, + 28, + 29, + 6, + 7, + 14, + 15, + 22, + 23, + 30, + 31, ], "8_4": [ - 0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23, 8, 9, 10, - 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31 - ] + 0, + 1, + 2, + 3, + 16, + 17, + 18, + 19, + 4, + 5, + 6, + 7, + 20, + 21, + 22, + 23, + 8, + 9, + 10, + 11, + 24, + 25, + 26, + 27, + 12, + 13, + 14, + 15, + 28, + 29, + 30, + 31, + ], } # permute_B_rows_for_mixed_gemm @@ -988,9 +117,9 @@ def preprocess_weights_for_mixed_gemm( num_rows = tensor.shape[1] num_cols = tensor.shape[2] - assert (sm_ >= 75) - assert (num_rows % B_ROWS_PER_MMA == 0) - assert (num_cols % MMA_SHAPE_N == 0) + assert sm_ >= 75 + assert num_rows % B_ROWS_PER_MMA == 0 + assert num_cols % MMA_SHAPE_N == 0 if do_weight_interleave and sm_ < 100: row_idx_list = [(row_idx // B_ROWS_PER_MMA) * B_ROWS_PER_MMA + @@ -1020,12 +149,16 @@ def preprocess_weights_for_mixed_gemm( rows_per_tile = 128 * 8 // BITS_PER_ELT_A elts_in_int32 = 32 // BITS_PER_ELT_B - assert (num_rows % elts_in_int32 == 0) - assert (num_rows % rows_per_tile == 0) + assert num_rows % elts_in_int32 == 0 + assert num_rows % rows_per_tile == 0 - tensor = tensor.reshape(num_experts, -1, interleave, - num_rows // rows_per_tile, - rows_per_tile * 4 // elts_in_int32) + tensor = tensor.reshape( + num_experts, + -1, + interleave, + num_rows // rows_per_tile, + rows_per_tile * 4 // elts_in_int32, + ) tensor = tensor.permute(0, 1, 3, 2, 4).reshape(original_shape) # add_bias_and_interleave_quantized_tensor_inplace @@ -1049,412 +182,3 @@ def preprocess_weights_for_mixed_gemm( raise NotImplementedError return tensor.squeeze(0).contiguous() - - -def get_weight_scale_interleave_factor(interleaved_dim: int, - group_size: int = 128) -> int: - # Calculate the weight_scale interleave factor for W4A8 groupwise MoE quant - # only Hopper w4a8 does interleave for weight scale, other arch or Hopper w4a16 default to 1 - factor = 1 - if get_sm_version() == 90: - if interleaved_dim % (4 * group_size) == 0: - factor = 4 - elif interleaved_dim % (2 * group_size) == 0: - factor = 2 - elif interleaved_dim % group_size == 0: - factor = 1 - else: - raise NotImplementedError( - f"Interleaved dimension must be a multiple of group_size ({group_size}), received {interleaved_dim}." - ) - return factor - - -def validate_group_size(layer): - # TODO: Remove this function and its usage after W4A8-AWQ with group_size = 64 is implemented. - W4A8_AWQ = 8 - if layer.quant_algo & W4A8_AWQ and layer.group_size == 64: - raise NotImplementedError( - "W4A8_AWQ with group_size = 64 is not implemented yet!") - - -def unpack_int32_into_int8(w_packed, autoawq_reorder=False): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - if autoawq_reorder: - w_unpacked = w_unpacked.view(-1, 8)[:, [0, 4, 1, 5, 2, 6, 3, 7]].view( - w_unpacked.shape) - return w_unpacked.contiguous() - - -def change_qkv_leading_dim(w, num_heads): - if w.dim() == 1: - w = w.reshape(num_heads, 3, -1) - w = w.transpose(0, 1).reshape(-1) - else: - shape = w.shape - head_dim = shape[1] // (3 * num_heads) - w = w.reshape(-1, num_heads, 3, head_dim) - w = w.transpose(1, 2).reshape(shape[0], -1) - return w - - -def pad_like(w, target_shape, value=0): - if w.shape != target_shape: - pad_dim = [] - for dim in range(len(target_shape)): - current_dim = -1 - dim - pad_dim.append(0) - pad_dim.append( - max(0, target_shape[current_dim] - w.shape[current_dim])) - res = F.pad(w, pad_dim, value=value) - return res - else: - return w - - -def postprocess_weight_only(tllm_key, weights, quant_mode, layer): - if weights.dim() > 2: - v = weights.transpose(-1, -2) - else: - v = weights.t() - - tp_dim = 1 if isinstance(layer, ColumnLinear) else 0 - if "weight" in tllm_key: - if layer.is_padded: - split_size = layer.out_features if tp_dim == 1 else layer.in_features - v = torch.split(v, split_size, tp_dim)[layer.tp_rank] - v = pad_like(v, (layer.in_features, layer.out_features)) - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.contiguous(), quant_mode) - return { - tllm_key: processed_torch_weights, - tllm_key.replace("weight", "per_channel_scale"): - torch_weight_scales, - } - else: - if layer.is_padded and tp_dim == 1: - weights = torch.split(weights, layer.out_features, - tp_dim)[layer.tp_rank] - weights = pad_like(weights, (layer.out_features, )) - return {tllm_key: weights} # Bias - - -def postprocess_weight_only_groupwise(tllm_key, weights, torch_dtype, layer, - **kwargs): - using_head_as_leading_dim = kwargs.get("using_head_as_leading_dim", False) - config = kwargs.get("config", None) - use_autoawq = kwargs.get("use_autoawq", None) - num_heads = config.num_attention_heads - USE_GPTQ = layer.prequant_scaling_factor is None and use_autoawq is None - USE_HF_AWQ = layer.prequant_scaling_factor is None and use_autoawq is not None - USE_MODELOPT_AWQ = layer.prequant_scaling_factor is not None - USE_INT8_WEIGHT = layer.quant_algo & 16 - - tp_dim = 1 if isinstance(layer, ColumnLinear) else 0 - is_qkv = layer.is_qkv if hasattr(layer, "is_qkv") else False - - if using_head_as_leading_dim: - assert config.num_attention_heads == config.num_key_value_heads, "using_head_as_leading_dim require head_size to be multiple of 3." - if tllm_key.endswith("weights_scaling_factor"): - # TODO: Remove reshaping after modelopt optimizes scale shape - if is_qkv: - for idx, w in enumerate(weights): - scales = w.to(torch_dtype) - scales = scales.reshape(-1, - layer.weights_scaling_factor.shape[0]).T - scales = scales.chunk(layer.tp_size, 1)[layer.tp_rank] - weights[idx] = scales - weights = torch.cat(weights, dim=1) - else: - scales = weights.to(torch_dtype) - scales_shape = [ - layer.weights_scaling_factor.shape[1], - layer.weights_scaling_factor.shape[0] - ] - scales_shape[1 - tp_dim] *= layer.tp_size - scales = scales.reshape(scales_shape).T - weights = scales.chunk(layer.tp_size, tp_dim)[layer.tp_rank] - if is_qkv and isinstance(weights, list) and len(weights) >= 3: - if USE_MODELOPT_AWQ: - if tllm_key.endswith("prequant_scaling_factor"): - weights = weights[0] - else: - weights = torch.cat(weights, dim=0) - elif len(weights) > 3: - weights = [ - torch.cat(weights[i::len(weights) // 3], dim=1) - for i in range(len(weights) // 3) - ] - - if tllm_key.endswith("bias"): - if is_qkv and isinstance(weights, list): - weights = torch.cat(weights) - if layer.is_padded: - weights = pad_like(weights, layer.bias.shape) - if using_head_as_leading_dim: - weights = change_qkv_leading_dim(weights, num_heads) - results = {tllm_key: weights.to(torch_dtype)} - elif tllm_key.endswith("weight"): - if not USE_INT8_WEIGHT: - # 4 bit quantization - if USE_GPTQ: - qweight = unpack_int32_into_int8(weights[0].T).T - 8 - elif USE_HF_AWQ: - qweight = unpack_int32_into_int8(weights[0], True) - 8 - else: - qweight = unpack_int32_into_int8(weights.T) - qweight -= (qweight >> 4) << 4 - qweight = qweight.view(torch.uint8) - elif USE_INT8_WEIGHT and USE_GPTQ: - # 8 bit quantization (only consider INT8 GPTQ here) - qweight = ( - weights[0].T.contiguous().view(torch.uint8).T.contiguous() - - 128).to(torch.int8) - else: - raise NotImplementedError( - "Unsupported quantization mode for weight.") - - if using_head_as_leading_dim: - qweight = change_qkv_leading_dim(qweight, num_heads) - if layer.is_padded: - qweight = torch.split(qweight, layer.out_features, - tp_dim)[layer.tp_rank] - qweight = pad_like(qweight, (layer.in_features, layer.out_features)) - # pack int8 tensor to packed int4 - if not USE_INT8_WEIGHT: - qweight = (qweight[:, 1::2] * 16 + qweight[:, ::2]).view(torch.int8) - weight_type = torch.int8 if USE_INT8_WEIGHT else torch.quint4x2 - qweight = preprocess_weights_for_mixed_gemm( - qweight, weight_type, torch.float16).view(torch_dtype) - results = {tllm_key: qweight} - - # scales and zeros for GPTQ and HF-AWQ - if USE_GPTQ or USE_HF_AWQ: - scales = weights[1].to(torch_dtype) - if USE_INT8_WEIGHT: - qzeros = weights[2].view(torch.uint8) - else: - qzeros = unpack_int32_into_int8(weights[2], USE_HF_AWQ) - if using_head_as_leading_dim: - scales = change_qkv_leading_dim(scales, num_heads) - qzeros = change_qkv_leading_dim(qzeros, num_heads) - if layer.is_padded: - scales = torch.split(scales, - layer.weights_scaling_factor.shape[tp_dim], - tp_dim)[layer.tp_rank] - scales = pad_like(scales, layer.weights_scaling_factor.shape, 1) - qzeros = torch.split(qzeros, - layer.weights_scaling_factor.shape[tp_dim], - tp_dim)[layer.tp_rank] - qzeros = pad_like(qzeros, layer.zero.shape, 7) - if USE_INT8_WEIGHT: - zeros_x_scales = (-qzeros + 128 - 1 * USE_GPTQ) * scales - else: - zeros_x_scales = (-qzeros + 8 - 1 * USE_GPTQ) * scales - zeros_x_scales = zeros_x_scales.to(torch_dtype) - results.update({ - tllm_key.replace("weight", "weights_scaling_factor"): - scales, - tllm_key.replace("weight", "zero"): - zeros_x_scales, - }) - elif tllm_key.endswith("weights_scaling_factor"): - # TODO: Remove reshaping after modelopt optimizes scale shape - if layer.is_padded: - raise NotImplementedError( - "Auto-padding is not Implemented for ModelOpt HF-AWQ.") - results = {tllm_key: weights} - elif tllm_key.endswith("prequant_scaling_factor"): - prequant_scale = weights.to(torch_dtype).reshape(1, -1) - if layer.is_padded and tp_dim == 1: - prequant_scale = torch.split(prequant_scale, - layer.prequant_scaling_factor.shape[1], - 1)[layer.tp_rank] - prequant_scale = pad_like(prequant_scale, - layer.prequant_scaling_factor.shape, 0) - results = {tllm_key: prequant_scale} - - return results - - -def postprocess_fp8_rowwise(tllm_key, weights, **kwargs): - if tllm_key.endswith("per_channel_scale"): - return {} - - config = kwargs.get("config", None) - weights, scales = weights[0::2], weights[1::2] - - if scales[0] is not None: - assert all(w.dtype == torch.float8_e4m3fn for w in weights) - weights = torch.cat(weights, dim=0) - scales = torch.cat([s.to(torch.float32).flatten() for s in scales]) - return { - tllm_key: weights, - tllm_key.replace("weight", "per_channel_scale"): scales - } - else: - x = torch.cat(weights, dim=0).to(torch.float32) - clamp_val = config.quantization.clamp_val - if clamp_val is not None: - # activation range bound. - x = x.clamp(clamp_val[0], clamp_val[1]) - xmax = x.abs().max(-1, keepdim=True).values - # minimum scaling factor. - torch_weight_scales = (xmax / 448.0).clamp(min=1.0 / (448.0 * 512.0)) - out = x / torch_weight_scales - torch_weight_scales = torch_weight_scales.reshape(-1) - out = torch.clamp(out, -448, 448) - processed_torch_weights = out.to(torch.float8_e4m3fn) - processed_torch_weights = processed_torch_weights.to( - torch.float8_e4m3fn) - return { - tllm_key: processed_torch_weights, - tllm_key.replace("weight", "per_channel_scale"): torch_weight_scales - } - - -def fp4_gemm(input: Tensor, - input_sf: Tensor, - weight: Tensor, - weight_sf: Tensor, - global_sf: Tensor, - output_dtype: str | trt.DataType, - scaling_vector_size: int = 16): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be fp4 - input_sf : Tensor (On GPU) - The input scaling factor tensor. Its shape is [batch_size, seq_len, input_dim / scaling_vector_size] or [num_tokens, input_dim / scaling_vector_size] for remove_input_padding, should be int32 (4 packed) - weight : Tensor (On GPU) - The weight tensor. Its shape is [output_dim, input_dim], should be fp4 - weight_sf : Tensor (On GPU) - The weight scaling factor tensor. Its shape is [output_dim, input_dim / scaling_vector_size], should be fp8 - global_sf : Tensor (On GPU) - The global scaling factor tensor. Its shape is [1,], should be float32, used as alpha of Gemm. - output_dtype: str - output data type - scaling_vector_size: int - scaling vector block size - ''' - if isinstance(output_dtype, str): - output_dtype = str_dtype_to_trt(output_dtype) - - fp4_gemm_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Fp4Gemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert fp4_gemm_plg_creator is not None - sv_vec_size = trt.PluginField("sv_vec_size", - np.array(scaling_vector_size, dtype=np.int32), - trt.PluginFieldType.INT32) - output_dtype = trt.PluginField("output_type_id", - np.array([int(output_dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([sv_vec_size, output_dtype]) - fp4_gemm_plug = fp4_gemm_plg_creator.create_plugin("fp4_gemm", pfc) - plug_inputs = [input, input_sf, weight, weight_sf, global_sf] - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, fp4_gemm_plug) - _add_plugin_info(layer, fp4_gemm_plg_creator, "fp4_gemm", pfc) - output = _create_tensor(layer.get_output(0), layer) - return output - - -def quantize_to_fp4_tensor(input: Tensor, sf_scale: Tensor): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be fp16 - sf_scale : Tensor (On GPU) - The global per-tensor scaling factor. Its shape is [1,], should be float32. - used to scale SF from input range to fp8 range (448.f / (MaxVal of input / 6.f)). - output : Tensor (On GPU) - The output tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be FP4 - output_sf : Tensor (On GPU) - The input scaling factor tensor. Its shape is [batch_size, seq_len, input_dim / scaling_vector_size] or [num_tokens, input_dim / scaling_vector_size] for remove_input_padding, should be FP8 - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizeToFP4', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pfc = trt.PluginFieldCollection([]) - quantize_plug = plg_creator.create_plugin("quantize_to_fp4_plugin", pfc) - - plug_inputs = [input.trt_tensor, sf_scale.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - _add_plugin_info(layer, plg_creator, "quantize_to_fp4_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - scales = _create_tensor(layer.get_output(1), layer) - return quantized, scales - - -def dynamic_quantize( - x: Tensor, - double_scale: Tensor, - axis: int = -1, - block_size: int = 16, - data_qtype: trt.DataType = trt.fp4, - scale_qtype: trt.DataType = trt.fp8) -> Tuple[Tensor, Tensor]: - ''' - Parameters: - x : Tensor (On GPU) - The input tensor. - double_scale : Tensor (On GPU) - The global per-tensor scaling factor. It should contain only 1 element. - axis : int - The axis to quantize. Default is -1 (the last axis). - block_size : int - The block size for quantization. Default is 16. - data_qtype : trt.DataType - The data type for quantized data. Default is FP4. - scale_qtype : trt.DataType - The data type for block scale. Default is FP8. - Returns: - A tuple of two tensors: quantized tensor and block scale tensor. - ''' - if axis < 0: - axis = len(x.shape) + axis - dynq = default_trtnet().add_dynamic_quantize(x.trt_tensor, axis, block_size, - data_qtype, scale_qtype) - dynq.set_input(1, double_scale.trt_tensor) - quantized = _create_tensor(dynq.get_output(0), dynq) - scale = _create_tensor(dynq.get_output(1), dynq) - return quantized, scale - - -def block_double_dequantize(x: Tensor, - scale: Tensor, - double_scale: Tensor, - dtype: trt.DataType | str = 'float16') -> Tensor: - ''' - Parameters: - x : Tensor (On GPU) - The input tensor. - scale : Tensor (On GPU) - The block scale tensor. - double_scale : Tensor (On GPU) - The global per-tensor scaling factor. It should contain only 1 element. - dtype : trt.DataType | str - The data type for dequantized data. Default is float32. - Returns: - The dequantized tensor. - ''' - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - dequantize_scale_layer = default_trtnet().add_dequantize( - scale.trt_tensor, double_scale.trt_tensor, dtype) - scale = _create_tensor(dequantize_scale_layer.get_output(0), - dequantize_scale_layer) - - dequantize_data_layer = default_trtnet().add_dequantize( - x.trt_tensor, scale.trt_tensor, dtype) - dequantize_data = _create_tensor(dequantize_data_layer.get_output(0), - dequantize_data_layer) - return dequantize_data diff --git a/tensorrt_llm/quantization/layers.py b/tensorrt_llm/quantization/layers.py deleted file mode 100644 index 28f4cbf2e8cf..000000000000 --- a/tensorrt_llm/quantization/layers.py +++ /dev/null @@ -1,3383 +0,0 @@ -# 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 math -from typing import Optional - -import numpy as np -import tensorrt as trt -import torch - -from .._common import default_net, precision -from .._utils import is_same_dtype, str_dtype_to_torch, trt_dtype_to_torch -from ..functional import (ACT2FN, AllReduceFusionOp, AllReduceParams, - AttentionMaskType, PositionEmbeddingType, - RotaryScalingType, Tensor, allgather, allreduce, cast, - concat, constant, div, embedding, gemm_allreduce, - generate_alibi_slopes, gpt_attention, matmul, mul, - shape, slice, softmax, split, where) -from ..layers import MropeParams, SpecDecodingParams -from ..layers.embedding import Embedding -from ..layers.linear import Linear, RowLinear -from ..module import Module -from ..parameter import Parameter -from .utils import fp4_utils - -# isort: off -from .functional import ( - block_double_dequantize, dequantize, dynamic_quantize, fp4_gemm, - quantize_to_fp4_tensor, fp8_rowwise_gemm, fp8_rowwise_rms_norm, - postprocess_fp8_rowwise, postprocess_weight_only, - postprocess_weight_only_groupwise, quantize, quantize_fp8_per_token, - quantize_per_token, quantize_tensor, validate_group_size, smooth_quant_gemm, - smooth_quant_layer_norm, smooth_quant_rms_norm, - weight_only_groupwise_quant_matmul, weight_only_quant_matmul, - qserve_gemm_per_group, qserve_gemm_per_channel, fp8_rowwise_layer_norm) -# isort: on -from .mode import GroupwiseQuantAlgo, QuantMode - - -class Quantize(Module): - """ - Quantize Layer - For per-tensor mode, the scaling factor is a scalar. - For per-channel mode, the scaling factor is a vector. - """ - - def __init__( - self, - output_dtype: str = 'int8', - scaling_factor_dtype: str = 'float32', - in_channels: int = -1, - axis=-1, - ) -> None: - super().__init__() - self.scaling_factor = Parameter(shape=(in_channels, ) if axis != -1 else - (), - dtype=scaling_factor_dtype) - self.output_dtype = output_dtype - self.axis = axis - - def forward(self, x): - return quantize(x, self.scaling_factor.value, self.output_dtype, - self.axis) - - -class QuantizePerToken(Module): - """ - Quantize Per Token and compute dynamic scales for SmoothQuant - """ - - def forward(self, x): - return quantize_per_token(x) - - -class Dequantize(Module): - """ - Dequantize Layer. - """ - - def __init__(self, axis: int = -1) -> None: - super().__init__() - self.scaling_factor = Parameter(shape=(), dtype='float32') - self.axis = axis - - def forward(self, input): - return dequantize(input, self.scaling_factor.value, self.axis) - - -class SmoothQuantLinear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - quant_mode=QuantMode(0), - prefer_managed_weight=True): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight) - - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant Linear has to have act+weight quantization mode set" - ) - - weights_dtype = dtype - if quant_mode.has_act_and_weight_quant(): - weights_dtype = "int8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype, - prefer_managed=self.prefer_managed_weight) - - if quant_mode.has_act_and_weight_quant(): - scale_shape = (1, self.out_features - ) if quant_mode.has_per_channel_scaling() else (1, 1) - self.per_channel_scale = Parameter(shape=scale_shape, - dtype="float32") - - if quant_mode.has_act_static_scaling(): - self.act_scale = Parameter(shape=(1, 1), dtype="float32") - - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantLinear now" - if self.quant_mode.has_act_static_scaling(): - per_token_scale = self.act_scale.value - else: - # If we are in SmoothQuant with dynamic activation scaling, - # input x has to be a tuple of int8 tensor and fp32 scaling factors - x, per_token_scale = x - x = smooth_quant_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_per_token_dynamic_scaling(), - self.quant_mode.has_per_channel_scaling(), - self.dtype) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - -SmoothQuantColumnLinear = SmoothQuantLinear - - -class SmoothQuantRowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - prefer_managed_weight=True, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight) - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant Linear has to have act+weight quantization mode set" - ) - weights_dtype = dtype - if quant_mode.has_act_and_weight_quant(): - weights_dtype = "int8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype, - prefer_managed=self.prefer_managed_weight) - self.smoother = Parameter(shape=(1, self.in_features), dtype="float32") - if quant_mode.has_act_and_weight_quant(): - scale_shape = (1, self.out_features - ) if quant_mode.has_per_channel_scaling() else (1, 1) - self.per_channel_scale = Parameter(shape=scale_shape, - dtype="float32") - - if quant_mode.has_act_static_scaling(): - self.act_scale = Parameter(shape=(1, 1), dtype="float32") - - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantRowLinear now" - if self.quant_mode.has_act_static_scaling(): - per_token_scale = self.act_scale.value - else: - x, per_token_scale = x - x = smooth_quant_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_per_token_dynamic_scaling(), - self.quant_mode.has_per_channel_scaling(), - self.dtype) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - -class SmoothQuantLayerNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-05, - elementwise_affine=True, - bias=True, - dtype=None, - quant_mode=QuantMode(0), - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant layer norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - if self.quant_mode.has_act_and_weight_quant(): - self.scale_to_int = Parameter(shape=(1, ), dtype=dtype) - else: - self.register_parameter('scale_to_int', None) - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None if self.scale_to_int is None else self.scale_to_int.value - return smooth_quant_layer_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - self.eps, - dynamic_act_scaling=self.quant_mode.has_per_token_dynamic_scaling()) - - -class SmoothQuantRmsNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-06, - elementwise_affine=True, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None, - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant Rms norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - if self.quant_mode.has_act_and_weight_quant(): - self.scale_to_int = Parameter(shape=(1, ), dtype=dtype) - else: - self.register_parameter('scale_to_int', None) - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None if self.scale_to_int is None else self.scale_to_int.value - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return smooth_quant_rms_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=self.quant_mode.has_per_token_dynamic_scaling()) - - -class QServeW4A8Linear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - quant_mode=QuantMode(0)): - assert dtype == "float16" # Currently the kernel only supports float16 output - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output) - - self.quant_mode = quant_mode - assert self.quant_mode.is_qserve_w4a8() - - # Only support g128 now. - if self.quant_mode.has_per_group_scaling(): - self.group_size = 128 - else: - self.group_size = -1 - - self.weight = Parameter(shape=(self.out_features, - self.in_features // 2), - dtype="int8") - - self.s1_scales = Parameter(shape=(self.out_features, ), dtype="float16") - - if self.group_size == -1: - self.s1_szeros = Parameter(shape=(self.out_features, ), - dtype="float16") - else: - self.s2_scales = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - self.s2_szeros = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - - def forward(self, x): - if self.group_size == -1: - x, per_token_scale, per_token_sum = x - x = qserve_gemm_per_channel(x, per_token_scale, per_token_sum, - self.weight.value, self.s1_scales.value, - self.s1_szeros.value) - - else: - x, per_token_scale = x - x = qserve_gemm_per_group(x, per_token_scale, self.weight.value, - self.s1_scales.value, - self.s2_scales.value, - self.s2_szeros.value, self.group_size) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - return x - - -QServeW4A8ColumnLinear = QServeW4A8Linear - - -class QServeW4A8RowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - assert dtype == "float16" # Currently the kernel only supports float16 output - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - self.quant_mode = quant_mode - assert self.quant_mode.is_qserve_w4a8() - # Only supports 128g now. - if self.quant_mode.has_per_group_scaling(): - self.group_size = 128 - else: - self.group_size = -1 - - self.weight = Parameter(shape=(self.out_features, - self.in_features // 2), - dtype="int8") - - self.s1_scales = Parameter(shape=(self.out_features, ), dtype="float16") - - if self.group_size == -1: - self.s1_szeros = Parameter(shape=(self.out_features, ), - dtype="float16") - else: - self.s2_scales = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - self.s2_szeros = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - - def forward(self, x, all_reduce_params=None): - if self.group_size == -1: - x, per_token_scale, per_token_sum = x - x = qserve_gemm_per_channel(x, per_token_scale, per_token_sum, - self.weight.value, self.s1_scales.value, - self.s1_szeros.value) - else: - x, per_token_scale = x - x = qserve_gemm_per_group(x, per_token_scale, self.weight.value, - self.s1_scales.value, - self.s2_scales.value, - self.s2_szeros.value, self.group_size) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - -class Fp8RowwiseRmsNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-06, - elementwise_affine=True, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None, - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Rms norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return fp8_rowwise_rms_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=self.quant_mode.has_fp8_rowwise()) - - -class Fp8RowwiseLinear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - quant_mode=QuantMode(0)): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output) - - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Linear has to have act+weight quantization mode set" - ) - - weights_dtype = dtype - if quant_mode.has_fp8_rowwise(): - weights_dtype = "fp8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype) - - if quant_mode.has_fp8_rowwise(): - self.per_channel_scale = Parameter(shape=(self.out_features, ), - dtype="float32") - - self.quant_mode = quant_mode - self.tllm_to_externel_key_dict = {"weight": ["weight", "weight_scale"]} - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantLinear now" - x, per_token_scale = x - x = fp8_rowwise_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_per_token_dynamic_scaling(), - self.quant_mode.has_per_channel_scaling()) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - return postprocess_fp8_rowwise(tllm_key, weights, **kwargs) - - -Fp8RowwiseColumnLinear = Fp8RowwiseLinear - - -class Fp8RowwiseRowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Linear has to have act+weight quantization mode set" - ) - weights_dtype = dtype - if quant_mode.has_fp8_rowwise(): - weights_dtype = "fp8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype) - if quant_mode.has_fp8_rowwise(): - self.per_channel_scale = Parameter(shape=(self.out_features, ), - dtype="float32") - - self.quant_mode = quant_mode - self.tllm_to_externel_key_dict = {"weight": ["weight", "weight_scale"]} - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantRowLinear now" - x, per_token_scale = x - x = fp8_rowwise_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_fp8_rowwise(), - self.quant_mode.has_fp8_rowwise()) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - return postprocess_fp8_rowwise(tllm_key, weights, **kwargs) - - -class WeightOnlyQuantLinear(Linear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - gather_output=True, - quant_mode=QuantMode.use_weight_only(), - transa=False, - transb=False, - is_qkv=False, - prefer_managed_weight=True, - ): - multiple = 64 * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - is_qkv=is_qkv, - prefer_managed_weight=prefer_managed_weight) - if quant_mode.is_int8_weight_only(): - self.weight_only_quant_mode = 1 - quant_type_size_in_bits = 8 - elif quant_mode.is_int4_weight_only(): - self.weight_only_quant_mode = 2 - quant_type_size_in_bits = 4 - # we use a fake tensor with data_type = int8 - self.weight = Parameter(shape=(self.in_features, - int(self.out_features * - quant_type_size_in_bits / 8)), - dtype="int8", - prefer_managed=self.prefer_managed_weight) - - scale_shape = (self.out_features, ) - self.per_channel_scale = Parameter(shape=scale_shape, dtype=dtype) - - self.transa = transa - self.transb = transb - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None): - # ootb has not supported int4 yet. - if self.weight_only_quant_mode == 2 and not default_net( - ).plugin_config.weight_only_quant_matmul_plugin: - raise TypeError( - "Int4 Weight-only Quant MatMul is only supported with plugin") - hidden_state = x - x = weight_only_quant_matmul(x, self.weight.value, - self.per_channel_scale.value, - self.weight_only_quant_mode, self.dtype, - self.transa, self.transb) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - weights = super().postprocess(tllm_key, weights, **kwargs)[tllm_key] - weights = weights.to(str_dtype_to_torch(self.dtype)) - return postprocess_weight_only( - tllm_key, weights, - torch.int8 if self.weight_only_quant_mode == 1 else torch.quint4x2, - self) - - -WeightOnlyQuantColumnLinear = WeightOnlyQuantLinear - - -class WeightOnlyQuantRowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - quant_mode=QuantMode.use_weight_only(), - prefer_managed_weight=True, - is_expert=False, - ): - multiple = 64 * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight, - is_expert=is_expert) - if quant_mode.is_int8_weight_only(): - self.weight_only_quant_mode = 1 - elif quant_mode.is_int4_weight_only(): - self.weight_only_quant_mode = 2 - #we use a fake tensor with data_type = int8 - self.weight = Parameter(shape=(self.in_features, - int(self.out_features / - self.weight_only_quant_mode)), - dtype="int8", - prefer_managed=prefer_managed_weight) - self.per_channel_scale = Parameter(shape=(self.out_features, ), - dtype=dtype) - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - hidden_state = x - x = weight_only_quant_matmul(x, self.weight.value, - self.per_channel_scale.value, - self.weight_only_quant_mode, self.dtype) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - if not self.is_expert: - x = allreduce(x, - self.tp_group, - all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias - else: - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias / self.tp_size - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - weights = weights.to(str_dtype_to_torch(self.dtype)) - return postprocess_weight_only( - tllm_key, weights, - torch.int8 if self.weight_only_quant_mode == 1 else torch.quint4x2, - self) - - -class WeightOnlyQuantEmbedding(Embedding): - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - dtype: Optional[str] = None, - tp_size: int = 1, - tp_group: Optional[list] = None, - sharding_dim: int = 0, - tp_rank: Optional[int] = None, - quant_mode=QuantMode.use_weight_only(), - ): - super().__init__( - num_embeddings, - embedding_dim, - dtype, # dtype, - tp_size, - tp_group, - sharding_dim, - tp_rank) - # only support int8 wo now - # TODO support int4 wo - self.quant_mode = quant_mode - self.per_token_scale = Parameter(shape=(self.num_embeddings, ), - dtype=dtype) - - if sharding_dim == 1: - self.weight = Parameter(shape=(self.num_embeddings, - self.embedding_dim // self.tp_size), - dtype="int8") - elif sharding_dim == 0: - self.weight = Parameter(shape=(math.ceil( - self.num_embeddings / self.tp_size), self.embedding_dim), - dtype="int8") - - def forward(self, x): - result = embedding(x, - self.weight.value, - tp_size=self.tp_size, - tp_group=self.tp_group, - sharding_dim=self.sharding_dim, - tp_rank=self.tp_rank, - per_token_scale=self.per_token_scale.value) - - return result - - -def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - w_unpacked = w_unpacked.view(-1, 8)[:, [0, 4, 1, 5, 2, 6, 3, 7]].view( - w_unpacked.shape) - return w_unpacked.contiguous() - - -class WeightOnlyGroupwiseQuantLinear(Linear): - - def __init__( - self, - in_features, - out_features, - group_size=128, - pre_quant_scale=False, - zero=False, - bias=False, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - gather_output=True, - use_w4a8_awq=False, - use_int8_weight=False, - is_qkv=False, - prefer_managed_weight=True, - ): - multiple = max((128 if use_w4a8_awq else 64), group_size) * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - is_qkv=is_qkv, - prefer_managed_weight=prefer_managed_weight) - - self.quant_algo = ( - use_int8_weight * GroupwiseQuantAlgo.INT8_WEIGHT + - use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA + - pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + - zero * GroupwiseQuantAlgo.ZERO + bias * GroupwiseQuantAlgo.BIAS) - self.group_size = group_size - # packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - pack_ratio = 2 if use_int8_weight else 4 - self.weight = Parameter(shape=(self.in_features, - self.out_features // pack_ratio), - dtype=dtype, - prefer_managed=self.prefer_managed_weight) - - scale_shape = (self.in_features // group_size, self.out_features) - self.weights_scaling_factor = Parameter(shape=scale_shape, dtype=dtype) - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.pre_quant_scale = pre_quant_scale - self.use_w4a8_awq = use_w4a8_awq - - if pre_quant_scale: - self.prequant_scaling_factor = Parameter(shape=(1, - self.in_features), - dtype=dtype) - else: - self.register_parameter('prequant_scaling_factor', None) - - if zero: - self.zero = Parameter(shape=scale_shape, dtype=dtype) - else: - self.register_parameter('zero', None) - - if use_w4a8_awq: - self.alpha = Parameter(shape=(1, ), dtype="float32") - else: - self.register_parameter('alpha', None) - - validate_group_size(self) - if pre_quant_scale: - self.tllm_to_externel_key_dict = { - "weights_scaling_factor": "weight_scale", - "prequant_scaling_factor": "input_quantizer._pre_quant_scale", - } # AWQ - else: - self.tllm_to_externel_key_dict = { - "weight": ["qweight", "scales", "qzeros"] - } # GPTQ - - def forward(self, x, lora_runtime_params=None): - pre_quant_scale = self.prequant_scaling_factor.value if self.prequant_scaling_factor else None - zero = self.zero.value if self.zero else None - bias = self.bias.value if self.bias else None - alpha = self.alpha if self.alpha else None - - hidden_state = x - x = weight_only_groupwise_quant_matmul( - x, pre_quant_scale, self.weight.value, - self.weights_scaling_factor.value, zero, bias, alpha, - self.quant_algo, self.group_size, self.dtype) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("zero") or ( - self.prequant_scaling_factor is None - and tllm_key.endswith("weights_scaling_factor")): - return {} - torch_dtype = str_dtype_to_torch(self.dtype) - return postprocess_weight_only_groupwise(tllm_key, weights, torch_dtype, - self, **kwargs) - - -WeightOnlyGroupwiseQuantColumnLinear = WeightOnlyGroupwiseQuantLinear - - -class WeightOnlyGroupwiseQuantRowLinear(RowLinear): - - def __init__(self, - in_features, - out_features, - group_size=128, - pre_quant_scale=False, - zero=False, - bias=False, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - use_w4a8_awq=False, - use_int8_weight=False, - prefer_managed_weight=True): - multiple = max((128 if use_w4a8_awq else 64), group_size) * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight) - - self.quant_algo = ( - use_int8_weight * GroupwiseQuantAlgo.INT8_WEIGHT + - use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA + - pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + - zero * GroupwiseQuantAlgo.ZERO + bias * GroupwiseQuantAlgo.BIAS) - self.group_size = group_size - # packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - pack_ratio = 2 if use_int8_weight else 4 - self.weight = Parameter(shape=(self.in_features, - self.out_features // pack_ratio), - dtype=dtype, - prefer_managed=self.prefer_managed_weight) - scale_shape = (self.in_features // group_size, self.out_features) - self.weights_scaling_factor = Parameter(shape=scale_shape, dtype=dtype) - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.pre_quant_scale = pre_quant_scale - self.use_w4a8_awq = use_w4a8_awq - - if pre_quant_scale: - self.prequant_scaling_factor = Parameter(shape=(1, - self.in_features), - dtype=dtype) - else: - self.register_parameter('prequant_scaling_factor', None) - - if zero: - self.zero = Parameter(shape=scale_shape, dtype=dtype) - else: - self.register_parameter('zero', None) - - if use_w4a8_awq: - self.alpha = Parameter(shape=(1, ), dtype="float32") - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - else: - self.register_parameter('alpha', None) - self.register_parameter('activation_scaling_factor', None) - - validate_group_size(self) - if pre_quant_scale: - self.tllm_to_externel_key_dict = { - "weights_scaling_factor": "weight_scale", - "prequant_scaling_factor": "input_quantizer._pre_quant_scale", - } # AWQ - else: - self.tllm_to_externel_key_dict = { - "weight": ["qweight", "scales", "qzeros"] - } # GPTQ - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - pre_quant_scale = self.prequant_scaling_factor.value if self.prequant_scaling_factor else None - zero = self.zero.value if self.zero else None - bias = self.bias.value if self.bias else None - alpha = self.alpha if self.alpha else None - activation_scaling_factor = self.activation_scaling_factor.value if self.activation_scaling_factor else None - - if self.alpha and x.dtype == trt.fp8: - x = dequantize(x, activation_scaling_factor, -1, self.dtype) - - hidden_state = x - x = weight_only_groupwise_quant_matmul( - x, pre_quant_scale, self.weight.value, - self.weights_scaling_factor.value, zero, bias, alpha, - self.quant_algo, self.group_size, self.dtype) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.tp_size > 1 and self.tp_group is not None: - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("zero") or ( - self.prequant_scaling_factor is None - and tllm_key.endswith("weights_scaling_factor")): - return {} - torch_dtype = str_dtype_to_torch(self.dtype) - return postprocess_weight_only_groupwise(tllm_key, weights, torch_dtype, - self, **kwargs) - - -class SmoothQuantMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act == 'swiglu' else ffn_hidden_size - self.fc = SmoothQuantColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = SmoothQuantRowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.hidden_act = hidden_act - self.quant_mode = quant_mode - self.dtype = dtype - - if self.quant_mode.has_act_static_scaling(): - self.quantization_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - else: - self.register_parameter('quantization_scaling_factor', None) - - def forward(self, hidden_states, lora_layer_params=None): - - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - value = cast(self.proj.smoother.value, inter.dtype) - inter = inter / value - if self.quant_mode.has_act_and_weight_quant(): - if self.quant_mode.has_act_static_scaling(): - # Avoid quantization layers as it breaks int8 plugins - inter = quantize_tensor(inter, - self.quantization_scaling_factor.value) - else: - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - inter = quantize_per_token(inter) - output = self.proj(inter) - return output - - -class Int8SmoothQuantRowLinear(RowLinear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - prefer_managed_weight=True): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(self.out_features, ), - dtype=trt.float32) - self.prequant_scaling_factor = Parameter(shape=(self.in_features, ), - dtype=dtype) - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.int8, - prefer_managed=self.prefer_managed_weight) - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - lora_hidden_state = x if lora_runtime_params is not None else None - if default_net().strongly_typed: - assert is_same_dtype( - x.dtype, - self.dtype), f"Got input type {x.dtype}, expecting {self.dtype}" - x = mul(x, self.prequant_scaling_factor.value) - - x = cast(x, self.activation_scaling_factor.value.dtype) - - quantized_out = quantize(x, self.activation_scaling_factor.value, - 'int8') - - dequantized_out = dequantize(quantized_out, - self.activation_scaling_factor.value, -1, - self.activation_scaling_factor.value.dtype) - - dequantized_out = cast(dequantized_out, self.dtype) - - w_deq_out = dequantize(self.weight.value, - self.weights_scaling_factor.value, 0, - self.weights_scaling_factor.value.dtype) - - w_deq_out = cast(w_deq_out, self.dtype) - return self.multiply_collect(dequantized_out, - w_deq_out, - gemm_plugin=None, - all_reduce_params=all_reduce_params, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - - -class Int8SmoothQuantLinear(Linear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - prefer_managed_weight=True, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - - self.weights_scaling_factor = Parameter(shape=(self.out_features, ), - dtype=trt.float32) - self.prequant_scaling_factor = Parameter(shape=(self.in_features, ), - dtype=dtype) - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.int8, - prefer_managed=self.prefer_managed_weight) - - def forward(self, x, lora_runtime_params=None): - lora_hidden_state = x if lora_runtime_params is not None else None - if default_net().strongly_typed: - assert is_same_dtype( - x.dtype, - self.dtype), f"Got input type {x.dtype}, expecting {self.dtype}" - x = mul(x, self.prequant_scaling_factor.value) - x = cast(x, self.activation_scaling_factor.value.dtype) - - quantized_out = quantize(x, self.activation_scaling_factor.value, - 'int8') - - dequantized_out = dequantize(quantized_out, - self.activation_scaling_factor.value, -1, - self.activation_scaling_factor.value.dtype) - - dequantized_out = cast(dequantized_out, self.dtype) - - w_deq_out = dequantize(self.weight.value, - self.weights_scaling_factor.value, 0, - self.weights_scaling_factor.value.dtype) - w_deq_out = cast(w_deq_out, self.dtype) - - return self.multiply_collect(dequantized_out, - w_deq_out, - gemm_plugin=None, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - - -class FP8Linear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - prefer_managed_weight=True, - is_qkv=False): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight, - is_qkv=is_qkv) - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype='fp8', - prefer_managed=self.prefer_managed_weight) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(1, ), dtype=trt.float32) - if self.is_qkv: - self.tllm_to_externel_key_dict = { - "activation_scaling_factor": "input_scale", - "weight": ["weight", "weight_scale"], - "weights_scaling_factor": "", - } - else: - self.tllm_to_externel_key_dict = { - "activation_scaling_factor": "input_scale", - "weights_scaling_factor": "weight_scale", - } - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None or default_net( - ).plugin_config.lora_plugin == self.dtype - - if default_net().strongly_typed: - assert default_net().plugin_config.user_buffer or is_same_dtype( - x.dtype, - self.dtype), f"Got input type {x.dtype}, expecting {self.dtype}" - - alpha = self.weights_scaling_factor.raw_value * self.activation_scaling_factor.raw_value - activation_scaling_factor = constant( - self.activation_scaling_factor.raw_value) - activation_scaling_factor = cast(activation_scaling_factor, self.dtype) - if x.dtype != trt.fp8: - quantized_out = quantize(x, activation_scaling_factor, 'fp8') - lora_hidden_state = x if lora_runtime_params is not None else None - else: - quantized_out = x - # TODO: add fp8 LoRA support - lora_hidden_state = dequantize( - x, activation_scaling_factor, -1, - self.dtype) if lora_runtime_params is not None else None - - weights_scaling_factor = cast(self.weights_scaling_factor.value, - self.dtype) - - if self.weight.value.dtype != trt.fp8: - w_quant_out = quantize(self.weight.value, weights_scaling_factor, - 'fp8') - else: - w_quant_out = self.weight.value - - gemm_plugin = default_net().plugin_config.gemm_plugin - - low_latency_gemm_plugin = default_net( - ).plugin_config.low_latency_gemm_plugin - if (low_latency_gemm_plugin == "fp8"): - return self.multiply_collect( - quantized_out, - w_quant_out, - gemm_plugin=None, - low_latency_gemm_plugin=low_latency_gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - elif gemm_plugin == 'fp8': - return self.multiply_collect( - quantized_out, - w_quant_out, - gemm_plugin=gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - else: - dequantized_out = dequantize(quantized_out, - activation_scaling_factor, -1, - self.dtype) - w_deq_out = dequantize(w_quant_out, weights_scaling_factor, -1, - self.dtype) - return self.multiply_collect( - dequantized_out, - w_deq_out, - gemm_plugin=None, - use_fp8=True, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - - def postprocess(self, tllm_key, weights, **kwargs): - if self.is_qkv: - if tllm_key.endswith("activation_scaling_factor"): - return max(weights).reshape(1, ).to(torch.float32) - elif tllm_key.endswith("weights_scaling_factor"): - return {} - elif tllm_key.endswith("weight"): - assert len(weights) == 6 - weight_scaling_factors = weights[1::2] - new_amax = max(weight_scaling_factors).reshape(1, ).to( - torch.float32) - for qkv_idx in range(3): - idx = qkv_idx * 2 - weights[idx] = weights[idx].view(torch.float8_e4m3fn).to( - torch.float32) - weights[idx] *= weights[idx + 1] - weights[idx] /= new_amax - weights[idx] = weights[idx].to(torch.float8_e4m3fn) - weights = torch.cat(weights[::2]) - scales = new_amax - return { - tllm_key: weights, - tllm_key.replace("weight", "weights_scaling_factor"): scales - } - else: - return super().postprocess(tllm_key, weights, **kwargs) - if tllm_key.endswith("scaling_factor"): - return weights.reshape(1, ).to(torch.float32) - elif tllm_key.endswith("weight"): - return weights.view(torch.float8_e4m3fn) - elif tllm_key.endswith("bias"): - return weights.to(trt_dtype_to_torch(self.bias.dtype)) - - -class FP8RowLinear(RowLinear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - prefer_managed_weight=True, - is_expert=False): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight, - is_expert=is_expert) - self.weight = Parameter( - shape=(self.out_features, self.in_features), - dtype="fp8", - prefer_managed=self.prefer_managed_weight, - ) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(1, ), dtype=trt.float32) - self.tllm_to_externel_key_dict = { - "activation_scaling_factor": "input_scale", - "weights_scaling_factor": "weight_scale", - } - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None or default_net( - ).plugin_config.lora_plugin == self.dtype - - alpha = self.weights_scaling_factor.raw_value * self.activation_scaling_factor.raw_value - activation_scaling_factor = cast(self.activation_scaling_factor.value, - self.dtype) - if x.dtype != trt.fp8: - quantized_out = quantize(x, activation_scaling_factor, 'fp8') - lora_hidden_state = x if lora_runtime_params is not None else None - else: - quantized_out = x - # TODO: add fp8 LoRA support - lora_hidden_state = dequantize( - x, activation_scaling_factor, -1, - self.dtype) if lora_runtime_params is not None else None - - weights_scaling_factor = cast(self.weights_scaling_factor.value, - self.dtype) - if self.weight.value.dtype != trt.fp8: - w_quant_out = quantize(self.weight.value, weights_scaling_factor, - 'fp8') - else: - w_quant_out = self.weight.value - - gemm_plugin = default_net().plugin_config.gemm_plugin - low_latency_gemm_plugin = default_net( - ).plugin_config.low_latency_gemm_plugin - gemm_allreduce_plugin = default_net( - ).plugin_config.gemm_allreduce_plugin - if gemm_allreduce_plugin: - ret = self.multiply_collect(quantized_out, - w_quant_out, - gemm_plugin=None, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - elif (low_latency_gemm_plugin == "fp8"): - ret = self.multiply_collect( - quantized_out, - w_quant_out, - gemm_plugin=None, - low_latency_gemm_plugin=low_latency_gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - elif gemm_plugin == 'fp8': - ret = self.multiply_collect(quantized_out, - w_quant_out, - gemm_plugin=gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - else: - dequantized_out = dequantize(quantized_out, - activation_scaling_factor, -1, - self.dtype) - w_deq_out = dequantize(w_quant_out, weights_scaling_factor, -1, - self.dtype) - ret = self.multiply_collect(dequantized_out, - w_deq_out, - gemm_plugin=None, - use_fp8=True, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - return ret - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("scaling_factor"): - return weights.reshape(1, ).to(torch.float32) - elif tllm_key.endswith("weight"): - return weights.view(torch.float8_e4m3fn) - elif tllm_key.endswith("bias"): - return weights.to(str_dtype_to_torch(self.bias.dtype)) - - -class Fp8RowwiseMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - clamp_val=None, - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act == 'swiglu' else ffn_hidden_size - self.fc = Fp8RowwiseColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = Fp8RowwiseRowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.bias = bias - self.dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - inter = quantize_fp8_per_token(inter, clamp_val) - output = self.proj(inter) - return output - - -class Fp8RowwiseGatedMLP(Fp8RowwiseMLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - clamp_val=None, - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode, - clamp_val=clamp_val) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - self.gate = Fp8RowwiseColumnLinear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states) - inter_x_gate = inter * gate - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - inter_x_gate = quantize_fp8_per_token(inter_x_gate, clamp_val) - output = self.proj(inter_x_gate) - return output - - -class Fp8RowwiseFusedGatedMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - clamp_val=None, - ): - super().__init__() - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.bias = bias - self.dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - - self.fused_fc = Fp8RowwiseColumnLinear(hidden_size, - ffn_hidden_size * 2, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = Fp8RowwiseRowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fused_fc(hidden_states) - - if self.hidden_act == 'silu': - inter = ACT2FN['swiglu'](inter) - elif self.hidden_act == 'gelu': - inter = ACT2FN['geglu'](inter) - else: - raise NotImplementedError( - f"Activation {self.hidden_act} not yet implemented for {self.__class__.__name__}." - ) - - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - inter = quantize_fp8_per_token(inter, clamp_val) - output = self.proj(inter) - return output - - -class Fp8RowwiseAttention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dense_bias=None, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - tp_group=None, - tp_size=1, - tp_rank=0, - scale_alibi_bias=False, - paged_kv_cache=False, - quant_mode=QuantMode(0), - clamp_val=None): - super().__init__() - self.local_layer_idx = local_layer_idx - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_kv_heads = num_kv_heads - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size // tp_size - self.max_position_embeddings = 0 if max_position_embeddings is None else max_position_embeddings - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = 1 - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = scale_alibi_bias - - self.position_embedding_type = position_embedding_type - self.paged_kv_cache = paged_kv_cache - - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.rotary_embedding_dim = 0 - - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes(self.num_attention_heads * - self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.dtype = dtype - - self.register_parameter('kv_cache_scaling_factor', None) - - self.qkv = Fp8RowwiseColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.dense = Fp8RowwiseRowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.use_lora = False - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - spec_decoding_params=None, - encoder_output=None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - ): - assert lora_layer_params is None, ( - f"LoRA is not supported by {self.__class__.__name__} (e.g., --use_fp8_rowwise). " - "If you need LoRA support, please use a non-quantized (e.g., bf16) attention implementation. " - "See https://github.com/NVIDIA/TensorRT-LLM/issues/2603 for details." - ) - qkv = self.qkv(hidden_states) - - alibi_slopes = None - if self.position_embedding_type == PositionEmbeddingType.alibi: - alibi_slopes = self.alibi_slopes.value - dtype = trt.float32 - if default_net().plugin_config.gpt_attention_plugin or default_net( - ).plugin_config.inflight_batching_gpt_attention_plugin: - dtype = hidden_states.dtype if self.quant_mode.has_act_static_scaling( - ) else hidden_states[0].dtype - if dtype == trt.int8: - dtype = trt.float16 - alibi_slopes = cast(alibi_slopes, dtype) - - if spec_decoding_params is None: - spec_decoding_params = SpecDecodingParams() - - assert default_net().plugin_config.gpt_attention_plugin - - assert attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - if use_cache: - assert kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - assert self.attention_mask_type == AttentionMaskType.causal, \ - 'Plugin only support masked MHA.' - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - if self.position_embedding_type.is_rope(): - rotary_inv_freq = attention_params.rotary_inv_freq - rotary_cos_sin = attention_params.embed_positions_for_gpt_attention - else: - rotary_inv_freq = None - rotary_cos_sin = None - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=kv_cache_params.get_first_past_key_value(), - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base, - rotary_embedding_scale_type=self.rotary_embedding_scale_type, - rotary_embedding_scale=self.rotary_embedding_scale, - rotary_embedding_max_positions=self.max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - host_runtime_perf_knobs=attention_params.host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress) - - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - context = quantize_fp8_per_token(context, clamp_val) - - context = self.dense( - context, - all_reduce_params=all_reduce_params, - ) - - if use_cache: - return (context, past_key_value) - - return context - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("kv_cache_scaling_factor") and weights is None: - return {tllm_key: torch.ones(1, )} - else: - return {tllm_key: weights} - - -class FP4Linear(Linear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - prefer_managed_weight=True, - is_qkv=False, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight, - is_qkv=is_qkv) - self.scaling_vector_size = 16 - assert self.in_features % self.scaling_vector_size == 0, \ - "Input features must be a multiple of 16 for FP4 GEMM" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.fp4) - self.weights_block_scaling_factor = Parameter( - shape=(self.out_features, - self.in_features // self.scaling_vector_size), - dtype=trt.fp8) - nrows = fp4_utils.pad_up(self.out_features, 128) - ncols = fp4_utils.pad_up(self.in_features // self.scaling_vector_size, - 4) - self.weights_block_scaling_factor_interleaved = Parameter(shape=(nrows, - ncols), - dtype=trt.fp8) - self.weights_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.activation_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - # alpha = 1.0 / (weight_global_scale * act_global_scale) - self.alpha = Parameter(shape=(1, ), dtype=trt.float32) - if self.is_qkv: - self.tllm_to_externel_key_dict = { - "weight": - ["weight", "weight_scale", "weight_scale_2", "input_scale"], - "weights_block_scaling_factor": - "", - "weights_block_scaling_factor_interleaved": - "", - "weights_global_scaling_factor": - "", - "activation_global_scaling_factor": - "", - "alpha": - "", - } - else: - self.tllm_to_externel_key_dict = { - "weight": ["weight"], - "weights_block_scaling_factor": "weight_scale", - "weights_block_scaling_factor_interleaved": "weight_scale", - "weights_global_scaling_factor": "weight_scale_2", - "activation_global_scaling_factor": "input_scale", - "alpha": ["weight_scale_2", "input_scale"], - } - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None, "lora is not supported on FP4Linear now" - if isinstance(x, (tuple, list)): - fp4_x, act_per_block_scale = x - else: - if default_net().plugin_config.gemm_plugin == 'nvfp4': - fp4_x, act_per_block_scale = quantize_to_fp4_tensor( - x, div(1, self.activation_global_scaling_factor.value)) - else: - fp4_x, act_per_block_scale = dynamic_quantize( - x, self.activation_global_scaling_factor.value) - if default_net().plugin_config.gemm_plugin == 'nvfp4': - x = fp4_gemm(fp4_x, act_per_block_scale, self.weight.value, - self.weights_block_scaling_factor_interleaved.value, - self.alpha.value, self.dtype) - else: - quant_w = self.weight.value - scale_w = self.weights_block_scaling_factor.value - dequant_w = block_double_dequantize( - quant_w, - scale_w, - self.weights_global_scaling_factor.value, - dtype=trt.float16) - dequant_x = block_double_dequantize( - fp4_x, - act_per_block_scale, - self.activation_global_scaling_factor.value, - dtype=trt.float16) - x = matmul(dequant_x, dequant_w, transb=True).cast(self.dtype) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if not any([ - tllm_key.endswith(suffix) - for suffix in self.tllm_to_externel_key_dict - ]): - return super().postprocess(tllm_key, weights, **kwargs) - if self.is_qkv: - if tllm_key.endswith("weight"): - assert len(weights) == 12 - qkv_weight = weights[0::4] - qkv_block_scale = weights[1::4] - qkv_global_scale = weights[2::4] - qkv_input_scale = weights[3::4] - - # Ckpt uses qkv max is guaranteed. So no need to re-quantize. - qkv_input_scale = max(qkv_input_scale).reshape(1, ).to( - torch.float32) - qkv_global_scale = max(qkv_global_scale).reshape(1, ).to( - torch.float32) - qkv_weight = torch.cat(qkv_weight) - qkv_block_scale = torch.cat(qkv_block_scale) - return { - tllm_key: - qkv_weight, - tllm_key.replace("weight", "weights_block_scaling_factor"): - qkv_block_scale, - tllm_key.replace( - 'weight', "weights_block_scaling_factor_interleaved"): - torch.ops.trtllm.block_scale_interleave( - qkv_block_scale.view( - torch.uint8).cpu().contiguous()).reshape( - qkv_block_scale.shape).view( - torch.float8_e4m3fn), - tllm_key.replace("weight", "weights_global_scaling_factor"): - qkv_global_scale, - tllm_key.replace("weight", "activation_global_scaling_factor"): - qkv_input_scale, - tllm_key.replace("weight", "alpha"): - qkv_input_scale * qkv_global_scale, - } - else: - return {} - else: - if tllm_key.endswith("weight"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor_interleaved"): - return torch.ops.trtllm.block_scale_interleave( - weights.view(torch.uint8).cpu().contiguous()).reshape( - weights.shape).view(torch.float8_e4m3fn) - elif tllm_key.endswith("weights_global_scaling_factor"): - return weights.float() - elif tllm_key.endswith('activation_global_scaling_factor'): - return weights.float() - elif tllm_key.endswith('alpha'): - weight_global_sf = weights[0].float() - act_global_sf = weights[1].float() - return act_global_sf * weight_global_sf - - -class FP4RowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - self.scaling_vector_size = 16 - assert self.in_features % self.scaling_vector_size == 0, \ - "Input features must be a multiple of 16 for FP4 GEMM" - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.fp4) - self.weights_block_scaling_factor = Parameter( - shape=(self.out_features, - self.in_features // self.scaling_vector_size), - dtype=trt.fp8) - nrows = fp4_utils.pad_up(self.out_features, 128) - ncols = fp4_utils.pad_up(self.in_features // self.scaling_vector_size, - 4) - self.weights_block_scaling_factor_interleaved = Parameter(shape=(nrows, - ncols), - dtype=trt.fp8) - self.weights_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.activation_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - # alpha = 1.0 / (weight_global_scale * act_global_scale) - self.alpha = Parameter(shape=(1, ), dtype=trt.float32) - - self.tllm_to_externel_key_dict = { - "weight": ["weight"], - "weights_block_scaling_factor": "weight_scale", - "weights_block_scaling_factor_interleaved": "weight_scale", - "weights_global_scaling_factor": "weight_scale_2", - "activation_global_scaling_factor": "input_scale", - "alpha": ["weight_scale_2", "input_scale"], - } - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None, "lora is not supported on FP4Linear now" - - if isinstance(x, (tuple, list)): - fp4_x, act_per_block_scale = x - else: - if default_net().plugin_config.gemm_plugin == "nvfp4": - fp4_x, act_per_block_scale = quantize_to_fp4_tensor( - x, div(1.0, self.activation_global_scaling_factor.value)) - else: - # WAR for FP8 output attention - if x.dtype == trt.fp8: - # Since the scale is NVFP4 scale, we need to make it back to fp8 scale - new_scale_factor = self.activation_global_scaling_factor.raw_value - new_scale_factor = constant(new_scale_factor * 6) - x = dequantize(x, new_scale_factor, 0, - new_scale_factor.dtype) - fp4_x, act_per_block_scale = dynamic_quantize( - x, self.activation_global_scaling_factor.value) - - if default_net().plugin_config.gemm_allreduce_plugin: - x = gemm_allreduce( - a=fp4_x, - b=self.weight.value, - a_sf=act_per_block_scale, - b_sf=self.weights_block_scaling_factor_interleaved.value, - transa=False, # row-major - transb=True, # col-major - alpha=self.alpha.value, - group=self.tp_group, # ranks participating - fp8_inputs_override=False) - else: - if default_net().plugin_config.gemm_plugin == "nvfp4": - x = fp4_gemm( - fp4_x, act_per_block_scale, self.weight.value, - self.weights_block_scaling_factor_interleaved.value, - self.alpha.value, self.dtype) - else: - quant_w = self.weight.value - scale_w = self.weights_block_scaling_factor.value - dequant_x = block_double_dequantize( - fp4_x, act_per_block_scale, - self.activation_global_scaling_factor.value, trt.float16) - dequant_w = block_double_dequantize( - quant_w, scale_w, self.weights_global_scaling_factor.value, - trt.float16) - x = matmul(dequant_x, dequant_w, transb=True).cast(self.dtype) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, - self.tp_group, - all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if not any([ - tllm_key.endswith(suffix) - for suffix in self.tllm_to_externel_key_dict - ]): - return super().postprocess(tllm_key, weights, **kwargs) - - if tllm_key.endswith("weight"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor_interleaved"): - return torch.ops.trtllm.block_scale_interleave( - weights.view(torch.uint8).cpu().contiguous()).reshape( - weights.shape).view(torch.float8_e4m3fn) - elif tllm_key.endswith("weights_global_scaling_factor"): - return weights.float() - elif tllm_key.endswith('activation_global_scaling_factor'): - return weights.float() - elif tllm_key.endswith('alpha'): - weight_global_sf = weights[0].float() - act_global_sf = weights[1].float() - return act_global_sf * weight_global_sf - - -class SmoothQuantGatedMLP(SmoothQuantMLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - self.gate = SmoothQuantColumnLinear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - if self.quant_mode.has_act_static_scaling(): - self.quantization_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - else: - self.register_parameter('quantization_scaling_factor', None) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states) - inter_x_gate = inter * gate - smoother = cast(self.proj.smoother.value, self.dtype) - inter_x_gate = inter_x_gate / smoother - if self.quant_mode.has_act_and_weight_quant(): - if self.quant_mode.has_act_static_scaling(): - # Avoid quantization layers as it breaks int8 plugins - inter_x_gate = quantize_tensor( - inter_x_gate, self.quantization_scaling_factor.value) - else: - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - inter_x_gate = quantize_per_token(inter_x_gate) - - output = self.proj(inter_x_gate) - return output - - -class SmoothQuantAttention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dense_bias=None, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - tp_group=None, - tp_size=1, - tp_rank=0, - scale_alibi_bias=False, - paged_kv_cache=False, - quant_mode=QuantMode(0)): - super().__init__() - self.local_layer_idx = local_layer_idx - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_kv_heads = num_kv_heads - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size // tp_size - self.max_position_embeddings = 0 if max_position_embeddings is None else max_position_embeddings - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = 1 - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = scale_alibi_bias - - self.position_embedding_type = position_embedding_type - self.paged_kv_cache = paged_kv_cache - - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.rotary_embedding_dim = 0 - - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes(self.num_attention_heads * - self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.dtype = dtype - - if self.quant_mode.has_act_static_scaling(): - self.quantization_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - else: - self.register_parameter('quantization_scaling_factor', None) - - qkv_quant_mode = quant_mode - if self.quant_mode.has_act_and_weight_quant(): - # We need to hijack quant_mode for QKV because QKV always uses per channel scaling - qkv_quant_mode = QuantMode.from_description( - True, True, quant_mode.has_per_token_dynamic_scaling(), True) - - self.register_parameter('kv_cache_scaling_factor', None) - - self.qkv = SmoothQuantColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=qkv_quant_mode) - - self.dense = SmoothQuantRowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.use_lora = False - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - spec_decoding_params=None, - mrope_params=None, - encoder_output=None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - ): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - qkv = self.qkv(hidden_states) - - alibi_slopes = None - if self.position_embedding_type == PositionEmbeddingType.alibi: - alibi_slopes = self.alibi_slopes.value - dtype = trt.float32 - if default_net().plugin_config.gpt_attention_plugin or default_net( - ).plugin_config.inflight_batching_gpt_attention_plugin: - dtype = hidden_states.dtype if self.quant_mode.has_act_static_scaling( - ) else hidden_states[0].dtype - if dtype == trt.int8: - dtype = trt.float16 - alibi_slopes = cast(alibi_slopes, dtype) - - if spec_decoding_params is None: - spec_decoding_params = SpecDecodingParams() - - if mrope_params is None: - mrope_params = MropeParams() - - if default_net().plugin_config.gpt_attention_plugin: - - assert attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - if use_cache: - assert kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - assert self.attention_mask_type == AttentionMaskType.causal, \ - 'Plugin only support masked MHA.' - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - if self.position_embedding_type.is_rope(): - rotary_inv_freq = attention_params.rotary_inv_freq - rotary_cos_sin = attention_params.embed_positions_for_gpt_attention - else: - rotary_inv_freq = None - rotary_cos_sin = None - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=kv_cache_params.get_first_past_key_value(), - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base, - rotary_embedding_scale_type=self.rotary_embedding_scale_type, - rotary_embedding_scale=self.rotary_embedding_scale, - rotary_embedding_max_positions=self.max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - mrope_rotary_cos_sin=mrope_params.mrope_rotary_cos_sin, - mrope_position_deltas=mrope_params.mrope_position_deltas, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - ) - else: - assert self.paged_kv_cache == False - - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - query, key, value = split(qkv, self.hidden_size, dim=2) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - past_key_value = kv_cache_params.get_first_past_key_value() - if past_key_value is not None: - - def dequantize_tensor(x, scale): - # Cast from int8 to dtype - casted_x = cast(x, self.dtype) - return casted_x * scale - - if self.quant_mode.has_int8_kv_cache(): - past_key_value = dequantize_tensor( - past_key_value, self.kv_dequantization_scale.value) - - # past_key_value [bs, 2, num_heads, max_seq_len, head_dim] - past_key, past_value = split(past_key_value, 1, dim=1) - - key_shape = concat([ - shape(past_key, 0), - shape(past_key, 2), - shape(past_key, 3), - shape(past_key, 4) - ]) - past_key = past_key.view(key_shape, zero_is_placeholder=False) - past_value = past_value.view(key_shape, - zero_is_placeholder=False) - key = concat([past_key, key], dim=2) - value = concat([past_value, value], dim=2) - - def merge_caches(): - key_inflated_shape = concat([ - shape(key, 0), 1, - shape(key, 1), - shape(key, 2), - shape(key, 3) - ]) - inflated_key = key.view(key_inflated_shape, - zero_is_placeholder=False) - inflated_value = value.view(key_inflated_shape, - zero_is_placeholder=False) - past_key_value = concat([inflated_key, inflated_value], dim=1) - return past_key_value - - if self.attention_mask_type == AttentionMaskType.causal: - query_length = shape(query, 2) - key_length = shape(key, 2) - starts = concat([0, 0, key_length - query_length, 0]) - sizes = concat([1, 1, query_length, key_length]) - buffer = constant( - np.expand_dims( - np.tril( - np.ones( - (self.max_position_embeddings, - self.max_position_embeddings))).astype(bool), - (0, 1))) - causal_mask = slice(buffer, starts, sizes) - - key = key.permute([0, 1, 3, 2]) - with precision("float32"): - attention_scores = matmul(query, key) - - if self.attention_mask_type == AttentionMaskType.causal: - attention_scores = where(causal_mask, attention_scores, - -10000.0) - - attention_scores = attention_scores / self.norm_factor - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([shape(context, 0), - shape(context, 1), self.hidden_size])) - - past_key_value = merge_caches() - - if use_cache and self.quant_mode.has_int8_kv_cache(): - past_key_value = quantize_tensor( - past_key_value, self.kv_quantization_scale.value) - value = cast(self.dense.smoother.value, context.dtype) - context = context / value - if self.quant_mode.has_act_and_weight_quant(): - if self.quant_mode.has_act_static_scaling(): - # Avoid quantization layers as it breaks int8 plugins - context = quantize_tensor( - context, self.quantization_scaling_factor.value) - else: - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - context = quantize_per_token(context) - - context = self.dense( - context, - all_reduce_params=all_reduce_params, - ) - - if use_cache: - return (context, past_key_value) - - return context - - -# TODO: Duplicates SmoothQuantRmsNorm -class QServeRmsNorm(Module): - - def __init__(self, - normalized_shape, - eps=1e-06, - elementwise_affine=False, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - assert quant_mode.is_qserve_w4a8() - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - if self.quant_mode.has_act_and_weight_quant(): - self.scale_to_int = Parameter(shape=(1, ), dtype=dtype) - else: - self.register_parameter('scale_to_int', None) - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None if self.scale_to_int is None else self.scale_to_int.value - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return smooth_quant_rms_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=True, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - - -# TODO: Mostly duplicates SmoothQuantMLP. -# TODO: MLP could represent GatedMLP if hidden_act=='swiglu'. -class QServeMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act == 'swiglu' else ffn_hidden_size - self.fc = QServeW4A8ColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = QServeW4A8RowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.hidden_act = hidden_act - self.quant_mode = quant_mode - self.dtype = dtype - - def forward(self, hidden_states): - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - inter = quantize_per_token( - inter, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - output = self.proj(inter) - return output - - -# TODO: Mostly duplicates SmoothQuantGatedMLP. -class QServeGatedMLP(QServeMLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - self.gate = QServeW4A8Linear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, "lora_layer_params not supported" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states) - - inter_x_gate = inter * gate - inter_x_gate = quantize_per_token( - inter_x_gate, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - - output = self.proj(inter_x_gate) - return output - - -# TODO: Duplicates SmoothQuantAttention. -class QServeAttention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dense_bias=None, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - tp_group=None, - tp_size=1, - tp_rank=0, - scale_alibi_bias=False, - paged_kv_cache=False, - quant_mode=QuantMode(0)): - super().__init__() - self.local_layer_idx = local_layer_idx - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_kv_heads = num_kv_heads - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size // tp_size - self.max_position_embeddings = 0 if max_position_embeddings is None else max_position_embeddings - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = 1 - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = scale_alibi_bias - - self.position_embedding_type = position_embedding_type - self.paged_kv_cache = paged_kv_cache - - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.rotary_embedding_dim = 0 - - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes(self.num_attention_heads * - self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.dtype = dtype - - # QServe does not use act static scaling - # if self.quant_mode.has_act_static_scaling(): - # self.quantization_scaling_factor = Parameter(shape=(1, ), - # dtype='float32') - # else: - # self.register_parameter('quantization_scaling_factor', None) - - qkv_quant_mode = quant_mode - if self.quant_mode.has_act_and_weight_quant(): - # We need to hijack quant_mode for QKV because QKV always uses per channel scaling - qkv_quant_mode = QuantMode.from_description( - True, True, quant_mode.has_per_token_dynamic_scaling(), True) - - self.register_parameter('kv_cache_scaling_factor', None) - - self.qkv = QServeW4A8ColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=qkv_quant_mode) - - self.dense = QServeW4A8RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.use_lora = False - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - spec_decoding_params=None, - encoder_output=None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - ): - assert lora_layer_params is None, "lora is not supported on SmoothQuantAttention now" - if default_net().plugin_config.qserve_gemm_plugin: - qkv = self.qkv(hidden_states) - else: - raise ValueError("qserve_gemm_plugin is not set") - - alibi_slopes = None - if self.position_embedding_type == PositionEmbeddingType.alibi: - alibi_slopes = self.alibi_slopes.value - dtype = trt.float32 - if default_net().plugin_config.gpt_attention_plugin or default_net( - ).plugin_config.inflight_batching_gpt_attention_plugin: - dtype = hidden_states.dtype if self.quant_mode.has_act_static_scaling( - ) else hidden_states[0].dtype - if dtype == trt.int8: - dtype = trt.float16 - alibi_slopes = cast(alibi_slopes, dtype) - - if spec_decoding_params is None: - spec_decoding_params = SpecDecodingParams() - - if default_net().plugin_config.gpt_attention_plugin: - - assert attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - if use_cache: - assert kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - assert self.attention_mask_type == AttentionMaskType.causal, \ - 'Plugin only support masked MHA.' - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - if self.position_embedding_type.is_rope(): - rotary_inv_freq = attention_params.rotary_inv_freq - rotary_cos_sin = attention_params.embed_positions_for_gpt_attention - else: - rotary_inv_freq = None - rotary_cos_sin = None - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=kv_cache_params.get_first_past_key_value(), - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base, - rotary_embedding_scale_type=self.rotary_embedding_scale_type, - rotary_embedding_scale=self.rotary_embedding_scale, - rotary_embedding_max_positions=self.max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress) - else: - assert self.paged_kv_cache == False - - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - query, key, value = split(qkv, self.hidden_size, dim=2) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - past_key_value = kv_cache_params.get_first_past_key_value() - if past_key_value is not None: - - def dequantize_tensor(x, scale): - # Cast from int8 to dtype - casted_x = cast(x, self.dtype) - return casted_x * scale - - if self.quant_mode.has_int8_kv_cache(): - past_key_value = dequantize_tensor( - past_key_value, self.kv_dequantization_scale.value) - - # past_key_value [bs, 2, num_heads, max_seq_len, head_dim] - past_key, past_value = split(past_key_value, 1, dim=1) - - key_shape = concat([ - shape(past_key, 0), - shape(past_key, 2), - shape(past_key, 3), - shape(past_key, 4) - ]) - past_key = past_key.view(key_shape, zero_is_placeholder=False) - past_value = past_value.view(key_shape, - zero_is_placeholder=False) - key = concat([past_key, key], dim=2) - value = concat([past_value, value], dim=2) - - def merge_caches(): - key_inflated_shape = concat([ - shape(key, 0), 1, - shape(key, 1), - shape(key, 2), - shape(key, 3) - ]) - inflated_key = key.view(key_inflated_shape, - zero_is_placeholder=False) - inflated_value = value.view(key_inflated_shape, - zero_is_placeholder=False) - past_key_value = concat([inflated_key, inflated_value], dim=1) - return past_key_value - - if self.attention_mask_type == AttentionMaskType.causal: - query_length = shape(query, 2) - key_length = shape(key, 2) - starts = concat([0, 0, key_length - query_length, 0]) - sizes = concat([1, 1, query_length, key_length]) - buffer = constant( - np.expand_dims( - np.tril( - np.ones( - (self.max_position_embeddings, - self.max_position_embeddings))).astype(bool), - (0, 1))) - causal_mask = slice(buffer, starts, sizes) - - key = key.permute([0, 1, 3, 2]) - with precision("float32"): - attention_scores = matmul(query, key) - - if self.attention_mask_type == AttentionMaskType.causal: - attention_scores = where(causal_mask, attention_scores, - -10000.0) - - attention_scores = attention_scores / self.norm_factor - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([shape(context, 0), - shape(context, 1), self.hidden_size])) - - past_key_value = merge_caches() - - if use_cache and self.quant_mode.has_int8_kv_cache(): - past_key_value = quantize_tensor( - past_key_value, self.kv_quantization_scale.value) - - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - context = quantize_per_token( - context, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - - context = self.dense(context, all_reduce_params=all_reduce_params) - - if use_cache: - return (context, past_key_value) - - return context - - -class Fp8RowwiseLayerNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-05, - elementwise_affine=True, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None, - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Layer norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return fp8_rowwise_layer_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=self.quant_mode.has_fp8_rowwise()) diff --git a/tensorrt_llm/quantization/quantize.py b/tensorrt_llm/quantization/quantize.py deleted file mode 100644 index a04a223914f2..000000000000 --- a/tensorrt_llm/quantization/quantize.py +++ /dev/null @@ -1,611 +0,0 @@ -import fnmatch -from typing import Union - -import torch - -from .._utils import get_init_params -from ..layers import (MLP, Attention, ColumnLinear, Embedding, GatedMLP, - LayerNorm, RmsNorm, RowLinear) -from ..layers.moe import MixtureOfExperts -from ..models.modeling_utils import LayerQuantConfig, QuantConfig -from ..parameter import Parameter - -# isort: off -from .layers import ( - FP4Linear, FP4RowLinear, FP8Linear, FP8RowLinear, Fp8RowwiseAttention, - Fp8RowwiseGatedMLP, Fp8RowwiseLayerNorm, Fp8RowwiseMLP, Fp8RowwiseRmsNorm, - Int8SmoothQuantLinear, Int8SmoothQuantRowLinear, QServeAttention, - QServeGatedMLP, QServeMLP, QServeRmsNorm, SmoothQuantAttention, - SmoothQuantGatedMLP, SmoothQuantLayerNorm, SmoothQuantMLP, - SmoothQuantRmsNorm, WeightOnlyGroupwiseQuantColumnLinear, - WeightOnlyGroupwiseQuantRowLinear, WeightOnlyQuantColumnLinear, - WeightOnlyQuantEmbedding, WeightOnlyQuantRowLinear) -# isort: on -from .mode import W8A8_SQ_PLUGIN_LIST, QuantAlgo, QuantMode - - -def quantize_layers( - model, - quant_config: QuantConfig, - quant_map, - preprocess_init_params=None, -): - exclude_modules = quant_config.exclude_modules - if exclude_modules is None: - exclude_modules = [ - '*lm_head', - '*router', - '*vocab_embedding', - '*position_embedding', - '*block_embedding', - '*shared_expert_gate', - ] - - for name, module, parent in model.named_modules_with_parent(): - module_name = name.rsplit('.', 1)[-1] - is_excluded = False - quant_cls = None - - # handle exclusion - for exclude_module in exclude_modules: - if fnmatch.fnmatchcase(name, exclude_module): - is_excluded = True - break - - # MoE modules are quantized on their constructor, so they must always - # be re-created with the appropriate quant_mode. When excluded, - # re-create with quant_mode 0. - # We need to handle it specially, we may want to redesign MoE implementation - if isinstance(module, MixtureOfExperts): - quant_cls = type(module) - elif not is_excluded: - for cls in quant_map: - if isinstance(module, cls): - quant_cls = quant_map[cls] - break - - if quant_cls: - init_params = get_init_params(module, quant_cls) - if isinstance(module, MixtureOfExperts): - if is_excluded: - quant_mode = QuantMode(0) - else: - quant_mode = quant_config.quant_mode - init_params["quant_mode"] = quant_mode - - # Auto-detect pre_quant_scale based on quant_algo - # For AWQ-based quantization methods that use pre_quant_scale - if quant_config.quant_algo in [ - QuantAlgo.W4A16_AWQ, QuantAlgo.NVFP4_AWQ, - QuantAlgo.W4A8_AWQ - ]: - init_params["pre_quant_scale"] = True - if "bias" in init_params and not isinstance(module, - MixtureOfExperts): - init_params["bias"] = init_params["bias"] is not None - if isinstance(module, ColumnLinear): - init_params[ - "out_features"] = module.out_features * module.tp_size - elif isinstance(module, RowLinear): - init_params["in_features"] = module.in_features * module.tp_size - if preprocess_init_params is not None: - preprocess_init_params(init_params, name, module) - quant_layer = quant_cls(**init_params) - if parent is not None: - setattr(parent, module_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_config.quant_mode) - return model - - -def weight_only_quantize(model, quant_config: QuantConfig, model_config=None): - assert quant_config.quant_mode.is_weight_only() - - try: - model_cfg = model.config - except AttributeError: - model_cfg = model_config - - quant_map = { - ColumnLinear: WeightOnlyQuantColumnLinear, - RowLinear: WeightOnlyQuantRowLinear, - Embedding: WeightOnlyQuantEmbedding, - } - - def preprocess_init_params(init_params, name, module): - init_params["quant_mode"] = quant_config.quant_mode - if isinstance(module, ColumnLinear): - module_name = name.rsplit('.', 1)[-1] - init_params["transb"] = module_name == "lm_head" - if "tp_rank" in init_params: - init_params["tp_rank"] = model_cfg.mapping.tp_rank - - model = quantize_layers( - model, - quant_config, - quant_map, - preprocess_init_params, - ) - return model - - -def weight_only_groupwise_quantize(model, - quant_config: QuantConfig, - model_config=None): - assert quant_config.quant_mode.is_weight_only() - - try: - model_cfg = model.config - except AttributeError: - model_cfg = model_config - - quant_map = { - ColumnLinear: WeightOnlyGroupwiseQuantColumnLinear, - RowLinear: WeightOnlyGroupwiseQuantRowLinear, - MixtureOfExperts: MixtureOfExperts, - } - - def preprocess_init_params(init_params, name, module): - init_params["group_size"] = quant_config.group_size - init_params["pre_quant_scale"] = quant_config.pre_quant_scale - init_params["zero"] = quant_config.has_zero_point - init_params[ - "use_w4a8_awq"] = quant_config.quant_algo == QuantAlgo.W4A8_AWQ - init_params[ - "use_int8_weight"] = quant_config.quant_algo == QuantAlgo.W8A16_GPTQ - if "tp_rank" in init_params: - init_params["tp_rank"] = model_cfg.mapping.tp_rank - - model = quantize_layers( - model, - quant_config, - quant_map, - preprocess_init_params, - ) - return model - - -def smooth_quantize_ootb( - model, - quant_config: QuantConfig, -): - quant_map = { - ColumnLinear: Int8SmoothQuantLinear, - RowLinear: Int8SmoothQuantRowLinear, - } - - model = quantize_layers( - model, - quant_config, - quant_map, - ) - return model - - -def smooth_quantize_plugin(model, quant_mode): - quant_map = { - RmsNorm: SmoothQuantRmsNorm, - LayerNorm: SmoothQuantLayerNorm, - GatedMLP: SmoothQuantGatedMLP, - MLP: SmoothQuantMLP, - Attention: SmoothQuantAttention, - } - for name, layer, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name in ['ln_f', 'ln_embed']: - continue - - quant_cls = None - for cls in quant_map: - if isinstance(layer, cls): - quant_cls = quant_map[cls] - break - - if quant_cls is None: - continue - - init_params = get_init_params(layer, quant_cls) - init_params["quant_mode"] = quant_mode - if isinstance(layer, Attention): - init_params[ - "num_attention_heads"] = layer.num_attention_heads * layer.tp_size - quant_layer = quant_cls(**init_params) - if parent is not None: - setattr(parent, layer_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_mode) - return model - - -def smooth_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_act_and_weight_quant() - if quant_config.quant_algo in W8A8_SQ_PLUGIN_LIST: - return smooth_quantize_plugin(model, quant_config.quant_mode) - else: - return smooth_quantize_ootb(model, quant_config) - - -def fp8_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_fp8_qdq() - - quant_map = { - ColumnLinear: FP8Linear, - RowLinear: FP8RowLinear, - MixtureOfExperts: MixtureOfExperts, - } - - model = quantize_layers( - model, - quant_config, - quant_map, - ) - return model - - -def fp8_rowwise_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_fp8_rowwise() - - quant_cls_map = { - RmsNorm: Fp8RowwiseRmsNorm, - LayerNorm: Fp8RowwiseLayerNorm, - GatedMLP: Fp8RowwiseGatedMLP, - MLP: Fp8RowwiseMLP, - Attention: Fp8RowwiseAttention, - } - - exclude_modules = quant_config.exclude_modules - if exclude_modules is None: - exclude_modules = [] - # Always exclude these modules for FP8 rowwise - exclude_modules = list( - set(exclude_modules + ['*ln_f', '*ln_embed', '*lm_head'])) - - def extract_layer_idx(name): - ss = name.split('.') - for s in ss: - if s.isdigit(): - return int(s) - return None - - # Meta's LLaMA 3.1 recipe: - # (1) Skip quantization for the first and last Transformer layers - # (2) Skip quantization for the Attention layers - if quant_config.use_meta_recipe: - exclude_modules.extend(['*input_layernorm', '*attention']) - - for name, layer, parent in model.named_modules_with_parent(): - module_name = name.rsplit('.', 1)[-1] - - if quant_config.use_meta_recipe: - local_layer_idx = extract_layer_idx(name) - mapping = model.config.mapping - layers_range = mapping.pp_layers(model.config.num_hidden_layers) - if mapping.is_first_pp_rank() and local_layer_idx == 0: - continue - if mapping.is_last_pp_rank( - ) and local_layer_idx == len(layers_range) - 1: - continue - - quant_cls = None - for cls in quant_cls_map: - if isinstance(layer, cls): - quant_cls = quant_cls_map[cls] - break - if quant_cls is None: - continue - - is_excluded = False - for exclude_module in exclude_modules: - if fnmatch.fnmatchcase(name, exclude_module): - is_excluded = True - break - if is_excluded: - continue - - init_params = get_init_params(layer, quant_cls) - init_params["quant_mode"] = quant_config.quant_mode - if isinstance(layer, Attention): - init_params[ - "num_attention_heads"] = layer.num_attention_heads * layer.tp_size - quant_layer = quant_cls(**init_params, clamp_val=quant_config.clamp_val) - if parent is not None: - setattr(parent, module_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_config.quant_mode) - return model - - -# TODO: These functions should be moved to ModelOpt. -def qserve_quantize_weight_per_group(linear_weight: torch.HalfTensor, - s1_scales: torch.FloatTensor, - s2_scales: torch.FloatTensor, - s2_szeros: torch.FloatTensor, - group_size: int) -> torch.CharTensor: - out_features = linear_weight.shape[0] - in_features = linear_weight.shape[1] - - # Step 1: Quantize the weights to int8 - linear_weight = linear_weight.div( - s1_scales.reshape(out_features, 1).to(linear_weight.device)) - linear_weight = linear_weight.round() - # assert linear_weight.min() >= -119 and linear_weight.max() <= 119, "Stage 1: Quantized weight out of range" # 119 is the "magic" number - assert (linear_weight.min() >= -128 and linear_weight.max() - <= 127), "Stage 1: Quantized weight out of range" - - # Step 2: Quantize the weights to int4 - linear_weight = linear_weight.reshape(out_features, - in_features // group_size, group_size) - s2_szeros = s2_szeros.reshape(out_features, in_features // group_size, - 1).to(torch.float16).to(linear_weight.device) - s2_scales = s2_scales.reshape(out_features, in_features // group_size, - 1).to(torch.float16).to(linear_weight.device) - linear_weight = linear_weight.add(s2_szeros).div(s2_scales).round() - assert (linear_weight.min() >= 0 and linear_weight.max() - <= 15), "Stage 2: Quantized weight out of range" - - qweight = linear_weight.reshape(out_features, in_features).to(torch.int8) - return qweight - - -def qserve_quantize_weight_per_channel( - linear_weight: torch.HalfTensor, s1_scales: torch.FloatTensor, - s1_szeros: torch.FloatTensor) -> torch.CharTensor: - out_features = linear_weight.shape[0] - in_features = linear_weight.shape[1] - - # Step 1: Quantize the weights to int4 - s1_scales = s1_scales.reshape(out_features, 1).to(linear_weight.device) - s1_szeros = s1_szeros.reshape(out_features, 1).to(linear_weight.device) - - qweight = linear_weight.add(s1_szeros).div(s1_scales).round() - assert (qweight.min() >= 0 - and qweight.max() <= 15), "Quantized weight out of range" - - return qweight.reshape(out_features, in_features).to(torch.int8) - - -# Pack the quantized weights, scales and zeros and apply the reordering required by QServe kernels. -# Return: processed [qweight, s1_scales, s2_scales, s2_zeros] -def qserve_pack_reorder_per_group(qweight: torch.CharTensor, - s1_scales: torch.FloatTensor, - s2_scales: torch.FloatTensor, - s2_szeros: torch.FloatTensor, group_size): - out_features = qweight.shape[0] - in_features = qweight.shape[1] - - outputs = [] - - s1_scales = s1_scales.reshape(out_features).to(torch.float16) - s2_szeros = s2_szeros.reshape(out_features, - in_features // group_size).to(torch.int8) - s2_scales = s2_scales.reshape(out_features, - in_features // group_size).to(torch.int8) - - # Step 3: Pack the quantized weights to real quantized weights - # ---- Repack the weight ---- # - assert qweight.dtype == torch.int8 - # pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4) - W_unpack_reorder = (qweight.reshape( - out_features // 32, - 2, - 2, - 8, - in_features // 32, - 2, - 4, - 4, - ).permute(0, 4, 3, 6, 1, 5, 2, 7).contiguous()) - W_unpack_reorder = (W_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, - 4).contiguous().to(torch.int8)) - # B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous() - # [16, 0, 17, 1, ...] - W_unpack_repacked = (W_unpack_reorder[..., 1] << 4) + W_unpack_reorder[..., - 0] - W_unpack_repacked = W_unpack_repacked.reshape(out_features // 32, - in_features // 32, 32, 16) - W_unpack_repacked = W_unpack_repacked.reshape(out_features, - in_features // 2) - - outputs.append(W_unpack_repacked) - - # for the last dimension, organize as 0, 8, 16, 24, 1, 9, 17, 25, ... following the requirement of tensor core gemm - # ---- Pack the scales ---- # - outputs.append(s1_scales.reshape(out_features)) - - s2_scales = (s2_scales.reshape(out_features, in_features // - group_size).transpose(0, 1).contiguous()) - s2_scales = s2_scales.reshape(in_features // group_size, out_features // 32, - 32) - s2_scales = (s2_scales.reshape(in_features // group_size, - out_features // 32, 4, - 8).transpose(-2, -1).contiguous()) - s2_scales = s2_scales.reshape(in_features // group_size, - out_features).contiguous() - outputs.append(s2_scales) - - # ---- Pack the zeros ---- # - s2_szeros = (s2_szeros.reshape(out_features, in_features // - group_size).transpose(0, 1).contiguous()) - s2_szeros = s2_szeros.reshape(in_features // group_size, out_features // 32, - 32) - s2_szeros = (s2_szeros.reshape(in_features // group_size, - out_features // 32, 4, - 8).transpose(-2, -1).contiguous()) - s2_szeros = (s2_szeros.reshape(in_features // group_size, - out_features).contiguous()) - - # (q - s2_zeros) * s2_scales = q * s2_scales - s2_zeros * s2_scales, - # We convert the s2_zeros -> -s2_zeros * s2_scales - s2_szeros = (-s2_szeros).int() # It has been pre-scaled in DeepCompressor - s2_szeros = s2_szeros.to(torch.int8) - - outputs.append(s2_szeros) - - return outputs - - -def qserve_pack_reorder_per_channel(qweight: torch.CharTensor, - s1_scales: torch.FloatTensor, - s1_szeros: torch.FloatTensor): - out_features = qweight.shape[0] - in_features = qweight.shape[1] - - outputs = [] - - # ---- Repack the weight ---- # - assert qweight.dtype == torch.int8 - # pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4) - W_unpack_reorder = (qweight.reshape( - out_features // 32, - 2, - 2, - 8, - in_features // 32, - 2, - 4, - 4, - ).permute(0, 4, 3, 6, 1, 5, 2, 7).contiguous()) - W_unpack_reorder = (W_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, - 4).contiguous()) - # B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous() - # [16, 0, 17, 1, ...] - W_unpack_repacked = (W_unpack_reorder[..., 1] << 4) + W_unpack_reorder[..., - 0] - W_unpack_repacked = W_unpack_repacked.reshape(out_features // 32, - in_features // 32, 32, 16) - W_unpack_repacked = W_unpack_repacked.reshape(out_features, in_features // - 2).contiguous() - - outputs.append(W_unpack_repacked) - - # ---- Pack the scales and zeros ---- # - s1_scales = s1_scales.reshape(out_features).contiguous() - outputs.append(s1_scales.half()) - - s1_szeros = s1_szeros.reshape(out_features).contiguous().half() - outputs.append(s1_szeros) - - return outputs - - -# TODO: Duplicates smooth_quantize and quantize_layers -def qserve_quantize(model, quant_config: QuantConfig): - quant_mode = quant_config.quant_mode - assert quant_config.quant_mode.is_qserve_w4a8() - - quant_map = { - RmsNorm: QServeRmsNorm, - LayerNorm: QServeRmsNorm, - GatedMLP: QServeGatedMLP, - MLP: QServeMLP, - Attention: QServeAttention, - } - - for name, layer, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name in ['ln_f', 'ln_embed']: - continue - - quant_cls = None - for cls in quant_map: - if isinstance(layer, cls): - quant_cls = quant_map[cls] - break - - if quant_cls is None: - continue - - init_params = get_init_params(layer, quant_cls) - init_params["quant_mode"] = quant_mode - if isinstance(layer, Attention): - init_params[ - "num_attention_heads"] = layer.num_attention_heads * layer.tp_size - quant_layer = quant_cls(**init_params) - if parent is not None: - setattr(parent, layer_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_mode) - return model - - -def fp4_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_nvfp4() - quant_map = { - ColumnLinear: FP4Linear, - RowLinear: FP4RowLinear, - MixtureOfExperts: MixtureOfExperts, - } - - model = quantize_layers( - model, - quant_config, - quant_map, - ) - return model - - -# Now consider the kv cache is enabled for all layers -def kv_cache_quantize(model): - for name, module in model.named_modules(): - if isinstance(module, - (Attention, SmoothQuantAttention, Fp8RowwiseAttention)): - # for dequant - module.kv_cache_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - # for quant - module.kv_cache_rcp_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - return model - - -def quantize(model, quant_config: Union[QuantConfig, LayerQuantConfig]): - - for name, module, parent in model.named_modules_with_parent(): - - if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION: - layer_quant_mode = quant_config.layer_quant_mode(name) - else: - layer_quant_mode = quant_config.layer_quant_mode - if layer_quant_mode == QuantMode(0): - continue - - layer_quant_cfg = quant_config._get_quant_cfg(name) - - if layer_quant_mode.has_fp8_qdq(): - module = fp8_quantize(module, layer_quant_cfg) - elif layer_quant_mode.has_fp8_rowwise(): - module = fp8_rowwise_quantize(module, layer_quant_cfg) - elif layer_quant_mode.is_qserve_w4a8(): - module = qserve_quantize(module, quant_config) - elif layer_quant_mode.has_nvfp4(): - module = fp4_quantize(module, layer_quant_cfg) - elif layer_quant_mode.has_act_and_weight_quant(): - module = smooth_quantize(module, layer_quant_cfg) - elif layer_quant_mode.is_weight_only(): - if layer_quant_mode.has_per_group_scaling(): - module = weight_only_groupwise_quantize(module, layer_quant_cfg, - model.config) - else: - module = weight_only_quantize(module, layer_quant_cfg, - model.config) - - if parent is not None: # for per layer - module_name = name.rsplit('.', 1)[-1] - setattr(parent, module_name, module) - else: # for all layer - model = module - break - - if quant_config.quant_mode.has_kv_cache_quant(): - model = kv_cache_quantize(model) - - setattr(model, 'quant_mode', quant_config.quant_mode) - return model diff --git a/tensorrt_llm/quantization/quantize_by_modelopt.py b/tensorrt_llm/quantization/quantize_by_modelopt.py index fbb50310f893..be9236aa6cef 100755 --- a/tensorrt_llm/quantization/quantize_by_modelopt.py +++ b/tensorrt_llm/quantization/quantize_by_modelopt.py @@ -29,7 +29,6 @@ from accelerate.hooks import remove_hook_from_module from datasets import load_dataset from modelopt.torch.utils import print_rank_0 -from safetensors.torch import load_file, save_file from torch import nn from torch.utils.data import DataLoader from transformers import (AutoConfig, AutoModelForCausalLM, AutoProcessor, @@ -37,7 +36,6 @@ from .._utils import get_hf_rope_theta, release_gc, str_dtype_to_torch from ..logger import logger -from ..mapping import Mapping from .image_processing import MllamaImageProcessor from .mode import QuantAlgo @@ -722,73 +720,6 @@ def calibrate_loop(): return model -def combine_medusa_weight(tp_size, pp_size, base_model_output_dir, - num_medusa_heads, num_medusa_layers, max_draft_len, - medusa_hidden_act, medusa_model_dir, - quant_medusa_head): - - with open(f"{medusa_model_dir}/config.json", "r") as fp: - medusa_config = json.load(fp) - - num_medusa_heads_from_config = medusa_config.get('medusa_num_heads', - num_medusa_heads) - num_medusa_layers = medusa_config.get('medusa_num_layers', - num_medusa_layers) - if num_medusa_heads is None: - num_medusa_heads = num_medusa_heads_from_config - - assert max_draft_len > 0, "should have max_draft_len > 0" - - world_size = tp_size * pp_size - # Process for each rank - for rank in range(world_size): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=tp_size, - pp_size=pp_size) - # 1. Load medusa weight for each rank - from tensorrt_llm.models.medusa.weight import load_medusa_hf - medusa_weights = load_medusa_hf(medusa_path=medusa_model_dir, - num_medusa_heads=num_medusa_heads, - num_medusa_layers=num_medusa_layers, - mapping=mapping, - dtype="float16") - # 2. Load base model safetensors (after quant) - base_model_weights = load_file( - f"{base_model_output_dir}/rank{rank}.safetensors") - - # 3. Combine and save weight - base_model_weights.update(medusa_weights) - save_file(base_model_weights, - f"{base_model_output_dir}/rank{rank}.safetensors") - - # 4. Add medusa config into config.json - with open(f"{base_model_output_dir}/config.json", 'r') as f: - base_model_config = json.load(f) - f.close() - - with open(f"{base_model_output_dir}/config.json", 'w') as f: - base_model_config['architecture'] = "MedusaForCausalLM" - base_model_config['quantization']['exclude_modules'] = [ - 'lm_head', - '*router', - '*vocab_embedding', - '*position_embedding', - '*block_embedding', - ] - if not quant_medusa_head: - base_model_config['quantization']['exclude_modules'].append( - '*medusa_heads*') - - base_model_config['max_draft_len'] = max_draft_len - base_model_config['num_medusa_heads'] = num_medusa_heads - base_model_config['num_medusa_layers'] = num_medusa_layers - json.dump(base_model_config, f, indent=4) - - torch.cuda.empty_cache() - logger.info("Combine medusa heads' weight, done.") - - def quantize_and_export(*, model_dir, device, @@ -1082,28 +1013,6 @@ def quantize_and_export(*, with open(f"{export_path}/config.json", "w") as f: json.dump(tensorrt_llm_config, f, indent=4) - # Workaround for combining medusa head - # TODO: move these integration into modelopt to avoid redundant reading and writing - if medusa_model_dir is not None: - combine_medusa_weight(tp_size, pp_size, export_path, - num_medusa_heads, num_medusa_layers, - max_draft_len, medusa_hidden_act, - medusa_model_dir, quant_medusa_head) - - # Workaround for mllama - if model_type == 'mllama': - from tensorrt_llm.models.mllama.config import MLLaMAConfig - config = MLLaMAConfig.from_hugging_face( - model_dir, - dtype=dtype, - ) - for key, value in config.to_dict().items(): - if key not in tensorrt_llm_config: - tensorrt_llm_config[key] = value - - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - end_time = time.time() logger.info( "Quantized model exported to {} \nTotal time used {:.2f} s.".format( diff --git a/tensorrt_llm/runtime/__init__.py b/tensorrt_llm/runtime/__init__.py index d6f018ff49aa..eaad2f1d3ab9 100644 --- a/tensorrt_llm/runtime/__init__.py +++ b/tensorrt_llm/runtime/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -12,6 +12,7 @@ # 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 os import sys from contextlib import contextmanager @@ -40,16 +41,7 @@ def temporary_sys_path(path: str) -> Iterator[None]: with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): import kv_cache_manager_v2 -from .enc_dec_model_runner import EncDecModelRunner -from .generation import SamplingConfig # autoflake: skip -from .generation import (ChatGLMGenerationSession, GenerationSession, - LogitsProcessor, LogitsProcessorList, ModelConfig, - QWenForCausalLMGenerationSession, StoppingCriteria, - StoppingCriteriaList, decode_words_list) -from .kv_cache_manager import GenerationSequence, KVCacheManager -from .model_runner import ModelRunner -from .multimodal_model_runner import MultimodalModelRunner -from .session import Session, TensorInfo +from .model_config import ModelConfig try: import tensorrt_llm.bindings # NOQA @@ -57,28 +49,8 @@ def temporary_sys_path(path: str) -> Iterator[None]: except ImportError: PYTHON_BINDINGS = False -if PYTHON_BINDINGS: - from .model_runner_cpp import ModelRunnerCpp - __all__ = [ 'ModelConfig', - 'GenerationSession', - 'GenerationSequence', - 'KVCacheManager', - 'SamplingConfig', - 'Session', - 'TensorInfo', - 'ChatGLMGenerationSession', - 'QWenForCausalLMGenerationSession', - 'decode_words_list', - 'LogitsProcessorList', - 'LogitsProcessor', - 'StoppingCriteriaList', - 'StoppingCriteria', - 'ModelRunner', - 'ModelRunnerCpp', - 'EncDecModelRunner', - 'MultimodalModelRunner', 'PYTHON_BINDINGS', 'kv_cache_manager_v2', ] diff --git a/tensorrt_llm/runtime/enc_dec_model_runner.py b/tensorrt_llm/runtime/enc_dec_model_runner.py deleted file mode 100644 index 44ce54f03a26..000000000000 --- a/tensorrt_llm/runtime/enc_dec_model_runner.py +++ /dev/null @@ -1,537 +0,0 @@ -import json -import time -from pathlib import Path - -# isort: off -import torch -import tensorrt as trt - -from .._deprecation import emit_engine_arch_deprecation -from .._utils import torch_to_numpy, trt_dtype_to_torch, mpi_world_size, mpi_rank -from ..logger import logger -from ..plugin.plugin import CustomAllReduceHelper -from .generation import ModelConfig, SamplingConfig, LoraManager, GenerationSession -from ..mapping import Mapping -from .session import Session -from ..models.modeling_utils import get_kv_cache_type_from_legacy - - -def get_engine_name(rank): - return 'rank{}.engine'.format(rank) - - -def read_config(config_path: Path): - with open(config_path, "r") as f: - config = json.load(f) - - builder_config = config['build_config'] - plugin_config = builder_config['plugin_config'] - pretrained_config = config['pretrained_config'] - lora_config = builder_config['lora_config'] - use_gpt_attention_plugin = plugin_config["gpt_attention_plugin"] - remove_input_padding = plugin_config["remove_input_padding"] - use_lora_plugin = plugin_config["lora_plugin"] - tp_size = pretrained_config['mapping']['tp_size'] - pp_size = pretrained_config['mapping']['pp_size'] - gpus_per_node = pretrained_config['mapping']['gpus_per_node'] - world_size = tp_size * pp_size - assert world_size == mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({mpi_world_size()})' - num_heads = pretrained_config["num_attention_heads"] - hidden_size = pretrained_config["hidden_size"] - head_size = pretrained_config["head_size"] - vocab_size = pretrained_config["vocab_size"] - max_batch_size = builder_config["max_batch_size"] - max_beam_width = builder_config["max_beam_width"] - num_layers = pretrained_config["num_hidden_layers"] - num_kv_heads = pretrained_config.get('num_kv_heads', num_heads) - - assert (num_heads % tp_size) == 0 - num_heads = num_heads // tp_size - hidden_size = hidden_size // tp_size - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - - cross_attention = pretrained_config["architecture"] == "DecoderModel" - skip_cross_kv = pretrained_config.get('skip_cross_kv', False) - has_position_embedding = pretrained_config["has_position_embedding"] - has_token_type_embedding = hasattr(pretrained_config, "type_vocab_size") - dtype = pretrained_config["dtype"] - - paged_kv_cache = plugin_config['paged_kv_cache'] - tokens_per_block = plugin_config['tokens_per_block'] - - gather_context_logits = builder_config.get('gather_context_logits', False) - gather_generation_logits = builder_config.get('gather_generation_logits', - False) - max_prompt_embedding_table_size = builder_config.get( - 'max_prompt_embedding_table_size', 0) - - kv_cache_type = get_kv_cache_type_from_legacy(True, paged_kv_cache) - language_adapter_config = pretrained_config.get("language_adapter_config", - None) - - model_config = ModelConfig( - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - head_size=head_size, - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - vocab_size=vocab_size, - num_layers=num_layers, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - cross_attention=cross_attention, - has_position_embedding=has_position_embedding, - has_token_type_embedding=has_token_type_embedding, - dtype=dtype, - gather_context_logits=gather_context_logits, - gather_generation_logits=gather_generation_logits, - max_prompt_embedding_table_size=max_prompt_embedding_table_size, - lora_plugin=use_lora_plugin, - lora_target_modules=lora_config.get('lora_target_modules'), - trtllm_modules_to_hf_modules=lora_config.get( - 'trtllm_modules_to_hf_modules'), - skip_cross_kv=skip_cross_kv, - language_adapter_config=language_adapter_config) - - return model_config, tp_size, pp_size, gpus_per_node, dtype - - -class EncDecModelRunner: - - def __init__(self, - engine_name, - engine_dir, - lora_dir=None, - lora_task_uids=None, - debug_mode=False, - skip_encoder=False, - stream: torch.cuda.Stream = None, - enable_context_fmha_fp32_acc: bool = None): - emit_engine_arch_deprecation("EncDecModelRunner") - # in multi-node setup, it's important to set_device at the very beginning so .to('cuda') refers to current device - # accordingly, all input & output tensors should be moved to current device - # otherwise, it's default to 'cuda:0' - self.runtime_rank = mpi_rank() - device_id = self.runtime_rank % torch.cuda.device_count() - torch.cuda.set_device(device_id) - self.device = torch.cuda.current_device() - self.skip_encoder = skip_encoder - self.lora_task_uids = lora_task_uids - self.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - - # when enc-dec runs by itself, stream can be None and we create new stream here - # when enc-dec has to run as a component in a bigger workflow (e.g., multimodal), earlier components in the workflow may have results in its stream, which we should pass that stream in to avoid unnecessary stream sync - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_dir = Path(engine_dir) - - def engine_setup(component): - # model config - config_path = engine_dir / component / "config.json" - logger.info(f"Using config path {config_path}") - model_config, tp_size, pp_size, gpus_per_node, dtype = read_config( - config_path) - - # MGMN config - world_size = tp_size * pp_size - runtime_rank = mpi_rank() - assert runtime_rank < world_size, "Runtime GPU rank exceeds MPI world size. Did you launch more MPI processes than required?" - runtime_mapping = Mapping(world_size, - runtime_rank, - tp_size=tp_size, - pp_size=pp_size, - gpus_per_node=gpus_per_node) - - # load engine - engine_fname = get_engine_name(runtime_rank) - with open(engine_dir / component / engine_fname, "rb") as f: - engine_buffer = f.read() - - return model_config, runtime_mapping, engine_buffer - - # Note: encoder and decoder doesn't necessarily have the same TP & PP config - - if not skip_encoder: - self.encoder_model_config, self.encoder_runtime_mapping, encoder_engine_buffer = engine_setup( - component='encoder') - - self.nccl_comm = None - if self.encoder_runtime_mapping.has_pp(): - # for Pipeline Parallelism in encoder - self.nccl_comm = torch.classes.trtllm.NcclCommunicatorOp( - self.encoder_runtime_mapping.world_size, - self.encoder_runtime_mapping.rank) - - # session setup - self.encoder_session = Session.from_serialized_engine( - encoder_engine_buffer) - - # encoder lora manager setup - if self.encoder_model_config.lora_plugin: - self.encoder_lora_manager = LoraManager( - mapping=self.encoder_runtime_mapping, - model_config=self.encoder_model_config, - ) - # TODO: this is only for bart - self.encoder_lora_manager.load_from_hf( - model_dirs=lora_dir, - model_config=self.encoder_model_config, - component='encoder', - ) - else: - self.encoder_lora_manager = None - else: - self.encoder_model_config, self.encoder_runtime_mapping, encoder_engine_buffer = None, None, None - self.nccl_comm, self.encoder_session = None, None - - self.decoder_model_config, self.decoder_runtime_mapping, decoder_engine_buffer = engine_setup( - component='decoder') - self.decoder_session = GenerationSession(self.decoder_model_config, - decoder_engine_buffer, - self.decoder_runtime_mapping, - debug_mode=debug_mode) - - # decoder lora manager setup - if self.decoder_model_config.lora_plugin: - self.decoder_lora_manager = LoraManager( - mapping=self.decoder_runtime_mapping, - model_config=self.decoder_model_config, - ) - # TODO: this is only for bart - self.decoder_lora_manager.load_from_hf( - model_dirs=lora_dir, - model_config=self.decoder_model_config, - component='decoder', - ) - else: - self.decoder_lora_manager = None - - @classmethod - def from_engine(cls, - engine_name, - engine_dir, - lora_dir=None, - lora_task_uids=None, - debug_mode=False, - skip_encoder=False, - stream=None, - enable_context_fmha_fp32_acc=None): - return cls(engine_name, - engine_dir, - lora_dir, - lora_task_uids, - debug_mode=debug_mode, - skip_encoder=skip_encoder, - stream=stream, - enable_context_fmha_fp32_acc=enable_context_fmha_fp32_acc) - - def process_input(self, - input_ids, - remove_input_padding=False, - pad_token_id=0, - prompt_tasks=None, - language_adapter_routings=None): - if remove_input_padding: - # in remove padding mode --> flatten input, calculate actual length and max length - # Note: 1st token should never be removed, even if it is pad_token_id - first_ids = input_ids[:, 0] - input_ids = input_ids[:, 1:] - input_lengths = 1 + (input_ids != pad_token_id).sum(dim=1).type( - torch.IntTensor).to(self.device) # [batch_size] - new_ids = [] - for i in range(len(input_ids)): - row = input_ids[i, :] - row = row[row != pad_token_id] - new_ids.append( - torch.cat( - (torch.IntTensor([first_ids[i]]).to(self.device), row))) - input_ids = torch.cat(new_ids) # [num_tokens] - if prompt_tasks is not None: - prompt_tasks = prompt_tasks[:input_ids.shape[0]] - else: - # in padding mode --> keep input, just calculate actual length and max length - # Note: 1st token should always count, even if it is pad_token_id. e.g., decoder start id in enc-dec models could be a single pad_token_id, we should count - input_lengths = torch.tensor( - 1 + (input_ids[:, 1:] != pad_token_id).sum(dim=1).type( - torch.IntTensor).to(self.device), - dtype=torch.int32, - device=self.device) - max_input_length = torch.max(input_lengths).item() - if language_adapter_routings is not None: - language_adapter_routings = language_adapter_routings.to( - self.device) - return input_ids, input_lengths, max_input_length, prompt_tasks, language_adapter_routings - - def encoder_run(self, - input_ids, - input_lengths, - max_input_length, - position_ids=None, - token_type_ids=None, - debug_mode=False, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_mask=None, - language_adapter_routings=None): - - # each engine has hidden_dim/TP, don't forget to multiply TP - hidden_size = self.encoder_model_config.hidden_size * self.encoder_runtime_mapping.tp_size - if input_ids.dim() == 1: - hidden_states_shape = (input_ids.shape[0], hidden_size - ) # [num_tokens,D] - else: - hidden_states_shape = (input_ids.shape[0], input_ids.shape[1], - hidden_size) # [BS,seqlen,D] - hidden_states_dtype = lambda name: trt_dtype_to_torch( - self.encoder_session.engine.get_tensor_dtype(name)) - - # input tensors. only first PP rank has id input, others are hidden_states input - inputs = {} - if self.encoder_runtime_mapping.is_first_pp_rank(): - inputs['input_ids'] = input_ids.contiguous() - if self.encoder_model_config.has_position_embedding: - if position_ids is None: - if self.encoder_model_config.remove_input_padding: - position_ids = [ - torch.arange(sample_length, - dtype=torch.int32, - device=input_ids.device) - for sample_length in torch_to_numpy(input_lengths) - ] - position_ids = torch.cat(position_ids) - else: - bsz, seq_len = input_ids.shape[:2] - position_ids = torch.arange( - seq_len, dtype=torch.int32, - device=input_ids.device).expand(bsz, -1) - inputs['position_ids'] = position_ids.contiguous() - if self.encoder_model_config.has_token_type_embedding: - inputs['token_type_ids'] = token_type_ids.contiguous() - - if self.encoder_model_config.max_prompt_embedding_table_size > 0: - inputs[ - 'prompt_embedding_table'] = prompt_embedding_table.contiguous( - ) - inputs['tasks'] = prompt_tasks.contiguous() - inputs['prompt_vocab_size'] = prompt_vocab_size.contiguous() - else: - # just need a placeholder, engine will call NCCL to recv and fill data from previous rank - inputs['hidden_states_input'] = torch.empty( - hidden_states_shape, - dtype=hidden_states_dtype('hidden_states_input'), - device=self.device).contiguous() - if attention_mask is not None and not self.encoder_model_config.gpt_attention_plugin: - inputs['attention_mask'] = attention_mask.contiguous() - - inputs['input_lengths'] = input_lengths - # use shape info to pass max length info in remove padding mode - inputs['max_input_length'] = torch.empty( - (max_input_length, ), - dtype=hidden_states_dtype('max_input_length'), - device=self.device).contiguous() - - if self.encoder_runtime_mapping.tp_size > 1: - ipc_buffers, all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.encoder_runtime_mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.encoder_runtime_mapping.tp_size)) - inputs['all_reduce_workspace'] = all_reduce_workspace - - if self.encoder_model_config.lora_plugin: - inputs.update( - self.encoder_lora_manager.input_buffers( - self.lora_task_uids, - self.encoder_runtime_mapping, - self.encoder_model_config.num_layers, - )) - batch_size = input_lengths.size(0) - inputs['host_request_types'] = torch.IntTensor([0] * - batch_size).to('cpu') - if self.encoder_model_config.remove_input_padding: - inputs['host_context_lengths'] = input_lengths.to('cpu') - if language_adapter_routings is not None: - inputs['language_adapter_routings'] = language_adapter_routings - # Note: runtime.Session's run() method will set input/output tensor address, here we only need to provide tensor shape - self.encoder_session.set_shapes(inputs) - - # output tensors. only last PP rank final encoder output, others are intermediate hidden_states output. Need broadcast later - outputs = {} - if self.encoder_runtime_mapping.is_last_pp_rank(): - outputs['encoder_output'] = torch.empty( - hidden_states_shape, - dtype=hidden_states_dtype('encoder_output'), - device=self.device).contiguous() - else: - outputs['hidden_states_output'] = torch.empty( - hidden_states_shape, - dtype=hidden_states_dtype('hidden_states_output'), - device=self.device).contiguous() - - # ------------------------------------------- - if debug_mode: - engine = self.encoder_session.engine - context = self.encoder_session.context - # setup debugging buffer for the encoder - for i in range(self.encoder_session.engine.num_io_tensors): - name = engine.get_tensor_name(i) - if engine.get_tensor_mode( - name - ) == trt.TensorIOMode.OUTPUT and name not in outputs.keys(): - dtype = engine.get_tensor_dtype(name) - shape = context.get_tensor_shape(name) - outputs[name] = torch.zeros(tuple(shape), - dtype=trt_dtype_to_torch(dtype), - device=self.device) - context.set_tensor_address(name, outputs[name].data_ptr()) - # ------------------------------------------- - - # TRT session run - # Note: need cuda stream ID, not a torch Stream - ok = self.encoder_session.run(inputs, outputs, self.stream.cuda_stream) - assert ok, "Runtime execution failed" - self.stream.synchronize() - - # Tensor Parallelism is handled by model/engine definition - # But we need to broadcast among PP group at the end of encoder's Pipeline Parallelism - # After this, all ranks should recv the encoder output, and world might be re-configured using decoder's TP-PP config - def pp_communicate_encoder_output(encoder_output): - if self.encoder_runtime_mapping.is_last_pp_rank(): - for pp_rank in self.encoder_runtime_mapping.pp_group: - if pp_rank != self.encoder_runtime_mapping.rank: - self.nccl_comm.send(encoder_output, pp_rank) - return encoder_output - else: - self.nccl_comm.recv(encoder_output, - self.encoder_runtime_mapping.pp_group[-1]) - return encoder_output - - if self.encoder_runtime_mapping.has_pp(): - # use hidden_states output buffer to receive output as the shapes are same - encoder_output_buf = outputs[ - 'encoder_output'] if self.encoder_runtime_mapping.is_last_pp_rank( - ) else outputs['hidden_states_output'] - encoder_output = pp_communicate_encoder_output(encoder_output_buf) - else: - encoder_output = outputs['encoder_output'] - - return encoder_output - - def generate(self, - encoder_input_ids, - decoder_input_ids, - max_new_tokens, - num_beams=1, - pad_token_id=None, - eos_token_id=None, - bos_token_id=None, - debug_mode=False, - return_dict=False, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_mask=None, - time_encoder=False, - return_encoder_output=False, - encoder_language_adapter_routings=None, - decoder_language_adapter_routings=None): - ## ensure all externally provided tensors are on the correct device. - encoder_input_ids = encoder_input_ids.to(self.device) - decoder_input_ids = decoder_input_ids.to(self.device) - - if attention_mask is not None: - attention_mask = torch.tensor(attention_mask, - dtype=torch.int32, - device=self.device) - - ## encoder run - encoder_remove_input_padding = self.encoder_model_config.remove_input_padding if self.encoder_model_config else self.decoder_model_config.remove_input_padding - - encoder_input_ids, encoder_input_lengths, encoder_max_input_length, prompt_tasks, encoder_language_adapter_routings = self.process_input( - encoder_input_ids, encoder_remove_input_padding, pad_token_id, - prompt_tasks, encoder_language_adapter_routings) - - if not self.skip_encoder: - logger.info(f"Rank {self.runtime_rank} Running encoder engine ...") - if time_encoder: - tik = time.time() - encoder_output = self.encoder_run( - encoder_input_ids, - encoder_input_lengths, - encoder_max_input_length, - debug_mode=debug_mode, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size, - attention_mask=attention_mask, - language_adapter_routings=encoder_language_adapter_routings) - if time_encoder: - tok = time.time() - print(f"TRT-LLM Encoder time {(tok-tik)*1000}ms") - else: - encoder_output = prompt_embedding_table - if encoder_input_ids.dim() > 1: - encoder_output = encoder_output.unsqueeze(0) - - ## decoder run - logger.info(f"Rank {self.runtime_rank} Running decoder engine ...") - decoder_input_ids, decoder_input_lengths, decoder_max_input_length, _, decoder_language_adapter_routings = self.process_input( - decoder_input_ids, self.decoder_model_config.remove_input_padding, - pad_token_id, None, decoder_language_adapter_routings) - # `cross_attention_mask` in context phase [batch_size, query_len, encoder_input_len] - # where query_len happens to be 1 in current cases, but not necessarily always, and - # `cross_attention_mask` in generation phase [batch_size, 1, encoder_input_len] where - # the query_len is always 1 since we have kv cache. But we use - # cross_attention_mask[:, step, :] during generation - cross_attention_mask = None - if attention_mask is not None: - cross_attention_mask = torch.tensor(attention_mask, - dtype=torch.int32, - device=self.device).reshape( - attention_mask.shape[0], 1, - attention_mask.shape[1]) - cross_attention_mask = cross_attention_mask.repeat( - [1, decoder_max_input_length + max_new_tokens, 1]) - - # generation config - sampling_config = SamplingConfig(end_id=eos_token_id, - pad_id=pad_token_id, - num_beams=num_beams, - min_length=1, - return_dict=return_dict) - sampling_config.update(output_cum_log_probs=return_dict, - output_log_probs=return_dict) - - # decoder autoregressive generation - self.decoder_session.setup( - decoder_input_lengths.size(0), - decoder_max_input_length, - max_new_tokens, - num_beams, - max_attention_window_size=None, - encoder_max_input_length=encoder_max_input_length, - lora_manager=self.decoder_lora_manager, - lora_uids=self.lora_task_uids, - enable_context_fmha_fp32_acc=self.enable_context_fmha_fp32_acc) - - output = self.decoder_session.decode( - decoder_input_ids, - decoder_input_lengths, - sampling_config, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - return_dict=return_dict, - cross_attention_mask=cross_attention_mask, - language_adapter_routings=decoder_language_adapter_routings) - - if return_dict and return_encoder_output: - output['encoder_output'] = encoder_output - - return output diff --git a/tensorrt_llm/runtime/generation.py b/tensorrt_llm/runtime/generation.py deleted file mode 100755 index 6a49587d3054..000000000000 --- a/tensorrt_llm/runtime/generation.py +++ /dev/null @@ -1,4785 +0,0 @@ -# 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 copy -import math -import os -import platform -from collections import Counter -from dataclasses import dataclass, field -from functools import reduce, wraps -from pathlib import Path -from typing import Dict, Iterable, List, Optional, Sequence, Set, Union - -import numpy as np - -# isort: off -import torch -import tensorrt as trt -# isort: on -try: - from cuda.bindings import runtime as cudart -except ImportError: - from cuda import cudart - -from tensorrt_llm.runtime.memory_pools.memory_pools_allocator import \ - MemoryPoolsAllocator -from tensorrt_llm.runtime.memory_pools.pools_kv_cache_manager import \ - PoolsKVCacheManager -from tensorrt_llm.runtime.redrafter_utils import * - -from .._utils import (binding_layer_type_to_str, binding_to_str_dtype, - pad_vocab_size, str_dtype_to_torch, torch_to_numpy, - trt_dtype_to_torch) -from ..bindings import ipc_nvls_allocate, ipc_nvls_free -from ..layers import LanguageAdapterConfig -from ..llmapi.kv_cache_type import KVCacheType -from ..logger import logger -from ..lora_manager import LoraManager -from ..mapping import Mapping -from ..plugin.plugin import CustomAllReduceHelper -from ..quantization import QuantMode -from .kv_cache_manager import GenerationSequence, KVCacheUpdater -from .session import _scoped_stream - -# When variable is set, this will disable torch.cuda.set_device(...) calls -# Useful in situations where device is already assigned by another library, i.e., megatron. -DISABLE_TORCH_DEVICE_SET = os.environ.get("DISABLE_TORCH_DEVICE_SET", False) - - -def decode_words_list(word_dict: List[List[str]], - tokenizer=None, - add_special_tokens=False): - ''' - format of word_dict - len(word_dict) should be same to batch_size - word_dict[i] means the words for batch i - len(word_dict[i]) >= 1, which means it must contain at least 1 string - For example, word_dict[2] = [" I am happy", " I am sad"]. - ''' - assert tokenizer != None, "need to set tokenizer" - - decoded_words_batch = [] - for word_dict_item in word_dict: - decoded_words_request = [] - - for item in word_dict_item: - if isinstance(item, bytes): - item = [item.decode()] - - ids = tokenizer.encode(item, add_special_tokens=add_special_tokens) - - if len(ids) == 0: - continue - - decoded_words_request.append(ids) - decoded_words_batch.append(decoded_words_request) - - return decoded_words_batch - - -def to_word_list_format(word_dict: List[List[List[int]]]): - ''' - format of word_dict - len(word_dict) should be same to batch_size - word_dict[i] means the words for batch i - len(word_dict[i]) >= 1, which means it must contain at least 1 word - For example, word_dict[2] = [[1, 267], [534]] has two words. - ''' - - flat_ids = [] - offsets = [] - for word_dict_item in word_dict: - items_flat_ids = [] - items_offsets = [] - - for ids in word_dict_item: - items_flat_ids += ids - items_offsets.append(len(ids)) - - flat_ids.append(np.array(items_flat_ids)) - offsets.append(np.cumsum(np.array(items_offsets))) - - pad_to = max(1, max(len(ids) for ids in flat_ids)) - - for i, (ids, offs) in enumerate(zip(flat_ids, offsets)): - flat_ids[i] = np.pad(ids, (0, pad_to - len(ids)), constant_values=0) - offsets[i] = np.pad(offs, (0, pad_to - len(offs)), constant_values=-1) - - return np.array([flat_ids, offsets], dtype="int32").transpose((1, 0, 2)) - - -def _prepare_input_ids(tensors: Sequence[torch.Tensor]): - tensors = [torch.flatten(t) for t in tensors] - data = torch.concat(tensors) - row_lengths = [t.size(0) for t in tensors] - row_lengths = torch.tensor(row_lengths, - dtype=torch.int32, - device=data.device) - return (data, row_lengths) - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -def _update_cuda_graph_instance(instance, graph): - err = cudart.cudaGraphExecUpdate(instance, graph) - if err != cudart.cudaError_t.cudaSuccess: - # When updating cuda graph failed, destroy and instantiate one. - CUASSERT(cudart.cudaGraphExecDestroy(instance)) - instance = CUASSERT(cudart.cudaGraphInstantiate(graph, 0))[0] - return instance - - -def _prepare_attention_mask(input_ids: torch.Tensor, - pad_id: Optional[int] = None): - is_pad_id_in_inputs = (pad_id is not None) and (pad_id in input_ids) - if input_ids is not None and is_pad_id_in_inputs: - mask = input_ids.ne(pad_id).int() - # for enc-dec models, pad_id could be the start token and should be always counted - # as valid token rather than padded token, so we force its mask to be 1. - # This doesn't impact the existing behavior - mask[:, 0] = 1 - return mask - else: - return torch.ones(input_ids.shape, - dtype=torch.int32, - device=input_ids.device) - - -def _tile_beam_width(tensor: torch.Tensor, num_beams: int): - new_shape = np.array(tensor.shape) - new_shape[0] = new_shape[0] * num_beams - - tile_size = np.ones(new_shape.shape, dtype=np.int32) - tile_size = np.insert(tile_size, 1, num_beams) - - new_tensor = torch.unsqueeze(tensor, 1) - new_tensor = new_tensor.tile(tile_size.tolist()) - new_tensor = new_tensor.reshape(new_shape.tolist()) - return new_tensor - - -class _Profiler(trt.IProfiler): - - def __init__(self): - super().__init__() - self.results = [] - - def report_layer_time(self, layer_name, ms): - self.results.append((layer_name, ms)) - - -def _contiguous_tile_beam_width(tensor: torch.Tensor, size: int, - num_beams: int): - new_shape = list(tensor.shape) - new_shape[0] *= num_beams - - numel = tensor.numel() - new_tensor = torch.empty(num_beams * numel, - device=tensor.device, - dtype=tensor.dtype) - - # Take the first 'size' values to tile and skip the others. - vals = tensor.view(-1)[:size] - for i in range(num_beams): - new_tensor[i * size:(i + 1) * size] = vals - - return new_tensor.view(new_shape) - - -class _Runtime(object): - runtime_rank: int - runtime: trt.Runtime - engine: trt.ICudaEngine - ctx_context: trt.IExecutionContext - context_0: trt.IExecutionContext - context_1: trt.IExecutionContext - profiler: _Profiler - engine_inspector: trt.EngineInspector - cuda_graph_instances: List[cudart.cudaGraphExec_t] - input_tensor_names: Set[str] - output_tensor_names: Set[str] - - def __init__(self, engine_buffer, mapping: Mapping): - self.address = None - self.device_memory_size = 0 - self.__prepare(mapping, engine_buffer) - - def _serialize_engine(self) -> trt.IHostMemory: - return self.engine.serialize() - - def __create_and_setup_context(self, address, size, profile_idx, - stream) -> trt.IExecutionContext: - context = self.engine.create_execution_context_without_device_memory() - assert context is not None, "Failed to create an execution context with the provided device memory!" - context.set_device_memory(address, size) - context.set_optimization_profile_async(profile_idx, stream) - # If nvtx verbosity is DETAILED, change it to LAYER_NAMES_ONLY for inference performance - if context.nvtx_verbosity == trt.ProfilingVerbosity.DETAILED: - context.nvtx_verbosity = trt.ProfilingVerbosity.LAYER_NAMES_ONLY - return context - - def _set_profiler(self): - if self.profiler is not None: - return - assert self.context_0 is not None - assert self.context_1 is not None - self.profiler = _Profiler() - self.context_0.profiler = self.profiler - self.context_0.enqueue_emits_profile = False - self.context_1.profiler = self.profiler - self.context_1.enqueue_emits_profile = False - if self.engine.num_optimization_profiles == 2: - assert self.ctx_context is not None - self.ctx_context.profiler = self.profiler - self.ctx_context.enqueue_emits_profile = False - - def __prepare(self, mapping: Mapping, engine_buffer): - self.runtime_rank = mapping.rank - local_rank = self.runtime_rank % mapping.gpus_per_node - if DISABLE_TORCH_DEVICE_SET: - CUASSERT(cudart.cudaSetDevice(torch.cuda.current_device())) - else: - torch.cuda.set_device(local_rank) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.runtime = trt.Runtime(logger.trt_logger) - self.engine = self.runtime.deserialize_cuda_engine(engine_buffer) - assert self.engine is not None - - self.input_tensor_names = set() - self.output_tensor_names = set() - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT: - self.output_tensor_names.add(name) - else: - self.input_tensor_names.add(name) - - self.profiler = None - self.engine_inspector = self.engine.create_engine_inspector() - # cuda graph ping-pong instances - self.cuda_graph_instances = [None for _ in range(2)] - if not self.engine.streamable_weights_size: - # engine does not have weight streaming enabled - self.__prepare_execution_contexts() - else: - self.engine.weight_streaming_budget_v2 = 0 # avoid OOM when print engine info - - if logger.level == "verbose": - self.__print_engine_info() - - def __prepare_execution_contexts(self): - self.context_0 = None - self.context_1 = None - self.ctx_context = None - - # The device_memory_size_v2 stores the memory required by the largest profile. - # When weight streaming is enable, it must be queried after the weight streaming budget set. - if self.address: - if self.device_memory_size != self.engine.device_memory_size_v2: - self.device_memory_size = self.engine.device_memory_size_v2 - CUASSERT(cudart.cudaFree(self.address)) - address = CUASSERT(cudart.cudaMalloc( - self.device_memory_size))[0] - self.address = address - else: - self.device_memory_size = self.engine.device_memory_size_v2 - address = CUASSERT(cudart.cudaMalloc(self.device_memory_size))[0] - self.address = address - - with _scoped_stream() as stream: - if self.engine.num_optimization_profiles == 1: - # At step = 0, context_1 is active - # At step = 1, context_0 is active - # At step = 2, context_1 is active - self.context_0 = self.__create_and_setup_context( - self.address, self.device_memory_size, 0, stream) - self.context_1 = self.__create_and_setup_context( - self.address, self.device_memory_size, 0, stream) - self.ctx_context = self.context_1 - elif self.engine.num_optimization_profiles == 2: - # At step = 0, ctx_context is active - # At step = 1, context_0 is active - # At step = 2, context_1 is active - self.ctx_context = self.__create_and_setup_context( - self.address, self.device_memory_size, 0, stream) - self.context_0 = self.__create_and_setup_context( - self.address, self.device_memory_size, 1, stream) - self.context_1 = self.__create_and_setup_context( - self.address, self.device_memory_size, 1, stream) - else: - logger.error( - f"Number of optimization profiles: {self.engine.num_optimization_profiles}" - ) - raise NotImplementedError( - "Python runtime only support 1 or 2 optimization profiles, " - "set --multiple_profiles=disable when calling trtllm-build " - "to disable the feature.") - - def __print_engine_info(self) -> None: - engine = self.engine - context = engine.create_execution_context( - trt.ExecutionContextAllocationStrategy.USER_MANAGED) - n_op = engine.num_optimization_profiles - max_name_width = 0 # Maximum Width of tensor Name - max_shape_width = 0 # Maximum Width of tensor Shape - tensor_name_list = [ - engine.get_tensor_name(i) for i in range(engine.num_io_tensors) - ] - - # Get information of engine input / output - tid = {} # Tensor Information Dictionary - for name in tensor_name_list: - item = dict() - max_name_width = max(max_name_width, len(name)) - item["mode"] = 'I' if engine.get_tensor_mode( - name) == trt.TensorIOMode.INPUT else 'O' - item["location"] = 'GPU' if engine.get_tensor_location( - name) else 'CPU' - item["data_type"] = str(engine.get_tensor_dtype(name))[9:] - item["build_shape"] = str(engine.get_tensor_shape(name)) - item["profile_list"] = [[] for _ in range(n_op)] - if item["mode"] == "I": - for k in range(n_op): - if item["location"] == "GPU": - shape = engine.get_tensor_profile_shape(name, k) - else: - shape = engine.get_tensor_profile_value(k, name) - item["profile_list"][k].extend(shape) - max_shape_width = max(max_shape_width, - *[len(str(s)) for s in shape]) - tid[name] = item - # Set input shape to get output shape - for k in range(n_op): - for j in range(3): # Min, Opt, Max - for name in tid.keys(): - if tid[name]["mode"] == "I": - if tid[name]["location"] == "GPU": - context.set_input_shape( - name, tid[name]["profile_list"][k][j]) - else: - context.set_tensor_address( - name, - tid[name]["profile_list"][k][j].ctypes.data) - elif tid[name]["mode"] == "O": - assert context.all_binding_shapes_specified and context.all_shape_inputs_specified - shape = context.get_tensor_shape(name) - tid[name]["profile_list"][k].append(shape) - - # Print information of engine input / output - logger.debug("Information of engine input / output.") - logger.debug(f"{'='*(max_name_width + max_shape_width + 24)}") - logger.debug( - f"{'Name':^{max_name_width}}|I/O|Location|DataType|{'Shape':^{max_shape_width}}|" - ) - logger.debug(f"{'-'*(max_name_width + max_shape_width + 24)}") - for name in tensor_name_list: - item = tid[name] - info = f"{name:<{max_name_width}}|{item['mode']:^3s}|{item['location']:^8s}|{item['data_type']:^8s}|" - info += f"{item['build_shape']:^{max_shape_width}}|" - logger.debug(info) - logger.debug(f"{'='*(max_name_width + max_shape_width + 24)}") - # Print information of optimization profile - logger.debug("Information of optimization profile.") - for k in range(n_op): - logger.debug(f"Optimization Profile {k}:") - logger.debug(f"{'='*(max_name_width + max_shape_width * 3 + 4)}") - logger.debug( - f"{'Name':^{max_name_width}}|{'Min':^{max_shape_width}}|{'Opt':^{max_shape_width}}|{'Max':^{max_shape_width}}|" - ) - logger.debug(f"{'-'*(max_name_width + max_shape_width * 3 + 4)}") - for name in tensor_name_list: - item = tid[name] - info = f"{name:<{max_name_width}}|" - info += f"{str(item['profile_list'][k][0]):^{max_shape_width}}|" - info += f"{str(item['profile_list'][k][1]):^{max_shape_width}}|" - info += f"{str(item['profile_list'][k][2]):^{max_shape_width}}|" - logger.debug(info) - logger.debug(f"{'='*(max_name_width + max_shape_width * 3 + 4)}") - - def print_context_info(self, context, context_index) -> None: - n_io = self.engine.num_io_tensors - max_name_width = 0 # Maximum Width of tensor Name - max_shape_width = 0 # Maximum Width of tensor Shape - tensorInfo = {} - for i in range(n_io): - name = self.engine.get_tensor_name(i) - b_input = self.engine.get_tensor_mode( - name) == trt.TensorIOMode.INPUT - shape = str(self.engine.get_tensor_shape(name)) - tensorInfo[i] = [name, b_input, shape] - max_name_width = max(max_name_width, len(name)) - max_shape_width = max(max_shape_width, len(shape)) - # Shape input tensor is not used in TRT-LLM yet - - logger.debug(f"Information of context input / output.") - logger.debug(f"Using Optimization Profile: {context_index}") - logger.debug(f"{'='*(max_name_width + max_shape_width + 6)}") - logger.debug( - f"{'Name':^{max_name_width}}|I/O|{'Shape':^{max_shape_width}}|") - logger.debug(f"{'-'*(max_name_width + max_shape_width + 6)}") - for i in range(n_io): - name, b_input, shape = tensorInfo[i] - info = f"{name:<{max_name_width}}|{'I' if b_input else 'O':^3s}|{shape:^{max_shape_width}}|" - logger.debug(info) - logger.debug(f"{'='*(max_name_width + max_shape_width + 6)}") - - def _set_shape(self, context: trt.IExecutionContext, - shape_dict: Dict[str, List[int]]): - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - if name not in shape_dict: - # shape and buffer can be set by calling _set_tensors API - continue - if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: - ok = context.set_input_shape(name, shape_dict[name]) - dtype = self.engine.get_tensor_dtype(name) - logger.debug( - f"setting input tensor {name} with shape {shape_dict[name]} and type {dtype}" - ) - if not ok: - raise ValueError( - f"Couldn't assign {name} with shape {shape_dict[name]}, " - f"engine supports [min, opt, max] = {self.engine.get_tensor_profile_shape(name, context.active_optimization_profile)}" - ) - - def _set_buffer(self, context: trt.IExecutionContext, - buffer_dict: Dict[str, torch.Tensor]): - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - if name not in buffer_dict.keys(): - dtype = self.engine.get_tensor_dtype(name) - shape = context.get_tensor_shape(name) - buffer_dict[name] = torch.zeros(tuple(shape), - dtype=trt_dtype_to_torch(dtype), - device='cuda') - assert buffer_dict[name].is_contiguous( - ), f"{name} is not contiguous()" - context.set_tensor_address(name, buffer_dict[name].data_ptr()) - - def _set_tensors(self, context: trt.IExecutionContext, - tensors: Dict[str, "RuntimeTensor"]): - for name in self.input_tensor_names: - # it's allowed to call set_tensors multi times with different tensors - # each time only set some of the engine tensors, so it is valid to skip the ones not in the current given tensors dict - if name not in tensors: - continue - - tensor = tensors[name] - if context.get_tensor_address(name) != tensor.data: - context.set_tensor_address(name, tensor.data) - - if list(context.get_tensor_shape(name)) != tensor.shape: - context.set_input_shape(name, tensor.shape) - - for name in self.output_tensor_names: - if name not in tensors: - dtype = self.engine.get_tensor_dtype(name) - shape = context.get_tensor_shape(name) - tensors[name] = RuntimeTensor.from_torch( - name, - torch.zeros(tuple(shape), - dtype=trt_dtype_to_torch(dtype), - device='cuda')) - t = tensors[name] - # output's shape is inference by TRT, no need to set the shape here - context.set_tensor_address(t.name, t.data) - - def _set_weight_streaming(self, gpu_weights_percent): - if not self.engine.streamable_weights_size: - assert gpu_weights_percent == 1, "Engine built without weight streaming. Cannot set gpu_weights_percent to a value other than 1." - return - - assert self.engine is not None - self.context_0 = None - self.context_1 = None - self.ctx_context = None - - min = 0 - max = self.engine.streamable_weights_size - budget = int(gpu_weights_percent * max) - self.engine.weight_streaming_budget_v2 = budget - assert self.engine.weight_streaming_budget_v2 == budget, "Failed to set weight streaming budget!" - logger.info( - f"Set gpu weights percent to {gpu_weights_percent}, which is {budget} bytes. Valid range: {min} bytes ~ {max} bytes." - ) - - try: - self.__prepare_execution_contexts() - except: - free_mem = torch.cuda.mem_get_info()[0] - if free_mem < budget: - print( - f"Failed to create context. Possibly out of memory: Memory budget is {budget} bytes but only {free_mem} bytes are available on the GPU." - ) - raise - - def _check_tensors(self, context: trt.IExecutionContext) -> None: - tensors = [] - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - ptr = context.get_tensor_address(name) - if ptr == 0: - raise RuntimeError(f"Engine I/O tensor {name} is unbound") - shp = list(context.get_tensor_shape(name)) - if any([s < 0 for s in shp]): # skip if shape is not available - continue - dt = self.engine.get_tensor_dtype(name) - tdt = trt_dtype_to_torch(dt) - sz = torch.tensor([], dtype=tdt).element_size() * np.prod(shp) - tensors.append((ptr, ptr + sz, name, shp, sz)) - tensors.sort() # sort by start address - starts, ends, names, _, _ = zip(*tensors) - starts = torch.tensor(starts) - ends = torch.tensor(ends) - overalps = (torch.nonzero((starts[1:] < ends[:-1]).int()) + 1).squeeze() - if overalps.ndim == 0: - # unsqueeze if there is a single value so it became scalar - overalps = torch.unsqueeze(overalps, 0) - if overalps.numel() > 0: - assert overalps.ndim == 1 - for i in list(overalps): - left_name = names[i] - right_name = names[i - 1] - if "key_value" in left_name and "key_value" in right_name: # kv - left_names = left_name.split("_") - right_names = right_name.split("_") - if left_names[-1] == right_names[-1]: # same kv layer - assert (left_names[0] == "past" and right_names[0] == "present") or ( - left_names[0] == "present" and right_names[0] == "past"), \ - f"Overlap found between {tensors[i]} and {tensors[i-1]}" - continue - logger.warning( - f"TENSOR BUFFER OVERLAP DETECTED: {tensors[i]} and {tensors[i-1]} !!!" - ) - return - - def _insert_step_to_profiler(self, step: int): - if not self.profiler: - raise RuntimeError("Profiler is disable") - self.profiler.results.append(("step", step)) - - def _is_profiling(self): - return self.profiler is not None - - def _run(self, - context: trt.IExecutionContext, - stream: Union[int, torch.cuda.Stream] = None) -> bool: - if stream is None: - stream = torch.cuda.current_stream().cuda_stream - elif isinstance(stream, torch.cuda.Stream): - stream = stream.cuda_stream - ok = context.execute_async_v3(stream) - return ok - - def __del__(self): - try: - if self.address is not None: - cudart.cudaFree(self.address) - except TypeError: - pass - - @property - def context_mem_size(self) -> int: - return self.engine.device_memory_size_v2 - - -@dataclass -class ModelConfig: - max_batch_size: int - max_beam_width: int - vocab_size: int - num_layers: int - num_heads: int - num_kv_heads: int - hidden_size: int - gpt_attention_plugin: bool - gemm_allreduce_plugin: str = None - remove_input_padding: bool = False - model_name: str = "" - kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS - cross_attention: bool = False - head_size: int = None - has_position_embedding: bool = True - has_token_type_embedding: bool = False - tokens_per_block: int = 32 - max_prompt_embedding_table_size: int = 0 - quant_mode: QuantMode = QuantMode(0) - gather_context_logits: bool = False - gather_generation_logits: bool = False - dtype: str = "" - lora_plugin: bool = False - lora_target_modules: List[str] = field(default_factory=list) - trtllm_modules_to_hf_modules: dict = None - skip_cross_kv: bool = False - num_medusa_heads: int = 0 - max_medusa_tokens: int = 0 - paged_state: bool = True - mamba_conv1d_plugin: bool = True - conv_kernel: int = 0 - layer_types: List[str] = field(default_factory=list) - rnn_hidden_size: int = 0 - rnn_head_size: int = 0 - rnn_conv_dim_size: int = 0 - state_size: int = 0 - state_dtype: str = "" - gpu_weights_percent: float = 1.0 - # ReDrafter - redrafter_num_beams: int = 0 - redrafter_draft_len_per_beam: int = 0 - num_kv_heads_per_layer: Optional[List[int]] = None - num_kv_heads_per_cross_attn_layer: Optional[List[int]] = None - skip_cross_attn_blocks: bool = False - # language adapter - language_adapter_config: Optional[LanguageAdapterConfig] = None - - @classmethod - def from_model_config_cpp(cls, model_config_cpp) -> 'ModelConfig': - """Create a partially initialized ModelConfig instance from a given ModelConfig CPP binding instance. - - Note that each of these classes have fields that don't exist in the other, so the created ModelConfigPython - won't have all of its fields initialized. - """ - return cls( - max_batch_size=model_config_cpp.max_batch_size, - max_beam_width=model_config_cpp.max_beam_width, - vocab_size=model_config_cpp.vocab_size, - num_layers=model_config_cpp.num_layers(), - num_heads=model_config_cpp.num_heads, - num_kv_heads=model_config_cpp.num_kv_heads(0), - hidden_size=model_config_cpp.hidden_size, - remove_input_padding=model_config_cpp.use_packed_input, - kv_cache_type=model_config_cpp.kv_cache_type, - cross_attention=model_config_cpp.use_cross_attention, - head_size=model_config_cpp.head_size, - max_prompt_embedding_table_size=model_config_cpp. - max_prompt_embedding_table_size, - quant_mode=QuantMode(model_config_cpp.quant_mode.value), - gather_context_logits=model_config_cpp.compute_context_logits, - gather_generation_logits=model_config_cpp.compute_generation_logits, - gpt_attention_plugin=model_config_cpp.use_gpt_attention_plugin, - dtype=binding_to_str_dtype(model_config_cpp.data_type), - num_kv_heads_per_layer=model_config_cpp.num_kv_heads_per_layer, - tokens_per_block=model_config_cpp.tokens_per_block, - lora_plugin=model_config_cpp.use_lora_plugin, - layer_types=[ - binding_layer_type_to_str(lt) - for lt in model_config_cpp.layer_types - ], - ) - - -@dataclass -class SamplingConfig: - end_id: int - pad_id: int - - max_new_tokens: int = field(default=20) - num_beams: int = field(default=1) - num_return_sequences: Optional[int] = field(default=None) - max_attention_window_size: Optional[int] = field(default=None) - sink_token_length: Optional[int] = field(default=None) - output_sequence_lengths: bool = field(default=False) - return_dict: bool = field(default=False) - stop_words_list: Optional[Union[list, np.ndarray, - torch.Tensor]] = field(default=None) - bad_words_list: Optional[Union[list, np.ndarray, - torch.Tensor]] = field(default=None) - - temperature: Union[float, torch.Tensor] = field(default=1.0) - top_k: Union[int, torch.Tensor] = field(default=1) - top_p: Union[float, torch.Tensor] = field(default=0.0) - top_p_decay: Optional[torch.Tensor] = field(default=None) # float - top_p_min: Optional[torch.Tensor] = field(default=None) # float - top_p_reset_ids: Optional[torch.Tensor] = field(default=None) # int - random_seed: Union[int, torch.Tensor] = field(default=None) - - length_penalty: Union[float, torch.Tensor] = field(default=1.0) - early_stopping: Union[int, torch.Tensor] = field(default=1) - repetition_penalty: Union[float, torch.Tensor] = field(default=1.0) - min_length: Union[int, torch.Tensor] = field(default=1) - presence_penalty: Union[float, torch.Tensor] = field(default=0.0) - frequency_penalty: Union[float, torch.Tensor] = field(default=0.0) - prompt_ignore_length: Union[int, torch.Tensor] = field(default=0) - use_beam_hyps: bool = field(default=True) - - # None here means user didn't set it, and dynamicDecodeOp.cpp take optional value - # The real default value is set in dynamicDecodeOp.cpp when it's None - beam_search_diversity_rate: Union[float, torch.Tensor] = field(init=False, - default=0.0) - output_cum_log_probs: bool = field(init=False, default=False) - output_log_probs: bool = field(init=False, default=False) - no_repeat_ngram_size: Union[int, torch.Tensor] = field(init=False, - default=None) - min_p: Union[float, torch.Tensor] = field(default=0.0) - - def update(self, **kwargs): - unused_kwargs = dict() - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - unused_kwargs[key] = value - return unused_kwargs - - -class LogitsProcessor: - """ - Base class for all logit processors that can be applied during generation. - """ - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> torch.Tensor: - raise NotImplementedError( - f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." - ) - - -class LogitsProcessorList(list, LogitsProcessor): - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> torch.Tensor: - for processor in self: - scores = processor(step, input_ids, scores) - return scores - - -class StoppingCriteria: - """ - Base class for all stopping criteria that can be applied during generation. - """ - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> bool: - raise NotImplementedError("StoppingCriteria needs to be subclassed") - - -class StoppingCriteriaList(list, StoppingCriteria): - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> bool: - return any(criteria(step, input_ids, scores) for criteria in self) - - -class RuntimeTensor: - - def __init__(self): - self._name = "" - # shape is the one sent to TRT, the actual torch tensor can be larger than the shape - # this is useful when allocating a big KV cache tensor at the beginning and incremental seq length dim of TRT engine's input tensor - self._shape = None - self._torch_tensor = None - # Used when pointer specified - self._data_ptr = None - self._dtype = None - - @staticmethod - def from_pointer(name: str, pointer, shape, - str_dtype: str) -> 'RuntimeTensor': - t = RuntimeTensor() - t._name = name - t._data_ptr = pointer - t._shape = shape - t._dtype = str_dtype_to_torch(str_dtype) - return t - - @staticmethod - def from_torch( - name: str, - data: torch.Tensor, - override_shape: Optional[Iterable] = None) -> 'RuntimeTensor': - assert (isinstance(data, torch.Tensor)), f"data {name} is {type(data)}" - t = RuntimeTensor() - t._name = name - # need to hold the torch tensor for memory life time - t._torch_tensor = data.contiguous() - t._dtype = t._torch_tensor.dtype - t._data_ptr = t._torch_tensor.data_ptr() - torch_shape = list(data.size()) - if override_shape is not None: - t._shape = override_shape - assert isinstance(override_shape, list) or isinstance( - override_shape, tuple) - assert all([lambda x: x >= 0 for x in override_shape - ]), f"Expect all dimensions >=0, got {override_shape}" - - def volume_func(dims): - return reduce(lambda x, y: x * y, dims, 1) - assert volume_func(override_shape) <= volume_func(torch_shape), \ - f"Override the shape to be larger than the underlying torch Tensor, got {override_shape}, torch tensor shape {torch_shape}" - else: - t._shape = torch_shape - return t - - def to_torch(self) -> torch.Tensor: - if self._torch_tensor is None: - raise RuntimeError( - 'RuntimeTensor cannot be converted to torch tensor as constructed from pointer' - ) - return self._torch_tensor - - @property - def shape(self) -> Iterable[int]: - return self._shape - - @property - def data(self): - return self._data_ptr - - @property - def name(self) -> str: - return self._name - - @property - def dtype(self) -> torch.dtype: - return self._dtype - - -class GenerationSession(object): - - _model_config: ModelConfig - mapping: Mapping - runtime: _Runtime - device: torch.device - batch_size: int - buffer_allocated: bool - debug_mode: bool - quant_mode: QuantMode - cuda_graph_mode: bool - dtype: trt.DataType - debug_tensors_to_save: None - num_draft_tokens: int = 0 - medusa_topks: List[int] = None - medusa_paths: List[List[int]] = None - medusa_tree_ids: List[int] = None - medusa_position_offsets: List[int] = None - medusa_temperature: float = 0.0 - - def __init__(self, - model_config: ModelConfig, - engine_buffer, - mapping: Mapping, - debug_mode=False, - debug_tensors_to_save=None, - cuda_graph_mode=False, - stream: torch.cuda.Stream = None): - assert isinstance(model_config, ModelConfig) - self._model_config = model_config - self.mapping = mapping - self.runtime = _Runtime(engine_buffer, mapping) - if DISABLE_TORCH_DEVICE_SET: - self.device = torch.device(f'cuda:{torch.cuda.current_device()}') - else: - self.device = torch.device( - f'cuda:{self.runtime.runtime_rank % mapping.gpus_per_node}') - torch.cuda.set_device(self.device) - # dynamic_decoder currently use torch's current stream, so must let TRT enqueue use same stream here - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - self.debug_mode = debug_mode - self.debug_tensors_to_save = debug_tensors_to_save - - self.cuda_graph_mode = cuda_graph_mode - # Optional inputs for dynamic decoder - self.top_p_decay = None - self.top_p_min = None - self.top_p_reset_ids = None - # TODO: in tensorrt_llm/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp it's T, can be float or half? - self.embedding_bias_opt = None - - self.buffer = None - self.buffer_allocated = False - - self.vocab_size_padded = pad_vocab_size(self.vocab_size, - self.mapping.tp_size) - if len(model_config.layer_types) == 0: - self.layer_types = ['attention'] * model_config.num_layers - else: - layer_types = model_config.layer_types - layer_types = layer_types * (model_config.num_layers // - len(layer_types)) - layer_types = layer_types + layer_types[0:(model_config.num_layers % - len(layer_types))] - self.layer_types = layer_types - self.num_attn_layers = \ - self.layer_types[self.first_layer:self.last_layer].count('attention') - self.has_attn_layers = self.num_attn_layers > 0 - self.has_rnn_layers = 'recurrent' in self.layer_types[ - self.first_layer:self.last_layer] - - self.attn_to_general_idx = {} - self.general_to_attn_idx = {} - attn_layer_idx = 0 - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - self.attn_to_general_idx[attn_layer_idx] = i - self.general_to_attn_idx[i] = attn_layer_idx - attn_layer_idx += 1 - - # Cyclic KV cache buffer names. - if self.attn_to_general_idx: - self.kv_cache_buffer_names = [ - f'present_key_value_{layer_idx}' - for _, layer_idx in self.attn_to_general_idx.items() - ] + [f'1_present_key_value_{self.attn_to_general_idx[0]}'] - else: - self.kv_cache_buffer_names = [] - - if self.paged_kv_cache: - logger.warning( - "The paged KV cache in Python runtime is experimental. For performance and correctness, please, use C++ runtime." - ) - - if self.mapping.has_pp(): - self.nccl_comm = torch.classes.trtllm.NcclCommunicatorOp( - self.mapping.world_size, self.mapping.rank) - - if self.mapping.is_last_pp_rank(): - self.decoder_logits_dtype = self._tensor_dtype('logits') - if self.decoder_logits_dtype not in [torch.float16, torch.float32]: - logger.warning( - "Logits dtype not supported by decoder. Falling back to float32. You may want to change the logits dtype to float16 in your model definition." - ) - self.decoder_logits_dtype = torch.float32 - self.dynamic_decoder = torch.classes.trtllm.DynamicDecodeOp( - model_config.max_batch_size, model_config.max_beam_width, - self.vocab_size, self.vocab_size_padded, self.mapping.tp_size, - self.mapping.pp_size, self.decoder_logits_dtype) - - expected_tensor_names = [] - if self.mapping.tp_size > 1: - self.ipc_buffers, self.all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.mapping.tp_size)) - - self.gather_tree = torch.ops.tensorrt_llm.gather_tree - - if self.mapping.is_first_pp_rank(): - expected_tensor_names += ['input_ids'] - else: - expected_tensor_names += ['hidden_states_input'] - - if self.mapping.is_last_pp_rank(): - expected_tensor_names += ['logits'] - if not model_config.gather_context_logits or self.has_rnn_layers: - expected_tensor_names += ['last_token_ids'] - else: - expected_tensor_names += ['hidden_states_output'] - - if self.has_attn_layers: - if model_config.has_position_embedding and self.mapping.is_first_pp_rank( - ): - expected_tensor_names += ['position_ids'] - if model_config.has_token_type_embedding and self.mapping.is_first_pp_rank( - ): - expected_tensor_names += ['token_type_ids'] - - if self.use_kv_cache: - expected_tensor_names += ['cache_indirection'] - - if self.paged_kv_cache and self.has_attn_layers: - expected_tensor_names += [f'kv_cache_block_offsets'] - expected_tensor_names += [f'host_kv_cache_block_offsets'] - expected_tensor_names += [f'host_kv_cache_pool_pointers'] - expected_tensor_names += [f'host_kv_cache_pool_mapping'] - if self.cross_attention: - expected_tensor_names += [f'cross_kv_cache_block_offsets'] - expected_tensor_names += [f'host_cross_kv_cache_block_offsets'] - expected_tensor_names += [f'host_cross_kv_cache_pool_pointers'] - expected_tensor_names += [f'host_cross_kv_cache_pool_mapping'] - expected_tensor_names += [f'cross_attention_mask'] - expected_tensor_names += [f'cross_attention_packed_mask'] - else: - # Refer to gpt_attention() inside functional.py - if self.use_kv_cache and not self.paged_kv_cache: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - expected_tensor_names += [ - f'past_key_value_{i}', f'present_key_value_{i}' - ] - if model_config.cross_attention: - if model_config.gpt_attention_plugin: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - expected_tensor_names += [ - f'cross_present_key_value_{i}', - f'cross_past_key_value_{i}' - ] - expected_tensor_names += [ - 'cross_attention_mask', - ] - expected_tensor_names += [f'cross_attention_packed_mask'] - else: - expected_tensor_names += [ - 'cross_attention_mask', - ] - - if self.paged_state and self.has_rnn_layers: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'recurrent': - expected_tensor_names += [ - f'conv_state_ptr_{i}', f'rnn_state_ptr_{i}' - ] - expected_tensor_names += ['slot_mapping'] - else: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'recurrent': - expected_tensor_names += [ - f'past_conv_state_{i}', f'present_conv_state_{i}', - f'past_rnn_state_{i}', f'present_rnn_state_{i}' - ] - - if model_config.gpt_attention_plugin and self.has_attn_layers: - if self.use_kv_cache: - expected_tensor_names += [ - 'sequence_length', 'host_past_key_value_lengths' - ] - - expected_tensor_names += [ - 'context_lengths', 'host_request_types', - 'host_sink_token_length', 'host_runtime_perf_knobs', - 'host_context_progress' - ] - expected_tensor_names += [f'host_max_attention_window_sizes'] - if model_config.remove_input_padding: - expected_tensor_names.append('host_context_lengths') - else: - if self.has_rnn_layers: - expected_tensor_names += ['host_request_types'] - if model_config.mamba_conv1d_plugin and model_config.remove_input_padding: - expected_tensor_names.append('host_context_lengths') - if self.has_attn_layers: - expected_tensor_names += ['attention_mask'] - - if model_config.max_prompt_embedding_table_size > 0: - expected_tensor_names += [ - 'prompt_embedding_table', 'tasks', 'prompt_vocab_size' - ] - - if model_config.cross_attention: - expected_tensor_names += [ - 'encoder_output', - 'encoder_input_lengths', - 'encoder_max_input_length', - 'cross_kv_cache_gen', - ] - if model_config.skip_cross_attn_blocks: - expected_tensor_names += ['skip_cross_attn_blocks'] - self.skip_cross_kv = model_config.skip_cross_kv - if self.skip_cross_kv: - expected_tensor_names += ['cross_kv_reuse'] - - if self.mapping.tp_size > 1: - expected_tensor_names += ['all_reduce_workspace'] - - self.lora_target_modules = model_config.lora_target_modules - self.missing_qkv_modules = LoraManager.get_missing_qkv_modules( - self.lora_target_modules) - if model_config.lora_plugin: - for lora_module in (self.lora_target_modules + - self.missing_qkv_modules): - for i in range(self.first_layer, self.last_layer): - expected_tensor_names += [ - f'{lora_module}_lora_ranks_{i}', - f'{lora_module}_lora_weights_pointers_{i}' - ] - if self.cross_attention and self.remove_input_padding: - expected_tensor_names += ['host_encoder_input_lengths'] - - if model_config.num_medusa_heads > 0: - expected_tensor_names += [ - 'spec_decoding_generation_lengths', - 'spec_decoding_position_offsets', 'spec_decoding_packed_mask', - 'spec_decoding_use', 'medusa_logits' - ] - - if self.is_redrafter_mode: - expected_tensor_names += get_redrafter_tensor_names() - - # language adapter - if model_config.language_adapter_config: - expected_tensor_names += ['language_adapter_routings'] - - found_tensor_names = [ - self.runtime.engine.get_tensor_name(i) - for i in range(self.runtime.engine.num_io_tensors) - ] - for name in found_tensor_names: - if name.startswith("allreduce_ub_") or name.startswith( - "gemm_allreduce"): - expected_tensor_names += [name] - if not self.debug_mode and set(expected_tensor_names) != set( - found_tensor_names): - logger.error( - f"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}" - ) - logger.error( - f"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}" - ) - logger.error(f"Expected tensor names: {expected_tensor_names}") - logger.error(f"Found tensor names: {found_tensor_names}") - raise RuntimeError( - "Tensor names in engine are not the same as expected, to use this GenerationSession, " - "you need to use PretrainedModel.prepare_inputs to create TRT Network inputs." - ) - if self.debug_mode: - self.debug_tensors = list( - set(found_tensor_names) - set(expected_tensor_names)) - if self.debug_tensors_to_save is None: - self.debug_tensors_to_save = self.debug_tensors - logger.info(f"Debug tensors found: {self.debug_tensors}") - logger.info(f"Debug tensors to save: {self.debug_tensors_to_save}") - - def __del__(self): - try: - if self.use_gemm_allreduce_plugin: - assert self.gemm_allreduce_output_handle is not None - ipc_nvls_free(self.gemm_allreduce_output_handle) - except TypeError: - pass - - @property - def context_mem_size(self) -> int: - return self.runtime.context_mem_size - - @property - def vocab_size(self): - return self._model_config.vocab_size - - @property - def num_layers(self): - assert self._model_config.num_layers % self.mapping.pp_size == 0, \ - f"num_layers {self._model_config.num_layers} must be a multiple of pipeline parallelism size {self.mapping.pp_size}" - return self._model_config.num_layers // self.mapping.pp_size - - @property - def first_layer(self): - return self.num_layers * self.mapping.pp_rank - - @property - def last_layer(self): - return self.first_layer + self.num_layers - - @property - def num_heads(self): - return self._model_config.num_heads - - @property - def hidden_size(self): - # For linear layer in attention block - return self._model_config.hidden_size - - @property - def use_gpt_attention_plugin(self): - return self._model_config.gpt_attention_plugin - - @property - def use_mamba_conv1d_plugin(self): - return self._model_config.mamba_conv1d_plugin - - @property - def paged_kv_cache(self): - return self._model_config.kv_cache_type == KVCacheType.PAGED - - @property - def kv_cache_type(self): - return self._model_config.kv_cache_type - - @property - def use_kv_cache(self): - return self._model_config.kv_cache_type != KVCacheType.DISABLED - - @property - def tokens_per_block(self): - return self._model_config.tokens_per_block - - @property - def remove_input_padding(self): - return self._model_config.remove_input_padding - - def get_num_heads_kv(self, layer_idx: Optional[int] = None) -> int: - if layer_idx is None or self._model_config.num_kv_heads_per_layer is None: - return self._model_config.num_kv_heads - - if self._model_config.layer_types: - assert self._model_config.layer_types[ - layer_idx] == "attention", f"Layer {layer_idx} is not an attention layer" - - if self._model_config.num_kv_heads_per_layer: - return self._model_config.num_kv_heads_per_layer[layer_idx] - - return self._model_config.num_kv_heads - - @property - def head_size(self): - return self.hidden_size // self.num_heads if self._model_config.head_size is None else self._model_config.head_size - - @property - def max_prompt_embedding_table_size(self): - return self._model_config.max_prompt_embedding_table_size - - @property - def quant_mode(self): - return self._model_config.quant_mode - - @property - def gather_context_logits(self): - return self._model_config.gather_context_logits - - @property - def gather_generation_logits(self): - return self._model_config.gather_generation_logits - - @property - def dtype(self): - return str_dtype_to_torch(self._model_config.dtype) - - @property - def profiler(self): - return self.runtime.profiler - - @property - def engine_inspector(self): - return self.runtime.engine_inspector - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - @property - def cross_attention(self): - return self._model_config.cross_attention - - @property - def has_position_embedding(self): - return self._model_config.has_position_embedding - - @property - def has_token_type_embedding(self): - return self._model_config.has_token_type_embedding - - @property - def use_lora_plugin(self): - return self._model_config.lora_plugin - - @property - def use_gemm_allreduce_plugin(self): - return bool(self._model_config.gemm_allreduce_plugin) - - @property - def gemm_allreduce_plugin(self): - return self._model_config.gemm_allreduce_plugin - - @property - def is_medusa_mode(self): - return self.num_medusa_heads > 0 - - @property - def is_redrafter_mode(self): - return self._model_config.redrafter_num_beams > 0 and self._model_config.redrafter_draft_len_per_beam > 0 - - @property - def max_draft_tokens(self): - if self.is_redrafter_mode: - return self._model_config.redrafter_num_beams * self._model_config.redrafter_draft_len_per_beam - return self._model_config.max_medusa_tokens - - @property - def num_medusa_heads(self): - return self._model_config.num_medusa_heads - - @property - def paged_state(self): - return self._model_config.paged_state - - @property - def conv_kernel(self): - return self._model_config.conv_kernel - - @property - def rnn_hidden_size(self): - return self._model_config.rnn_hidden_size - - @property - def rnn_head_size(self): - return self._model_config.rnn_head_size - - @property - def rnn_conv_dim_size(self): - return self._model_config.rnn_conv_dim_size - - @property - def state_size(self): - return self._model_config.state_size - - @property - def state_dtype(self): - if self._model_config.state_dtype == "": - return str_dtype_to_torch(self._model_config.dtype) - return str_dtype_to_torch(self._model_config.state_dtype) - - def _capture_cuda_graph_and_instantiate(self, context, stream, step): - instance_idx = (step + 1) % 2 - if not self.has_attn_layers: - # Create two cuda graph once.If cuda graph has already existed, skip it. - if self.runtime.cuda_graph_instances[instance_idx] is not None: - return - # capture cuda graph - CUASSERT( - cudart.cudaStreamBeginCapture( - stream, - cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)) - context.execute_async_v3(stream) - next_graph = CUASSERT(cudart.cudaStreamEndCapture(stream))[0] - - if self.runtime.cuda_graph_instances[instance_idx] is not None: - self.runtime.cuda_graph_instances[ - instance_idx] = _update_cuda_graph_instance( - self.runtime.cuda_graph_instances[instance_idx], next_graph) - else: - self.runtime.cuda_graph_instances[instance_idx] = CUASSERT( - cudart.cudaGraphInstantiate(next_graph, 0))[0] - - # Pre-upload cuda graph to stream - CUASSERT( - cudart.cudaGraphUpload( - self.runtime.cuda_graph_instances[instance_idx], stream)) - - def __setup_decoder(self, input_ids: torch.Tensor, - sampling_config: SamplingConfig, - host_context_lengths: torch.Tensor): - '''Allocate buffers and setup the post-processing decoder kernel - ''' - batch_size = host_context_lengths.shape[0] - scfg = sampling_config # just to make a shorter name, no other meaning - if isinstance(scfg.top_k, torch.Tensor): - assert scfg.top_k.dtype == torch.int32, f"scfg.top_k.dtype ({scfg.top_k.dtype}) must be torch.int32" - assert scfg.top_k.shape[ - 0] == batch_size, f"scfg.top_k.shape[0] ({scfg.top_k.shape[0]}) must equal to batch_size ({batch_size})" - self.top_k = scfg.top_k - else: - self.top_k = torch.full([batch_size], scfg.top_k, dtype=torch.int32) - - if isinstance(scfg.top_p, torch.Tensor): - assert scfg.top_p.dtype == torch.float32, f"scfg.top_p.dtype ({scfg.top_p.dtype}) must be torch.float32" - assert scfg.top_p.shape[ - 0] == batch_size, f"scfg.top_p.shape[0] ({scfg.top_p.shape[0]}) must equal to batch_size ({batch_size})" - self.top_p = scfg.top_p - else: - self.top_p = torch.full([batch_size], - scfg.top_p, - dtype=torch.float32) - - if isinstance(scfg.temperature, torch.Tensor): - assert scfg.temperature.dtype == torch.float32, f"scfg.temperature.dtype ({scfg.temperature.dtype}) must be torch.float32" - assert scfg.temperature.shape[ - 0] == batch_size, f"scfg.temperature.shape[0] ({scfg.temperature.shape[0]}) must equal to batch_size ({batch_size})" - self.temperature = scfg.temperature - else: - self.temperature = torch.full([batch_size], - scfg.temperature, - dtype=torch.float32) - - if isinstance(scfg.repetition_penalty, torch.Tensor): - assert scfg.repetition_penalty.dtype == torch.float32, f"scfg.repetition_penalty.dtype ({scfg.repetition_penalty.dtype}) must be torch.float32" - assert scfg.repetition_penalty.shape[ - 0] == batch_size, f"scfg.repetition_penalty.shape[0] ({scfg.repetition_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.repetition_penalty = scfg.repetition_penalty - elif scfg.repetition_penalty == 1.0: - self.repetition_penalty = None - else: - self.repetition_penalty = torch.full([batch_size], - scfg.repetition_penalty, - dtype=torch.float32) - - if isinstance(scfg.length_penalty, torch.Tensor): - assert scfg.length_penalty.dtype == torch.float32, f"scfg.length_penalty.dtype ({scfg.length_penalty.dtype}) must be torch.float32" - assert scfg.length_penalty.shape[ - 0] == batch_size, f"scfg.length_penalty.shape[0] ({scfg.length_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.host_length_penalty = scfg.length_penalty - else: - self.host_length_penalty = torch.full([batch_size], - scfg.length_penalty, - dtype=torch.float32) - self.length_penalty = self.host_length_penalty.to(self.device) - - if isinstance(scfg.early_stopping, torch.Tensor): - assert scfg.early_stopping.dtype == torch.int32, f"scfg.early_stopping.dtype ({scfg.early_stopping.dtype}) must be torch.int32" - assert scfg.early_stopping.shape[ - 0] == batch_size, f"scfg.early_stopping.shape[0] ({scfg.early_stopping.shape[0]}) must equal to batch_size ({batch_size})" - self.host_early_stopping = scfg.early_stopping - else: - self.host_early_stopping = torch.full([batch_size], - scfg.early_stopping, - dtype=torch.int32) - - if isinstance(scfg.presence_penalty, torch.Tensor): - assert scfg.presence_penalty.dtype == torch.float32, f"scfg.presence_penalty.dtype ({scfg.presence_penalty.dtype}) must be torch.float32" - assert scfg.presence_penalty.shape[ - 0] == batch_size, f"scfg.presence_penalty.shape[0] ({scfg.presence_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.presence_penalty = scfg.presence_penalty - elif scfg.presence_penalty == 0.0: - self.presence_penalty = None - else: - self.presence_penalty = torch.full([batch_size], - scfg.presence_penalty, - dtype=torch.float32) - - if isinstance(scfg.frequency_penalty, torch.Tensor): - assert scfg.frequency_penalty.dtype == torch.float32, f"scfg.frequency_penalty.dtype ({scfg.frequency_penalty.dtype}) must be torch.float32" - assert scfg.frequency_penalty.shape[ - 0] == batch_size, f"scfg.frequency_penalty.shape[0] ({scfg.frequency_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.frequency_penalty = scfg.frequency_penalty - elif scfg.frequency_penalty == 0.0: - self.frequency_penalty = None - else: - self.frequency_penalty = torch.full([batch_size], - scfg.frequency_penalty, - dtype=torch.float32) - - if isinstance(scfg.prompt_ignore_length, torch.Tensor): - assert scfg.prompt_ignore_length.dtype == torch.int32, f"scfg.prompt_ignore_length.dtype ({scfg.prompt_ignore_length.dtype}) must be torch.int32" - assert scfg.prompt_ignore_length.shape[ - 0] == batch_size, f"scfg.prompt_ignore_length.shape[0] ({scfg.prompt_ignore_length.shape[0]}) must equal to batch_size ({batch_size})" - self.prompt_ignore_length = scfg.prompt_ignore_length - else: - self.prompt_ignore_length = torch.full([batch_size], - scfg.prompt_ignore_length, - dtype=torch.int32) - - if isinstance(scfg.min_length, torch.Tensor): - assert scfg.min_length.dtype == torch.int32, f"scfg.min_length.dtype ({scfg.min_length.dtype}) must be torch.int32" - assert scfg.min_length.shape[ - 0] == batch_size, f"scfg.min_length.shape[0] ({scfg.min_length.shape[0]}) must equal to batch_size ({batch_size})" - self.min_length = scfg.min_length - else: - self.min_length = torch.full([batch_size], - scfg.min_length, - dtype=torch.int32) - - if isinstance(scfg.beam_search_diversity_rate, torch.Tensor): - assert scfg.beam_search_diversity_rate.dtype == torch.float32, f"scfg.beam_search_diversity_rate.dtype ({scfg.beam_search_diversity_rate.dtype}) must be torch.float32" - assert scfg.beam_search_diversity_rate.shape[ - 0] == batch_size, f"scfg.beam_search_diversity_rate.shape[0] ({scfg.beam_search_diversity_rate.shape[0]}) must equal to batch_size ({batch_size})" - self.beam_search_diversity_rate = scfg.beam_search_diversity_rate - elif scfg.beam_search_diversity_rate is not None: - self.beam_search_diversity_rate = torch.full( - [batch_size], - scfg.beam_search_diversity_rate, - dtype=torch.float32) - else: - self.beam_search_diversity_rate = None - - if isinstance(scfg.random_seed, torch.Tensor): - assert scfg.random_seed.dtype == torch.int64, f"scfg.random_seed.dtype ({scfg.random_seed.dtype}) must be torch.int64" - assert scfg.random_seed.shape[ - 0] == batch_size, f"scfg.random_seed.shape[0] ({scfg.random_seed.shape[0]}) must equal to batch_size ({batch_size})" - self.random_seed = scfg.random_seed - elif scfg.random_seed is not None: - self.random_seed = torch.full([batch_size], - scfg.random_seed, - dtype=torch.int64) - else: - self.random_seed = None - - if isinstance(scfg.no_repeat_ngram_size, torch.Tensor): - assert scfg.no_repeat_ngram_size.dtype == torch.int32, f"scfg.no_repeat_ngram_size.dtype ({scfg.no_repeat_ngram_size.dtype}) must be torch.int32" - assert scfg.no_repeat_ngram_size.shape[ - 0] == batch_size, f"scfg.no_repeat_ngram_size.shape[0] ({scfg.no_repeat_ngram_size.shape[0]}) must equal to batch_size ({batch_size})" - self.no_repeat_ngram_size = scfg.no_repeat_ngram_size - elif scfg.no_repeat_ngram_size is not None: - self.no_repeat_ngram_size = torch.full([batch_size], - scfg.no_repeat_ngram_size, - dtype=torch.int32) - else: - self.no_repeat_ngram_size = None - - if isinstance(scfg.min_p, torch.Tensor): - assert scfg.min_p.dtype == torch.float32, f"scfg.min_p.dtype ({scfg.min_p.dtype}) must be torch.float32" - assert scfg.min_p.shape[ - 0] == batch_size, f"scfg.min_p.shape[0] ({scfg.min_p.shape[0]}) must equal to batch_size ({batch_size})" - self.min_p = scfg.min_p - elif scfg.min_p == 1.0: - self.min_p = None - else: - self.min_p = torch.full([batch_size], - scfg.min_p, - dtype=torch.float32) - - if self.mapping.is_last_pp_rank(): - self.dynamic_decoder.setup( - batch_size, - scfg.num_beams, - self.top_k, - self.top_p, - self.temperature, - self.repetition_penalty, - self.presence_penalty, - self.frequency_penalty, - self.prompt_ignore_length, - self.min_length, - self.host_length_penalty, - self.host_early_stopping, - self.beam_search_diversity_rate, - self.random_seed, - self.top_p_decay, - self.top_p_min, - self.top_p_reset_ids, - self.no_repeat_ngram_size, - self.min_p, - scfg.output_log_probs, - scfg.num_beams > 1 or scfg.output_cum_log_probs, - ) - - assert scfg.end_id is not None, "end_id cannot be none" - assert scfg.pad_id is not None, 'pad_id cannot be none' - self.end_ids = torch.full((batch_size, ), - scfg.end_id, - dtype=torch.int32, - device=self.device) - max_context_length = host_context_lengths.max() - - # setup output ids buffer - if input_ids.dim() == 1: - # input_ids only have one dimension, which means remove_padding is enabled - split_ids_list = list( - torch.split(input_ids.unsqueeze(0), - host_context_lengths.numpy().tolist(), - dim=1)) - padded_input_ids = torch.nested.to_padded_tensor( - torch.nested.nested_tensor(split_ids_list, - dtype=torch.int32, - device='cuda'), - scfg.pad_id).reshape(batch_size, max_context_length) - else: - padded_input_ids = input_ids - if scfg.num_beams > 1: - tiled_input_ids = _tile_beam_width(padded_input_ids, scfg.num_beams) - tiled_input_ids = tiled_input_ids.reshape(batch_size, - scfg.num_beams, - max_context_length) - tiled_input_ids.permute(2, 0, 1) # TODO: delete? - self.output_ids = torch.cat( - (tiled_input_ids, - torch.full((batch_size, scfg.num_beams, - self.max_seq_length - max_context_length), - scfg.end_id, - dtype=padded_input_ids.dtype, - device=padded_input_ids.device)), - axis=-1) - else: - self.output_ids = torch.cat( - (padded_input_ids, - torch.full( - (batch_size, self.max_seq_length - max_context_length), - scfg.end_id, - dtype=padded_input_ids.dtype, - device=padded_input_ids.device)), - axis=-1) - - # Note: we still allocate max_seq_length size of parent ids (not max_attention_window_size). - self.parent_ids = torch.zeros( - (batch_size, scfg.num_beams, self.max_seq_length), - dtype=torch.int32, - device=self.device) - - if self.is_redrafter_mode: - self.new_tokens = torch.zeros([ - batch_size, self._model_config.redrafter_draft_len_per_beam + 1 - ], - dtype=torch.int32, - device=self.device) - self.accept_lengths = torch.ones([batch_size], - dtype=torch.int32, - device=self.device) - self.buffer["redrafter_inverted_temperature"] = torch.reciprocal( - self.temperature).to(device=self.device, dtype=self.dtype) - elif self.is_medusa_mode: - self.new_tokens = torch.zeros( - [batch_size, self.num_medusa_heads + 1], - dtype=torch.int32, - device=self.device) - self.medusa_output_tokens = torch.zeros( - [batch_size, self.num_draft_tokens], - dtype=torch.int32, - device=self.device) - self.generation_input_ids = torch.zeros( - [batch_size, self.num_draft_tokens + 1], - dtype=torch.int32, - device=self.device) - self.accept_lengths = torch.ones([batch_size], - dtype=torch.int32, - device=self.device) - if self.medusa_temperature != 0: - self.medusa_output_logits = torch.empty( - [batch_size, self.num_medusa_heads, self.vocab_size_padded], - dtype=self._tensor_dtype('logits'), - device=self.device) - elif scfg.num_beams > 1: - self.new_tokens = torch.zeros([batch_size, scfg.num_beams, 1], - dtype=torch.int32, - device=self.device) - else: - self.new_tokens = torch.zeros([batch_size, 1], - dtype=torch.int32, - device=self.device) - - if scfg.num_beams > 1 or scfg.output_cum_log_probs: - self.cum_log_probs = torch.full((batch_size, scfg.num_beams), - -1e20, - dtype=torch.float32, - device=self.device) - self.cum_log_probs[:, 0] = 0.0 - else: - self.cum_log_probs = None - - if scfg.output_log_probs: - self.log_probs = torch.zeros( - (batch_size, scfg.num_beams, self.max_seq_length), - dtype=torch.float32, - device=self.device) - self.log_probs_tiled = torch.zeros( - (self.max_seq_length, self._model_config.max_batch_size, - scfg.num_beams), - dtype=torch.float32, - device=self.device) - else: - self.log_probs = None - self.log_probs_tiled = None - - self.finished = torch.zeros((batch_size, scfg.num_beams), - dtype=torch.uint8, - device=self.device) - - if scfg.use_beam_hyps: - self.beam_hyps_output_ids_cba = torch.full( - size=[batch_size, scfg.num_beams * 2, self.max_seq_length], - fill_value=scfg.end_id, - dtype=torch.int32, - device=self.device) - self.beam_hyps_seq_len_cba = torch.zeros( - [batch_size, scfg.num_beams * 2], - dtype=torch.int32, - device=self.device) - self.beam_hyps_cum_log_probs_cba = torch.zeros( - [batch_size, scfg.num_beams * 2], - dtype=torch.float, - device=self.device) - self.beam_hyps_normed_scores_cba = torch.zeros( - [batch_size, scfg.num_beams * 2], - dtype=torch.float, - device=self.device) - self.beam_hyps_log_probs_cba = torch.zeros( - [batch_size, scfg.num_beams * 2, self.max_seq_length], - dtype=torch.float, - device=self.device) - self.beam_hyps_min_normed_scores = torch.zeros([batch_size], - dtype=torch.float, - device=self.device) - self.beam_hyps_num_beams = torch.zeros([batch_size], - dtype=torch.int32, - device=self.device) - self.beam_hyps_is_done = torch.zeros([batch_size], - dtype=torch.bool, - device=self.device) - else: - self.beam_hyps_output_ids_cba = None - self.beam_hyps_seq_len_cba = None - self.beam_hyps_cum_log_probs_cba = None - self.beam_hyps_normed_scores_cba = None - self.beam_hyps_log_probs_cba = None - self.beam_hyps_min_normed_scores = None - self.beam_hyps_num_beams = None - self.beam_hyps_is_done = None - - self.cross_kv_reuse = None - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.runtime.engine.get_tensor_dtype(name)) - return dtype - - def _init_medusa(self, medusa_choices: List[List[int]]): - from tensorrt_llm.runtime.medusa_utils import (_medusa_setup, - expand_choices_if_needed) - medusa_choices = expand_choices_if_needed(medusa_choices) - self.num_draft_tokens = len(medusa_choices) - assert self.num_draft_tokens > 0 and self.num_draft_tokens <= self.max_draft_tokens - medusa_info = _medusa_setup(medusa_choices, self.num_medusa_heads) - self.medusa_topks = medusa_info.medusa_topks - self.medusa_mask = medusa_info.medusa_mask[1:, 1:].to( - torch.bool - ) # convert to bool, original mask includes true token as well - - # Expand medusa position offsets to number of batch size in order to be compatible with the new Medusa. - target_shape = list(medusa_info.medusa_packed_mask.unsqueeze(0).shape) - target_shape[0] = self.batch_size - # Note: spec_decoding_packed_mask has no paddings in the first dimension. - self.spec_decoding_packed_mask = medusa_info.medusa_packed_mask.unsqueeze( - 0).expand(target_shape).reshape(-1, target_shape[-1]).cuda() - self.spec_decoding_use = medusa_info.medusa_spec_decoding_use - - self.medusa_paths = medusa_info.medusa_paths - self.medusa_tree_ids = medusa_info.medusa_tree_ids - - # Expand medusa position offsets to number of batch size in order to be compatible with the new Medusa. - target_shape = list( - medusa_info.medusa_position_offsets.unsqueeze(0).shape) - target_shape[0] = self.batch_size - # Note: medusa_position_offsets still keeps the paddings in order to get max_gen_input_length from the shape info. - self.spec_decoding_position_offsets = medusa_info.medusa_position_offsets.unsqueeze( - 0).expand(target_shape).int().cuda() - # Fixed sequence lengths currently. - # Support variable sequence lengths later. - self.spec_decoding_generation_lengths = (torch.ones( - (self.batch_size)) * (self.num_draft_tokens + 1)).int().cuda() - if not self.use_gpt_attention_plugin: - medusa_fp_mask = torch.zeros_like(self.medusa_mask, - dtype=torch.float32) - medusa_fp_mask[torch.logical_not(self.medusa_mask)] = float('-inf') - self.medusa_mask = medusa_fp_mask - return - - def _get_num_paged_blocks(self, max_attention_window_size, - sink_token_length): - bubble_len = 0 - if sink_token_length % self.tokens_per_block > 0: - bubble_len += (self.tokens_per_block - - sink_token_length % self.tokens_per_block) - max_blocks_per_seq = math.ceil( - (max_attention_window_size + bubble_len) / self.tokens_per_block) - num_blocks = self.batch_size * self.beam_width * max_blocks_per_seq - - return num_blocks, max_blocks_per_seq - - def setup(self, - batch_size: int, - max_context_length: int, - max_new_tokens: int, - beam_width: int = 1, - max_attention_window_size: Optional[int] = None, - sink_token_length: Optional[int] = None, - encoder_max_input_length: Optional[int] = None, - lora_manager: LoraManager = None, - lora_uids: List[str] = None, - medusa_choices: List[List[int]] = None, - multi_block_mode: bool = True, - enable_context_fmha_fp32_acc: bool = None): - # Store these params related to buffer size to check against - # the input shape with the params given in decode() - self.batch_size = batch_size - self.max_context_length = max_context_length - self.max_new_tokens = max_new_tokens - self.max_seq_length = max_context_length + max_new_tokens - if medusa_choices is not None or self.is_redrafter_mode: - self.max_seq_length += self.max_draft_tokens - self.beam_width = beam_width - self.encoder_max_input_length = encoder_max_input_length - self.multi_block_mode = multi_block_mode - self.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - if max_attention_window_size is None: - self.max_attention_window_size = self.max_seq_length - logger.debug( - "The max_attention_window_size is not set, we will use max_seq_length by default." - ) - self.host_max_attention_window_sizes = torch.ones( - (self.num_attn_layers, ), - dtype=torch.int32) * self.max_attention_window_size - - elif isinstance(max_attention_window_size, int): - if max_attention_window_size > self.max_seq_length: - logger.warning( - "The value of max_attention_window_size should ideally not exceed max_seq_length. " - "Therefore, it has been adjusted to match the value of max_seq_length." - ) - self.max_attention_window_size = min(max_attention_window_size, - self.max_seq_length) - self.host_max_attention_window_sizes = torch.ones( - (self.num_attn_layers, ), - dtype=torch.int32) * self.max_attention_window_size - - elif isinstance(max_attention_window_size, (torch.Tensor, list)): - if isinstance(max_attention_window_size, list): - max_attention_window_size = torch.tensor( - max_attention_window_size, dtype=torch.int32) - self.max_attention_window_size = int( - torch.max(max_attention_window_size).item()) - attn_win_size_len = max_attention_window_size.shape[0] - num_total_attn_layers = self.layer_types.count('attention') - if attn_win_size_len < num_total_attn_layers: - repeat_num = num_total_attn_layers // attn_win_size_len - remain_num = num_total_attn_layers % attn_win_size_len - warning_info = "The size of max_attention_window_size tensor/list is less than num_attn_layers, " \ - + "and it will be repeated to num_attn_layers. So the actual max_attention_window_size " \ - + f"is {max_attention_window_size.tolist()} * {repeat_num}" - warning_info += f" + {max_attention_window_size.tolist()[0:remain_num]}. " if remain_num > 0 else ". " - warning_info += "Note that num_attn_layers is the number of total attention layers." - logger.warning(warning_info) - elif attn_win_size_len > num_total_attn_layers: - logger.error( - "The size of max_attention_window_size tensor/list is larger than num_attn_layers! " - "Note that num_attn_layers is the number of total attention layers." - ) - assert False - if self.max_attention_window_size > self.max_seq_length: - logger.warning( - "The value of max_attention_window_size should ideally not exceed max_seq_length. " - "Therefore, it has been adjusted to match the value of max_seq_length." - ) - self.max_attention_window_size = min(self.max_attention_window_size, - self.max_seq_length) - max_attention_window_size = torch.minimum( - max_attention_window_size.to(torch.int32), - torch.IntTensor([self.max_seq_length] * attn_win_size_len)) - self.host_max_attention_window_sizes = torch.ones( - (self.num_attn_layers, ), dtype=torch.int32) - for i in range(self.num_attn_layers): - self.host_max_attention_window_sizes[ - i] = max_attention_window_size[ - (self.layer_types[0:self.first_layer].count('attention') - + i) % attn_win_size_len] - else: - assert False, "invalid max_attention_window_size!" - - if sink_token_length is None: - self.sink_token_length = 0 - self.host_sink_token_length = torch.zeros((1, ), dtype=torch.int32) - elif isinstance(sink_token_length, int): - self.sink_token_length = sink_token_length - self.host_sink_token_length = torch.ones( - (1, ), dtype=torch.int32) * self.sink_token_length - else: - assert False, "invalid sink_token_length!" - - self.lora_manager = lora_manager - if medusa_choices is not None: - self._init_medusa(medusa_choices) - - self.buffer = {} - if self.mapping.is_last_pp_rank(): - if self.is_redrafter_mode: - init_allocate_redrafter_tensors(self, batch_size) - self.buffer['logits'] = torch.empty( - (batch_size, self.max_draft_tokens + 1, - self.vocab_size_padded) - if not self.gather_context_logits else - (batch_size, max_context_length, self.vocab_size_padded), - dtype=self._tensor_dtype('logits'), - device=self.device) - elif self.is_medusa_mode: - self.buffer['logits'] = torch.empty( - (batch_size, self.num_draft_tokens + 1, - self.vocab_size_padded) - if not self.gather_context_logits else - (batch_size, max_context_length, self.vocab_size_padded), - dtype=self._tensor_dtype('logits'), - device=self.device) - medusa_logits_shape = (self.num_medusa_heads, batch_size, - (self.num_draft_tokens + 1), - self.vocab_size_padded) - if self.remove_input_padding: - medusa_logits_shape = (self.num_medusa_heads, batch_size * - (self.num_draft_tokens + 1), - self.vocab_size_padded) - - self.buffer['medusa_logits'] = torch.empty( - medusa_logits_shape if not self.gather_context_logits else - (self.num_medusa_heads, batch_size, max_context_length, - self.vocab_size_padded), - dtype=self._tensor_dtype('medusa_logits'), - device=self.device) - else: - self.buffer['logits'] = torch.empty( - (batch_size, self.vocab_size_padded) - if not self.gather_context_logits else - (batch_size, max_context_length, self.vocab_size_padded), - dtype=self._tensor_dtype('logits'), - device=self.device) - - if self.cross_attention: - # use shape info to pass max length info in remove padding mode - self.buffer['encoder_max_input_length'] = torch.empty( - (encoder_max_input_length, ), - dtype=self._tensor_dtype('encoder_max_input_length'), - device=self.device) - - if self.quant_mode.has_kv_cache_quant(): - # Since torch does not support fp8 now, using int8 here. - kv_cache_type = torch.int8 - else: - if self.use_kv_cache and self.has_attn_layers: - first_atten_layer = self.layer_types[ - self.first_layer:self.last_layer].index( - 'attention') + self.first_layer - kv_cache_type = self.dtype if self.paged_kv_cache else self._tensor_dtype( - f'present_key_value_{first_atten_layer}') - else: - kv_cache_type = None - - if self.use_kv_cache: - if self.paged_kv_cache and self.has_attn_layers: - num_blocks, _ = self._get_num_paged_blocks( - self.max_attention_window_size, self.sink_token_length) - self._memory_pool_allocator = MemoryPoolsAllocator( - num_blocks=num_blocks, - tokens_per_block=self.tokens_per_block, - head_size=self.head_size) - if self._model_config.num_kv_heads_per_layer is None: - num_kv_heads_per_layer = MemoryPoolsAllocator.prepare_num_kv_heads_per_layer( - self.get_num_heads_kv(), self.num_attn_layers) - else: - num_kv_heads_per_layer = self._model_config.num_kv_heads_per_layer - - self._memory_pool_allocator.allocate(kv_cache_type, - num_kv_heads_per_layer) - - if self.cross_attention: # As for now we enable cross paged kv and self paged kv to share the same tokens_per_block - cross_num_blocks, _ = self._get_num_paged_blocks( - self.encoder_max_input_length, sink_token_length=0) - - num_kv_heads_per_layer = MemoryPoolsAllocator.prepare_num_kv_heads_per_layer( - self.get_num_heads_kv(), self.num_attn_layers) - - self._cross_memory_pool_allocator = MemoryPoolsAllocator( - num_blocks=cross_num_blocks, - tokens_per_block=self.tokens_per_block, - head_size=self.head_size) - if self._model_config.num_kv_heads_per_cross_attn_layer is None: - num_kv_heads_per_cross_attn_layer = MemoryPoolsAllocator.prepare_num_kv_heads_per_layer( - self.get_num_heads_kv(), self.num_attn_layers) - else: - num_kv_heads_per_cross_attn_layer = self._model_config.num_kv_heads_per_cross_attn_layer - - self._cross_memory_pool_allocator.allocate( - kv_cache_type, num_kv_heads_per_cross_attn_layer) - - elif self.has_attn_layers: - - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - cache_shape = ( - batch_size, - 2, - self.get_num_heads_kv(i), - self.max_attention_window_size, - self.head_size, - ) - self.buffer[f'present_key_value_{i}'] = torch.empty( - cache_shape, - dtype=kv_cache_type, - device=self.device) - - if self.cross_attention: - cross_cache_shape = ( - batch_size, - 2, - self.get_num_heads_kv(), - self.encoder_max_input_length, - self.head_size, - ) - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - self.buffer[ - f'cross_present_key_value_{i}'] = torch.empty( - cross_cache_shape, - dtype=kv_cache_type, - device=self.device) - - if self.use_gpt_attention_plugin: - self.sequence_length_buffer = torch.ones((batch_size, ), - dtype=torch.int32, - device=self.device) - else: - # Without plugin, we need extra kv cache buffers. - # Because we don't support inplace update, so we need separate buffer for inputs and outputs. - # We can do reuse between different layers' inputs and outputs, i.e. current layer's output can - # reuse previous layer's input memory. But this need one extra buffer as the guard. - if self.use_kv_cache and self.has_attn_layers: # Not applicable to cross KV buffers as it's constant - i = self.attn_to_general_idx[0] - trt_dtype = self.runtime.engine.get_tensor_dtype( - f'present_key_value_{i}') - - if trt_dtype == trt.fp8: - # PyTorch doesn't support fp8 datatype, use int8 instead of it because int8 datatype size is same with fp8. - # TODO: Remove this section when PyTorch support fp8 datatype - dtype = torch.int8 - else: - dtype = self._tensor_dtype(f'present_key_value_{i}') - self.buffer[f'1_present_key_value_{i}'] = torch.empty( - cache_shape, dtype=dtype, device=self.device) - - if self.use_mamba_conv1d_plugin: - conv_state_shape = ( - batch_size, - self.conv_kernel - 1, - self.rnn_conv_dim_size, - ) - else: - conv_state_shape = ( - batch_size, - self.rnn_conv_dim_size, - self.conv_kernel - 1, - ) - - if self.rnn_head_size > 1: - rnn_state_shape = ( - batch_size, - self.rnn_hidden_size // self.rnn_head_size, - self.state_size, - self.rnn_head_size, - ) - else: - rnn_state_shape = ( - batch_size, - self.state_size, - self.rnn_hidden_size, - ) - - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'recurrent': - dtype = self.dtype - self.buffer[f'present_conv_state_{i}'] = torch.empty( - conv_state_shape, dtype=dtype, device=self.device) - self.buffer[f'1_present_conv_state_{i}'] = torch.empty( - conv_state_shape, dtype=dtype, device=self.device) - self.buffer[f'present_rnn_state_{i}'] = torch.empty( - rnn_state_shape, dtype=self.state_dtype, device=self.device) - if self.paged_state: - conv_state_ptr = torch.tensor( - [self.buffer[f'present_conv_state_{i}'].data_ptr()], - dtype=torch.int64, - device='cpu') - rnn_state_ptr = torch.tensor( - [self.buffer[f'present_rnn_state_{i}'].data_ptr()], - dtype=torch.int64, - device='cpu') - self.buffer[f'conv_state_ptr_{i}'] = conv_state_ptr - self.buffer[f'rnn_state_ptr_{i}'] = rnn_state_ptr - - if self.use_lora_plugin and self.lora_manager is not None: - lora_uids = lora_uids or ["-1"] - self.buffer.update( - self.lora_manager.input_buffers( - lora_uids, - self.mapping, - self._model_config.num_layers, - )) - - if self.use_gemm_allreduce_plugin: - max_num_tokens = max(batch_size * beam_width, - batch_size * self.max_seq_length) - M = max_num_tokens - N = self.hidden_size - self.gemm_allreduce_output_size = M * N - itemsize = str_dtype_to_torch(self.gemm_allreduce_plugin).itemsize - alloc_bytes = self.gemm_allreduce_output_size * itemsize - self.gemm_allreduce_output_handle = ipc_nvls_allocate( - alloc_bytes, set(self.mapping.tp_group)) - logger.debug(f'Allocated NVLS IPC memory: {alloc_bytes} bytes') - - if self.is_medusa_mode: - self.buffer[ - 'spec_decoding_packed_mask'] = self.spec_decoding_packed_mask - self.buffer[ - 'spec_decoding_position_offsets'] = self.spec_decoding_position_offsets - self.buffer[ - 'spec_decoding_generation_lengths'] = self.spec_decoding_generation_lengths - self.buffer['spec_decoding_use'] = self.spec_decoding_use - self.buffer_allocated = True - if self.is_medusa_mode: - return self.num_draft_tokens - - def _allocate_empty_kv_cache_pools(self, kv_cache_type, num_blocks): - # Layers are homogeneous, use old kv cache shape - unique_cache_pools = [] - if self._model_config.num_kv_heads_per_layer is None: - cache_shape = ( - num_blocks, - self.num_attn_layers, - 2, - self.get_num_heads_kv(), - self.tokens_per_block, - self.head_size, - ) - unique_cache_pools.append( - torch.empty(cache_shape, - dtype=kv_cache_type, - device=self.device)) - - # Layers are not homogeneous, use new kv cache shape - else: - kv_heads_unique_counter = Counter( - self._model_config.num_kv_heads_per_layer) - for kv_head, num_layers in kv_heads_unique_counter.items(): - cache_shape = ( - num_blocks, - num_layers, - 2, - kv_head, - self.tokens_per_block, - self.head_size, - ) - unique_cache_pools.append( - torch.empty(cache_shape, - dtype=kv_cache_type, - device=self.device)) - - return unique_cache_pools - - def _get_context_shape_buffer( - self, - input_ids: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths: torch.Tensor, - position_ids: torch.Tensor, - last_token_ids: torch.Tensor, - attention_mask: torch.Tensor, - cross_attention_mask: torch.Tensor, - cache_indirection: torch.Tensor, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_block_offsets: torch.Tensor, - cross_kv_cache_block_offsets: torch.Tensor = None, - host_cross_kv_cache_block_offsets: torch.Tensor = None, - hidden_states_input: torch.Tensor = None, - prompt_embedding_table: torch.Tensor = None, - tasks: torch.Tensor = None, - prompt_vocab_size: torch.Tensor = None, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - host_runtime_perf_knobs: torch.Tensor = None, - host_context_progress: torch.Tensor = None, - skip_cross_attn_blocks: torch.Tensor = None, - language_adapter_routings: torch.Tensor = None, - ) -> Dict[str, RuntimeTensor]: - tensors = {} - - def sym(x, name): - return RuntimeTensor.from_torch(name, x) - - def add_tensor_from_pointer(pointer, name, shape, str_dtype): - return tensors.update({ - name: - RuntimeTensor.from_pointer(name, pointer, shape, str_dtype) - }) - - def add_tensor(x, name): - return tensors.update({name: sym(x, name)}) - - def add_tensor_with_shape(x, name, shape): - return tensors.update( - {name: RuntimeTensor.from_torch(name, x, override_shape=shape)}) - - def add_tensor_with_bs(x, name, bs): - # this assumes dim0 to be bs and only overrides dim0 with given bs - shape = list(x.shape) - shape[0] = bs - return tensors.update( - {name: RuntimeTensor.from_torch(name, x, override_shape=shape)}) - - if self.has_attn_layers: - if self.use_gpt_attention_plugin: - add_tensor(context_lengths, 'context_lengths') - assert host_runtime_perf_knobs != None, "gpt_attention_plugin needs to set host_runtime_perf_knobs" - add_tensor(host_runtime_perf_knobs, 'host_runtime_perf_knobs') - add_tensor(host_context_progress, 'host_context_progress') - add_tensor(cache_indirection, 'cache_indirection') - - if self.has_position_embedding: - add_tensor(position_ids, 'position_ids') - - if self.cross_attention: - # in context phase, need to generate cross kv cache, set to True - add_tensor(torch.ones(1, dtype=torch.bool, device=self.device), - 'cross_kv_cache_gen') - if self._model_config.skip_cross_attn_blocks: - add_tensor(skip_cross_attn_blocks, 'skip_cross_attn_blocks') - if self.skip_cross_kv: - if self.cross_kv_reuse is None: - # see Attention's self.qkv output dim - cross_kv_out_dim = 2 * self.get_num_heads_kv( - ) * self.head_size - cross_kv_shape = encoder_output.shape[:-1] + ( - cross_kv_out_dim, ) - cross_kv_reuse = torch.empty(cross_kv_shape, - dtype=encoder_output.dtype, - device=encoder_output.device) - self.cross_kv_reuse = cross_kv_reuse - add_tensor(self.cross_kv_reuse, 'cross_kv_reuse') - add_tensor(encoder_output, 'encoder_output') - add_tensor(encoder_input_lengths, 'encoder_input_lengths') - if language_adapter_routings is not None: - add_tensor(language_adapter_routings, - 'language_adapter_routings') - add_tensor(self.buffer['encoder_max_input_length'], - 'encoder_max_input_length') - if not self.use_gpt_attention_plugin: - add_tensor(cross_attention_mask, 'cross_attention_mask') - else: - if cross_attention_mask != None: - # cross-attention packed mask (used by fmha). - cross_attention_packed_mask = torch.ops.tensorrt_llm.pack_fmha_mask_by_input( - cross_attention_mask, context_lengths, - encoder_input_lengths, 1.0) - add_tensor(cross_attention_mask, 'cross_attention_mask') - add_tensor(cross_attention_packed_mask, - 'cross_attention_packed_mask') - else: - # create a full 1 cross_attention_mask because it is necessary - batch_size = context_lengths.shape[0] - cross_attention_mask = torch.ones( - (np.asarray(input_ids.shape).prod(), - np.asarray(list(encoder_output.shape)[:-1]).prod()), - dtype=torch.bool, - device=self.device) - add_tensor(cross_attention_mask, "cross_attention_mask") - cross_attention_packed_mask = torch.ops.tensorrt_llm.pack_fmha_mask_by_input( - cross_attention_mask, context_lengths, - encoder_input_lengths, 1.0) - add_tensor(cross_attention_packed_mask, - "cross_attention_packed_mask") - - if self.mapping.has_pp(): - hidden_size = self.hidden_size * self.mapping.tp_size - if input_ids.dim() == 2: - hidden_states_input = hidden_states_input.resize_( - input_ids.shape[0], input_ids.shape[1], hidden_size) - else: - hidden_states_input = hidden_states_input.resize_( - input_ids.shape[0], hidden_size) - - if self.mapping.is_last_pp_rank(): - if self.is_redrafter_mode: - set_redrafter_ctx_tensors(self, add_tensor, add_tensor_with_bs) - add_tensor(self.buffer['logits'], 'logits') - if self.is_medusa_mode: - add_tensor(self.buffer['medusa_logits'], 'medusa_logits') - - if not self.gather_context_logits or self.has_rnn_layers: - add_tensor(last_token_ids, 'last_token_ids') - else: - add_tensor(hidden_states_input, 'hidden_states_output') - - if self.mapping.is_first_pp_rank(): - add_tensor(input_ids, 'input_ids') - else: - add_tensor(hidden_states_input, 'hidden_states_input') - - if prompt_embedding_table is not None: - add_tensor(prompt_embedding_table, 'prompt_embedding_table') - - if self.remove_input_padding: - tasks_generation = torch.concat([ - torch.full([context_lengths[b].item()], - tasks[b].item(), - dtype=torch.int32) - for b in range(context_lengths.size(0)) - ]).cuda() - else: - tasks_generation = tasks.unsqueeze(-1) - add_tensor(tasks_generation, 'tasks') - add_tensor(prompt_vocab_size, 'prompt_vocab_size') - - if self.paged_kv_cache and self.has_attn_layers: - buffer = kv_cache_block_offsets.contiguous() - shape = kv_cache_block_offsets.shape - shape = [shape[0], shape[1] * shape[2], *shape[3:]] - add_tensor_with_shape(buffer, f'kv_cache_block_offsets', shape) - add_tensor_with_shape(host_kv_cache_block_offsets, - f'host_kv_cache_block_offsets', shape) - pool_pointers = f'host_kv_cache_pool_pointers' - pool_mapping = f'host_kv_cache_pool_mapping' - add_tensor(self.buffer[pool_pointers], pool_pointers) - add_tensor(self.buffer[pool_mapping], pool_mapping) - if self.cross_attention: - cross_buffer = cross_kv_cache_block_offsets.contiguous() - cross_shape = cross_kv_cache_block_offsets.shape - cross_shape = [ - cross_shape[0], cross_shape[1] * cross_shape[2], - *cross_shape[3:] - ] - add_tensor_with_shape(cross_buffer, - f'cross_kv_cache_block_offsets', - cross_shape) - add_tensor_with_shape(host_cross_kv_cache_block_offsets, - f'host_cross_kv_cache_block_offsets', - cross_shape) - cross_pool_pointers = f'host_cross_kv_cache_pool_pointers' - cross_pool_mapping = f'host_cross_kv_cache_pool_mapping' - add_tensor(self.buffer[cross_pool_pointers], - cross_pool_pointers) - add_tensor(self.buffer[cross_pool_mapping], cross_pool_mapping) - - batch_size = context_lengths.shape[0] - if self.use_kv_cache and not self.paged_kv_cache: - for idx in range(self.first_layer, self.last_layer): - if not self.use_gpt_attention_plugin and self.layer_types[ - idx] == 'attention': - kv_cache_shape = (batch_size, 2, - self.get_num_heads_kv( - self.general_to_attn_idx[idx]), 0, - self.head_size) - # for empty tensor, TRT does not really use the tensor data, so any dtype is fine - kv_cache_buffer = torch.zeros((1, ), - dtype=torch.float32, - device=self.device) - add_tensor_with_shape(kv_cache_buffer, - f'past_key_value_{idx}', - kv_cache_shape) - present = f'present_key_value_{idx}' - add_tensor(self.buffer[present], present) - - if self.cross_attention: - cross_kv_cache_shape = (batch_size, 2, - self.get_num_heads_kv(), 0, - self.head_size) - # for empty tensor, TRT does not really use the tensor data, so any dtype is fine - cross_kv_cache_buffer = torch.zeros((1, ), - dtype=torch.float32, - device=self.device) - add_tensor_with_shape(cross_kv_cache_buffer, - f'cross_past_key_value_{idx}', - cross_kv_cache_shape) - cross_present = f'cross_present_key_value_{idx}' - add_tensor(self.buffer[cross_present], cross_present) - elif self.layer_types[idx] == 'attention': - key_value_cache = self.buffer[f'present_key_value_{idx}'] - # when plugin is used, past_ket_value tensor does not need to be empty tensor - # because plugin does not care, and does not use this shape. - add_tensor(key_value_cache, f'past_key_value_{idx}') - add_tensor(key_value_cache, f'present_key_value_{idx}') - - if self.cross_attention: - cross_cache_buffer = self.buffer[ - f'cross_present_key_value_{idx}'] - add_tensor(cross_cache_buffer, - f'cross_past_key_value_{idx}') - add_tensor(cross_cache_buffer, - f'cross_present_key_value_{idx}') - - for idx in range(self.first_layer, self.last_layer): - if self.layer_types[idx] != 'recurrent': - continue - if self.paged_state: - add_tensor(self.buffer[f'conv_state_ptr_{idx}'], - f'conv_state_ptr_{idx}') - add_tensor(self.buffer[f'rnn_state_ptr_{idx}'], - f'rnn_state_ptr_{idx}') - else: - # conv state - dtype = self._tensor_dtype(f'present_conv_state_{idx}') - if self.use_mamba_conv1d_plugin: - conv_state_shape = (batch_size, self.conv_kernel - 1, - self.rnn_conv_dim_size) - else: - conv_state_shape = (batch_size, self.rnn_conv_dim_size, - self.conv_kernel - 1) - - conv_state = torch.zeros(conv_state_shape, - dtype=dtype, - device=self.device) - add_tensor(conv_state, f'past_conv_state_{idx}') - present = f'present_conv_state_{idx}' - add_tensor(self.buffer[present], present) - # rnn state - rnn_state = self.buffer[f'present_rnn_state_{idx}'] - add_tensor(rnn_state, f'past_rnn_state_{idx}') - add_tensor(rnn_state, f'present_rnn_state_{idx}') - - if self.paged_state and self.has_rnn_layers: - slot_mapping = torch.arange(0, - batch_size, - device='cuda', - dtype=torch.int32) - add_tensor(slot_mapping, 'slot_mapping') - - if self.use_gpt_attention_plugin and self.has_attn_layers: - # context request - host_request_types = torch.zeros_like(context_lengths, - device='cpu').int() - self.sequence_length_buffer = context_lengths.detach().clone() - if self.is_redrafter_mode: - device_request_types = torch.zeros_like( - context_lengths, device=self.device).int() - add_tensor(device_request_types, 'device_request_types') - add_tensor_with_shape(self.sequence_length_buffer, - 'sequence_length', (batch_size, )) - - # field 0: past_key_value_length, field 1: is_context (deprecated). changed to [0], otherwise affects batch padded input mode - add_tensor_with_shape(host_context_lengths.clone(), - 'host_past_key_value_lengths', (batch_size, )) - add_tensor_with_shape(self.host_sink_token_length, - 'host_sink_token_length', (1, )) - add_tensor(host_request_types, 'host_request_types') - add_tensor_with_shape(self.host_max_attention_window_sizes, - f'host_max_attention_window_sizes', - (self.num_attn_layers, )) - if self.remove_input_padding: - add_tensor(host_context_lengths, 'host_context_lengths') - else: - if self.has_rnn_layers: - host_request_types = torch.zeros_like(context_lengths, - device='cpu').int() - add_tensor(host_request_types, 'host_request_types') - if self.remove_input_padding: - add_tensor(host_context_lengths, 'host_context_lengths') - if self.has_attn_layers: - add_tensor(attention_mask, 'attention_mask') - - if self.mapping.tp_size > 1: - add_tensor(self.all_reduce_workspace, 'all_reduce_workspace') - if self.use_gemm_allreduce_plugin: - found_tensor_names = [ - self.runtime.engine.get_tensor_name(i) - for i in range(self.runtime.engine.num_io_tensors) - ] - for name in found_tensor_names: - if name.startswith("gemm_allreduce_uc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.uc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_mc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.mc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_ipc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.get_ipc_ptrs(), - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - - if self.use_lora_plugin: - for idx in range(self.num_layers): - for lora_module in (self.lora_target_modules + - self.missing_qkv_modules): - layer_idx = idx + self.first_layer - lora_ranks = f'{lora_module}_lora_ranks_{layer_idx}' - add_tensor(self.buffer[lora_ranks], lora_ranks) - lora_weights = f'{lora_module}_lora_weights_pointers_{layer_idx}' - add_tensor(self.buffer[lora_weights], lora_weights) - if self.cross_attention and self.remove_input_padding: - add_tensor(encoder_input_lengths.to('cpu'), - 'host_encoder_input_lengths') - if self.is_medusa_mode: - # Medusa mask and position offsets are fixed for the whole session. - add_tensor(self.buffer['spec_decoding_packed_mask'], - 'spec_decoding_packed_mask') - add_tensor(self.buffer['spec_decoding_position_offsets'], - 'spec_decoding_position_offsets') - add_tensor(self.buffer['spec_decoding_generation_lengths'], - 'spec_decoding_generation_lengths') - add_tensor(self.buffer['spec_decoding_use'], 'spec_decoding_use') - - return tensors - - def _get_next_step_shape_buffer( - self, - batch_size: int, - beam_width: int, - max_context_length: int, - step: int, - context_lengths: torch.Tensor, - host_context_lengths: torch.Tensor, - position_ids: torch.Tensor, - last_token_ids: torch.Tensor, - attention_mask: torch.Tensor, - cross_attention_mask: torch.Tensor, - cache_indirection: torch.Tensor, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_block_offsets: torch.Tensor, - cross_kv_cache_block_offsets: torch.Tensor = None, - host_cross_kv_cache_block_offsets: torch.Tensor = None, - hidden_states_input: torch.Tensor = None, - prompt_embedding_table: torch.Tensor = None, - tasks: torch.Tensor = None, - prompt_vocab_size: torch.Tensor = None, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - host_runtime_perf_knobs: torch.Tensor = None, - host_context_progress: torch.Tensor = None, - skip_cross_attn_blocks: torch.Tensor = None, - language_adapter_routings: torch.Tensor = None, - ): - torch.cuda.nvtx.range_push("_get_next_step_shape_buffer") - tensors = {} # Dict[str, RuntimeTensor] - - def add_tensor_from_pointer(pointer, name, shape, str_dtype): - return tensors.update({ - name: - RuntimeTensor.from_pointer(name, pointer, shape, str_dtype) - }) - - def sym(x, name): - return RuntimeTensor.from_torch(name, x) - - def add_tensor(x, name): - return tensors.update({name: sym(x, name)}) - - def add_tensor_with_shape(x, name, shape): - return tensors.update( - {name: RuntimeTensor.from_torch(name, x, override_shape=shape)}) - - context_lengths_local = context_lengths.clone() - host_context_lengths_local = host_context_lengths.clone() - if self.has_attn_layers: - if self.use_gpt_attention_plugin: - add_tensor(context_lengths_local, 'context_lengths') - assert host_runtime_perf_knobs != None, "gpt_attention_plugin needs to set host_runtime_perf_knobs" - add_tensor(host_runtime_perf_knobs, 'host_runtime_perf_knobs') - add_tensor(host_context_progress, 'host_context_progress') - add_tensor(cache_indirection, 'cache_indirection') - if self.has_position_embedding: - add_tensor(position_ids, 'position_ids') - - if self.mapping.has_pp(): - hidden_size = self.hidden_size * self.mapping.tp_size - shape = (batch_size * beam_width, - hidden_size) if self.remove_input_padding else ( - batch_size * beam_width, 1, hidden_size) - hidden_states_input = hidden_states_input.resize_(*shape) - - if self.mapping.is_last_pp_rank(): - add_tensor(self.buffer['logits'], 'logits') - if self.is_medusa_mode: - add_tensor(self.buffer['medusa_logits'], 'medusa_logits') - - if not self.gather_context_logits or self.has_rnn_layers: - add_tensor(last_token_ids, 'last_token_ids') - else: - add_tensor(hidden_states_input, 'hidden_states_output') - - if self.mapping.is_first_pp_rank(): - if self.is_redrafter_mode: - input_ids_shape = (self.host_total_gen_token, ) - else: - input_ids_shape = ( - batch_size * beam_width * (self.num_draft_tokens + 1), - ) if self.remove_input_padding else (batch_size * beam_width, - self.num_draft_tokens + 1) - if self.is_redrafter_mode: - add_tensor_with_shape(self.buffer['flat_tokens'], 'input_ids', - input_ids_shape) - elif self.is_medusa_mode: - add_tensor_with_shape(self.generation_input_ids, 'input_ids', - input_ids_shape) - else: - add_tensor_with_shape(self.new_tokens, 'input_ids', - input_ids_shape) - else: - add_tensor(hidden_states_input, 'hidden_states_input') - - if self.cross_attention: - if self.use_gpt_attention_plugin: - # disable (or minimize) cross qkv computation at generation phase - if self.skip_cross_kv: - # disable - encoder_output_shape = encoder_output.shape - add_tensor(self.cross_kv_reuse, 'cross_kv_reuse') - else: - # minimize - # use TensorRT Empty Tensor to skip redundant computation - # 0 for generation phase, >0 for context phase - encoder_output_shape = list(encoder_output.shape) - if self.remove_input_padding: - encoder_output_shape[-2] = 0 - else: - encoder_output_shape = [1, 0, encoder_output.shape[-1]] - else: - # OOTB path doesn't have kv cache for now, so this encoder_output is - # a must-have input. We just use the encoder_output - encoder_output_shape = encoder_output.shape - - # in generation phase, cross kv cache is already filled during context phase, set to False - add_tensor(torch.zeros(1, dtype=torch.bool, device=self.device), - 'cross_kv_cache_gen') - if self._model_config.skip_cross_attn_blocks: - add_tensor(skip_cross_attn_blocks, 'skip_cross_attn_blocks') - add_tensor_with_shape(encoder_output, 'encoder_output', - encoder_output_shape) - add_tensor(encoder_input_lengths, 'encoder_input_lengths') - if language_adapter_routings is not None: - add_tensor(language_adapter_routings, - 'language_adapter_routings') - add_tensor(self.buffer['encoder_max_input_length'], - 'encoder_max_input_length') - if not self.use_gpt_attention_plugin: - add_tensor(cross_attention_mask, 'cross_attention_mask') - else: - if cross_attention_mask != None: - cross_attention_mask = _tile_beam_width( - cross_attention_mask, beam_width) - # Empty packed mask is passed in the generation phase as it is not used. - cross_attention_packed_mask = torch.empty( - (batch_size, - (cross_attention_mask.shape[1] + 31) // 32), - dtype=torch.int32, - device=self.device) - add_tensor(cross_attention_mask, 'cross_attention_mask') - add_tensor(cross_attention_packed_mask, - 'cross_attention_packed_mask') - else: - # create a full 1 cross_attention_mask because it is necessary in generation phase - add_tensor( - torch.ones((batch_size, - np.asarray(list( - encoder_output.shape)[:-1]).prod()), - dtype=torch.bool, - device=self.device), "cross_attention_mask") - # Empty packed mask is passed in the generation phase as it is not used. - add_tensor( - torch.empty((batch_size, 1), - dtype=torch.int32, - device=self.device), - "cross_attention_packed_mask") - - if self.paged_kv_cache and self.has_attn_layers: - shape = kv_cache_block_offsets.shape - shape = [shape[0], shape[1] * shape[2], *shape[3:]] - add_tensor_with_shape(kv_cache_block_offsets, - f'kv_cache_block_offsets', shape) - add_tensor_with_shape(host_kv_cache_block_offsets, - f'host_kv_cache_block_offsets', shape) - pool_pointers = f'host_kv_cache_pool_pointers' - pool_mapping = f'host_kv_cache_pool_mapping' - add_tensor(self.buffer[pool_pointers], pool_pointers) - add_tensor(self.buffer[pool_mapping], pool_mapping) - if self.cross_attention: - cross_shape = cross_kv_cache_block_offsets.shape - cross_shape = [ - cross_shape[0], cross_shape[1] * cross_shape[2], - *cross_shape[3:] - ] - add_tensor_with_shape(cross_kv_cache_block_offsets, - f'cross_kv_cache_block_offsets', - cross_shape) - add_tensor_with_shape(host_cross_kv_cache_block_offsets, - f'host_cross_kv_cache_block_offsets', - cross_shape) - cross_pool_pointers = f'host_cross_kv_cache_pool_pointers' - cross_pool_mapping = f'host_cross_kv_cache_pool_mapping' - add_tensor(self.buffer[cross_pool_pointers], - cross_pool_pointers) - add_tensor(self.buffer[cross_pool_mapping], cross_pool_mapping) - - if prompt_embedding_table is not None: - add_tensor(prompt_embedding_table, 'prompt_embedding_table') - - if self.remove_input_padding: - gen_tasks = tasks - else: - gen_tasks = tasks.unsqueeze(-1) - add_tensor(gen_tasks, 'tasks') - add_tensor(prompt_vocab_size, 'prompt_vocab_size') - - if not self.paged_kv_cache: - for attn_idx, layer_idx in self.attn_to_general_idx.items(): - if not self.use_gpt_attention_plugin: - next_shape = (batch_size * beam_width, 2, - self.get_num_heads_kv(), - max_context_length + step, self.head_size) - # We will make current layer's output KV-cache overwrite previous layers input KV-cache - # buffer id: ... 5, 6, 7, 8, 9, ... - # layer n: out in - # layer n+1: out in - # layer n+2 out in - # And when finish a step, we will make every layer's in/out buffer index subtract 1 in - # a circular buffer way to make sure current outputs become next step's inputs. - num_buffers = self.num_attn_layers + 1 - input_idx = (attn_idx - (step % num_buffers)) % num_buffers - output_idx = (input_idx - 1) % num_buffers - input_name = self.kv_cache_buffer_names[input_idx] - output_name = self.kv_cache_buffer_names[output_idx] - - add_tensor_with_shape(self.buffer[input_name], - f'past_key_value_{layer_idx}', - next_shape) - add_tensor(self.buffer[output_name], - f'present_key_value_{layer_idx}') - else: - key_value_cache = self.buffer[ - f'present_key_value_{layer_idx}'] - add_tensor(key_value_cache, f'past_key_value_{layer_idx}') - add_tensor(key_value_cache, - f'present_key_value_{layer_idx}') - - if self.cross_attention: - cross_cache_buffer = self.buffer[ - f'cross_present_key_value_{layer_idx}'] - add_tensor(cross_cache_buffer, - f'cross_past_key_value_{layer_idx}') - add_tensor(cross_cache_buffer, - f'cross_present_key_value_{layer_idx}') - - for idx in range(self.first_layer, self.last_layer): - if self.layer_types[idx] != 'recurrent': - continue - if self.paged_state: - add_tensor(self.buffer[f'conv_state_ptr_{idx}'], - f'conv_state_ptr_{idx}') - add_tensor(self.buffer[f'rnn_state_ptr_{idx}'], - f'rnn_state_ptr_{idx}') - else: - # conv state - if self.use_mamba_conv1d_plugin: - conv_state_shape = (batch_size, self.conv_kernel - 1, - self.rnn_conv_dim_size) - else: - conv_state_shape = (batch_size, self.rnn_conv_dim_size, - self.conv_kernel - 1) - if step % 2: - add_tensor_with_shape( - self.buffer[f'1_present_conv_state_{idx}'], - f'past_conv_state_{idx}', conv_state_shape) - add_tensor(self.buffer[f'present_conv_state_{idx}'], - f'present_conv_state_{idx}') - else: - add_tensor_with_shape( - self.buffer[f'present_conv_state_{idx}'], - f'past_conv_state_{idx}', conv_state_shape) - add_tensor(self.buffer[f'1_present_conv_state_{idx}'], - f'present_conv_state_{idx}') - # rnn state - rnn_state = self.buffer[f'present_rnn_state_{idx}'] - add_tensor(rnn_state, f'past_rnn_state_{idx}') - add_tensor(rnn_state, f'present_rnn_state_{idx}') - - if self.paged_state and self.has_rnn_layers: - slot_mapping = torch.arange(0, - batch_size, - device='cuda', - dtype=torch.int32) - add_tensor(slot_mapping, 'slot_mapping') - - if self.use_gpt_attention_plugin and self.has_attn_layers: - # generation requests - host_request_types = torch.ones_like(context_lengths, - device='cpu').int() - if self.is_redrafter_mode: - torch.cuda.nvtx.range_push("device_request_types") - device_request_types = torch.ones_like( - context_lengths, device=self.device).int() - add_tensor(device_request_types, 'device_request_types') - torch.cuda.nvtx.range_pop() - if self.is_medusa_mode or self.is_redrafter_mode: - host_past_key_value_lengths = self.sequence_length_buffer.cpu() - else: - # previous [past_kv_length, is_context] has been deprecated. only past_kv_length should be given here - # Note we should use max_context_length here to align to max -- but isn't this done in attn plugin's max_element() already? - host_past_key_value_lengths = torch.tensor( - [max_context_length + step] * (batch_size * beam_width), - dtype=torch.int32, - device='cpu') - add_tensor(host_past_key_value_lengths, - 'host_past_key_value_lengths') - add_tensor(host_request_types, 'host_request_types') - # Sequence lengths are not used in the context phase actually. - sequence_length = self.sequence_length_buffer - - add_tensor_with_shape(sequence_length, 'sequence_length', - (batch_size * beam_width, )) - add_tensor_with_shape(self.host_sink_token_length, - 'host_sink_token_length', (1, )) - add_tensor_with_shape(self.host_max_attention_window_sizes, - f'host_max_attention_window_sizes', - (self.num_attn_layers, )) - if self.remove_input_padding: - add_tensor(host_context_lengths_local, 'host_context_lengths') - else: - if self.has_rnn_layers: - host_request_types = torch.ones_like(context_lengths, - device='cpu').int() - add_tensor(host_request_types, 'host_request_types') - if self.remove_input_padding: - add_tensor(host_context_lengths_local, - 'host_context_lengths') - if self.has_attn_layers: - add_tensor(attention_mask, 'attention_mask') - - if self.mapping.tp_size > 1: - add_tensor(self.all_reduce_workspace, 'all_reduce_workspace') - if self.use_gemm_allreduce_plugin: - found_tensor_names = [ - self.runtime.engine.get_tensor_name(i) - for i in range(self.runtime.engine.num_io_tensors) - ] - for name in found_tensor_names: - if name.startswith("gemm_allreduce_uc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.uc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_mc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.mc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_ipc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.get_ipc_ptrs(), - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - - # Since we are using a ping-pong context design and the lora weight remains constant within the same request, - # it is only necessary to set the lora weight for the first two steps. - if self.use_lora_plugin and step < 2: - for idx in range(self.num_layers): - layer_idx = idx + self.first_layer - for lora_module in (self.lora_target_modules + - self.missing_qkv_modules): - lora_ranks = f'{lora_module}_lora_ranks_{layer_idx}' - add_tensor(self.buffer[lora_ranks], lora_ranks) - lora_module = f'{lora_module}_lora_weights_pointers_{layer_idx}' - add_tensor(self.buffer[lora_module], lora_module) - if self.cross_attention and self.remove_input_padding: - add_tensor(encoder_input_lengths.to('cpu'), - 'host_encoder_input_lengths') - - if self.is_medusa_mode: - # Spec Decoding mask and position offsets are fixed for the whole session for Medusa. - add_tensor(self.buffer['spec_decoding_packed_mask'], - 'spec_decoding_packed_mask') - add_tensor(self.buffer['spec_decoding_position_offsets'], - 'spec_decoding_position_offsets') - add_tensor(self.buffer['spec_decoding_generation_lengths'], - 'spec_decoding_generation_lengths') - add_tensor(self.buffer['spec_decoding_use'], 'spec_decoding_use') - - if self.is_redrafter_mode: - set_redrafter_gen_tensors(self, batch_size, add_tensor, - add_tensor_with_shape) - torch.cuda.nvtx.range_pop() - - return tensors - - def _prepare_context_inputs(self, batch_size, context_lengths, - host_context_lengths, use_gpt_attention_plugin, - remove_input_padding, **kwargs): - - last_token_ids = context_lengths.detach().clone() - if (self.is_medusa_mode - or self.is_redrafter_mode) and not remove_input_padding: - # For Medusa, last_token_ids should contain the actual indices - last_token_ids = last_token_ids - 1 # sub 1 from context_lengths for indices - last_token_ids = last_token_ids.reshape([batch_size, -1]) - if (use_gpt_attention_plugin - or self.has_rnn_layers) and remove_input_padding: - last_token_ids = torch.cumsum(last_token_ids, dim=0).int() - ret = {'last_token_ids': last_token_ids} - - if use_gpt_attention_plugin: - max_context_length = kwargs.pop('max_context_length') - if remove_input_padding: - position_ids = torch.concat([ - torch.arange(0, - host_context_lengths[i], - dtype=torch.int32, - device='cuda') for i in range(batch_size) - ]) - else: - position_ids = torch.tensor(range(max_context_length), - dtype=torch.int32, - device='cuda').reshape( - [1, - -1]).expand([batch_size, -1]) - - perf_knob_tensor_size = 16 - context_runtime_perf_knobs = torch.tensor([-1] * - perf_knob_tensor_size, - dtype=torch.int64) - if self.multi_block_mode: - context_runtime_perf_knobs[0] = 1 # multi_block_mode - if self.enable_context_fmha_fp32_acc: - context_runtime_perf_knobs[ - 1] = 1 # enable_context_fmha_fp32_acc - ret['host_runtime_perf_knobs'] = context_runtime_perf_knobs - else: - if self.has_attn_layers: - input_ids = kwargs.pop('input_ids') - pad_id = kwargs.pop('pad_id', None) - attention_mask = _prepare_attention_mask(input_ids, pad_id) - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids.int() - ret['attention_mask'] = attention_mask - - if self.has_position_embedding and self.has_attn_layers: - ret['position_ids'] = position_ids - - if self.is_redrafter_mode: - self.buffer['position_ids_base'] = context_lengths.clone() - # NOTE: Generate random tensors using torch - redrafter_prepare_random_tensors(self, batch_size, initialize=True) - - return ret - - def _prepare_generation_inputs(self, batch_size, context_lengths, - use_gpt_attention_plugin, - remove_input_padding, **kwargs): - torch.cuda.nvtx.range_push("_prepare_generation_inputs") - - step = kwargs.pop('step') - last_token_ids = torch.ones_like(context_lengths) - if use_gpt_attention_plugin and (self.is_medusa_mode - or self.is_redrafter_mode): - if remove_input_padding: - if self.is_medusa_mode: - # For Medusa, last_token_ids should be [bs * seq] and should contain the actual indices (starts from 1) - last_token_ids = torch.ones(batch_size * - (self.num_draft_tokens + 1), - dtype=torch.int32, - device=context_lengths.device) - elif self.is_redrafter_mode: - torch.cuda.nvtx.range_push("last_token_ids_1s") - # update last_token_ids here (buffers already swapped) - last_token_ids = torch.ones(self.host_total_gen_token, - dtype=torch.int32, - device=context_lengths.device) - torch.cuda.nvtx.range_pop() - else: - # For Medusa, last_token_ids should be [bs, seq] and should contain the actual indices (starts from 0) - last_token_ids = torch.arange(self.num_draft_tokens + 1, - dtype=torch.int32, - device=context_lengths.device) - last_token_ids = last_token_ids.expand([batch_size, -1]) - if (use_gpt_attention_plugin - or self.has_rnn_layers) and remove_input_padding: - torch.cuda.nvtx.range_push("last_token_ids_cumsum") - last_token_ids = torch.cumsum(last_token_ids, dim=0).int() - torch.cuda.nvtx.range_pop() - ret = {'last_token_ids': last_token_ids} - - if use_gpt_attention_plugin: - if self.is_redrafter_mode: - torch.cuda.nvtx.range_push("position_ids_update") - # set position_ids - # buffers are swapped but sequence_length is not updated at this point - - if step != 0: - self.buffer['position_ids_base'] += self.buffer[ - 'num_accepted_tokens'] - position_ids = self.buffer['packed_position_ids'].view( - -1)[:self.host_total_gen_token] - if step == 0: - position_ids -= 1 - - torch.cuda.nvtx.range_pop() - else: - position_ids = context_lengths + step - if not remove_input_padding: - position_ids = torch.unsqueeze(position_ids, 1) - - perf_knob_tensor_size = 16 - gen_runtime_perf_knobs = torch.tensor([-1] * perf_knob_tensor_size, - dtype=torch.int64) - if self.multi_block_mode: - gen_runtime_perf_knobs[0] = 1 # multi_block_mode - if self.enable_context_fmha_fp32_acc: - gen_runtime_perf_knobs[1] = 1 # enable_context_fmha_fp32_acc - ret['host_runtime_perf_knobs'] = gen_runtime_perf_knobs - elif self.has_attn_layers: - attention_mask = kwargs.pop('attention_mask') - num_beams = kwargs.pop('num_beams') - attention_mask = torch.cat((attention_mask, - attention_mask.new_ones( - (batch_size * num_beams, 1))), - dim=-1).contiguous() - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids[:, -1].unsqueeze(-1) - position_ids = position_ids.int() - ret['attention_mask'] = attention_mask - - if self.has_position_embedding and self.has_attn_layers: - ret['position_ids'] = position_ids - if self.is_redrafter_mode: - # buffers are already swapped - # convert spec_decoding_mask to spec_decoding_packed_mask - redrafter_convert_spec_decoding_mask_to_packed_mask( - self, self.buffer['spec_decoding_generation_lengths']) - # NOTE: Generate random tensors using torch - redrafter_prepare_random_tensors(self, batch_size) - torch.cuda.nvtx.range_pop() - - return ret - - def _prepare_cross_attention_mask(self, batch_size, context_lengths, - cross_attention_mask): - cross_attention_mask_for_context = [] - cross_attention_mask_for_gen = [] - max_decoder_input_length = torch.max(context_lengths).item() - for batch_idx in range(batch_size): - decoder_input_length = context_lengths[batch_idx].item() - local_mask_for_context = cross_attention_mask[ - batch_idx][:decoder_input_length, :] - local_mask_for_gen = cross_attention_mask[batch_idx][ - decoder_input_length:, :] - if not self.use_gpt_attention_plugin: - local_mask_for_context = local_mask_for_context.unsqueeze(0) - if not self.remove_input_padding: - local_mask_for_context = torch.nn.functional.pad( - local_mask_for_context, - (0, 0, 0, - (max_decoder_input_length - decoder_input_length)), - "constant", False) - local_mask_for_gen = torch.nn.functional.pad( - local_mask_for_gen, - (0, 0, 0, - (max_decoder_input_length - decoder_input_length)), - "constant", False) - cross_attention_mask_for_context.append(local_mask_for_context) - # add additional dimension for batch size. - cross_attention_mask_for_gen.append(local_mask_for_gen.unsqueeze(0)) - - return torch.concat(cross_attention_mask_for_context), torch.concat( - cross_attention_mask_for_gen) - - def pp_communicate_new_tokens(self, should_stop, cache_indir, - sequence_length): - if self.mapping.is_last_pp_rank(): - for pg in self.mapping.pp_group: - if pg == self.mapping.rank: - continue - should_stop = should_stop.to(self.device) - self.nccl_comm.send(should_stop, pg) - self.nccl_comm.send(cache_indir, pg) - self.nccl_comm.send(sequence_length, pg) - self.nccl_comm.send(self.new_tokens, self.mapping.pp_group[0]) - else: - should_stop = torch.zeros(1, dtype=torch.bool, device=self.device) - self.nccl_comm.recv(should_stop, self.mapping.pp_group[-1]) - self.nccl_comm.recv(cache_indir, self.mapping.pp_group[-1]) - self.nccl_comm.recv(sequence_length, self.mapping.pp_group[-1]) - if self.mapping.is_first_pp_rank(): - self.nccl_comm.recv(self.new_tokens, self.mapping.pp_group[-1]) - return should_stop - - def pp_communicate_final_output_ids(self, final_output_ids, batch_size, - beam_width): - if self.mapping.is_last_pp_rank(): - self.nccl_comm.send(final_output_ids, self.mapping.pp_group[0]) - elif self.mapping.is_first_pp_rank(): - final_output_ids = torch.zeros( - (batch_size, beam_width, self.max_seq_length), - dtype=torch.int32, - device=self.device) - self.nccl_comm.recv(final_output_ids, self.mapping.pp_group[-1]) - return final_output_ids - - def finalize_decoder(self, - context_lengths, - batch_size, - beam_width, - scfg, - in_progress=False): - final_output_ids = None - if self.mapping.is_last_pp_rank(): - # output shape of self.gather_tree: [batch_size, beam_width, output_len] - beam_hyps_args = [ - self.beam_hyps_output_ids_cba, self.beam_hyps_seq_len_cba, - self.beam_hyps_cum_log_probs_cba, - self.beam_hyps_normed_scores_cba, self.beam_hyps_log_probs_cba, - self.beam_hyps_min_normed_scores, self.beam_hyps_num_beams, - self.beam_hyps_is_done - ] - - if scfg.use_beam_hyps and in_progress: - # self.gather_tree modifies these args. - # In streaming mode, this results in incorrect decoding in the following steps. - beam_hyps_args = copy.deepcopy(beam_hyps_args) - - final_output_ids = self.gather_tree( - self.sequence_length_buffer, self.output_ids, self.parent_ids, - self.end_ids, context_lengths, self.cum_log_probs, - self.log_probs, self.log_probs_tiled, *beam_hyps_args, - self.finished, self.length_penalty, batch_size, beam_width, - self.max_seq_length, scfg.use_beam_hyps) - - # Communicate ranks in Pipeline Parallelism - if self.mapping.has_pp(): - final_output_ids = self.pp_communicate_final_output_ids( - final_output_ids, batch_size, beam_width) - - return final_output_ids - - def find_best_medusa_path(self, - batch_size, - input_ids: torch.Tensor, - next_logits, - temp=0): - assert input_ids.shape[-1] == self.num_draft_tokens + 1 - best_path = [0] * batch_size - best_path_len = [1] * batch_size - next_tokens = [None] * batch_size - zero_pad = torch.zeros((batch_size, 1), - dtype=input_ids.dtype, - device=input_ids.device) - input_ids = torch.cat((input_ids, zero_pad), dim=-1) - if temp == 0: - new_tokens_raw = torch.argmax( - next_logits, dim=-1 - ) # TODO: can be done by treating [bs, nT, vocab] as [bs*nT, vocab] and using decoderOp? - new_tokens = torch.cat((new_tokens_raw, zero_pad), dim=-1) - input_paths = [ - input_ids[b, self.medusa_paths] for b in range(batch_size) - ] - new_paths = [ - new_tokens[b, self.medusa_paths] for b in range(batch_size) - ] - for b in range(batch_size): - equality = input_paths[b][:, 1:] == new_paths[b][:, :-1] - paths_correct_len = torch.cumprod(equality.int(), - dim=1).sum(dim=1) - best_path_len[b] = paths_correct_len.max().item() + 1 - if best_path_len[b] > 1: - best_path[b] = torch.argmax(paths_correct_len) - next_tokens[b] = new_paths[b][ - best_path[b]][:best_path_len[b]].clone() - - return best_path, best_path_len, next_tokens - - def filter_medusa_logits(self, batch_size, best_path, best_path_lengths, - medusa_logits): - """ - medusa_logits is of shape [nMH, bs, nMT+1, vocab] - - Returns [nMH, bs, vocab] - """ - filtered_logits = torch.empty( - (self.num_medusa_heads, batch_size, self.vocab_size_padded), - dtype=medusa_logits.dtype, - device=medusa_logits.device) - medusa_logits = medusa_logits.view(self.num_medusa_heads, batch_size, - self.num_draft_tokens + 1, -1) - for b in range(batch_size): - idx = self.medusa_paths[best_path[b], best_path_lengths[b] - 1] - filtered_logits[:, b, ...] = medusa_logits[:, b, idx, ...] - return filtered_logits - - def get_next_medusa_tokens(self, batch_size, next_medusa_logits): - next_medusa_tokens = [ - torch.zeros((batch_size, 1), - dtype=torch.int32, - device=next_medusa_logits.device) - ] # dummy token for now, TODO: update tree_ids and remove this - for i in range(self.num_medusa_heads): - medusa_token = torch.topk(next_medusa_logits[i, :, :], - self.medusa_topks[i], - dim=-1).indices - next_medusa_tokens.append(medusa_token) - next_medusa_tokens = torch.cat(next_medusa_tokens, dim=-1) - return next_medusa_tokens - - def locate_accepted_draft_tokens(self, batch_size, best_path, best_path_len, - draft_paths): - torch.cuda.nvtx.range_push("locate_accepted_draft_tokens") - best_path_len_tensor = best_path_len if isinstance( - best_path_len, torch.Tensor) else torch.tensor( - best_path_len, dtype=torch.int, device='cuda') - accepted_draft_token_counts = torch.maximum( - best_path_len_tensor - 1, - torch.tensor([0], device=best_path_len_tensor.device)) - accepted_draft_token_offsets = torch.zeros(batch_size + 1, - dtype=torch.int32, - device='cuda') - accepted_draft_token_offsets[1:] = torch.cumsum( - accepted_draft_token_counts, dim=0) - accepted_draft_token_offsets_cpu = accepted_draft_token_offsets.to( - 'cpu') - packed_accepted_draft_tokens_indices = torch.empty( - accepted_draft_token_offsets_cpu[batch_size], - dtype=torch.int32, - device='cuda') - for seq_idx in range(batch_size): - cur_draft_paths = draft_paths if self.is_medusa_mode else draft_paths[ - seq_idx] - seq_start = accepted_draft_token_offsets_cpu[seq_idx] - seq_end = accepted_draft_token_offsets_cpu[seq_idx + 1] - seq_accepted_draft_count = seq_end - seq_start - best_path_idx = best_path[seq_idx].cpu() if isinstance( - best_path[seq_idx], torch.Tensor) else best_path[seq_idx] - seq_accepted_token_indices = cur_draft_paths[ - best_path_idx, 1:1 + seq_accepted_draft_count] - packed_accepted_draft_tokens_indices[ - seq_start:seq_end] = seq_accepted_token_indices - 1 - # print("KV offsets & indices", accepted_draft_token_offsets, - # packed_accepted_draft_tokens_indices,) - torch.cuda.nvtx.range_pop() - return accepted_draft_token_offsets, packed_accepted_draft_tokens_indices - - def update_output_ids_by_offset(self, new_generated_ids, offsets): - # output_ids [batch_size, padded_input_length] - # new_generated_ids [batch_size, padded_accepted_length] - # offsets [batch_size] - # FIXME: using fused kernel to update the padded output ids. - batch_size = self.output_ids.shape[0] - for b in range(batch_size): - self.output_ids[b, offsets[b]:( - offsets[b] + self.accept_lengths[b] - )] = new_generated_ids[b][:self.accept_lengths[b]] - return - - def next_medusa_input_ids(self): - # self.new_tokens [batch_size, padded_accepted_length] - # self.accept_lengths [batch_size] - # self.medusa_new_tokens [batch_size, num_draft_tokens] - # FIXME: using fused kernel to generate the new medusa input ids. - batch_size = self.new_tokens.shape[0] - for b in range(batch_size): - self.generation_input_ids[b, 0] = self.new_tokens[ - b, self.accept_lengths[b] - 1] - self.generation_input_ids[b, 1:] = self.medusa_output_tokens[b, :] - - def reorder_kv_cache_for_beam_search( - self, - batch_size: int, - beam_width: int, - max_context_length: int, - step: int, - ): - if self.use_gpt_attention_plugin: - # Do nothing. - return - - # WAR: This degrades the latency performance in beam search - # due to memcpy. Recommend to use gpt attention plugin instead. - assert self.buffer is not None - assert self.parent_ids.shape[:2] == (batch_size, beam_width) - - cache_shape = (batch_size * beam_width, 2, self.get_num_heads_kv(), - max_context_length + step, self.head_size) - - import functools - numel = functools.reduce(lambda x, y: x * y, cache_shape) - - # attention layer num + 1 extra buffer. - num_buffers = self.num_attn_layers + 1 - for i in self.attn_to_general_idx: - # Cyclic buffers, an output becomes the next step's input. - input_idx = (i - (step % num_buffers)) % num_buffers - presents = self.buffer[self.kv_cache_buffer_names[input_idx]] - presents = presents.view(-1)[:numel].view(*cache_shape) - # parent_ids = (batch, beam, max_seq_len) - parent_ids = self.parent_ids[..., - max_context_length + step].view(-1) - - for batch_beam in range(batch_size * beam_width): - batch = batch_beam // beam_width - if parent_ids[batch_beam] != batch_beam % beam_width: - # Update past kv cache to parent beam's cache. - src_bbid = batch * beam_width + parent_ids[batch_beam] - presents[batch_beam, ...] = presents[src_bbid, ...] - - # OPTIMIZE: need to optimize this early-stop workflow. - def early_stop_criteria(self, batch_size, step, should_stop): - for b in range(batch_size): - if self.medusa_should_stop[b]: - self.accept_lengths[b] = 0 - continue - # output sequence length criteria. - prev_total_output_length = self.total_accept_lengths[b] - # end id criteria. - end_id_mask = self.new_tokens[ - b, :self.accept_lengths[b]] == self.end_ids[b] - should_stop_with_end_id = torch.any(end_id_mask) - self.medusa_should_stop[b] = self.medusa_should_stop[b] or ( - prev_total_output_length + self.accept_lengths[b] - >= self.max_new_tokens) or should_stop_with_end_id - # update accept lengths for the current step. - if (prev_total_output_length + self.accept_lengths[b] - >= self.max_new_tokens): - self.accept_lengths[b] = min( - self.max_new_tokens - prev_total_output_length, - self.accept_lengths[b]) - if should_stop_with_end_id: - # get the position of first end_id. - end_id_pos = (end_id_mask).nonzero(as_tuple=True)[0] - self.accept_lengths[b] = min(end_id_pos[0] + 1, - self.accept_lengths[b]) - self.total_accept_lengths[b] += self.accept_lengths[b] - - should_stop[0] = should_stop[0] or (step == self.max_new_tokens - - 1) or torch.all( - self.medusa_should_stop) - return should_stop - - def medusa_decode_and_verify(self, step, batch_size, logits): - medusa_logits = self.buffer['medusa_logits'] - best_path = None - best_path_lengths = None - if step == 0: - # logits buffer is of shape [bs, medusa_tokens+1, vocab] - # but during context phase, we get only [bs, 1, vocab] but contiguous - logits = logits.view(-1)[:batch_size * logits.shape[-1]].view( - batch_size, -1) - next_main_token_logits = logits.to(self.decoder_logits_dtype) - next_main_token = torch.argmax(next_main_token_logits, - dim=-1, - keepdim=True) - self.new_tokens = next_main_token - # NOTE: only one token's medusa logit will be written in. - medusa_logits = medusa_logits.view(self.num_draft_tokens + 1, - -1)[0, ...] - next_medusa_logits = medusa_logits.reshape( - self.num_medusa_heads, batch_size, - -1).to(self.decoder_logits_dtype) - next_medusa_tokens = self.get_next_medusa_tokens( - batch_size, next_medusa_logits) - self.medusa_output_tokens = next_medusa_tokens[:, - self.medusa_tree_ids[ - -self. - num_draft_tokens:]] - self.accept_lengths = torch.ones([batch_size], - dtype=torch.int32, - device=self.device) - else: - next_token_logits = logits.to(self.decoder_logits_dtype) - - best_path, best_path_lengths, next_main_tokens = self.find_best_medusa_path( - batch_size, self.generation_input_ids.view(batch_size, -1), - next_token_logits.view(batch_size, self.num_draft_tokens + 1, - -1)) - self.accept_lengths = torch.tensor(best_path_lengths, - device=self.device) - self.new_tokens = torch.nested.to_padded_tensor( - torch.nested.nested_tensor(next_main_tokens, dtype=torch.int32), - self.end_ids[0]) #FIXME end id padding. - next_medusa_logits = self.filter_medusa_logits( - batch_size, best_path, best_path_lengths, medusa_logits) - next_medusa_tokens = self.get_next_medusa_tokens( - batch_size, next_medusa_logits) - - self.medusa_output_tokens = next_medusa_tokens[:, - self.medusa_tree_ids[ - -self. - num_draft_tokens:]] - return best_path, best_path_lengths - - def process_logits_including_draft(self, step, batch_size, logits, - next_step_buffer): - """ - 1. Process logits to tokens and validate (Medusa) or process outputs (ReDrafter) - 2. Extract early stop criteria here : self.accept_length - 3. Update output ids : needs self.new_tokens and past_sequence_length - 4. Get next input_ids : self.[new_tokens, accept_lengths, medusa_output_tokens] - 5. Update KV cache : self.[sequence_length, num_draft_tokens] - 6. Update sequence_length_buffer and past_kv_length - """ - should_stop = torch.tensor([False], dtype=bool) - if self.is_medusa_mode: - # NOTE: this function call also updates self.[accept_lengths, new_tokens, medusa_output_tokens] - best_path, best_path_lengths = self.medusa_decode_and_verify( - step, batch_size, logits) - last_draft_paths = self.medusa_paths - # print(best_path, self.new_tokens, self.medusa_output_tokens) - last_draft_tokens_len = self.num_draft_tokens if step > 0 else 0 - cur_draft_tokens_len = self.num_draft_tokens - elif self.is_redrafter_mode: - # buffers are swapped at this point - last_draft_tokens = self.buffer['next_draft_tokens'] - new_draft_tokens = self.buffer['draft_tokens'] - last_draft_paths = self.buffer["next_draft_indices"] - last_draft_tokens_len = self.buffer[ - 'next_spec_decoding_generation_lengths'] - 1 if step > 0 else 0 - cur_draft_tokens_len = self.buffer[ - 'spec_decoding_generation_lengths'] - 1 - - best_path, best_path_lengths = process_redrafter_outputs( - self, step, batch_size, last_draft_tokens, new_draft_tokens) - # NOTE: stop criteria - torch.cuda.nvtx.range_push("early_stop_check") - if step == 0: - self.total_accept_lengths = self.accept_lengths.clone() - self.medusa_should_stop = torch.eq(self.new_tokens.reshape(-1), - self.end_ids) - should_stop[0] = torch.equal( - self.new_tokens.reshape(-1), - self.end_ids) or (step == self.max_new_tokens - 1) - else: - should_stop = self.early_stop_criteria(batch_size, step, - should_stop) - torch.cuda.nvtx.range_pop() - # NOTE: self.accept_lengths are the lengths of accepted tokens in the current step - # NOTE: self.sequence_length_buffer = num_past_kv_cache (accepted) + accept_lengths - torch.cuda.nvtx.range_push("update_output_ids") - self.update_output_ids_by_offset( - self.new_tokens, - self.sequence_length_buffer - last_draft_tokens_len) - torch.cuda.nvtx.range_pop() - - if step != self.max_new_tokens - 1 and not should_stop.item(): - if self.is_medusa_mode: - self.next_medusa_input_ids() - if step != 0: - assert best_path is not None and best_path_lengths is not None - accepted_draft_token_offsets, packed_accepted_draft_tokens_indices = self.locate_accepted_draft_tokens( - batch_size, best_path, best_path_lengths, last_draft_paths) - # update the KV cache - torch.cuda.nvtx.range_push("kv_update") - self.kv_cache_updater.update( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - self.sequence_length_buffer, last_draft_tokens_len) - torch.cuda.nvtx.range_pop() - - self.sequence_length_buffer += self.accept_lengths + cur_draft_tokens_len - last_draft_tokens_len - else: - self.sequence_length_buffer += cur_draft_tokens_len + 1 - - # NOTE: set the accepted tokens for the last step. - if should_stop.item(): - # remove num_draft_tokens for next generation. - # Runtime: denotes kv cache length start positions. - # Output: denotes the length of sequence length (input ids + output ids) - self.sequence_length_buffer += self.accept_lengths - last_draft_tokens_len - - if next_step_buffer is not None: - next_step_buffer['host_past_key_value_lengths'].to_torch().copy_( - self.sequence_length_buffer) - - return should_stop - - def handle_per_step( - self, - *, - cache_indirections: list, - step: int, - batch_size: int, - max_context_length: int, - beam_width: int, - input_ids: torch.Tensor, - hidden_states: torch.Tensor, - scfg: SamplingConfig, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_block_offsets: torch.Tensor, - cross_kv_cache_block_offsets: torch.Tensor, - host_cross_kv_cache_block_offsets: torch.Tensor, - prompt_embedding_table: torch.Tensor, - tasks: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths, - attention_mask: torch.Tensor, - cross_attention_mask_for_context: torch.Tensor, - cross_attention_mask_for_gen: torch.Tensor, - prompt_vocab_size: torch.Tensor, - ite: int, - sequence_limit_lengths: torch.Tensor, - sequence_lengths: torch.Tensor, - next_step_tensors: Dict[str, RuntimeTensor], - stop_words_data, - bad_words_data, - encoder_output: torch.Tensor, - encoder_input_lengths: torch.Tensor, - stopping_criteria: StoppingCriteria, - logits_processor: LogitsProcessor, - output_generation_logits: bool, - **kwargs, - ): - if self.debug_mode: - print( - f"=================================== STEP {step} ==================================" - ) - if step % 2: - context = self.runtime.context_0 - this_src_cache_indirection = cache_indirections[1] - this_tgt_cache_indirection = cache_indirections[0] - next_src_cache_indirection = cache_indirections[0] - else: - context = self.runtime.context_1 - this_src_cache_indirection = cache_indirections[0] - this_tgt_cache_indirection = cache_indirections[1] - next_src_cache_indirection = cache_indirections[1] - - position_ids_raw = kwargs.get('position_ids', None) - skip_cross_attn_blocks = kwargs.get('skip_cross_attn_blocks', None) - language_adapter_routings = kwargs.get('language_adapter_routings', - None) - if step == 0: - model_inputs = self._prepare_context_inputs( - batch_size=batch_size, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - use_gpt_attention_plugin=self.use_gpt_attention_plugin, - remove_input_padding=self.remove_input_padding, - max_context_length=max_context_length, - input_ids=input_ids, - pad_id=scfg.pad_id, - eos_id=scfg.end_id) - - if position_ids_raw is None: - # default iota position ids - position_ids = model_inputs.get('position_ids', None) - else: - # user input position ids - if self.remove_input_padding: - position_ids = torch.cat(position_ids_raw, dim=0) - else: - padded_position_ids = torch.nn.utils.rnn.pad_sequence( - position_ids_raw, batch_first=True, padding_value=0) - position_ids = padded_position_ids - last_token_ids = model_inputs.get('last_token_ids') - attention_mask = model_inputs.get('attention_mask', None) - context_runtime_perf_knobs = model_inputs.get( - 'host_runtime_perf_knobs', None) - host_context_progress = torch.tensor([0], dtype=torch.int64) - - if self.paged_kv_cache and self.has_attn_layers: - host_kv_cache_block_offsets = self.pools_kv_cache_manager.get_block_offsets( - beam_width=1) - kv_cache_block_offsets = host_kv_cache_block_offsets.to('cuda') - if self.cross_attention: - host_cross_kv_cache_block_offsets = self.cross_pools_kv_cache_manager.get_block_offsets( - beam_width=1) - cross_kv_cache_block_offsets = host_cross_kv_cache_block_offsets.to( - 'cuda') - - ctx_tensors = self._get_context_shape_buffer( - input_ids, - context_lengths, - host_context_lengths, - position_ids, - last_token_ids, - attention_mask, - cross_attention_mask_for_context, - this_src_cache_indirection, - kv_cache_block_offsets, - host_kv_cache_block_offsets, - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets, - hidden_states, - prompt_embedding_table, - tasks, - prompt_vocab_size, - encoder_output, - encoder_input_lengths, - host_runtime_perf_knobs=context_runtime_perf_knobs, - host_context_progress=host_context_progress, - skip_cross_attn_blocks=skip_cross_attn_blocks, - language_adapter_routings=language_adapter_routings, - ) - - context = self.runtime.ctx_context - self.runtime._set_tensors(context, ctx_tensors) - if self.debug_mode: - self.debug_buffer = { - name: tensor.to_torch() - for name, tensor in ctx_tensors.items() - } - if self.cuda_graph_mode: - # context mode, clean cuda graph instances - self.runtime.cuda_graph_instances = [None for _ in range(2)] - - if self.debug_mode and False: # TODO: after TRT bug is fixed - self.runtime._check_tensors(context) - # dynamic_decoder currently use torch's current stream, so must let TRT enqueue use same stream here - stream = torch.cuda.current_stream().cuda_stream - instance_idx = step % 2 - if self.cuda_graph_mode and self.runtime.cuda_graph_instances[ - instance_idx] is not None: - # launch cuda graph - CUASSERT( - cudart.cudaGraphLaunch( - self.runtime.cuda_graph_instances[instance_idx], stream)) - ok = True - else: - ok = self.runtime._run(context, stream) - - if not ok: - raise RuntimeError(f"Executing TRT engine failed step={step}!") - - # TODO: remove this Windows WAR after https://nvbugs/4460474 is fixed. - if platform.system() == "Windows" or self.debug_mode: - torch.cuda.synchronize() - - context_logits = None - if self.mapping.is_last_pp_rank(): - if step == 0 and self.gather_context_logits: - assert not self.is_medusa_mode and not self.is_redrafter_mode - context_logits = self.buffer['logits'].detach().clone() - # gather last token of context - if self.remove_input_padding: - # reshape self.buffer['logits'] from [bs, max_context_length, vocab] - # to [1, bs * max_context_length, vocab] - # Note that the data are put in the buffer without padding although - # the allocated buffer has padding. - self.buffer['logits'] = self.buffer['logits'].reshape( - [1, -1, self.vocab_size_padded]) - self.buffer['logits'] = torch.index_select( - self.buffer['logits'], 1, - last_token_ids - 1).view(batch_size, - self.vocab_size_padded) - else: - last_token_ids = last_token_ids.reshape(batch_size, 1, 1) - last_token_ids = last_token_ids.expand( - batch_size, 1, self.vocab_size_padded) - 1 - self.buffer['logits'] = torch.gather( - self.buffer['logits'], - dim=1, - index=last_token_ids.to(dtype=torch.int64)).view( - batch_size, self.vocab_size_padded) - - if step == 0 and beam_width > 1: - assert not self.is_medusa_mode and not self.is_redrafter_mode - assert not self.has_rnn_layers - # these tiled tensors are returned by handle_per_step(), so they can relay to the next generation calls - if not self.use_gpt_attention_plugin: - attention_mask = _tile_beam_width(attention_mask, beam_width) - context_lengths = _tile_beam_width(context_lengths, beam_width) - host_context_lengths = _tile_beam_width(host_context_lengths, - beam_width) - if encoder_input_lengths is not None: - encoder_input_lengths = _tile_beam_width( - encoder_input_lengths, beam_width) - - if tasks is not None: - tasks = _tile_beam_width(tasks, beam_width) - - # Move tiling before logit computing of context - if not self.paged_kv_cache: - for key in self.buffer: - # Note: this tiles both self attn cache and cross attn - # cache! both names contain "present_key_value" - if "present_key_value" in key: - if self.use_gpt_attention_plugin: - self.buffer[key] = _tile_beam_width( - self.buffer[key], beam_width) - else: - # In the OOTB path, KV cache should be contiguously - # tiled since TRT engine allocates past_kv cache of - # length context_length, i.e., we need a buffer of - # shape (batch * beam, 2, heads, context_length, head_size). - b, _, h, _, d = self.buffer[key].shape - numel = 2 * b * h * (max_context_length + step) * d - self.buffer[key] = _contiguous_tile_beam_width( - self.buffer[key], numel, beam_width) - - if self.mapping.is_last_pp_rank(): - self.buffer['logits'] = _tile_beam_width( - self.buffer['logits'], beam_width) - - generation_logits = None - if self.mapping.is_last_pp_rank(): - if self.gather_generation_logits or output_generation_logits: - generation_logits = self.buffer['logits'].detach().clone() - - # Initialize sequence_lengths (no paddings) for the generation phase. - if step == 0 and not self.is_medusa_mode and not self.is_redrafter_mode: # Medusa/ReDrafter has its own logic - self.sequence_length_buffer = context_lengths.detach().clone() - - if self.is_redrafter_mode: - # to simplify some processing logic, always swap buffers after execution - exchange_redrafter_buffers(self) - - # NOTE: handle next step. - if not step == self.max_new_tokens - 1: - # Set shape and address for the next step - model_inputs = self._prepare_generation_inputs( - batch_size=batch_size, - context_lengths=context_lengths, - use_gpt_attention_plugin=self.use_gpt_attention_plugin, - remove_input_padding=self.remove_input_padding, - step=step, - num_beams=beam_width, - attention_mask=attention_mask, - ) - - if position_ids_raw is None: - position_ids = model_inputs.get('position_ids', None) - else: - position_ids = torch.cat( - [p[-1:] + step + 1 for p in position_ids_raw], dim=0) - if not self.remove_input_padding: - position_ids = torch.unsqueeze(position_ids, 1) - last_token_ids = model_inputs.get('last_token_ids') - attention_mask = model_inputs.get('attention_mask', None) - gen_runtime_perf_knobs = model_inputs.get('host_runtime_perf_knobs', - None) - host_context_progress = torch.tensor([0], dtype=torch.int64) - - # Prepare for the next step, and always allocate 1 token slot. - if self.paged_kv_cache and self.has_attn_layers: - # Iterate to the next step in KV cache manager. - # Increase number of tokens for all unfinished sequences. - # And allocate new blocks if needed. - # We set this to False for all sequences, since we use only length criterion to stop now - # OPTIMIZE: find a better of adding multiple tokens for paged kv cache. - torch.cuda.nvtx.range_push("paged_kv_alloc") - if self.is_redrafter_mode and self.max_draft_tokens > 0: - add_token_count = (self.max_draft_tokens + - 1) * 2 if step == 0 else torch.max( - self.accept_lengths).item() - assert add_token_count > 0 - for _ in range(add_token_count): - self.pools_kv_cache_manager.step([False] * batch_size) - if self.is_medusa_mode and self.num_draft_tokens > 0: - # Allocate kv cache token slots for next step. - # Make sure there are always > (num_draft_tokens + 1) free token slots. - # Allocate (num_draft_tokens + 1) * 2 for safety as we don't know the current step or next step's accepted lengths. - add_token_count = (self.num_draft_tokens + - 1) * 2 if step == 0 else torch.max( - self.accept_lengths).item() - assert add_token_count > 0 - for _ in range(add_token_count): - self.pools_kv_cache_manager.step([False] * batch_size) - else: - self.pools_kv_cache_manager.step([False] * batch_size) - torch.cuda.nvtx.range_pop() - torch.cuda.nvtx.range_push("paged_kv_post_alloc") - host_kv_cache_block_offsets = self.pools_kv_cache_manager.get_block_offsets( - beam_width) - kv_cache_block_offsets = host_kv_cache_block_offsets.to('cuda') - if self.cross_attention: - host_cross_kv_cache_block_offsets = self.cross_pools_kv_cache_manager.get_block_offsets( - beam_width) - cross_kv_cache_block_offsets = host_cross_kv_cache_block_offsets.to( - 'cuda') - torch.cuda.nvtx.range_pop() - - next_context = self.runtime.context_1 if step % 2 else self.runtime.context_0 - cross_attention_mask_step = None - if cross_attention_mask_for_gen is not None: - # cross_attention_mask_for_gen shape [batch_size, max_output_length, max_encoder_input_length] - decode_step = step - if decode_step == 0: - decode_step += 1 - if self.use_gpt_attention_plugin: - cross_attention_mask_step = cross_attention_mask_for_gen[:, ( - decode_step - 1), :] - else: - cross_attention_mask_step = cross_attention_mask_for_gen[:, ( - decode_step - 1):decode_step, :] - next_step_tensors = self._get_next_step_shape_buffer( - batch_size, - beam_width, - max_context_length, - step, - context_lengths, - host_context_lengths, - position_ids, - last_token_ids, - attention_mask, - cross_attention_mask_step, - next_src_cache_indirection, - kv_cache_block_offsets, - host_kv_cache_block_offsets, - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets, - hidden_states, - prompt_embedding_table, - tasks, - prompt_vocab_size, - encoder_output, - encoder_input_lengths, - host_runtime_perf_knobs=gen_runtime_perf_knobs, - host_context_progress=host_context_progress, - skip_cross_attn_blocks=skip_cross_attn_blocks, - language_adapter_routings=language_adapter_routings) - - # there are some tensors created inside the _get_next_step_shape_buffer, not owned by any object - # needs to pro-long the life time of the tensors inside the next_step_tensors array - # otherwise, it maybe released before the next step actually enqueued - # one way to prolong it is to return the list, and destroy it in next step by assigning new values - torch.cuda.nvtx.range_push("_set_tensors") - self.runtime._set_tensors(next_context, next_step_tensors) - torch.cuda.nvtx.range_pop() - - if logger.level == "verbose": - self.runtime.print_context_info( - next_context, int(next_context == self.runtime.context_1)) - - if self.cuda_graph_mode: - self._capture_cuda_graph_and_instantiate( - next_context, stream, step) - - should_stop = None - logits = None - if self.mapping.is_last_pp_rank(): - logits = self.buffer['logits'] - if self.is_redrafter_mode: - should_stop = self.process_logits_including_draft( - step, batch_size, logits, next_step_tensors) - elif logits is not None: - if self.is_medusa_mode: - should_stop = self.process_logits_including_draft( - step, batch_size, logits, next_step_tensors) - else: - if logits_processor is not None: - final_output_ids = self.finalize_decoder( - context_lengths, - batch_size, - beam_width, - scfg, - in_progress=True) - # keep the shape as same as huggingface stopping_criteria - final_output_ids_ = final_output_ids.reshape( - -1, final_output_ids.size(-1)) - logits = logits_processor(step, final_output_ids_, - logits) - self.buffer['logits'] = logits - # [batch_size x beam_width, vocab_size_padded] -> [batch_size, beam_width, vocab_size_padded] - next_token_logits = logits.reshape( - (batch_size, beam_width, - -1)).to(self.decoder_logits_dtype) - decode_step = step + max_context_length - - stop_words_list_ptrs, stop_words_lens, max_stop_words_len = stop_words_data - bad_words_list_ptrs, bad_words_lens, max_bad_words_len = bad_words_data - - should_stop = self.dynamic_decoder.forward( - next_token_logits, decode_step, max_context_length, - self.max_attention_window_size, self.sink_token_length, - ite, batch_size, self.end_ids, self.embedding_bias_opt, - context_lengths, sequence_limit_lengths, - stop_words_list_ptrs, stop_words_lens, - max_stop_words_len, bad_words_list_ptrs, bad_words_lens, - max_bad_words_len, this_src_cache_indirection, - self.output_ids, self.new_tokens, self.finished, - self.finished, self.sequence_length_buffer, - self.cum_log_probs, self.log_probs, - self.log_probs_tiled, self.parent_ids, - this_tgt_cache_indirection, - self.beam_hyps_output_ids_cba, - self.beam_hyps_seq_len_cba, - self.beam_hyps_cum_log_probs_cba, - self.beam_hyps_normed_scores_cba, - self.beam_hyps_log_probs_cba, - self.beam_hyps_min_normed_scores, - self.beam_hyps_num_beams, self.beam_hyps_is_done, - scfg.use_beam_hyps) - - if not self.use_gpt_attention_plugin: - self.reorder_kv_cache_for_beam_search( - batch_size, beam_width, max_context_length, step) - - if stopping_criteria is not None and not should_stop.item(): - final_output_ids = self.finalize_decoder( - context_lengths, - batch_size, - beam_width, - scfg, - in_progress=True) - # keep the shape as same as huggingface stopping_criteria - final_output_ids_ = final_output_ids.reshape( - -1, final_output_ids.size(-1)) - should_stop[0] = stopping_criteria( - step, final_output_ids_, logits) - - if self.runtime._is_profiling(): - if not context.report_to_profiler(): - logger.warning("Runtime report to profiler failed.") - self.runtime._insert_step_to_profiler(step) - - if self.mapping.has_pp(): - should_stop = self.pp_communicate_new_tokens( - should_stop, this_tgt_cache_indirection, - self.sequence_length_buffer) - - if self.paged_kv_cache and self.has_attn_layers: - if (step >= self.max_new_tokens - 1) or (should_stop is not None - and should_stop.item()): - # Free all blocks in all sequences. - # With in-flight batching and while loop we'll free some sequences, when they are done - self.pools_kv_cache_manager.step([True] * batch_size) - if self.cross_attention: - self.cross_pools_kv_cache_manager.step([True] * batch_size) - - if self.debug_mode: - self.dump_debug_buffers(step) - - if next_step_tensors is not None: - self.debug_buffer = { - name: tensor.to_torch() - for name, tensor in next_step_tensors.items() - } - - return should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths - - def dump_debug_buffers(self, step: int) -> None: - if self.debug_tensors_to_save is not None: - # restricted written tensors according to filter - debug_tensor_names = copy.deepcopy(list(self.debug_buffer.keys())) - for k in debug_tensor_names: - if all([kk not in k for kk in self.debug_tensors_to_save]): - self.debug_buffer.pop(k) - - debug_dir = Path( - f"tllm_debug/PP_{self.mapping.pp_rank}/TP_{self.mapping.tp_rank}/CP_{self.mapping.cp_rank}" - ) - debug_dir.mkdir(parents=True, exist_ok=True) - - for name, t in self.debug_buffer.items(): - # convert tensor name to valid file name - print("Saving: ", name) - fname = name.replace("/", ".") - t = torch_to_numpy(t.float()) - np.save(debug_dir / f"{fname}-step{step}.npy", t) - - txt_format = "%d" if t.dtype in [np.int32, np.int8] else '%.18e' - np.savetxt( - debug_dir / f"{fname}-step{step}.txt", - t.reshape(-1, t.shape[-1]), # savetxt accepts 2 dims only - fmt=txt_format) - - def decode_regular(self, - *, - batch_size: int, - scfg: SamplingConfig, - sequence_lengths: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths, - max_context_length: int, - beam_width: int, - cache_indirections: list, - input_ids: torch.Tensor, - hidden_states: torch.Tensor, - prompt_embedding_table: torch.Tensor, - tasks: torch.Tensor, - prompt_vocab_size: torch.Tensor, - ite: int, - sequence_limit_lengths: torch.Tensor, - stop_words_data, - bad_words_data, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - return_dict: bool = False, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - stopping_criteria: StoppingCriteria = None, - logits_processor: LogitsProcessor = None, - cross_attention_mask: List[torch.Tensor] = None, - **kwargs): - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - attention_mask = None - outputs_context_logits = None - outputs_generation_logits = [] - - def get_outputs_dict(output_ids, num_steps=self.max_new_tokens): - outputs = {} - outputs['output_ids'] = output_ids - if scfg.output_log_probs: - outputs['log_probs'] = self.log_probs - if scfg.output_cum_log_probs: - outputs['cum_log_probs'] = self.cum_log_probs - if output_sequence_lengths: - outputs[ - 'sequence_lengths'] = self.sequence_length_buffer.reshape( - [batch_size, beam_width]) - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - if self.gather_generation_logits or output_generation_logits: - outputs['generation_logits'] = outputs_generation_logits - if self.is_medusa_mode or self.is_redrafter_mode: - outputs['steps_to_finish'] = num_steps - if self.is_medusa_mode: - outputs['medusa_output_tokens'] = self.medusa_output_tokens - outputs['accept_lengths'] = self.accept_lengths - if self.medusa_temperature != 0.0: - outputs['medusa_output_logits'] = self.medusa_output_logits - return outputs - - benchmark_profiler = kwargs.get('benchmark_profiler', None) - generation_phase_step_count = 0 - - if benchmark_profiler is not None and benchmark_profiler.is_recording_perf_profile: - self.runtime._set_profiler() - - def profile_fn(benchmark_profiler_obj, step_count): - if benchmark_profiler_obj is not None: - benchmark_profiler_obj.record_cuda_event('last_token') - benchmark_profiler_obj.record_elapsed_time( - 'first_token', 'last_token', 'generation_time') - benchmark_profiler_obj.add_aux_info('generation_step_count', - step_count) - - # prepare cross attention mask. - cross_attention_mask_for_context = None - cross_attention_mask_for_gen = None - if cross_attention_mask is not None: - cross_attention_mask_for_context, cross_attention_mask_for_gen = self._prepare_cross_attention_mask( - batch_size, context_lengths, cross_attention_mask) - if self.use_gpt_attention_plugin: - # When we use plugin, the data type of cross_attention_mask is bool. - # When we don't use plugin, the data type of cross_attention_mask is int32 - cross_attention_mask_for_context = cross_attention_mask_for_context.to( - torch.bool) - cross_attention_mask_for_gen = cross_attention_mask_for_gen.to( - torch.bool) - - next_step_tensors = None - for step in range(0, self.max_new_tokens): - - should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths = self.handle_per_step( - cache_indirections=cache_indirections, - step=step, - batch_size=batch_size, - max_context_length=max_context_length, - beam_width=beam_width, - input_ids=input_ids, - hidden_states=hidden_states, - scfg=scfg, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - attention_mask=attention_mask, - cross_attention_mask_for_context= - cross_attention_mask_for_context, - cross_attention_mask_for_gen=cross_attention_mask_for_gen, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - sequence_lengths=sequence_lengths, - next_step_tensors=next_step_tensors, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - output_generation_logits=output_generation_logits, - **kwargs, - ) - if step == 0: - if benchmark_profiler is not None: - benchmark_profiler.record_cuda_event('first_token') - else: - generation_phase_step_count = generation_phase_step_count + 1 - - if self.mapping.is_last_pp_rank(): - if step == 0 and self.gather_context_logits: - outputs_context_logits = context_logits - if self.gather_generation_logits or output_generation_logits: - outputs_generation_logits.append(generation_logits) - - if should_stop is not None and should_stop.item(): - profile_fn(benchmark_profiler, generation_phase_step_count) - if self.is_medusa_mode or self.is_redrafter_mode: - # just hack away for now - final_output_ids = self.output_ids.clone().unsqueeze(1) - final_output_ids = final_output_ids[:, :, :self. - max_seq_length - - self.max_draft_tokens] - else: - final_output_ids = self.finalize_decoder( - context_lengths, batch_size, beam_width, scfg) - - if self.mapping.is_first_pp_rank(): - if return_dict: - return get_outputs_dict(final_output_ids, step + 1) - else: - return final_output_ids - elif self.mapping.is_last_pp_rank(): - outputs = {} - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - if self.gather_generation_logits or output_generation_logits: - outputs['generation_logits'] = outputs_generation_logits - return outputs - else: - return None - - assert not self.is_medusa_mode and not self.is_redrafter_mode, "the custom decoder doesn't support medusa/redrafter." - - profile_fn(benchmark_profiler, generation_phase_step_count) - - final_output_ids = self.finalize_decoder(context_lengths, batch_size, - beam_width, scfg) - if self.mapping.is_first_pp_rank(): - if return_dict: - return get_outputs_dict(final_output_ids) - else: - return final_output_ids - elif self.mapping.is_last_pp_rank(): - outputs = {} - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - if self.gather_generation_logits or output_generation_logits: - outputs['generation_logits'] = outputs_generation_logits - return outputs - else: - return None - - def decode_stream(self, - *, - batch_size: int, - scfg: SamplingConfig, - sequence_lengths: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths, - max_context_length: int, - beam_width: int, - cache_indirections: list, - input_ids: torch.Tensor, - hidden_states: torch.Tensor, - prompt_embedding_table: torch.Tensor, - tasks: torch.Tensor, - prompt_vocab_size: torch.Tensor, - ite: int, - sequence_limit_lengths: torch.Tensor, - stop_words_data, - bad_words_data, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - return_dict: bool = False, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - stopping_criteria: StoppingCriteria = None, - logits_processor: LogitsProcessor = None, - cross_attention_mask: List[torch.Tensor] = None, - **kwargs): - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - attention_mask = None - outputs_context_logits = None - - def get_outputs_dict(output_ids): - outputs = {} - outputs['output_ids'] = output_ids - if output_sequence_lengths: - outputs[ - 'sequence_lengths'] = self.sequence_length_buffer.reshape( - [batch_size, beam_width]) - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - return outputs - - # prepare cross attention mask. - cross_attention_mask_for_context = None - cross_attention_mask_for_gen = None - if cross_attention_mask is not None: - cross_attention_mask_for_context, cross_attention_mask_for_gen = self._prepare_cross_attention_mask( - batch_size, context_lengths, cross_attention_mask) - - next_step_tensors = None - for step in range(0, self.max_new_tokens): - - should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths = self.handle_per_step( - cache_indirections=cache_indirections, - step=step, - batch_size=batch_size, - max_context_length=max_context_length, - beam_width=beam_width, - input_ids=input_ids, - hidden_states=hidden_states, - scfg=scfg, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - attention_mask=attention_mask, - cross_attention_mask_for_context= - cross_attention_mask_for_context, - cross_attention_mask_for_gen=cross_attention_mask_for_gen, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - sequence_lengths=sequence_lengths, - next_step_tensors=next_step_tensors, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - output_generation_logits=output_generation_logits, - ) - if step == 0: - outputs_context_logits = context_logits - if should_stop is not None: - - final_output_ids = self.finalize_decoder(context_lengths, - batch_size, - beam_width, - scfg, - in_progress=True) - - if self.mapping.is_first_pp_rank(): - if return_dict: - yield get_outputs_dict(final_output_ids) - else: - yield final_output_ids - else: - yield None - - if should_stop.item(): - return - - final_output_ids = self.finalize_decoder(context_lengths, batch_size, - beam_width, scfg) - if self.mapping.is_first_pp_rank(): - if return_dict: - yield get_outputs_dict(final_output_ids) - else: - yield final_output_ids - else: - yield None - - def decode_batch(self, - input_ids: Sequence[torch.Tensor], - sampling_config: SamplingConfig, - streaming: bool = False, - **kwargs): - input_ids, context_lengths = _prepare_input_ids(input_ids) - return self.decode(input_ids, - context_lengths, - sampling_config, - streaming=streaming, - **kwargs) - - # As dynamic_decoder uses torch's current stream, we must ensure it runs on the same stream that - # dynamic_decoder was set up with - @cuda_stream_guard - def decode(self, - input_ids: torch.Tensor, - context_lengths: torch.Tensor, - sampling_config: SamplingConfig, - prompt_embedding_table: torch.Tensor = None, - tasks: torch.Tensor = None, - prompt_vocab_size: torch.Tensor = None, - stop_words_list=None, - bad_words_list=None, - streaming: bool = False, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - return_dict: bool = False, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - stopping_criteria: StoppingCriteria = None, - logits_processor: LogitsProcessor = None, - cross_attention_mask: List[torch.Tensor] = None, - **kwargs): - scfg = sampling_config - batch_size = context_lengths.size(0) - beam_width = scfg.num_beams - max_context_length = torch.max(context_lengths).item() - host_context_lengths = context_lengths.cpu() - assert batch_size == self.batch_size, \ - "Given batch size is different from the one used in setup()," \ - "rerun the setup function with the new batch size to avoid buffer overflow." - assert max_context_length <= self.max_context_length, \ - "Given input length is large then the one used in setup()," \ - "rerun the setup function with the new max_context_length to avoid buffer overflow." - assert beam_width == self.beam_width, \ - "Given beam width is different from the one used in setup()," \ - "rerun the setup function with the new beam width to avoid buffer overflow." - assert self.sink_token_length <= torch.min(context_lengths).item(), \ - "Given sink token length is larger than shortest context length," \ - "rerun the setup function with a smaller sink token length." - ite = 0 # index of local batches, will always be 0 if pp_size = 1 - - if self.remove_input_padding and input_ids.dim() == 2: - assert input_ids.shape[ - 0] == 1, "Packed 2D input must have shape [1, ]" - input_ids = input_ids.squeeze(0) - - self.__setup_decoder(input_ids, scfg, host_context_lengths) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - sequence_limit_lengths = torch.full((batch_size, 1), - self.max_seq_length, - dtype=torch.int32, - device=self.device) - - # Sequence_lengths for the dynamic decoder still has the input paddings. - sequence_lengths = torch.full((batch_size * beam_width, 1), - max_context_length, - dtype=torch.int32, - device=self.device) - - cache_indirections = [ - torch.full(( - batch_size, - beam_width, - self.max_attention_window_size, - ), - 0, - dtype=torch.int32, - device=self.device), - torch.full(( - batch_size, - beam_width, - self.max_attention_window_size, - ), - 0, - dtype=torch.int32, - device=self.device) - ] # ping-pong buffers - - hidden_states = None - if self.mapping.has_pp(): - max_num_tokens = max(batch_size * beam_width, - batch_size * self.max_seq_length) - hidden_size = self.hidden_size * self.mapping.tp_size - hidden_states = torch.zeros((1, max_num_tokens, hidden_size)) - - # Init KV cache block manager - if self.paged_kv_cache and self.has_attn_layers: - num_blocks, max_blocks_per_seq = self._get_num_paged_blocks( - self.max_attention_window_size, self.sink_token_length) - - self.buffer[ - f'host_kv_cache_pool_pointers'] = self._memory_pool_allocator.get_kv_cache_pool_pointers( - ) - self.buffer[ - f'host_kv_cache_pool_mapping'] = self._memory_pool_allocator.pool_mapping - - self.pools_kv_cache_manager = PoolsKVCacheManager( - self._memory_pool_allocator.pools_metadata, - max_blocks_per_seq, - num_blocks, - self.tokens_per_block, - self.head_size, - max_attention_window_size=self.max_attention_window_size, - beam_width=beam_width, - sink_token_len=self.sink_token_length) - - if self.cross_attention: - cross_num_blocks, max_cross_blocks_per_seq = self._get_num_paged_blocks( - self.encoder_max_input_length, sink_token_length=0) - self.buffer[ - f'host_cross_kv_cache_pool_pointers'] = self._cross_memory_pool_allocator.get_kv_cache_pool_pointers( - ) - self.buffer[ - f'host_cross_kv_cache_pool_mapping'] = self._cross_memory_pool_allocator.pool_mapping - - self.cross_pools_kv_cache_manager = PoolsKVCacheManager( - self._cross_memory_pool_allocator.pools_metadata, - max_cross_blocks_per_seq, - cross_num_blocks, - self.tokens_per_block, - self.head_size, - max_attention_window_size=self.encoder_max_input_length, - beam_width=beam_width, - sink_token_len=self.sink_token_length) - - # Add sequences to the manager - for bi in range(batch_size): - generation_sequence = GenerationSequence(seq_idx=bi, - batch_idx=bi) - self.pools_kv_cache_manager.add_sequence( - generation_sequence, max_context_length) - if self.cross_attention: - cross_generation_sequence = GenerationSequence(seq_idx=bi, - batch_idx=bi) - self.cross_pools_kv_cache_manager.add_sequence( - cross_generation_sequence, - self.encoder_max_input_length, - always_share_across_beam=True) - # cross attention paged kv cache should always share the context blocks across beams - # due to the fact that we are not adding new key/value cache to cross kv in generation - - if self.is_medusa_mode or self.is_redrafter_mode: - if self.quant_mode.has_kv_cache_quant(): - # Since torch does not support fp8 now, using int8 here. - kv_cache_type = torch.int8 - else: - kv_cache_type = self.dtype if self.paged_kv_cache else self._tensor_dtype( - f'present_key_value_{self.first_layer}') - self.history_max_seq_length = [max_context_length] - self.kv_cache_updater = KVCacheUpdater() - assert not self.cross_attention - assert self.use_gpt_attention_plugin - - if self.paged_kv_cache: - self.kv_cache_updater.init_paged_kv_cache( - self.num_layers, self.get_num_heads_kv(), self.head_size, - kv_cache_type, self.pools_kv_cache_manager, - self.buffer[f'host_kv_cache_pool_pointers']) - else: - past_key_value_list = [ - self.buffer[f'present_key_value_{i}'] - for i in range(self.first_layer, self.last_layer) - ] - self.kv_cache_updater.init_linear_kv_cache( - self.num_layers, self.get_num_heads_kv(), self.head_size, - kv_cache_type, past_key_value_list) - - stop_words_lens = None - stop_words_list_ptrs = None - max_stop_words_len = 0 - if stop_words_list is not None: - stop_words_list = torch.from_numpy(stop_words_list).contiguous().to( - 'cuda') - max_stop_words_len = stop_words_list.shape[2] - stop_words_lens = torch.full((batch_size, ), - max_stop_words_len, - dtype=torch.int32).to('cuda') - stop_words_list_ptrs = torch.zeros((batch_size), dtype=torch.int64) - for bi in range(batch_size): - stop_words_list_ptrs[bi] = stop_words_list.data_ptr( - ) + bi * 2 * max_stop_words_len * stop_words_list.element_size( - ) - stop_words_list_ptrs = stop_words_list_ptrs.to('cuda') - stop_words_data = (stop_words_list_ptrs, stop_words_lens, - max_stop_words_len) - - bad_words_lens = None - bad_words_list_ptrs = None - max_bad_words_len = 0 - if bad_words_list is not None: - bad_words_list = torch.from_numpy(bad_words_list).contiguous().to( - 'cuda') - max_bad_words_len = bad_words_list.shape[2] - bad_words_lens = torch.full((batch_size, ), - max_bad_words_len, - dtype=torch.int32).to('cuda') - bad_words_list_ptrs = torch.zeros((batch_size), dtype=torch.int64) - for bi in range(batch_size): - bad_words_list_ptrs[bi] = bad_words_list.data_ptr( - ) + bi * 2 * max_bad_words_len * bad_words_list.element_size() - bad_words_list_ptrs = bad_words_list_ptrs.to('cuda') - bad_words_data = (bad_words_list_ptrs, bad_words_lens, - max_bad_words_len) - - # start context phase - if streaming: - return self.decode_stream( - batch_size=batch_size, - scfg=scfg, - sequence_lengths=sequence_lengths, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_context_length, - beam_width=beam_width, - cache_indirections=cache_indirections, - input_ids=input_ids, - hidden_states=hidden_states, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - output_generation_logits=output_generation_logits, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - output_sequence_lengths=output_sequence_lengths, - return_dict=return_dict, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - cross_attention_mask=cross_attention_mask, - **kwargs, - ) - else: - return self.decode_regular( - batch_size=batch_size, - scfg=scfg, - sequence_lengths=sequence_lengths, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_context_length, - beam_width=beam_width, - cache_indirections=cache_indirections, - input_ids=input_ids, - hidden_states=hidden_states, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - return_dict=return_dict, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - cross_attention_mask=cross_attention_mask, - **kwargs, - ) - - -class ChatGLMGenerationSession(GenerationSession): - - def __init__( - self, - model_config: ModelConfig, - engine_buffer, - mapping: Mapping, - debug_mode=False, - debug_tensors_to_save=None, - cuda_graph_mode=False, - stream: torch.cuda.Stream = None, - ): - - super().__init__( - model_config, - engine_buffer, - mapping, - debug_mode, - debug_tensors_to_save, - cuda_graph_mode, - stream, - ) - - self.mask_index_tensor = None - - def _prepare_context_inputs(self, batch_size, context_lengths, - use_gpt_attention_plugin, remove_input_padding, - **kwargs): - - max_context_length = kwargs.pop('max_context_length') - last_token_ids = context_lengths.detach().clone() - - if remove_input_padding: - input_lengths_acc = torch.cumsum(torch.cat( - [torch.IntTensor([0]).cuda(), context_lengths], dim=0), - dim=0) - position_ids = torch.zeros([2, input_lengths_acc[-1]], - dtype=torch.int32) - for i in range(batch_size): - position_ids[0, input_lengths_acc[i]:input_lengths_acc[ - i + 1]] = torch.arange(0, - context_lengths[i], - dtype=torch.int32) - position_ids[0, input_lengths_acc[i + 1] - - 1] = context_lengths[i] - 2 - position_ids[1, input_lengths_acc[i + 1] - 1] = 1 - position_ids = position_ids.int().cuda() - last_token_ids = torch.cumsum(last_token_ids, dim=0).int().cuda() - - # specialization for GLM series models - if kwargs["pad_id"] in [50256, 50259]: - if kwargs["pad_id"] == 50256: # glm_2b / glm_10b - mask_ids = [50260, 50264, 50263] - else: # glm_10b_chinese / glm_large_chinese - mask_ids = [50003, 50008, 50009] - - self.mask_index_tensor = \ - torch.zeros([batch_size], dtype=torch.int32) - position_ids = position_ids.cpu() - for i in range(batch_size): - length = context_lengths[i] - input_ids = kwargs["input_ids"][ - 0:context_lengths[i]] if i == 0 else kwargs[ - "input_ids"][sum(context_lengths[0:i] - ):sum(context_lengths[0:i]) + - length] - mask_index = [ - torch.where(input_ids == id)[0].int() for id in mask_ids - ] - tail_index = torch.Tensor([max_context_length]).int().cuda() - mask_index.append(tail_index) - mask_index = torch.cat(mask_index, dim=0).min() - self.mask_index_tensor[i] = int(mask_index) - position_ids[0][sum(context_lengths[0:i + 1]) - - 1] = int(mask_index) - position_ids = position_ids.cuda() - else: - position_ids = torch.zeros([batch_size, 2, max_context_length], - dtype=torch.int32) - position_ids[:, 0, :] = torch.arange(max_context_length) - - # specialization for GLM series models - if kwargs["pad_id"] in [50256, 50259]: - if kwargs["pad_id"] == 50256: # glm_2b / glm_10b - mask_ids = [50260, 50264, 50263] - else: # glm_10b_chinese / glm_large_chinese - mask_ids = [50003, 50008, 50009] - - self.mask_index_tensor = \ - torch.zeros([batch_size], dtype=torch.int32) - for i in range(batch_size): - length = context_lengths[i] - input_ids = kwargs["input_ids"][i] - mask_index = [ - torch.where(input_ids == id)[0].int() for id in mask_ids - ] - tail_index = torch.Tensor([max_context_length]).int().cuda() - mask_index.append(tail_index) - mask_index = torch.cat(mask_index, dim=0).min() - position_ids[i, 0, length - 1] = int(mask_index) - position_ids[i, 1, length - 1] = 1 - self.mask_index_tensor[i] = int(mask_index) - else: - for i in range(batch_size): - length = context_lengths[i] - position_ids[i, 0, length - 1] = length - 2 - position_ids[i, 1, length - 1] = 1 - - position_ids = position_ids.cuda() - - perf_knob_tensor_size = 16 - context_runtime_perf_knobs = torch.tensor([-1] * perf_knob_tensor_size, - dtype=torch.int64) - - inputs = { - 'position_ids': position_ids, - 'last_token_ids': last_token_ids, - 'host_runtime_perf_knobs': context_runtime_perf_knobs - } - if not use_gpt_attention_plugin: - attention_mask = torch.zeros((batch_size, 1)) - inputs['attention_mask'] = attention_mask - return inputs - - def _prepare_generation_inputs(self, batch_size, context_lengths, - use_gpt_attention_plugin, - remove_input_padding, **kwargs): - - step = kwargs.pop('step') - num_beams = kwargs.pop('num_beams') - last_token_ids = torch.ones_like(context_lengths) - - if remove_input_padding: - - def _tile_beam_width_chatglm(tensor: torch.Tensor, num_beams: int): - new_shape = np.array(tensor.shape) - new_shape[1] = new_shape[1] * num_beams - tile_size = np.ones(new_shape.shape, dtype=np.int32) - tile_size = np.insert(tile_size, 2, num_beams) - new_tensor = torch.unsqueeze(tensor, 2) - new_tensor = new_tensor.tile(tile_size.tolist()) - new_tensor = new_tensor.reshape(new_shape.tolist()) - return new_tensor - - position_ids = torch.zeros([2, batch_size], dtype=torch.int32) - for i in range(batch_size): - position_ids[0, i] = context_lengths[i * num_beams] - 2 - position_ids[1, i] = step + 2 - position_ids = _tile_beam_width_chatglm(position_ids, num_beams) - position_ids = position_ids.int().cuda() - last_token_ids = torch.cumsum(last_token_ids, dim=0).int().cuda() - - if self.mask_index_tensor is not None: # specialization for GLM series models - position_ids = position_ids.cpu() - for i in range(batch_size): - position_ids[0][i] = self.mask_index_tensor[i] - position_ids = position_ids.cuda() - else: - data = [] - if self.mask_index_tensor is not None: # specialization for GLM series models - for i in range(batch_size): - data.append([[self.mask_index_tensor[i]], [step + 2]]) - else: - for i in range(batch_size): - data.append([[context_lengths[i * num_beams] - 2], - [step + 2]]) - position_ids = torch.tensor(data, dtype=torch.int32, device='cuda') - position_ids = _tile_beam_width(position_ids, num_beams) - - perf_knob_tensor_size = 16 - generation_runtime_perf_knobs = torch.tensor([-1] * - perf_knob_tensor_size, - dtype=torch.int64) - - inputs = { - 'position_ids': position_ids, - 'last_token_ids': last_token_ids, - 'host_runtime_perf_knobs': generation_runtime_perf_knobs - } - if not use_gpt_attention_plugin: - attention_mask = torch.zeros((batch_size, 1)) - inputs['attention_mask'] = attention_mask - return inputs - - -class QWenForCausalLMGenerationSession(GenerationSession): - - def __init__( - self, - model_config: ModelConfig, - engine_buffer, - mapping: Mapping, - debug_mode=False, - debug_tensors_to_save=None, - cuda_graph_mode=False, - stream: torch.cuda.Stream = None, - global_max_input_length: int = 2048, - global_max_output_length: int = 4096, - ): - super().__init__(model_config, - engine_buffer, - mapping, - debug_mode, - debug_tensors_to_save=debug_tensors_to_save, - cuda_graph_mode=cuda_graph_mode, - stream=stream) - self.global_max_input_length = global_max_input_length - self.global_max_output_length = global_max_output_length - - def generate( - self, - input_ids: torch.Tensor, - input_lengths: torch.Tensor, - sampling_config: SamplingConfig, - max_new_tokens: int, - runtime_rank: int = 0, - ): - max_input_length = torch.max(input_lengths).item() - max_new_tokens = min(max_new_tokens, - self.global_max_output_length - max_input_length) - # setup batch_size, max_input_length, max_output_len - self.setup(batch_size=input_lengths.size(0), - max_context_length=max_input_length, - max_new_tokens=max_new_tokens) - output_ids = self.decode(input_ids, input_lengths, sampling_config) - with torch.no_grad(): - torch.cuda.synchronize() - if runtime_rank == 0: - outputs = output_ids[:, 0, :] - return outputs diff --git a/tensorrt_llm/runtime/kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager.py deleted file mode 100644 index c2b6c3f9b1a1..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager.py +++ /dev/null @@ -1,473 +0,0 @@ -# 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 collections import defaultdict -from typing import List - -import torch - - -class Block(object): - - def __init__(self, block_idx: int): - self.idx = block_idx - self.ref_count = 0 - - def add_link(self): - self.ref_count += 1 - - def remove_link(self): - self.ref_count -= 1 - - def has_link(self) -> bool: - return self.ref_count > 0 - - def is_shared(self) -> bool: - return self.ref_count > 1 - - -class GenerationSequence(object): - - def __init__(self, seq_idx, batch_idx): - self.seq_idx = seq_idx - self.batch_idx = batch_idx - - def get_batch_idx(self) -> int: - """ - Returns idx of sequence in batch - """ - return self.batch_idx - - def get_seq_idx(self) -> int: - """ - Returns sequence idx - """ - return self.seq_idx - - def __eq__(self, another): - return hasattr(another, 'seq_idx') and self.seq_idx == another.seq_idx and \ - hasattr(another, 'batch_idx') and self.batch_idx == another.batch_idx - - def __hash__(self): - return self.seq_idx - - -class BlocksManager(object): - _sizeof = { - torch.float32: 4, - torch.float16: 2, - torch.bfloat16: 2, - torch.int8: 1 - } - - def __init__(self, - *, - num_layers: int, - num_blocks: int, - block_size: int, - max_blocks_per_seq: int = 128, - beam_width: int = 1): - """ - If layers are homogeneous then the expected block pool shape is: [num_blocks, num_layers, 2, block_size] - Otherwise, the expected block pool shape is: [num_blocks, 2, block_size] - """ - - self.max_blocks_per_seq = max_blocks_per_seq - - self.num_layers = num_layers - self.num_blocks = num_blocks - self.block_size = block_size - self.beam_width = beam_width - - self.free_blocks = [] - for bi in range(num_blocks): - self.free_blocks.append(Block(bi)) - - beam_width = self.beam_width - # Here use beam_width instead of self.beam_width to remove cyclic reference between self and - # self.allocated_blocks by preventing capture self, which may cause memory leak. - self.allocated_blocks = defaultdict( - lambda: [[] for _ in range(beam_width)]) - - def has_free_block(self) -> bool: - """ - Returns True if we have at least 1 free block - """ - return len(self.free_blocks) > 0 - - def allocate(self, - owner: GenerationSequence, - share_across_beam: bool = False): - """ - Add block to owner and increase ref count - """ - # Add blocks for whole beam width - block = None - for bi in range(self.beam_width): - if not self.has_free_block(): - raise RuntimeError("Can't allocate new block for KV cache") - - # Use the same block for all seqs in beam if share_across_beam - if block is None or share_across_beam == False: - block = self.free_blocks.pop(0) - # Add one reference to the block - block.add_link() - self.allocated_blocks[owner][bi].append(block) - - def replace_shared_block(self, owner: GenerationSequence, block_idx: int): - """ - Replace the shared block. - Free the shared block, and allocate blocks with share_across_beam=False - """ - if not self.allocated_blocks[owner][0][block_idx].is_shared(): - return - - # Free shared block - for bi in range(self.beam_width): - block = self.allocated_blocks[owner][bi][block_idx] - block.remove_link() - if not block.has_link(): - self.free_blocks.append(block) - - # Allocate new block - for bi in range(self.beam_width): - if not self.has_free_block(): - raise RuntimeError("Can't allocate new block for KV cache") - block = self.free_blocks.pop(0) - block.add_link() - self.allocated_blocks[owner][bi][block_idx] = block - return - - def free(self, owner: GenerationSequence): - """ - Unlink all blocks of given owner. - Moves blocks with ref_count == 0 to free. - Removes owner from allocated blocks. - """ - for bi in range(self.beam_width): - for block in self.allocated_blocks[owner][bi]: - # Move block to free if no one refers to it - block.remove_link() - - # Move block to free if no one refers to it - if not block.has_link(): - self.free_blocks.append(block) - # Remove owner from allocated blocks - self.allocated_blocks.pop(owner) - - def get_number_blocks(self, owner: GenerationSequence) -> int: - """ - Returns number of blocks allocated to the sequence owner - """ - return len(self.allocated_blocks[owner][0]) - - def get_k_or_v_block_offset(self, block_idx, field_idx): - """ - Get offset in memory pool to K or V block. field_idx should be 0 (K) or 1 (V). - """ - return block_idx * self.num_layers * 2 + field_idx - - def get_offset_array(self, beam_width: int) -> torch.Tensor: - """ - Returns array of [batch size, beam_width, 2, max_blocks_per_seq] of offsets - to the allocated blocks in memory pool - """ - assert (beam_width <= self.beam_width) - - def create_nested_list(dims): - """Recursive function to generate nested list.""" - if len(dims) == 1: - return [0 for _ in range(dims[0])] - return [create_nested_list(dims[1:]) for _ in range(dims[0])] - - offset_array = create_nested_list( - (len(self.allocated_blocks), beam_width, 2, - self.max_blocks_per_seq)) - - k_idx = 0 - v_idx = 1 - for owner, beams_blocks in self.allocated_blocks.items(): - for bi in range(beam_width): - for block_linear_idx, block in enumerate(beams_blocks[bi]): - for x_idx in [k_idx, v_idx]: - offset_array[owner.get_batch_idx()][bi][x_idx][ - block_linear_idx] = self.get_k_or_v_block_offset( - block.idx, x_idx) - - self.offset_array = torch.tensor(offset_array, dtype=torch.int32) - return self.offset_array - - def get_continuous_caches(self, memory_pool: torch.Tensor) -> torch.Tensor: - """ - Returns continuous KV caches. - Used only for debug purposes. - """ - assert self.beam_width == 1 - - pool = memory_pool.flatten() - continuous_kv_cache = torch.zeros(len(self.allocated_blocks), - 2, - self.max_blocks_per_seq * - self.block_size, - dtype=pool.dtype, - device="cuda") - k_idx = 0 - v_idx = 1 - for owner, beam_blocks in self.allocated_blocks.items(): - for bi in range(self.beam_width): - for block_linear_idx, block in enumerate(beam_blocks[bi]): - # The batch index. - batch_idx = owner.get_batch_idx() - # The first index in the sequence. - block_offset = block_linear_idx * self.block_size - - for x_idx in [k_idx, v_idx]: - x_start = self.get_k_or_v_block_offset( - block.idx, x_idx) * self.block_size - - continuous_kv_cache[batch_idx][ - x_idx][block_offset:block_offset + - self.block_size] = pool[x_start:x_start + - self.block_size] - - return continuous_kv_cache - - -class KVCacheManager(object): - - def __init__(self, - *, - num_layers: int, - num_blocks: int, - block_size: int, - tokens_per_block: int, - max_blocks_per_seq: int, - max_attention_window_size: int, - sink_token_len: int, - beam_width: int = 1, - use_one_more_block: bool = False): - - self.blocks_manager = BlocksManager( - num_layers=num_layers, - num_blocks=num_blocks, - block_size=block_size, - max_blocks_per_seq=max_blocks_per_seq, - beam_width=beam_width) - - self.tokens_per_block = tokens_per_block - self.max_attention_window_size = max_attention_window_size - self.sink_token_len = sink_token_len - self.beam_width = beam_width - - # The sink tokens are not stored into the same block with other tokens. - # Need to add the bubble after the sink tokens. - if sink_token_len % tokens_per_block == 0: - self.bubble_len = 0 - else: - self.bubble_len = tokens_per_block - sink_token_len % tokens_per_block - - # Token num in the sink blocks - self.sink_block_token_num = self.sink_token_len + self.bubble_len - - # Max token num in the cache - self.max_token_num = self.max_attention_window_size + self.bubble_len - if use_one_more_block: - self.max_token_num += self.tokens_per_block - - self.lens = [] - self.sequences = [] - - def step(self, finished: List[bool]): - """ - Iterate to the next generation step. - Add new blocks where needed and clear finished sequences. - """ - for seq in self.sequences: - batch_idx = seq.get_batch_idx() - # Enable cyclic kv cache when it exceeds the max_token_num - cyclic_token_num = self.max_token_num - self.sink_block_token_num - next_token_idx_in_cache = self.sink_block_token_num + \ - (self.lens[batch_idx] - self.sink_block_token_num) % cyclic_token_num - if not finished[batch_idx] and ( - next_token_idx_in_cache % self.tokens_per_block == 0 or - (next_token_idx_in_cache - self.sink_block_token_num) % - cyclic_token_num == 0): - if self.lens[batch_idx] < self.max_token_num: - self.blocks_manager.allocate(seq) - elif self.beam_width > 1: - # Get next block index - next_block_idx = next_token_idx_in_cache // self.tokens_per_block - # Replace the shared block with the unshared ones - self.blocks_manager.replace_shared_block( - seq, next_block_idx) - - self.lens[batch_idx] += 1 - - # Remove finished sequences - for fi in range(len(finished)): - if finished[fi]: - self.blocks_manager.free(self.sequences[fi]) - self.lens = [l for l, f in zip(self.lens, finished) if not f] - - # Remap sequence ids - new_sequences = [] - batch_idx = 0 - for seq, finish in zip(self.sequences, finished): - if not finish: - seq.batch_idx = batch_idx - new_sequences.append(seq) - batch_idx += 1 - self.sequences = new_sequences - - def add_sequence(self, - sequence: GenerationSequence, - context_len: int, - always_share_across_beam: bool = False): - """ - Add sequence to the manager and allocate minimum amount of blocks for context - """ - seq_len = context_len + self.bubble_len - self.lens.append(seq_len) - self.sequences.append(sequence) - - # Enable cyclic kv cache when inputLength exceeds maxAttentionWindow. - # Note that currently cyclic kv cache doesn't work with shared kv cache of different beams. - enable_cyclic_kv_cache = seq_len >= self.max_token_num - - # Get the final token index in kv cache - final_token_kv_index = self.sink_block_token_num + ( - (seq_len - 1 - self.sink_block_token_num) % - (self.max_token_num - self.sink_block_token_num)) - - # Get block index that with shareAmongBeams=False. - unshared_block_idx = -1 - if (not enable_cyclic_kv_cache or self.beam_width > 1 - or final_token_kv_index % self.tokens_per_block > 0): - unshared_block_idx = final_token_kv_index // self.tokens_per_block + 1 if ( - final_token_kv_index + 1 - ) % self.tokens_per_block == 0 else final_token_kv_index // self.tokens_per_block - - # Get context block num. - # Allocate one more block if there are tokens that can't be shared across beams. - seq_len = min(seq_len, self.max_token_num) - context_blocks = seq_len // self.tokens_per_block - if seq_len % self.tokens_per_block > 0: - context_blocks += 1 - - # Allocate blocks - for i in range(context_blocks): - self.blocks_manager.allocate( - sequence, - share_across_beam=True if always_share_across_beam else - (i != unshared_block_idx)) - - def get_block_offsets(self, beam_width: int) -> torch.Tensor: - """ - Returns array of offsets into memory pools - """ - return self.blocks_manager.get_offset_array(beam_width) - - -class KVCacheUpdater: - - def __init__(self): - self.use_paged_kv_cache = None - self.num_layers = None - self.num_kv_heads = None - self.head_dim = None - self.elt_size = None - self.past_key_value_list = None - self.max_kv_cache_length = None - self.kv_cache_manager = None - self.host_kv_cache_pool_pointers = None - - def init_linear_kv_cache(self, num_layers, num_kv_heads, head_dim, - kv_cache_type, past_key_value_list): - self.use_paged_kv_cache = False - self.num_layers = num_layers - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim - self.past_key_value_list = past_key_value_list - self.elt_size = torch.zeros(1, dtype=kv_cache_type).element_size() - self.max_kv_cache_length = past_key_value_list[0].shape[3] - - def init_paged_kv_cache(self, num_layers, num_kv_heads, head_dim, - kv_cache_type, kv_cache_manager, - host_kv_cache_pool_pointers): - self.use_paged_kv_cache = True - self.num_layers = num_layers - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim - self.kv_cache_manager = kv_cache_manager - self.host_kv_cache_pool_pointers = host_kv_cache_pool_pointers - self.elt_size = torch.zeros(1, dtype=kv_cache_type).element_size() - - def update(self, accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, sequence_length_buffer, - rewind_tokens): - assert isinstance(rewind_tokens, torch.Tensor) or isinstance( - rewind_tokens, int) - rewind_tokens_tensor = rewind_tokens if isinstance( - rewind_tokens, torch.Tensor) else None - rewind_tokens_count = rewind_tokens if isinstance(rewind_tokens, - int) else 0 - assert self.use_paged_kv_cache is not None - if self.use_paged_kv_cache: - if self.kv_cache_manager.has_single_pool(): - kv_cache_manager = self.kv_cache_manager.get_single_kv_cache_manager( - ) - else: - raise RuntimeError( - "Currently, using KVCacheUpdater with more then single memory pool is not supported" - ) - - host_kv_cache_block_offsets = kv_cache_manager.get_block_offsets(1) - kv_cache_block_offsets = host_kv_cache_block_offsets.to('cuda') - torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - sequence_length_buffer, - True, - self.num_layers, - self.num_kv_heads, - self.head_dim * self.elt_size, - rewind_tokens_count, - kv_cache_manager.max_attention_window_size, - rewind_tokens_tensor, - None, - self.host_kv_cache_pool_pointers, - kv_cache_block_offsets, - kv_cache_manager.blocks_manager.max_blocks_per_seq, - kv_cache_manager.tokens_per_block, - None, - ) - else: - torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - sequence_length_buffer, - False, - self.num_layers, - self.num_kv_heads, - self.head_dim * self.elt_size, - rewind_tokens_count, - self.max_kv_cache_length, - rewind_tokens_tensor, - self.past_key_value_list, - None, - None, - None, - None, - None, - ) diff --git a/tensorrt_llm/runtime/medusa_utils.py b/tensorrt_llm/runtime/medusa_utils.py deleted file mode 100644 index ffd5126c7013..000000000000 --- a/tensorrt_llm/runtime/medusa_utils.py +++ /dev/null @@ -1,182 +0,0 @@ -import copy -from argparse import Namespace -from functools import cmp_to_key -from typing import List - -import numpy as np -import torch - -from tensorrt_llm.logger import logger - - -def path_sorter(a, b): - for i in range(min(len(a), len(b))): - if a[i] != b[i]: - return -1 if a[i] < b[i] else 1 - return 0 # shouldn't reach - - -path_sorting_key = cmp_to_key(path_sorter) - - -def expand_choices_if_needed(medusa_choices: List[List[int]]): - """ - Do a simple check to see if the given choices are path-only or vanilla. - """ - assert len(medusa_choices) > 0 - for c in medusa_choices: - if len(c) > 1: - try: - _ = medusa_choices.index( - [c[0]]) # find the first parent of current path - logger.debug( - "Detected vanilla-style of Medusa choices. No need to expand." - ) - return medusa_choices # if found, just return assuming it is already expanded - except ValueError: - logger.debug( - "Detected path-only style of Medusa choices. Expanding ...") - break - expanded_choices = set() - for c in medusa_choices: - cur = () - for n in c: - cur = (*cur, n) - expanded_choices.add(cur) - expanded_choices = [list(c) for c in expanded_choices] - return expanded_choices - - -def get_packed_mask(num_medusa_tokens, medusa_mask, max_medusa_tokens=None): - max_medusa_tokens = num_medusa_tokens if max_medusa_tokens is None else max_medusa_tokens - num_packed_masks = (max_medusa_tokens + 1 + 32 - 1) // 32 - medusa_packed_mask = torch.zeros((num_medusa_tokens + 1, num_packed_masks), - dtype=torch.int32) - for token_idx in range(num_medusa_tokens + 1): - if token_idx == 0: - medusa_packed_mask[0, 0] = 1 - else: - mask_list = medusa_mask[token_idx - 1, :].tolist() - # insert 1 as there is one extra new token from the original lm head. - mask_list.insert(0, True) - # convert binary bits into 4 int32_t - mask_str_list = [str(int(val)) for val in mask_list] - mask_str_list.reverse() - - for mask_idx in range(num_packed_masks): - if mask_idx * 32 >= len(mask_str_list): - break - mask_32bits_str = ''.join(mask_str_list[-(mask_idx + 1) * 32: - (-mask_idx * 32 - 1)] + - [mask_str_list[(-mask_idx * 32 - 1)]]) - valid_num_bits = len(mask_32bits_str) - first_bit1 = mask_32bits_str[0] == '1' - mask_31bits_str = mask_32bits_str[1:] - mask_31bits = 0 if mask_31bits_str == "" else int( - mask_31bits_str, 2) - if valid_num_bits == 32: - mask_32bits = mask_31bits - first_bit1 * (2**( - valid_num_bits - 1)) - else: - mask_32bits = mask_31bits + first_bit1 * (2**( - valid_num_bits - 1)) - medusa_packed_mask[token_idx, mask_idx] = mask_32bits - return medusa_packed_mask - - -def choices_2_paths(num_medusa_heads, choices): - paths = {} - all_paths = {} - level_counts = [0] * num_medusa_heads - choices.sort(key=len, reverse=True) - for c in choices: - k = ":".join([str(ci) for ci in c]) - if k not in all_paths: - paths[k] = c - for i in range(len(c)): - k = ":".join([str(ci) for ci in c[:i + 1]]) - if k not in all_paths: - all_paths[k] = c[:i + 1] - level_counts[i] += 1 - return list(paths.values()), level_counts, paths, all_paths - - -def get_medusa_topks(num_medusa_heads, paths): - medusa_topks = [0] * num_medusa_heads - for p in paths: - for i, k in enumerate(p): - medusa_topks[i] = max(medusa_topks[i], k + 1) - return medusa_topks - - -def get_medusa_tree(num_medusa_heads, medusa_topks, level_counts, paths): - cum_topks = np.cumsum([0] + medusa_topks) - cum_level_counts = np.cumsum([0] + level_counts) - tree_paths = copy.deepcopy(paths) - medusa_tree_ids = list(np.arange(medusa_topks[0])) - medusa_position_offsets = [0] * medusa_topks[0] - for i in range(1, num_medusa_heads): - last_prefix = "-1" - last = -1 - c = -1 - for pi, p in enumerate(paths): - if i < len(p): - prefix_str = ":".join([str(k) for k in p[:i]]) - if last_prefix != prefix_str or last != p[i]: - # new path - medusa_position_offsets.append(i) - medusa_tree_ids.append(p[i] + cum_topks[i]) - last_prefix = prefix_str - last = p[i] - c += 1 - tree_paths[pi][i] = cum_level_counts[i] + c - return medusa_tree_ids, medusa_position_offsets, tree_paths - - -def get_medusa_mask(medusa_tree_ids, medusa_paths): - medusa_mask = torch.zeros((len(medusa_tree_ids), len(medusa_tree_ids))) - medusa_mask[:, 0] = 1 - for p in medusa_paths: - for i, idx in enumerate(p): - if idx < 0: - continue - for j in range(i + 1): - medusa_mask[idx, p[j]] = 1 - return medusa_mask - - -def _medusa_setup(choices_or_paths, num_medusa_heads=None): - choices = copy.deepcopy(choices_or_paths) - sorted_choices = sorted(choices, key=path_sorting_key) - if num_medusa_heads is None: - num_medusa_heads = max([len(c) for c in sorted_choices]) - paths, level_counts, _, _ = choices_2_paths(num_medusa_heads, - sorted_choices) - paths = sorted(paths, key=path_sorting_key) - medusa_topks = get_medusa_topks(num_medusa_heads, paths) - medusa_tree_ids, medusa_position_offsets, tree_paths = get_medusa_tree( - num_medusa_heads, medusa_topks, level_counts, paths) - - num_medusa_tokens = len(medusa_tree_ids) - # now do the padding before converting to torch.Tensor - medusa_paths = [] - for p in tree_paths: - medusa_paths.append( - torch.tensor([-1] + p + ([-2] * (num_medusa_heads - len(p))))) - medusa_topks = torch.tensor(medusa_topks) - medusa_paths = torch.stack(medusa_paths) + 1 - medusa_tree_ids = torch.tensor([-1] + medusa_tree_ids) + 1 - medusa_position_offsets = torch.tensor([-1] + medusa_position_offsets) + 1 - medusa_mask = get_medusa_mask(medusa_tree_ids, medusa_paths) - medusa_packed_mask = get_packed_mask(num_medusa_tokens, medusa_mask[1:, 1:]) - medusa_spec_decoding_use = torch.tensor([1], device="cpu") - - return Namespace( - medusa_mask=medusa_mask.cuda(), - medusa_packed_mask=medusa_packed_mask.cuda(), - medusa_topks=medusa_topks.cuda(), - medusa_paths=medusa_paths.cuda(), - medusa_tree_ids=medusa_tree_ids.cuda(), - medusa_position_offsets=medusa_position_offsets.cuda(), - medusa_spec_decoding_use=medusa_spec_decoding_use, - ) diff --git a/tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py b/tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py deleted file mode 100644 index 719aca5f8444..000000000000 --- a/tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py +++ /dev/null @@ -1,80 +0,0 @@ -from itertools import chain -from typing import List - -import torch - -import tensorrt_llm -from tensorrt_llm.runtime.memory_pools.pool import Pool - - -class MemoryPoolsAllocator(object): - - def __init__(self, num_blocks, tokens_per_block, head_size): - self._pools_metadata = [] - self._pool_pointers = [] - self._pool_mapping = None - - self._num_blocks = num_blocks - self._tokens_per_block = tokens_per_block - self._head_size = head_size - - def allocate(self, dtype, num_kv_heads_per_layer: List[int], device="cuda"): - if isinstance(dtype, str): - dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype) - - # LayerCachePoolLocator{.indexOfPool, .layerIdxInCachePool}" - layers_mapping = [[-1, -1]] * len(num_kv_heads_per_layer) - unique_nkvh = sorted(set(num_kv_heads_per_layer)) - for index_of_pool, kv_head in enumerate(unique_nkvh): - layers = [ - layer for layer, nkvh in enumerate(num_kv_heads_per_layer) - if nkvh == kv_head - ] - - num_layers = len(layers) - cache_shape = ( - self._num_blocks, - num_layers, - 2, - kv_head, - self._tokens_per_block, - self._head_size, - ) - self._pool_pointers.append( - torch.empty(cache_shape, dtype=dtype, device=device)) - self._pools_metadata.append( - Pool(num_kv_heads=kv_head, num_layers=num_layers)) - - for layer_idx_in_cache_pool, layer in enumerate(layers): - layers_mapping[layer] = [index_of_pool, layer_idx_in_cache_pool] - - assert -1 not in set(chain(*layers_mapping)) - self._pool_mapping = torch.tensor(layers_mapping, dtype=torch.int32) - - def get_kv_cache_pool_pointers(self): - return self._get_primarmy_secondary_pool_pointers() - - def _get_primarmy_secondary_pool_pointers(self): - assert len(self._pool_pointers - ) >= 1, "pool pointers haven't been initiated yet" - data_ptr_pointers = torch.tensor(list( - map(lambda x: x.data_ptr(), self._pool_pointers)), - dtype=torch.int64) - host_kv_cache_pool_pointers = torch.cat( - (data_ptr_pointers.view(-1, 1), - torch.zeros(len(self._pool_pointers), 1, dtype=torch.int64)), - dim=1) - - return host_kv_cache_pool_pointers - - @classmethod - def prepare_num_kv_heads_per_layer(cls, kv_head, num_layers): - return [kv_head] * num_layers - - @property - def pools_metadata(self): - return self._pools_metadata - - @property - def pool_mapping(self): - return self._pool_mapping diff --git a/tensorrt_llm/runtime/memory_pools/pool.py b/tensorrt_llm/runtime/memory_pools/pool.py deleted file mode 100644 index 63308ad0d3f0..000000000000 --- a/tensorrt_llm/runtime/memory_pools/pool.py +++ /dev/null @@ -1,7 +0,0 @@ -from dataclasses import dataclass - - -@dataclass -class Pool(object): - num_kv_heads: int - num_layers: int diff --git a/tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py b/tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py deleted file mode 100644 index 187117302419..000000000000 --- a/tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py +++ /dev/null @@ -1,59 +0,0 @@ -from typing import List - -import torch - -from tensorrt_llm.runtime.kv_cache_manager import (GenerationSequence, - KVCacheManager) -from tensorrt_llm.runtime.memory_pools.pool import Pool - - -class PoolsKVCacheManager(object): - - def __init__(self, pools_metadata: List[Pool], max_blocks_per_seq, - num_blocks, tokens_per_block, head_size, - max_attention_window_size, beam_width, sink_token_len) -> None: - self._num_pools = len(pools_metadata) - self._kv_cache_managers = [] - - for pool in pools_metadata: - block_size = pool.num_kv_heads * tokens_per_block * head_size - self._kv_cache_managers.append( - KVCacheManager( - num_layers=pool.num_layers, - num_blocks=num_blocks, - block_size=block_size, - tokens_per_block=tokens_per_block, - max_blocks_per_seq=max_blocks_per_seq, - max_attention_window_size=max_attention_window_size, - sink_token_len=sink_token_len, - beam_width=beam_width, - )) - - def add_sequence(self, - sequence: GenerationSequence, - context_len: int, - always_share_across_beam: bool = False): - for kv_cache_manager in self._kv_cache_managers: - kv_cache_manager.add_sequence(sequence, context_len, - always_share_across_beam) - - def step(self, finished: List[bool]): - for kv_cache_manager in self._kv_cache_managers: - kv_cache_manager.step(finished) - - def get_block_offsets(self, beam_width: int) -> torch.Tensor: - offsets = [] - for kv_cache_manager in self._kv_cache_managers: - block_offset = kv_cache_manager.get_block_offsets(beam_width) - offsets.append(block_offset) - - return torch.stack(offsets) - - def get_single_kv_cache_manager(self): - assert len(self._kv_cache_managers - ) == 1, f"More then one kv cache manager exists" - - return self._kv_cache_managers[0] - - def has_single_pool(self): - return len(self._kv_cache_managers) == 1 diff --git a/tensorrt_llm/runtime/model_config.py b/tensorrt_llm/runtime/model_config.py new file mode 100644 index 000000000000..31e2d084849b --- /dev/null +++ b/tensorrt_llm/runtime/model_config.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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 __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, List, Optional + +from .._utils import binding_layer_type_to_str, binding_to_str_dtype +from ..llmapi.kv_cache_type import KVCacheType +from ..quantization import QuantMode + + +@dataclass +class ModelConfig: + max_batch_size: int + max_beam_width: int + vocab_size: int + num_layers: int + num_heads: int + num_kv_heads: int + hidden_size: int + gpt_attention_plugin: bool + gemm_allreduce_plugin: str = None + remove_input_padding: bool = False + model_name: str = "" + kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS + cross_attention: bool = False + head_size: int = None + has_position_embedding: bool = True + has_token_type_embedding: bool = False + tokens_per_block: int = 32 + max_prompt_embedding_table_size: int = 0 + quant_mode: QuantMode = QuantMode(0) + gather_context_logits: bool = False + gather_generation_logits: bool = False + dtype: str = "" + lora_plugin: bool = False + lora_target_modules: List[str] = field(default_factory=list) + trtllm_modules_to_hf_modules: dict = None + skip_cross_kv: bool = False + num_medusa_heads: int = 0 + max_medusa_tokens: int = 0 + paged_state: bool = True + mamba_conv1d_plugin: bool = True + conv_kernel: int = 0 + layer_types: List[str] = field(default_factory=list) + rnn_hidden_size: int = 0 + rnn_head_size: int = 0 + rnn_conv_dim_size: int = 0 + state_size: int = 0 + state_dtype: str = "" + gpu_weights_percent: float = 1.0 + # ReDrafter + redrafter_num_beams: int = 0 + redrafter_draft_len_per_beam: int = 0 + num_kv_heads_per_layer: Optional[List[int]] = None + num_kv_heads_per_cross_attn_layer: Optional[List[int]] = None + skip_cross_attn_blocks: bool = False + # language adapter (typed as Optional[Any]; the concrete config type is + # not imported here) + language_adapter_config: Optional[Any] = None + + @classmethod + def from_model_config_cpp(cls, model_config_cpp) -> "ModelConfig": + """Create a partially initialized ModelConfig instance from a given ModelConfig CPP binding instance. + + Note that each of these classes have fields that don't exist in the other, so the created ModelConfigPython + won't have all of its fields initialized. + """ + return cls( + max_batch_size=model_config_cpp.max_batch_size, + max_beam_width=model_config_cpp.max_beam_width, + vocab_size=model_config_cpp.vocab_size, + num_layers=model_config_cpp.num_layers(), + num_heads=model_config_cpp.num_heads, + num_kv_heads=model_config_cpp.num_kv_heads(0), + hidden_size=model_config_cpp.hidden_size, + remove_input_padding=model_config_cpp.use_packed_input, + kv_cache_type=model_config_cpp.kv_cache_type, + cross_attention=model_config_cpp.use_cross_attention, + head_size=model_config_cpp.head_size, + max_prompt_embedding_table_size=model_config_cpp.max_prompt_embedding_table_size, + quant_mode=QuantMode(model_config_cpp.quant_mode.value), + gather_context_logits=model_config_cpp.compute_context_logits, + gather_generation_logits=model_config_cpp.compute_generation_logits, + gpt_attention_plugin=model_config_cpp.use_gpt_attention_plugin, + dtype=binding_to_str_dtype(model_config_cpp.data_type), + num_kv_heads_per_layer=model_config_cpp.num_kv_heads_per_layer, + tokens_per_block=model_config_cpp.tokens_per_block, + lora_plugin=model_config_cpp.use_lora_plugin, + layer_types=[binding_layer_type_to_str(lt) for lt in model_config_cpp.layer_types], + ) diff --git a/tensorrt_llm/runtime/model_runner.py b/tensorrt_llm/runtime/model_runner.py deleted file mode 100644 index 20c250f940eb..000000000000 --- a/tensorrt_llm/runtime/model_runner.py +++ /dev/null @@ -1,1001 +0,0 @@ -# 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 copy -import json -import math -from pathlib import Path -from typing import List, Optional, Tuple, Union - -import numpy as np -import tensorrt as trt -import torch - -from .. import profiler -from .._deprecation import emit_engine_arch_deprecation -from .._utils import mpi_comm, mpi_world_size, numpy_to_torch -from ..bindings import MpiComm -from ..bindings.executor import Executor -from ..builder import Engine, EngineConfig, get_engine_version -from ..llmapi.kv_cache_type import KVCacheType -from ..logger import logger -from ..mapping import Mapping -from ..quantization import QuantMode -from .generation import (DISABLE_TORCH_DEVICE_SET, ChatGLMGenerationSession, - GenerationSession, LogitsProcessor, LoraManager, - ModelConfig, QWenForCausalLMGenerationSession, - SamplingConfig, StoppingCriteria, to_word_list_format) - - -def get_engine_name(model: str, dtype: str, tp_size: int, pp_size: int, - rank: int) -> str: - """ - Get the serialized engine file name. - - Args: - model (str): - Model name, e.g., bloom, gpt. - dtype (str): - Data type, e.g., float32, float16, bfloat16, - tp_size (int): - The size of tensor parallel. - pp_size (int): - The size of pipeline parallel. - rank (int): - The rank id. - - Returns: - str: The serialized engine file name. - """ - if pp_size == 1: - return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) - return '{}_{}_tp{}_pp{}_rank{}.engine'.format(model, dtype, tp_size, - pp_size, rank) - - -def read_config(config_path: Path) -> Tuple[ModelConfig, dict]: - """ - Read the engine config file and create a ModelConfig instance, return the ModelConfig instance - and other config fields in a dict. - - Args: - config_path (Path): - The path of engine config file. - - Returns: - Tuple[ModelConfig, dict]: A ModelConfig instance and other config fields. - """ - with open(config_path, 'r') as f: - config = json.load(f) - return _builder_to_model_config(config) - - -def _builder_to_model_config(config: dict) -> Tuple[ModelConfig, dict]: - builder_config = config['builder_config'] - model_name = builder_config['name'] - dtype = builder_config['precision'] - tp_size = builder_config['tensor_parallel'] - pp_size = builder_config.get('pipeline_parallel', 1) - kv_cache_type = builder_config.get('kv_cache_type') - if kv_cache_type is not None: - kv_cache_type = KVCacheType(kv_cache_type) - world_size = tp_size * pp_size - assert world_size == mpi_world_size(), \ - f'Engine world size ({tp_size} * {pp_size}) != Runtime world size ({mpi_world_size()})' - - num_heads = builder_config['num_heads'] - assert num_heads % tp_size == 0, \ - f"The number of heads ({num_heads}) is not a multiple of tp_size ({tp_size})" - num_kv_heads = builder_config.get('num_kv_heads', num_heads) - # TODO: multi_query_mode should be removed - multi_query_mode = builder_config.get('multi_query_mode', False) - if multi_query_mode: - logger.warning( - "`multi_query_mode` config is deprecated. Please rebuild the engine." - ) - # num_kv_heads, if exists in config, should override multi_query_mode - if multi_query_mode and ('num_kv_heads' not in builder_config): - num_kv_heads = 1 - num_heads = num_heads // tp_size - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - head_size = builder_config.get('head_size', None) - - hidden_size = builder_config['hidden_size'] // tp_size - vocab_size = builder_config['vocab_size'] - num_layers = builder_config['num_layers'] - max_batch_size = builder_config['max_batch_size'] - max_beam_width = builder_config['max_beam_width'] - - cross_attention = builder_config.get('cross_attention', False) - has_position_embedding = builder_config.get('has_position_embedding', True) - has_token_type_embedding = builder_config.get('has_token_type_embedding', - False) - gather_context_logits = builder_config.get('gather_context_logits', False) - gather_generation_logits = builder_config.get('gather_generation_logits', - False) - max_prompt_embedding_table_size = builder_config.get( - 'max_prompt_embedding_table_size', 0) - quant_mode = QuantMode(builder_config.get('quant_mode', 0)) - lora_target_modules = builder_config.get('lora_target_modules') - lora_trtllm_modules_to_hf_modules = builder_config.get( - 'trtllm_modules_to_hf_modules') - max_medusa_token_len = builder_config.get('max_draft_len', 0) - num_medusa_heads = builder_config.get('num_medusa_heads', 0) - - skip_cross_attn_blocks = bool(config['pretrained_config'].get( - 'skip_cross_attn_blocks', False)) - - # ReDrafter - redrafter_num_beams = config['pretrained_config'].get( - 'redrafter_num_beams', 0) - redrafter_draft_len_per_beam = config['pretrained_config'].get( - 'redrafter_draft_len_per_beam', 0) - - plugin_config = config['plugin_config'] - use_gpt_attention_plugin = bool(plugin_config['gpt_attention_plugin']) - gemm_allreduce_plugin = plugin_config['gemm_allreduce_plugin'] - mamba_conv1d_plugin = bool(plugin_config['mamba_conv1d_plugin']) - remove_input_padding = plugin_config['remove_input_padding'] - paged_state = plugin_config['paged_state'] - tokens_per_block = plugin_config['tokens_per_block'] - lora_plugin = plugin_config.get('lora_plugin') - - model_config = ModelConfig( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - head_size=head_size, - gpt_attention_plugin=use_gpt_attention_plugin, - gemm_allreduce_plugin=gemm_allreduce_plugin, - mamba_conv1d_plugin=mamba_conv1d_plugin, - remove_input_padding=remove_input_padding, - model_name=model_name, - kv_cache_type=kv_cache_type, - paged_state=paged_state, - cross_attention=cross_attention, - has_position_embedding=has_position_embedding, - has_token_type_embedding=has_token_type_embedding, - tokens_per_block=tokens_per_block, - max_prompt_embedding_table_size=max_prompt_embedding_table_size, - quant_mode=quant_mode, - gather_context_logits=gather_context_logits, - gather_generation_logits=gather_generation_logits, - dtype=dtype, - lora_plugin=lora_plugin, - lora_target_modules=lora_target_modules, - trtllm_modules_to_hf_modules=lora_trtllm_modules_to_hf_modules, - num_medusa_heads=num_medusa_heads, - max_medusa_tokens=max_medusa_token_len, - skip_cross_attn_blocks=skip_cross_attn_blocks, - # ReDrafter - redrafter_num_beams=redrafter_num_beams, - redrafter_draft_len_per_beam=redrafter_draft_len_per_beam, - ) - - other_config = { - 'world_size': world_size, - 'tp_size': tp_size, - 'pp_size': pp_size, - 'max_batch_size': builder_config['max_batch_size'], - 'max_input_len': builder_config['max_input_len'], - 'max_output_len': builder_config['max_output_len'], - 'max_beam_width': builder_config['max_beam_width'] - } - return model_config, other_config - - -def _engine_config_to_model_config(engine_config: EngineConfig, - **kwargs) -> ModelConfig: - pretrained_config = engine_config.pretrained_config - build_config = engine_config.build_config - - tp_size = pretrained_config.mapping.tp_size - num_heads = pretrained_config.num_attention_heads // tp_size - num_kv_heads = pretrained_config.num_key_value_heads - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - hidden_size = pretrained_config.hidden_size // tp_size - head_size = pretrained_config.head_size - - rnn_config_items = [ - 'conv_kernel', 'layer_types', 'rnn_hidden_size', 'state_size', - 'state_dtype', 'rnn_head_size', 'rnn_conv_dim_size' - ] - rnn_configs_kwargs = {} - for item in rnn_config_items: - if hasattr(pretrained_config, item): - rnn_configs_kwargs[item] = getattr(pretrained_config, item) - - if not hasattr(build_config, 'kv_cache_type'): - logger.Warning( - 'Build config doesn\'t have kv_cache_type, you might need to rebuild your enigne.' - ) - - # TODO(oargov): this is a hack, make it prettier! - if hasattr(pretrained_config, "num_kv_heads_per_layer"): - pp_rank = pretrained_config.mapping.pp_rank - pp_size = pretrained_config.mapping.pp_size - layers_per_pp_rank = pretrained_config.num_hidden_layers // pp_size - first_local_layer = layers_per_pp_rank * pp_rank - first_layer_next_rank = first_local_layer + layers_per_pp_rank - layer_types = getattr(pretrained_config, "layer_types", ["attention"]) - num_attn_layers_lower_ranks = [ - layer_types[layer_idx % len(layer_types)] - for layer_idx in range(first_local_layer) - ].count("attention") - num_local_attn_layers = [ - layer_types[layer_idx % len(layer_types)] - for layer_idx in range(first_local_layer, first_layer_next_rank) - ].count("attention") - num_kv_heads_per_layer = pretrained_config.num_kv_heads_per_layer[ - num_attn_layers_lower_ranks:num_attn_layers_lower_ranks + - num_local_attn_layers] - num_kv_heads_per_layer = [(nheads + tp_size - 1) // tp_size - for nheads in num_kv_heads_per_layer] - - elif hasattr(pretrained_config, "get_layer_num_kv_heads"): - # each layer has a different number of kv heads - attention_layers = [ - layer_idx for layer_idx, layer_type in enumerate( - pretrained_config.layer_types) if layer_type == "attention" - ] if hasattr(pretrained_config, "layer_types") else list( - range(pretrained_config.num_hidden_layers)) - num_kv_heads_per_layer = [ - pretrained_config.get_layer_num_kv_heads(layer_idx) - if layer_idx in attention_layers else 0 - for layer_idx in range(pretrained_config.num_hidden_layers) - ] - else: - num_kv_heads_per_layer = None - - if hasattr(pretrained_config, "num_kv_heads_per_cross_attn_layer"): - num_kv_heads_per_cross_attn_layer = pretrained_config.num_kv_heads_per_cross_attn_layer - else: - num_kv_heads_per_cross_attn_layer = None - - return ModelConfig( - max_batch_size=build_config.max_batch_size, - max_beam_width=build_config.max_beam_width, - vocab_size=pretrained_config.vocab_size, - num_layers=pretrained_config.num_hidden_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - head_size=head_size, - gpt_attention_plugin=bool( - build_config.plugin_config.gpt_attention_plugin), - gemm_allreduce_plugin=build_config.plugin_config.gemm_allreduce_plugin, - mamba_conv1d_plugin=bool( - build_config.plugin_config.mamba_conv1d_plugin), - remove_input_padding=build_config.plugin_config.remove_input_padding, - paged_state=build_config.plugin_config.paged_state, - tokens_per_block=build_config.plugin_config.tokens_per_block, - quant_mode=pretrained_config.quant_mode, - gather_context_logits=build_config.gather_context_logits, - gather_generation_logits=build_config.gather_generation_logits, - dtype=pretrained_config.dtype, - max_prompt_embedding_table_size=build_config. - max_prompt_embedding_table_size, - lora_plugin=build_config.plugin_config.lora_plugin, - lora_target_modules=build_config.lora_config.lora_target_modules, - trtllm_modules_to_hf_modules=build_config.lora_config. - trtllm_modules_to_hf_modules, - max_medusa_tokens=pretrained_config.max_draft_len if hasattr( - pretrained_config, 'max_draft_len') else 0, - num_medusa_heads=pretrained_config.num_medusa_heads if hasattr( - pretrained_config, 'num_medusa_heads') else 0, - **rnn_configs_kwargs, - num_kv_heads_per_layer=num_kv_heads_per_layer, - num_kv_heads_per_cross_attn_layer=num_kv_heads_per_cross_attn_layer, - redrafter_num_beams=pretrained_config.redrafter_num_beams if hasattr( - pretrained_config, 'redrafter_num_beams') else 0, - redrafter_draft_len_per_beam=pretrained_config. - redrafter_draft_len_per_beam - if hasattr(pretrained_config, 'redrafter_draft_len_per_beam') else 0, - kv_cache_type=getattr(build_config, 'kv_cache_type', - KVCacheType.CONTINUOUS), - cross_attention=getattr(pretrained_config, 'cross_attention', False), - has_position_embedding=getattr(pretrained_config, - 'has_position_embedding', True), - skip_cross_attn_blocks=getattr(pretrained_config, - 'skip_cross_attn_blocks', False), - **kwargs) - - -class ModelRunnerMixin: - - def _check_inputs(self, batch_input_ids: List[torch.Tensor], - sampling_config: SamplingConfig): - batch_size = len(batch_input_ids) - if batch_size > self.max_batch_size: - raise RuntimeError( - f"Input batch size ({batch_size}) exceeds the engine or specified limit ({self.max_batch_size})" - ) - input_lengths = [x.size(0) for x in batch_input_ids] - max_length = max(input_lengths) - if max_length > self.max_input_len: - raise RuntimeError( - f"Maximum input length ({max_length}) exceeds the engine or specified limit ({self.max_input_len})" - ) - if max_length + sampling_config.max_new_tokens > self.max_seq_len: - raise RuntimeError( - f"Maximum input length ({max_length}) + maximum new tokens ({sampling_config.max_new_tokens}) exceeds the engine or specified limit ({self.max_seq_len})" - ) - if sampling_config.num_beams > self.max_beam_width: - raise RuntimeError( - f"Num beams ({sampling_config.num_beams}) exceeds the engine or specified limit ({self.max_beam_width})" - ) - - def _prepare_inputs(self, batch_input_ids: List[torch.Tensor], - pad_id: int) -> Tuple[torch.Tensor]: - # Cast to int32 - batch_input_ids = [x.type(torch.int32) for x in batch_input_ids] - input_lengths = [x.size(0) for x in batch_input_ids] - max_length = max(input_lengths) - - if self.remove_input_padding: - batch_input_ids = torch.concat(batch_input_ids) - else: - # Right padding for trt-llm - paddings = [ - torch.ones(max_length - l, dtype=torch.int32) * pad_id - for l in input_lengths - ] - batch_input_ids = [ - torch.cat([x, pad]) for x, pad in zip(batch_input_ids, paddings) - ] - batch_input_ids = torch.stack(batch_input_ids) - input_lengths = torch.tensor(input_lengths, dtype=torch.int32) - return batch_input_ids, input_lengths - - def _prepare_outputs(self, outputs: Optional[dict], - input_lengths: torch.Tensor) -> dict: - if outputs is not None: - batch_size = input_lengths.size(0) - if 'context_logits' in outputs: - if self.mapping.has_pp(): - # If pp size > 1, the context logits and generation logits are both in last pp - # Last pp rank send context logits and generation logits to rank 0 - if self.mapping.is_last_pp_rank(): - context_logits = outputs['context_logits'] - context_logits_host = context_logits.cpu() - mpi_comm().send(context_logits_host, dest=0) - elif self.mapping.is_first_pp_rank(): - context_logits_host = mpi_comm().recv( - source=self.mapping.prev_pp_rank() - ) # Prev pp rank of rank=0 is the last pp - context_logits = context_logits_host.to( - torch.device('cuda:0')) - outputs['context_logits'] = context_logits - - context_logits = outputs['context_logits'] - - context_logits_output = [] - if self.remove_input_padding: - if isinstance(self.session, Executor) and batch_size > 1: - # The starting position of the context logits buffer of each micro batch is separated - num_batches = self.mapping.pp_size - micro_batch_size = math.ceil(batch_size / - self.mapping.pp_size) - - for i in range(num_batches): - start_idx = i * micro_batch_size - end_idx = min(start_idx + micro_batch_size, - batch_size) - micro_context_logits = context_logits[ - start_idx:end_idx] - micro_input_lengths = input_lengths[ - start_idx:end_idx] - - micro_context_logits = micro_context_logits.flatten( - end_dim=-2) - seg_points = [0] + micro_input_lengths.cumsum( - dim=0).tolist() - context_logits_output += [ - micro_context_logits[s:e] - for s, e in zip(seg_points[:-1], seg_points[1:]) - ] - else: - context_logits = context_logits.flatten(end_dim=-2) - - seg_points = [0] + input_lengths.cumsum(dim=0).tolist() - context_logits_output = [ - context_logits[s:e] - for s, e in zip(seg_points[:-1], seg_points[1:]) - ] - else: - context_logits_output = [ - context_logits[bidx, :input_lengths[bidx]] - for bidx in range(batch_size) - ] - - assert len(context_logits_output) == batch_size - outputs['context_logits'] = context_logits_output - - if 'generation_logits' in outputs: - if self.mapping.has_pp(): - if self.mapping.is_last_pp_rank(): - generation_logits = outputs['generation_logits'] - if isinstance(generation_logits, list): - generation_logits_host = [ - logits.cpu() for logits in generation_logits - ] - else: - generation_logits_host = generation_logits.cpu() - mpi_comm().send(generation_logits_host, dest=0) - elif self.mapping.is_first_pp_rank(): - generation_logits_host = mpi_comm().recv( - source=self.mapping.prev_pp_rank() - ) # Prev pp rank of rank=0 is the last pp - if isinstance(generation_logits_host, list): - generation_logits = [ - logits.to(torch.device('cuda:0')) - for logits in generation_logits_host - ] - else: - generation_logits = generation_logits_host.to( - torch.device('cuda:0')) - outputs['generation_logits'] = generation_logits - - if isinstance(self.session, GenerationSession): - # Convert logits format to be same as GptSession - generation_logits = torch.stack( - outputs['generation_logits'], dim=1) - batch_x_beam, max_gen_len, voc_size = generation_logits.size( - ) - num_beams = batch_x_beam // batch_size - generation_logits = generation_logits.view( - batch_size, num_beams, max_gen_len, voc_size) - outputs['generation_logits'] = generation_logits - - return outputs - - def _prepare_embedding_table(self, prompt_table: Union[str, torch.Tensor]): - if isinstance(prompt_table, str): - prompt_table_data = numpy_to_torch( - np.load(prompt_table)).to(dtype=self.dtype) - else: - assert isinstance( - prompt_table, - torch.Tensor), "Prompt table should be str or torch.Tensor" - prompt_table_data = prompt_table.to(dtype=self.dtype) - torch.cuda.current_stream().synchronize() - - return prompt_table_data - - def _prepare_ptuning(self, prompt_table: Union[str, torch.Tensor], - tasks: str, batch_size: int): - if self.max_prompt_embedding_table_size == 0: - return {} - - if prompt_table is not None: - prompt_table_data = self._prepare_embedding_table(prompt_table) - if len(prompt_table_data.size()) == 3: - _, task_vocab_size, hidden_size = prompt_table_data.size() - elif len(prompt_table_data.size()) == 2: - task_vocab_size, hidden_size = prompt_table_data.size() - task_vocab_size = torch.tensor([task_vocab_size], dtype=torch.int32) - prompt_table_data = prompt_table_data.view(-1, hidden_size) - else: - prompt_table_data = torch.empty( - [1, self.hidden_size * self.mapping.tp_size], dtype=self.dtype) - task_vocab_size = torch.zeros([1], dtype=torch.int32) - if tasks is not None: - tasks = torch.tensor([int(t) for t in tasks.split(',')], - dtype=torch.int32) - assert tasks.size(0) == batch_size, \ - f"Number of supplied tasks ({tasks.size(0)}) must match input batch size ({batch_size})" - else: - tasks = torch.zeros([batch_size], dtype=torch.int32) - - if isinstance(self.session, GenerationSession): - return { - 'prompt_embedding_table': prompt_table_data.cuda(), - 'tasks': tasks.cuda(), - 'prompt_vocab_size': task_vocab_size.cuda() - } - else: - return { - 'embedding_table': prompt_table_data.cuda(), - 'tasks': tasks.cuda(), - 'vocab_size': task_vocab_size.cuda() - } - - -class ModelRunner(ModelRunnerMixin): - """ - An interface class that wraps GenerationSession and provides generation methods. - """ - - def __init__( - self, - session: GenerationSession, - max_batch_size: int, - max_input_len: int, - max_seq_len: int, - max_beam_width: int, - kv_cache_type: KVCacheType, - lora_manager: Optional[LoraManager] = None, - ) -> None: - """ - Create a ModelRunner instance. - You are recommended to use the from_dir method to load the engine and create a ModelRunner instance. - - Args: - session (GenerationSession): - The TensorRT session created from an engine. - max_batch_size (int): - The maximum batch size allowed for the input. - max_input_len (int): - The maximum input length allowed for the input. - max_seq_len (int): - The maximum sequence length (input + new tokens). - max_beam_width (int): - The maximum beam width. - lora_manager (LoraManager): - The LoRA manager to handle LoRA weights. - """ - emit_engine_arch_deprecation("ModelRunner") - self.session = session - self.max_batch_size = max_batch_size - self.max_input_len = max_input_len - self.max_seq_len = max_seq_len - self.max_beam_width = max_beam_width - self.lora_manager = lora_manager - self.kv_cache_type = kv_cache_type - self.enable_context_fmha_fp32_acc = False - self.multi_block_mode = True - - @classmethod - def from_engine( - cls, - engine: Engine, - *, - max_output_len: Optional[int], - lora_dir: Optional[List[str]], - rank: int, - debug_mode: bool, - lora_ckpt_source: str, - medusa_choices: List[List[int]], - stream: torch.cuda.Stream, - gpu_weights_percent: float, - enable_context_fmha_fp32_acc: Optional[bool], - multi_block_mode: Optional[bool], - ) -> 'ModelRunner': - model_config = _engine_config_to_model_config( - engine.config, gpu_weights_percent=gpu_weights_percent) - - if model_config.kv_cache_type == KVCacheType.DISABLED: - assert max_output_len == 1 or max_output_len is None, 'Disabled KV cache is intended for context phase only now.' - - pretrained_config = engine.config.pretrained_config - build_config = engine.config.build_config - max_batch_size = build_config.max_batch_size - max_input_len = build_config.max_input_len - max_seq_len = build_config.max_seq_len - max_beam_width = build_config.max_beam_width - if 'GLM' in pretrained_config.architecture and pretrained_config.chatglm_version in [ - 'glm', 'chatglm' - ]: - session_cls = ChatGLMGenerationSession - else: - session_cls = GenerationSession - engine_buffer = engine.engine - runtime_mapping = pretrained_config.mapping - - if medusa_choices is not None: - assert session_cls == GenerationSession, "Medusa is only supported by GenerationSession" - - assert model_config.max_medusa_tokens > 0, \ - "medusa_chioce is specified but model_config.max_medusa_tokens is 0." - - if MpiComm.size() > runtime_mapping.gpus_per_node: - assert MpiComm.local_size() == runtime_mapping.gpus_per_node - if not DISABLE_TORCH_DEVICE_SET: - torch.cuda.set_device(rank % runtime_mapping.gpus_per_node) - session = session_cls(model_config, - engine_buffer, - runtime_mapping, - debug_mode=debug_mode, - stream=stream) - if session.runtime.engine.streamable_weights_size: - session.runtime._set_weight_streaming(gpu_weights_percent) - - if session.use_lora_plugin: - lora_manager = LoraManager(mapping=runtime_mapping, - model_config=model_config) - if lora_dir is not None: - lora_manager.load_from_ckpt(lora_dir, - model_config=model_config, - ckpt_source=lora_ckpt_source) - else: - lora_manager = None - - runner = cls(session=session, - max_batch_size=max_batch_size, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - max_beam_width=max_beam_width, - kv_cache_type=model_config.kv_cache_type, - lora_manager=lora_manager) - runner.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - runner.multi_block_mode = multi_block_mode - return runner - - @classmethod - def from_dir( - cls, - engine_dir: str, - *, - max_output_len: Optional[int] = None, - lora_dir: Optional[List[str]] = None, - rank: int = 0, - debug_mode: bool = False, - lora_ckpt_source: str = "hf", - medusa_choices: List[List[int]] = None, - stream: torch.cuda.Stream = None, - gpu_weights_percent: float = 1, - enable_context_fmha_fp32_acc: Optional[bool] = None, - multi_block_mode: Optional[bool] = None, - fail_fast_on_attention_window_too_large: bool = False, - ) -> 'ModelRunner': - """ - Create a ModelRunner instance from an engine directory. - - Args: - engine_dir (str): - The directory that contains the serialized engine files and config files. - max_output_len (Optional[int]): - max_output_len, this arg might be available only when loading time, generate will still to check when disable_kv_cache is enabled. - lora_dir (Optional[List[str]]): - The directories that contain LoRA weights. - rank (int): - The runtime rank id. - debug_mode (bool): - Whether or not to turn on the debug mode. - medusa_choices (List[List[int]]): - Medusa choices to use when in Medusa decoding - stream (torch.cuda.Stream): - Stream to use. - multi_block_mode (bool): - Whether to distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel. - fail_fast_on_attention_window_too_large (bool): - Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache. - Note: This parameter is only applicable to C++ runtime (ModelRunnerCpp). - Returns: - ModelRunner: An instance of ModelRunner. - """ - engine_version = get_engine_version(engine_dir) - profiler.start('load tensorrt_llm engine') - # the old engine format - if engine_version is None: - engine_dir = Path(engine_dir) - config_path = engine_dir / "config.json" - model_config, other_config = read_config(config_path) - world_size = other_config.pop('world_size') - tp_size = other_config.pop('tp_size') - pp_size = other_config.pop('pp_size') - max_batch_size = other_config.pop('max_batch_size') - max_input_len = other_config.pop('max_input_len') - max_output_len = other_config.pop('max_output_len') - max_beam_width = other_config.pop('max_beam_width') - runtime_mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=tp_size, - pp_size=pp_size) - - engine_name = get_engine_name(model_config.model_name, - model_config.dtype, tp_size, pp_size, - rank) - serialize_path = engine_dir / engine_name - - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - - if model_config.model_name in ('chatglm_6b', 'glm_10b'): - session_cls = ChatGLMGenerationSession - elif model_config.model_name == 'qwen': - session_cls = QWenForCausalLMGenerationSession - else: - session_cls = GenerationSession - - if medusa_choices is not None: - assert model_config.max_medusa_tokens > 0, \ - "medusa_choice is specified but model_config.max_medusa_tokens is 0." - - if not DISABLE_TORCH_DEVICE_SET: - torch.cuda.set_device(rank % runtime_mapping.gpus_per_node) - session = session_cls(model_config, - engine_buffer, - runtime_mapping, - debug_mode=debug_mode, - stream=stream) - if session.use_lora_plugin: - lora_manager = LoraManager(mapping=runtime_mapping, - model_config=model_config) - if lora_dir is not None: - lora_manager.load_from_ckpt(lora_dir, - model_config=model_config, - ckpt_source=lora_ckpt_source) - else: - lora_manager = None - - if session.runtime.engine.streamable_weights_size: - session.runtime._set_weight_streaming(gpu_weights_percent) - - profiler.stop('load tensorrt_llm engine') - loading_time = profiler.elapsed_time_in_sec( - "load tensorrt_llm engine") - logger.info(f'Load engine takes: {loading_time} sec') - - runner = cls(session=session, - max_batch_size=max_batch_size, - max_input_len=max_input_len, - max_seq_len=max_input_len + max_output_len, - max_beam_width=max_beam_width, - kv_cache_type=KVCacheType.CONTINUOUS, - lora_manager=lora_manager) - runner.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - runner.multi_block_mode = multi_block_mode - return runner - else: - # the new engine format - engine = Engine.from_dir(engine_dir, rank) - if lora_dir is None: - config_lora_dir = engine.config.build_config.lora_config.lora_dir - if len(config_lora_dir) > 0: - lora_dir = [ - f"{engine_dir}/{dir}" for dir in config_lora_dir - ] - lora_ckpt_source = engine.config.build_config.lora_config.lora_ckpt_source - - runner = ModelRunner.from_engine( - engine=engine, - max_output_len=max_output_len, - lora_dir=lora_dir, - rank=rank, - debug_mode=debug_mode, - lora_ckpt_source=lora_ckpt_source, - medusa_choices=medusa_choices, - stream=stream, - gpu_weights_percent=gpu_weights_percent, - enable_context_fmha_fp32_acc=enable_context_fmha_fp32_acc, - multi_block_mode=multi_block_mode, - ) - profiler.stop('load tensorrt_llm engine') - loading_time = profiler.elapsed_time_in_sec( - "load tensorrt_llm engine") - logger.info(f'Load engine takes: {loading_time} sec') - return runner - - @property - def dtype(self) -> torch.dtype: - return self.session.dtype - - @property - def vocab_size(self) -> int: - return self.session.vocab_size - - @property - def vocab_size_padded(self) -> int: - return self.session.vocab_size_padded - - @property - def hidden_size(self) -> int: - return self.session.hidden_size - - @property - def num_heads(self) -> int: - return self.session.num_heads - - @property - def num_layers(self) -> int: - return self.session.num_layers - - @property - def max_sequence_length(self) -> int: - return self.max_seq_len - - @property - def remove_input_padding(self) -> bool: - return self.session.remove_input_padding - - @property - def use_lora_plugin(self) -> bool: - return self.session.use_lora_plugin - - @property - def max_prompt_embedding_table_size(self) -> int: - return self.session.max_prompt_embedding_table_size - - @property - def mapping(self) -> Mapping: - return self.session.mapping - - @property - def gather_context_logits(self) -> bool: - return self.session.gather_context_logits - - @property - def gather_generation_logits(self) -> bool: - return self.session.gather_generation_logits - - def generate(self, - batch_input_ids: List[torch.Tensor], - position_ids: List[torch.Tensor] = None, - sampling_config: Optional[SamplingConfig] = None, - prompt_table: Optional[Union[str, torch.Tensor]] = None, - prompt_tasks: Optional[str] = None, - lora_uids: Optional[list] = None, - streaming: bool = False, - output_generation_logits: bool = False, - stopping_criteria: Optional[StoppingCriteria] = None, - logits_processor: Optional[LogitsProcessor] = None, - medusa_choices: Optional[List[List[int]]] = None, - encoder_max_input_length: int = None, - encoder_input_features: List[torch.Tensor] = None, - encoder_output_lengths: List[torch.Tensor] = None, - cross_attention_masks: List[torch.Tensor] = None, - **kwargs) -> Union[torch.Tensor, dict]: - """ - Generates sequences of token ids. - The generation-controlling parameters are set in the sampling_config; it will be set to a default one if not passed. - You can override any sampling_config's attributes by passing corresponding parameters. - - Args: - batch_input_ids (List[torch.Tensor]): - A list of input id tensors. Each tensor is of shape (sequence_length, ). - sampling_config (SamplingConfig): - The sampling configuration to be used as base parametrization for the generation call. - The passed **kwargs matching the sampling_config's attributes will override them. - If the sampling_config is not provided, a default will be used. - prompt_table (str or torch.Tensor): - The file path of prompt table (.npy format, exported by nemo_prompt_convert.py) or the prompt table itself. - prompt_tasks (str): - The prompt tuning task ids for the input batch, in format of comma-separated list (e.g., 0,3,1,0). - lora_uids (list): - The uids of LoRA weights for the input batch. Use -1 to disable the LoRA module. - streaming (bool): - Whether or not to use streaming mode for generation. - stopping_criteria (StoppingCriteria): - Custom stopping criteria. - logits_processor (LogitsProcessor): - Custom logits processors. - medusa_choices (List[List[int]]): - Medusa decoding choices. - kwargs (Dict[str, Any]: - Ad hoc parametrization of sampling_config. - The passed **kwargs matching the sampling_config's attributes will override them. - Returns: - torch.Tensor or dict: - If return_dict=False, the method returns generated output_ids. - If return_dict=True, the method returns a dict of output_ids, - sequence_lengths (if sampling_config.output_sequence_lengths=True), - context_logits and generation_logits (if self.gather_context_logits=True - and self.gather_generation_logits=True, respectively). - """ - # Use sampling_config like HF's generation_config - if sampling_config is None: - sampling_config = SamplingConfig(end_id=None, pad_id=None) - else: - sampling_config = copy.deepcopy(sampling_config) - sampling_config.update(**kwargs) - - # To prevent numerical overflow when the temperature is set to 0.0 - # Modify generation.SamplingConfig - if isinstance(sampling_config.temperature, - float) and sampling_config.temperature == 0.0: - logger.warning( - "Convert `temperature=0.0` to `temperature=1.0` and `top_k=1` to prevent overflow." - ) - sampling_config.temperature = 1.0 - sampling_config.top_k = 1 - - self._check_inputs(batch_input_ids, sampling_config) - - if kwargs.get('num_return_sequences', None) is not None: - raise ValueError( - 'num_return_sequences will be ignored since ' - 'num_return_sequences > 1 is not supported on python runtime. ' - 'Please use C++ runtime.') - - batch_size = len(batch_input_ids) - batch_input_ids, input_lengths = self._prepare_inputs( - batch_input_ids, sampling_config.pad_id) - - def maybe_convert_to_words_list_format( - words_list: Optional[Union[list, np.ndarray, torch.Tensor]] - ) -> Optional[np.ndarray]: - if words_list is None or isinstance(words_list, np.ndarray): - return words_list - elif isinstance(words_list, torch.Tensor): - return words_list.numpy() - elif isinstance(words_list, list): - return to_word_list_format(words_list) - else: - raise TypeError( - f"Unexpected words_list type={type(words_list)}. Only list, np.ndarray, and torch.Tensor are supported." - ) - - if cross_attention_masks is not None: - encoder_input_features = torch.concat(encoder_input_features) - encoder_output_lengths = torch.concat(encoder_output_lengths) - - sampling_config.bad_words_list = maybe_convert_to_words_list_format( - sampling_config.bad_words_list) - sampling_config.stop_words_list = maybe_convert_to_words_list_format( - sampling_config.stop_words_list) - - if not self.kv_cache_type and sampling_config.max_new_tokens > 1: - raise RuntimeError( - 'Disabled KV cache is intended for context phase only now.') - - self.session.setup( - batch_size=batch_size, - max_context_length=input_lengths.max().item(), - max_new_tokens=sampling_config.max_new_tokens, - beam_width=sampling_config.num_beams, - max_attention_window_size=sampling_config.max_attention_window_size, - sink_token_length=sampling_config.sink_token_length, - lora_manager=self.lora_manager, - lora_uids=lora_uids, - medusa_choices=medusa_choices, - enable_context_fmha_fp32_acc=self.enable_context_fmha_fp32_acc, - multi_block_mode=self.multi_block_mode, - encoder_max_input_length=encoder_max_input_length, - ) - - batch_input_ids = batch_input_ids.cuda() - input_lengths = input_lengths.cuda() - other_kwargs = self._prepare_ptuning(prompt_table, prompt_tasks, - batch_size) - other_kwargs['skip_cross_attn_blocks'] = kwargs.get( - 'skip_cross_attn_blocks', None) - outputs = self.session.decode( - batch_input_ids, - input_lengths, - sampling_config, - stop_words_list=sampling_config.stop_words_list, - bad_words_list=sampling_config.bad_words_list, - output_sequence_lengths=sampling_config.output_sequence_lengths, - output_generation_logits=output_generation_logits, - return_dict=sampling_config.return_dict, - streaming=streaming, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - position_ids=position_ids, - encoder_output=encoder_input_features, - encoder_input_lengths=encoder_output_lengths, - cross_attention_mask=cross_attention_masks, - **other_kwargs) - if sampling_config.return_dict: - if streaming: - outputs = (self._prepare_outputs(curr_outputs, input_lengths) - for curr_outputs in outputs) - else: - outputs = self._prepare_outputs(outputs, input_lengths) - return outputs - - def serialize_engine(self) -> trt.IHostMemory: - """ - Serialize the engine. - - Returns: - bytes: The serialized engine. - """ - return self.session.runtime._serialize_engine() diff --git a/tensorrt_llm/runtime/model_runner_cpp.py b/tensorrt_llm/runtime/model_runner_cpp.py deleted file mode 100644 index 3186a47f697a..000000000000 --- a/tensorrt_llm/runtime/model_runner_cpp.py +++ /dev/null @@ -1,1223 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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 copy -import os -from os.path import join -from pathlib import Path -from typing import Dict, List, Optional, Union - -import torch - -from .. import profiler -from .._deprecation import emit_engine_arch_deprecation -from .._utils import maybe_pin_memory, mpi_broadcast -from ..bindings import DataType, GptJsonConfig, ModelConfig, WorldConfig -from ..bindings import executor as trtllm -from ..bindings.executor import (DecodingMode, ExternalDraftTokensConfig, - OrchestratorConfig, ParallelConfig) -from ..builder import EngineConfig -from ..layers import MropeParams -from ..llmapi.kv_cache_type import KVCacheType -from ..logger import logger -from ..mapping import Mapping -from .generation import LogitsProcessor, LoraManager -from .generation import ModelConfig as ModelConfigPython -from .generation import SamplingConfig, StoppingCriteria -from .model_runner import ModelRunnerMixin, _engine_config_to_model_config - -_bindings_dtype_to_torch_dtype_dict = { - DataType.FLOAT: torch.float, - DataType.HALF: torch.half, - DataType.INT8: torch.int8, - DataType.INT32: torch.int32, - DataType.BOOL: torch.bool, - DataType.UINT8: torch.uint8, - DataType.BF16: torch.bfloat16, - DataType.INT64: torch.int64 -} - -SamplingConfigType = Union[SamplingConfig, trtllm.SamplingConfig] - - -def _world_config_to_mapping(world_config: WorldConfig): - return Mapping(world_size=world_config.size, - rank=world_config.rank, - gpus_per_node=world_config.gpus_per_node, - tp_size=world_config.tensor_parallelism, - pp_size=world_config.pipeline_parallelism, - cp_size=world_config.context_parallelism) - - -class ModelRunnerCpp(ModelRunnerMixin): - """ - An interface class that wraps Executor and provides generation methods. - """ - - def __init__(self, - executor: trtllm.Executor, - max_batch_size: int, - max_input_len: int, - max_seq_len: int, - max_beam_width: int, - model_config: ModelConfig, - world_config: WorldConfig, - use_kv_cache: bool, - lora_manager: Optional[LoraManager] = None) -> None: - emit_engine_arch_deprecation("ModelRunnerCpp") - self.session = executor - self.max_batch_size = max_batch_size - self.max_input_len = max_input_len - self.max_seq_len = max_seq_len - self.max_beam_width = max_beam_width - self.model_config = model_config - self.mapping = _world_config_to_mapping(world_config) - self.world_config = world_config - self.use_kv_cache = use_kv_cache - self.lora_manager = lora_manager - - @classmethod - def from_dir( - cls, - engine_dir: str, - *, - lora_dir: Optional[str] = None, - rank: int = 0, - max_batch_size: Optional[int] = None, - max_input_len: Optional[int] = None, - max_output_len: Optional[int] = None, - max_beam_width: Optional[int] = None, - max_attention_window_size: Optional[list[int]] = None, - sink_token_length: Optional[int] = None, - kv_cache_free_gpu_memory_fraction: Optional[float] = None, - cross_kv_cache_fraction: Optional[float] = None, - medusa_choices: list[list[int]] | None = None, - eagle_choices: list[list[int]] | None = None, - eagle_posterior_threshold: float | None = None, - eagle_use_dynamic_tree: bool = False, - eagle_dynamic_tree_max_top_k: Optional[int] = None, - lookahead_config: list[int] | None = None, - debug_mode: bool = False, - lora_ckpt_source: str = "hf", - use_gpu_direct_storage: bool = False, - gpu_weights_percent: float = 1, - max_tokens_in_paged_kv_cache: int | None = None, - kv_cache_enable_block_reuse: bool = False, - enable_chunked_context: bool = False, - is_enc_dec: bool = False, - multi_block_mode: bool = True, - enable_context_fmha_fp32_acc: Optional[bool] = None, - cuda_graph_mode: Optional[bool] = None, - logits_processor_map: Optional[Dict[str, LogitsProcessor]] = None, - device_ids: List[int] | None = None, - is_orchestrator_mode: bool = False, - use_runtime_defaults: bool = True, - gather_generation_logits: bool = False, - use_variable_beam_width_search: bool = False, - mm_embedding_offloading: bool = False, - fail_fast_on_attention_window_too_large: bool = False, - normalize_log_probs: bool = False, - ) -> 'ModelRunnerCpp': - """ - Create a ModelRunnerCpp instance from an engine directory. - - Args: - engine_dir (str): - The directory that contains the serialized engine files and config files. - lora_dir (str): - The directory that contains LoRA weights. - rank (int): - The runtime rank id. - max_batch_size (int): - The runtime batch size limit. If max_batch_size is not None, it should not - be larger than the engine's max_batch_size; otherwise, the engine's max_batch_size - will be used. - max_input_len (int): - The runtime input length limit. If max_input_len is not None, it should not - be larger than the engine's max_input_len; otherwise, the engine's max_input_len - will be used. - max_output_len (int): - The runtime output length limit. If max_output_len is not None, it should not - be larger than the engine's max_output_len; otherwise, the engine's max_output_len - will be used. - max_beam_width (int): - The runtime beam width limit. If max_beam_width is not None, it should not - be larger than the engine's max_beam_width; otherwise, the engine's max_beam_width - will be used. - max_attention_window_size (List[int]): - The attention window size that controls the sliding window attention / cyclic kv cache behavior. - sink_token_length (int) : - The sink token length, default=0. - kv_cache_free_gpu_memory_fraction (float) : - Free GPU memory fraction that KV cache used. - cross_kv_cache_fraction (float) : - KV Cache fraction reserved for cross attention, should only be used with enc-dec models. - debug_mode (bool): - Whether or not to turn on the debug mode. - medusa_choices (List[List[int]]): - Medusa choices to use when in Medusa decoding. - eagle_choices (List[List[int]]): - Eagle choices to use when in Eagle-1 decoding. - eagle_posterior_threshold float: - Minimum token probability threshold for typical acceptance. - Value different from None enables typical acceptance in Eagle. - eagle_use_dynamic_tree bool: - Whether to use Eagle-2, which is dynamic tree. - eagle_dynamic_tree_max_top_k int: - The maximum number of draft tokens to expand for each node in Eagle-2. - lora_ckpt_source (str): - Source of checkpoint. Should be one of ['hf', 'nemo']. - max_tokens_in_paged_kv_cache (int): - Maximum amount of tokens configured in kv cache. - kv_cache_enable_block_reuse (bool): - Enables block reuse in kv cache. - enable_chunked_context (bool): - Enables chunked context. - is_enc_dec (bool): - Whether the model is encoder-decoder architecture. - multi_block_mode (bool): - Whether to distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel. - enable_context_fmha_fp32_acc (bool): - Enable FMHA runner FP32 accumulation. - cuda_graph_mode (bool): - Whether to use cuda graph for inference. - logits_processor_map (Dict[str, LogitsProcessor]) - A map of logits processor functions indexed by names. A name can be provided later to - the generate() function to specify which logits processor to run. - device_ids (List[int]): - Device indices to run the Executor on. - is_orchestrator_mode (bool): - The mode to run the model-runner, Leader mode by default. - gather_generation_logits (bool): - Enable gathering generation logits. - fail_fast_on_attention_window_too_large (bool): - Whether to fail fast if the attention window(s) are too large to fit even a single sequence in the KVCache. - Returns: - ModelRunnerCpp: An instance of ModelRunnerCpp. - """ - extended_runtime_perf_knob_config = trtllm.ExtendedRuntimePerfKnobConfig( - ) - if multi_block_mode is not None: - extended_runtime_perf_knob_config.multi_block_mode = multi_block_mode - if enable_context_fmha_fp32_acc is not None: - extended_runtime_perf_knob_config.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - if cuda_graph_mode is not None: - extended_runtime_perf_knob_config.cuda_graph_mode = cuda_graph_mode - - model_config = None - is_multimodal = {'vision', 'llm'}.issubset({ - name - for name in os.listdir(engine_dir) - if os.path.isdir(os.path.join(engine_dir, name)) - }) - encoder_path = None - decoder_path = None - - if is_enc_dec: - if is_multimodal: - encoder_path = join(engine_dir, 'vision') - decoder_path = join(engine_dir, 'llm') - else: - encoder_path = join(engine_dir, 'encoder') - decoder_path = join(engine_dir, 'decoder') - - encoder_config_path = Path(encoder_path) / "config.json" - encoder_json_config = GptJsonConfig.parse_file(encoder_config_path) - encoder_model_config = encoder_json_config.model_config - decoder_config_path = Path(decoder_path) / "config.json" - decoder_json_config = GptJsonConfig.parse_file(decoder_config_path) - decoder_model_config = decoder_json_config.model_config - - json_config = decoder_json_config - model_config = decoder_model_config - engine_dir = decoder_path - - if max_input_len is None and not is_multimodal: - max_input_len = encoder_model_config.max_input_len - else: - config_path = Path(engine_dir) / "config.json" - json_config = GptJsonConfig.parse_file(config_path) - model_config = json_config.model_config - - use_kv_cache = KVCacheType.from_cpp( - model_config.kv_cache_type) != KVCacheType.DISABLED - if not model_config.use_cross_attention: - assert cross_kv_cache_fraction is None, "cross_kv_cache_fraction should only be used with enc-dec models." - - if not use_kv_cache: - assert max_output_len == 1 or max_output_len is None, 'Disabled KV cache is intended for context phase only now.' - - # Note: Parallel configuration will be fetched automatically from trtllm.Executor constructor - # by inspecting the json file. These lines serve the purpose of serving vocab_size_padded and - # num_layers properties. - # MPI world size must be 1 in Orchestrator mode - if is_orchestrator_mode: - tp_size = 1 - pp_size = 1 - cp_size = 1 - # Check the count of devices equal to tp_size of engine - # assert len(device_ids) == json_config.tensor_parallelism - else: - tp_size = json_config.tensor_parallelism - pp_size = json_config.pipeline_parallelism - cp_size = json_config.context_parallelism - gpus_per_node = json_config.gpus_per_node - world_config = WorldConfig.mpi(tensor_parallelism=tp_size, - pipeline_parallelism=pp_size, - context_parallelism=cp_size, - gpus_per_node=gpus_per_node) - assert rank == world_config.rank - - engine_config = EngineConfig.from_json_file(f"{engine_dir}/config.json") - if model_config.use_lora_plugin and rank == 0: - mapping = _world_config_to_mapping(world_config) - lora_manager = LoraManager( - mapping=mapping, - model_config=ModelConfigPython.from_model_config_cpp( - model_config)) - if lora_dir is None: - config_lora_dir = engine_config.build_config.lora_config.lora_dir - if len(config_lora_dir) > 0: - lora_dir = [ - f"{engine_dir}/{dir}" for dir in config_lora_dir - ] - lora_ckpt_source = engine_config.build_config.lora_config.lora_ckpt_source - - if lora_dir is not None: - runtime_model_config = _engine_config_to_model_config( - engine_config, gpu_weights_percent=gpu_weights_percent) - # For Executor, only rank 0 can enqueue requests, and should hold all lora weights - lora_manager.load_from_ckpt(lora_dir, - model_config=runtime_model_config, - ckpt_source=lora_ckpt_source) - else: - raise RuntimeError( - f"LoRA weights are unspecified and also unavailable in the engine_dir ({engine_dir})." - ) - - max_lora_rank = engine_config.build_config.lora_config.max_lora_rank - num_lora_modules = engine_config.pretrained_config.num_hidden_layers * \ - len(lora_manager.lora_target_modules + lora_manager.missing_qkv_modules) - num_lora_adapters = min(lora_manager.num_lora_adapters, 8) - peft_cache_config = trtllm.PeftCacheConfig( - num_device_module_layer=max_lora_rank * num_lora_modules * - num_lora_adapters, - num_host_module_layer=max_lora_rank * num_lora_modules * - num_lora_adapters, - ) - else: - lora_manager = None - peft_cache_config = trtllm.PeftCacheConfig() - - if world_config.size > 1: - peft_cache_config = mpi_broadcast(peft_cache_config, 0) - - profiler.start('load tensorrt_llm engine') - - kv_cache_config = trtllm.KvCacheConfig( - free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction, - max_attention_window=max_attention_window_size, - sink_token_length=sink_token_length, - max_tokens=max_tokens_in_paged_kv_cache, - enable_block_reuse=kv_cache_enable_block_reuse, - cross_kv_cache_fraction=cross_kv_cache_fraction, - runtime_defaults=json_config.runtime_defaults - if use_runtime_defaults else None, - ) - - decoding_config = trtllm.DecodingConfig() - if medusa_choices is not None: - decoding_config.medusa_choices = medusa_choices - if multi_block_mode is not None: - multi_block_mode = False # Medusa doesn't support multi-block mode. - - if eagle_choices is not None or eagle_posterior_threshold is not None or eagle_use_dynamic_tree: - greedy_sampling = eagle_posterior_threshold is None - decoding_config.eagle_config = trtllm.EagleConfig( - eagle_choices, greedy_sampling, eagle_posterior_threshold, - eagle_use_dynamic_tree, eagle_dynamic_tree_max_top_k) - if multi_block_mode is not None: - logger.warning( - f'Multi block mode is not supported for EAGLE. Disabling it.' - ) - multi_block_mode = False # Eagle doesn't support multi-block mode. - - if lookahead_config is not None: - [w, n, g] = lookahead_config - decoding_config.lookahead_decoding_config = trtllm.LookaheadDecodingConfig( - w, n, g) - - if use_variable_beam_width_search: - decoding_config.decoding_mode = DecodingMode.BeamSearch( - ).useVariableBeamWidthSearch(True) - - if max_batch_size is None: - max_batch_size = model_config.max_batch_size - else: - assert max_batch_size <= model_config.max_batch_size - if max_input_len is None: - max_input_len = model_config.max_input_len - # NOTE: remove assertion here for temp fix, - # model_config.max_input_len is not the upper bound of input length. - # If runtime max_input_len is not properly set, - # C++ runtime will throw an error when fetching new requests - if max_output_len is None or is_enc_dec: - max_seq_len = model_config.max_seq_len - else: - max_seq_len = max_input_len + max_output_len - assert max_seq_len <= model_config.max_seq_len - if max_beam_width is None: - max_beam_width = model_config.max_beam_width - else: - assert max_beam_width <= model_config.max_beam_width - - debug_config = None - if debug_mode: - # To debug specific tensors, add tensor names in the following list - # if none provided, all input and output tensors will be dumped - # if not none, it will disable all input/output dump - debug_tensor_names: List[str] = [ - ] # modify this list for specific tensor dump - debug_config = trtllm.DebugConfig( - debug_input_tensors=True, - debug_output_tensors=True, - debug_tensor_names=debug_tensor_names) - - trtllm_config = trtllm.ExecutorConfig( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - kv_cache_config=kv_cache_config, - decoding_config=decoding_config, - peft_cache_config=peft_cache_config, - debug_config=debug_config, - use_gpu_direct_storage=use_gpu_direct_storage, - gpu_weights_percent=gpu_weights_percent, - gather_generation_logits=gather_generation_logits, - normalize_log_probs=normalize_log_probs, - ) - trtllm_config.enable_chunked_context = enable_chunked_context - trtllm_config.extended_runtime_perf_knob_config = extended_runtime_perf_knob_config - trtllm_config.mm_embedding_offloading = mm_embedding_offloading - trtllm_config.fail_fast_on_attention_window_too_large = fail_fast_on_attention_window_too_large - if is_orchestrator_mode: - communication_mode = trtllm.CommunicationMode.ORCHESTRATOR - path = str(Path(__file__).parent.parent / 'bin' / 'executorWorker') - orchestrator_config = OrchestratorConfig(True, path) - else: - communication_mode = trtllm.CommunicationMode.LEADER - orchestrator_config = None - - trtllm_config.parallel_config = ParallelConfig( - trtllm.CommunicationType.MPI, - communication_mode, - device_ids=device_ids, - orchestrator_config=orchestrator_config) - - # LogitsPostProcessor in Orchestrator mode is not supported yet. - if not is_orchestrator_mode: - logits_proc_config = trtllm.LogitsPostProcessorConfig() - if logits_processor_map is not None: - logits_proc_config.processor_map = logits_processor_map - trtllm_config.logits_post_processor_config = logits_proc_config - - if is_enc_dec: - executor = trtllm.Executor(encoder_path, decoder_path, - trtllm.ModelType.ENCODER_DECODER, - trtllm_config) - else: - executor = trtllm.Executor(Path(engine_dir), - trtllm.ModelType.DECODER_ONLY, - trtllm_config) - - profiler.stop('load tensorrt_llm engine') - - loading_time = profiler.elapsed_time_in_sec("load tensorrt_llm engine") - logger.info(f'Load engine takes: {loading_time} sec') - - return cls(executor, - max_batch_size=max_batch_size, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - max_beam_width=max_beam_width, - model_config=model_config, - world_config=world_config, - use_kv_cache=use_kv_cache, - lora_manager=lora_manager) - - def _check_inputs(self, batch_input_ids: List[List[int]], - encoder_input_ids: Optional[List[List[int]]], - sampling_config: trtllm.SamplingConfig, max_new_tokens): - batch_size = len(encoder_input_ids) if encoder_input_ids else len( - batch_input_ids) - if batch_size > self.max_batch_size: - raise RuntimeError( - f"Input batch size ({batch_size}) exceeds the engine or specified limit ({self.max_batch_size})" - ) - input_lengths = [ - len(x) for x in encoder_input_ids - ] if encoder_input_ids else [len(x) for x in batch_input_ids] - max_length = max(input_lengths) - if max_length > self.max_input_len: - raise RuntimeError( - f"Maximum input length ({max_length}) exceeds the engine or specified limit ({self.max_input_len})" - ) - if encoder_input_ids: - decoder_max_length = max([len(x) for x in batch_input_ids]) - if decoder_max_length + max_new_tokens > self.max_seq_len: - raise RuntimeError( - f"Decoder prefix tokens ({decoder_max_length}) + maximum new tokens ({max_new_tokens}) exceeds the engine or specified limit ({self.max_seq_len})" - ) - else: - if max_length + max_new_tokens > self.max_seq_len: - raise RuntimeError( - f"Maximum input length ({max_length}) + maximum new tokens ({max_new_tokens}) exceeds the engine or specified limit ({self.max_seq_len})" - ) - if sampling_config.beam_width > self.max_beam_width: - raise RuntimeError( - f"Num beams ({sampling_config.beam_width}) exceeds the engine or specified limit ({self.max_beam_width})" - ) - - @property - def dtype(self) -> torch.dtype: - bindings_dtype = self.model_config.data_type - return _bindings_dtype_to_torch_dtype_dict[bindings_dtype] - - @property - def vocab_size(self) -> int: - return self.model_config.vocab_size - - @property - def vocab_size_padded(self) -> int: - return self.model_config.vocab_size_padded(self.world_config.size) - - @property - def hidden_size(self) -> int: - return self.model_config.hidden_size - - @property - def num_heads(self) -> int: - return self.model_config.num_heads - - @property - def num_layers(self) -> int: - return self.model_config.num_layers( - self.world_config.pipeline_parallelism, - self.world_config.pipeline_parallel_rank, - ) - - @property - def max_sequence_length(self) -> int: - return self.max_seq_len - - @property - def remove_input_padding(self) -> bool: - return self.model_config.use_packed_input - - @property - def max_prompt_embedding_table_size(self) -> int: - return self.model_config.max_prompt_embedding_table_size - - @property - def gather_context_logits(self) -> bool: - return self.model_config.compute_context_logits - - @property - def gather_generation_logits(self) -> bool: - return self.model_config.compute_generation_logits - - def generate( - self, - batch_input_ids: List[torch.Tensor], - *, - position_ids: List[torch.Tensor] = None, - encoder_input_ids: List[torch.Tensor] = None, - encoder_input_features: List[ - torch.Tensor] = None, # TODO: add to doc string - encoder_output_lengths: List[int] = None, - cross_attention_masks: List[ - torch.Tensor] = None, # TODO: add to doc string - mrope_params: Optional[MropeParams] = None, - sampling_config: Optional[SamplingConfig] = None, - lora_uids: Optional[list] = None, - lookahead_config: list[int] | None = None, - streaming: bool = False, - stopping_criteria: Optional[StoppingCriteria] = None, - logits_processor_names: list[str] | None = None, - max_new_tokens: int = 1, - end_id: int | None = None, - pad_id: int | None = None, - bad_words_list: list[list[int]] | None = None, - stop_words_list: list[list[int]] | None = None, - return_dict: bool = False, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - output_log_probs: bool = False, - output_cum_log_probs: bool = False, - prompt_table: Optional[Union[str, torch.Tensor]] = None, - prompt_tasks: Optional[str] = None, - input_token_extra_ids: List[List[int]] = None, - return_all_generated_tokens: bool = False, - language_adapter_uids: Optional[List[int]] = None, - mm_embedding_offloading: bool = False, - **kwargs) -> Union[torch.Tensor, dict]: - """ - Generates sequences of token ids. - The generation-controlling parameters are set in the sampling_config; it will be set to a default one if not passed. - You can override any sampling_config's attributes by passing corresponding parameters. - - Args: - batch_input_ids (List[torch.Tensor]): - A list of input id tensors. Each tensor is of shape (sequence_length, ). - position_ids (List[torch.Tensor]): - A list of position id tensors. Each tensor is of shape (sequence_length, ). - encoder_input_ids (List[torch.Tensor]): - A list of encoder input id tensors for encoder-decoder models (optional). Each tensor is of shape (sequence_length, ). - encoder_input_features: (List[torch.Tensor]): - A list of encoder input feature tensors for multimodal encoder-decoder models (optional). Each tensor is of shape (sequence_length, feature_dim). - encoder_output_lengths: (List[int]): - A list of encoder output lengths (optional) if encoder output has different length from encoder input (due to convolution down-sampling, etc.) - sampling_config (SamplingConfig): - The sampling configuration to be used as base parametrization for the generation call. - The passed **kwargs matching the sampling_config's attributes will override them. - If the sampling_config is not provided, a default will be used. - prompt_table (str or torch.Tensor): - The file path of prompt table (.npy format, exported by nemo_prompt_convert.py) or the prompt table itself. - prompt_tasks (str): - The prompt tuning task ids for the input batch, in format of comma-separated list (e.g., 0,3,1,0). - input_token_extra_ids (List[List[int]]): - Input token extra ids for using p-tuning and KV Cache reuse together - lora_uids (list): - The uids of LoRA weights for the input batch. Use -1 to disable the LoRA module. - streaming (bool): - Whether or not to use streaming mode for generation. - stopping_criteria (StoppingCriteria): - Custom stopping criteria. - logits_processor_names (List[str]): - Custom logits processor names. - return_all_generated_tokens (bool): - Whether the full output is returned at each streaming step - kwargs (Dict[str, Any]: - Ad hoc parametrization of sampling_config. - The passed **kwargs matching the sampling_config's attributes will override them. - Returns: - torch.Tensor or dict: - If return_dict=False, the method returns generated output_ids. - If return_dict=True, the method returns a dict of output_ids, - sequence_lengths (if sampling_config.output_sequence_lengths=True), - context_logits and generation_logits (if self.gather_context_logits=True and - self.gather_generation_logits=True, respectively). - """ - # TODO: Check if these can be supported now and support them - if stopping_criteria is not None: - raise RuntimeError( - "Stopping criteria is not supported in C++ session.") - - if not self.use_kv_cache and max_new_tokens > 1: - raise RuntimeError( - 'Disabled KV cache is intended for context phase only now.') - - # If we are in a multi-gpu scenario, only rank 0 continues - if not self.session.can_enqueue_requests(): - return [] - - # Convert tensor input to plain lists - batch_input_ids_list = [a.tolist() for a in batch_input_ids] - encoder_input_ids_list = [a.tolist() for a in encoder_input_ids - ] if encoder_input_ids else None - - if sampling_config is None: - # Convert from old API of SamplingConfig - # Note: Due to a Python3.10 bug one cannot use inspect on it currently - accepted_parameters = [ - "num_beams", - "top_k", - "top_p", - "top_p_min", - "top_p_reset_ids", - "top_p_decay", - "temperature", - "min_tokens", - "beam_search_diversity_rate", - "repetition_penalty", - "presence_penalty", - "frequency_penalty", - "prompt_ignore_length", - "length_penalty", - "early_stopping", - "no_repeat_ngram_size", - "random_seed", - "num_return_sequences", - "min_p", - "beam_width_array", - ] - rename_params = {"num_beams": "beam_width", "random_seed": "seed"} - sampling_params = { - k: v - for k, v in kwargs.items() if k in accepted_parameters - } - for k, v in rename_params.items(): - if k in sampling_params: - sampling_params[v] = sampling_params.pop(k) - if "top_p" in sampling_params and sampling_params["top_p"] == 0.0: - sampling_params["top_p"] = None - - # TODO: improve usage of SamplingConfig. For example, - # construct SamplingConfig for each request, rather than one for the whole batch. - # Here we use beam width array for each request for Variable-Beam-Width-Search. - batch_size = len(batch_input_ids) - use_sampling_config_for_each_request = False - # Just placeholder for non-Variable-Beam-Width-Search - sampling_config_list = [None] * batch_size - if "beam_width_array" in sampling_params and sampling_params[ - "beam_width_array"] is not None and len( - sampling_params["beam_width_array"]) == batch_size: - use_sampling_config_for_each_request = True - sp_copy = copy.deepcopy(sampling_params) - for i in range(batch_size): - bwa = sampling_params["beam_width_array"][i] - sp_copy["beam_width_array"] = bwa - sp_copy["beam_width"] = max(bwa) - sampling_config_list[i] = trtllm.SamplingConfig(**sp_copy) - # Just placeholder for Variable-Beam-Width-Search and for `self._check_inputs` - max_beam_width = max(sc.beam_width - for sc in sampling_config_list) - sampling_params["beam_width"] = max_beam_width - sampling_params["beam_width_array"] = [max_beam_width] * 8 - sampling_config = trtllm.SamplingConfig(**sampling_params) - else: - sampling_config = copy.deepcopy(sampling_config) - - self._check_inputs(batch_input_ids_list, encoder_input_ids_list, - sampling_config, max_new_tokens) - - output_config = trtllm.OutputConfig( - return_context_logits=self.gather_context_logits, - return_generation_logits=self.gather_generation_logits - or output_generation_logits, - return_log_probs=output_log_probs, - ) - - prompt_tuning_configs = self._prepare_ptuning_executor( - batch_input_ids_list, - prompt_table, - prompt_tasks, - input_token_extra_ids, - mm_embedding_offloading=mm_embedding_offloading) - mrope_configs = self._prepare_mrope_executor(batch_input_ids_list, - mrope_params) - - stop_words_list = self._prepare_words_list(stop_words_list, - len(batch_input_ids_list)) - bad_words_list = self._prepare_words_list(bad_words_list, - len(batch_input_ids_list)) - logits_processor_names = self._prepare_names_list( - logits_processor_names, len(batch_input_ids_list)) - - lora_configs = self._prepare_lora_configs(lora_uids, - len(batch_input_ids_list)) - request_lookahead_config = None - if lookahead_config is not None: - [w, n, g] = lookahead_config - request_lookahead_config = trtllm.LookaheadDecodingConfig(w, n, g) - skip_cross_attn_blocks = kwargs.get('skip_cross_attn_blocks', None) - - # Draft-Target-Model speculative decoding - if "draft_tokens_list" in kwargs.keys() and kwargs[ - "draft_tokens_list"] is not None and "draft_logits_list" in kwargs.keys( - ) and kwargs["draft_logits_list"] is not None: - # Use logits to accept - external_draft_tokens_configs = [ - ExternalDraftTokensConfig(draft_tokens, draft_logits) - for draft_tokens, draft_logits in zip( - kwargs["draft_tokens_list"], kwargs["draft_logits_list"]) - ] - is_draft_target_model = True - elif "draft_tokens_list" in kwargs.keys( - ) and kwargs["draft_tokens_list"] is not None: - # Use tokens to accept - external_draft_tokens_configs = [ - ExternalDraftTokensConfig(draft_tokens) - for draft_tokens in kwargs["draft_tokens_list"] - ] - is_draft_target_model = True - else: - external_draft_tokens_configs = [None] * len(batch_input_ids_list) - is_draft_target_model = False - - if language_adapter_uids is None: - language_adapter_uids = [None] * len(batch_input_ids_list) - - requests = [ - trtllm.Request( - input_token_ids=input_ids, - encoder_input_token_ids=encoder_input_ids_list[i] - if encoder_input_ids is not None else None, - encoder_output_length=encoder_output_lengths[i] - if encoder_output_lengths is not None else None, - encoder_input_features=encoder_input_features[i].contiguous() - if encoder_input_features is not None else None, - position_ids=position_ids[i].tolist() - if position_ids is not None else None, - cross_attention_mask=cross_attention_masks[i].contiguous() if - (cross_attention_masks is not None - and cross_attention_masks[i] is not None) else None, - max_tokens=max_new_tokens, - pad_id=pad_id, - end_id=end_id, - stop_words=stop_words, - bad_words=bad_words, - sampling_config=(sampling_config_each_request - if use_sampling_config_for_each_request else - sampling_config), - lookahead_config=request_lookahead_config, - streaming=streaming, - output_config=output_config, - prompt_tuning_config=prompt_tuning_config, - mrope_config=mrope_config, - lora_config=lora_config, - return_all_generated_tokens=return_all_generated_tokens, - logits_post_processor_name=logits_post_processor_name, - external_draft_tokens_config=external_draft_tokens_config, - skip_cross_attn_blocks=skip_cross_attn_blocks, - language_adapter_uid=language_adapter_uid, - ) for i, - (input_ids, stop_words, bad_words, prompt_tuning_config, - mrope_config, lora_config, logits_post_processor_name, - external_draft_tokens_config, language_adapter_uid, - sampling_config_each_request) in enumerate( - zip(batch_input_ids_list, stop_words_list, bad_words_list, - prompt_tuning_configs, mrope_configs, lora_configs, - logits_processor_names, external_draft_tokens_configs, - language_adapter_uids, sampling_config_list)) - ] - - request_ids = self.session.enqueue_requests(requests) - if not streaming: - return self._initialize_and_fill_output( - request_ids=request_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - streaming=streaming, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - else: - return self._stream( - request_ids=request_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - batch_input_ids_list=batch_input_ids_list, - streaming=streaming, - return_all_generated_tokens=return_all_generated_tokens, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - - def _prepare_words_list(self, words_list: List[List[List[int]]], - batch_size: int): - if words_list is None: - return [None] * batch_size - return words_list - - def _prepare_names_list(self, names_list: List[str], batch_size: int): - if names_list is None: - return [None] * batch_size - return names_list - - def _prepare_ptuning_executor(self, batch_input_ids_list, prompt_table, - prompt_tasks, input_token_extra_ids, - mm_embedding_offloading): - if input_token_extra_ids: - assert len(batch_input_ids_list) == len(input_token_extra_ids), \ - f"Batch size of input_token_extra_ids ({len(input_token_extra_ids)}) must be the same as input batch size ({len(batch_input_ids_list)})" - prompt_tuning_configs = len(batch_input_ids_list) * [None] - if prompt_table is not None: - if mm_embedding_offloading: - # CUDA Stream Overlapping Requirements: - # 1. Both memory copy stream and kernel execution stream must be non-default streams - # 2. For host<->device transfers (H2D/D2H), host memory MUST be page-locked (pinned) - # NOTE: pinning is skipped under Confidential Compute - # (see maybe_pin_memory() and prefer_pinned()) - prompt_table_data = maybe_pin_memory( - self._prepare_embedding_table(prompt_table)) - else: - prompt_table_data = self._prepare_embedding_table( - prompt_table).cuda() - if prompt_tasks is not None: - task_indices = [int(t) for t in prompt_tasks.split(',')] - assert len(task_indices) == len(batch_input_ids_list), \ - f"Number of supplied tasks ({len(task_indices)}) must match input batch size ({len(batch_input_ids_list)})" - prompt_tuning_configs = [ - trtllm.PromptTuningConfig( - embedding_table=prompt_table_data[task_indices[i]], - input_token_extra_ids=input_token_extra_ids[i] - if input_token_extra_ids else None) - for i in range(len(batch_input_ids_list)) - ] - else: - prompt_tuning_configs = [ - trtllm.PromptTuningConfig( - embedding_table=prompt_table_data[0], - input_token_extra_ids=input_token_extra_ids[i] - if input_token_extra_ids else None) - for i in range(len(batch_input_ids_list)) - ] - return prompt_tuning_configs - - # TODO: add multimodal input for TRT engine backend - - def _prepare_mrope_executor(self, batch_input_ids_list, mrope: MropeParams): - mrope_configs = len(batch_input_ids_list) * [None] - if mrope != None: - mrope_rotary_cos_sin = mrope.mrope_rotary_cos_sin - assert isinstance( - mrope_rotary_cos_sin, - torch.Tensor), "mrope_rotary_cos_sin should be torch.Tensor" - mrope_rotary_cos_sin_data = mrope_rotary_cos_sin.to( - dtype=torch.float32) - - mrope_position_deltas = mrope.mrope_position_deltas - assert isinstance( - mrope_position_deltas, - torch.Tensor), "mrope_position_deltas should be torch.Tensor" - mrope_position_deltas_data = mrope_position_deltas.to( - dtype=torch.int32) - - mrope_configs = [ - trtllm.MropeConfig( - mrope_rotary_cos_sin=mrope_rotary_cos_sin_data[i], - mrope_position_deltas=mrope_position_deltas_data[i]) - for i in range(len(batch_input_ids_list)) - ] - return mrope_configs - - def _prepare_lora_configs(self, lora_uids, batch_size): - if lora_uids is None: - return [None] * batch_size - assert len(lora_uids) == batch_size - return [ - trtllm.LoraConfig(task_id=int(uid), - weights=self.lora_manager.cpp_lora_weights[uid], - config=self.lora_manager.cpp_lora_config[uid]) - if int(uid) >= 0 else None for uid in lora_uids - ] - - def _get_num_sequences(self, sampling_config: SamplingConfigType): - num_beams = sampling_config.num_beams if isinstance( - sampling_config, SamplingConfig) else sampling_config.beam_width - num_sequences = sampling_config.num_return_sequences or num_beams - assert num_beams == 1 or num_sequences <= num_beams - return num_sequences - - def _initialize_and_fill_output( - self, - *, - request_ids, - end_id, - return_dict, - output_sequence_lengths, - output_generation_logits, - output_log_probs, - output_cum_log_probs, - batch_input_ids, - streaming, - sampling_config: SamplingConfigType, - is_draft_target_model: bool = False, - ): - num_sequences = self._get_num_sequences(sampling_config) - # (batch_size, num_sequences, sequence_len) - output_ids = [[[] for _ in range(num_sequences)] - for _ in range(len(request_ids))] - - all_responses = [] - finished_request_ids = set() - while finished_request_ids != set(request_ids): - responses = self.session.await_responses() - for response in responses: - if response.result.is_final: - finished_request_ids.add(response.request_id) - all_responses.extend(responses) - - return self._fill_output( - responses=all_responses, - output_ids=output_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - batch_input_ids_list=[], - streaming=streaming, - request_ids=request_ids, - return_all_generated_tokens=False, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - - def _stream( - self, - *, - request_ids, - end_id, - return_dict, - output_sequence_lengths, - output_generation_logits, - output_log_probs, - output_cum_log_probs, - batch_input_ids, - batch_input_ids_list, - streaming, - return_all_generated_tokens, - sampling_config: SamplingConfigType, - is_draft_target_model: bool = False, - ): - num_sequences = self._get_num_sequences(sampling_config) - # (batch_size, num_sequences, sequence_len) - output_ids = [[ - copy.deepcopy(batch_input_ids_list[batch_idx]) - for _ in range(num_sequences) - ] for batch_idx in range(len(request_ids))] - - finished_request_ids = set() - while finished_request_ids != set(request_ids): - responses = self.session.await_responses() - for response in responses: - if response.result.is_final: - finished_request_ids.add(response.request_id) - - yield self._fill_output( - responses=responses, - output_ids=output_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - batch_input_ids_list=batch_input_ids_list, - streaming=streaming, - request_ids=request_ids, - return_all_generated_tokens=return_all_generated_tokens, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - - def _fill_output( - self, - *, - responses, - output_ids, - end_id, - return_dict, - output_sequence_lengths, - output_generation_logits, - output_log_probs, - output_cum_log_probs, - batch_input_ids, - batch_input_ids_list, - streaming, - request_ids, - return_all_generated_tokens, - sampling_config: SamplingConfigType, - is_draft_target_model: bool, - ): - cuda_device = torch.device("cuda") - - batch_size = len(batch_input_ids) - num_sequences = len(output_ids[0]) - beam_width = getattr(sampling_config, 'num_beams', - getattr(sampling_config, 'beam_width')) - is_beam_search = beam_width > 1 - - def fill_output_ids(result_token_ids, batch_idx, seq_idx): - # Return shape = (batch_size, num_sequences, seq_len) - if return_all_generated_tokens: - output_ids[batch_idx][seq_idx] = ( - batch_input_ids_list[batch_idx] + result_token_ids) - else: - output_ids[batch_idx][seq_idx] += result_token_ids - - for response in responses: - if response.has_error(): - raise RuntimeError(response.error_msg) - - result = response.result - batch_idx = request_ids.index(response.request_id) - if is_beam_search: - for beam, output_tokens in enumerate(result.output_token_ids): - fill_output_ids(output_tokens, batch_idx, beam) - else: - fill_output_ids(result.output_token_ids[0], batch_idx, - result.sequence_index) - - if output_sequence_lengths: - sequence_lengths = [[len(token_ids) for token_ids in beams] - for beams in output_ids] - - if streaming: - output_ids = copy.deepcopy(output_ids) - - # Pad by end_id tokens (batch, num_sequences, max_seq_len). - for beams in output_ids: - for token_ids in beams: - token_ids += [end_id] * (self.max_seq_len - len(token_ids)) - output_ids = torch.tensor(output_ids, - dtype=torch.int32, - device=cuda_device) - - if return_dict: - outputs = {'output_ids': output_ids} - - input_lengths = torch.tensor([x.size(0) for x in batch_input_ids], - dtype=torch.int32, - device=cuda_device) - - if output_sequence_lengths: - outputs['sequence_lengths'] = torch.tensor(sequence_lengths, - dtype=torch.int32, - device=cuda_device) - - if self.gather_context_logits: - context_logits = None - max_input_len = input_lengths.max() - for response in responses: - result = response.result - logits = result.context_logits - if logits is None: - continue - input_len, vocab_size = logits.shape - if context_logits is None: - context_logits = torch.zeros( - (batch_size, max_input_len, vocab_size), - dtype=logits.dtype, - device=cuda_device) - if result.sequence_index == 0: - batch_idx = request_ids.index(response.request_id) - context_logits[batch_idx, :input_len, :] = logits - assert context_logits is not None - outputs['context_logits'] = context_logits - - if self.gather_generation_logits or output_generation_logits: - gen_logits = None - if is_draft_target_model: - # Put the outputs in a list rather than a tensor since their - # length may vary among requests in a batch - gen_logits = [ - a.result.generation_logits.cuda() for a in responses - if a.result.generation_logits is not None - ] - else: - # The shape of generation logits - # (num_sequences, seq_len, vocab_size) in non-streaming - # (seq_len, num_sequences, vocab_size) in streaming - seq_dim = 0 if streaming else 1 - max_out_len = max( - response.result.generation_logits.size(seq_dim) - for response in responses - if response.result.generation_logits is not None) - vocab_size = responses[0].result.generation_logits.size(-1) - if not streaming: - gen_shape = (num_sequences, max_out_len, vocab_size) - elif streaming and return_all_generated_tokens: - gen_shape = (max_out_len, num_sequences, vocab_size) - else: - # streaming and not return_all_generated_tokens - gen_shape = (1, num_sequences, vocab_size) - logits_dtype = responses[0].result.generation_logits.dtype - gen_logits = torch.zeros((batch_size, *gen_shape), - dtype=logits_dtype, - device=cuda_device) - - for response in responses: - logits = response.result.generation_logits - if logits is None: - continue - seq_len = logits.size(seq_dim) - - batch_idx = request_ids.index(response.request_id) - seq_idx = response.result.sequence_index - if streaming: - if is_beam_search: - # WAR: gen_logits contains all beams, clipping - # the first n beams as a postprocessing. - gen_logits[batch_idx, :seq_len, - ...] = logits[:, :num_sequences, :] - else: - gen_logits[batch_idx, :seq_len, seq_idx, - ...] = logits[:, 0, :] - else: - if is_beam_search: - gen_logits[batch_idx, :, :seq_len, ...] = logits - else: - gen_logits[batch_idx, seq_idx, :seq_len, - ...] = logits[0] - outputs['generation_logits'] = gen_logits - - if output_log_probs: - max_log_probs_len = max( - len(lprobs) for response in responses - for lprobs in response.result.log_probs) - log_probs = torch.zeros( - (batch_size, num_sequences, max_log_probs_len), - dtype=torch.float32) - for response in responses: - batch_idx = request_ids.index(response.request_id) - if is_beam_search: - for beam_idx, lprobs in enumerate( - response.result.log_probs): - log_probs[batch_idx, - beam_idx, :len(lprobs)] = torch.tensor( - lprobs) - else: - seq_idx = response.result.sequence_index - lprobs = response.result.log_probs[0] - log_probs[batch_idx, - seq_idx, :len(lprobs)] = torch.tensor(lprobs) - assert isinstance(log_probs, torch.Tensor) - outputs['log_probs'] = log_probs.to(cuda_device) - - if output_cum_log_probs: - cum_log_probs = torch.zeros((batch_size, num_sequences), - dtype=torch.float32) - for response in responses: - if response.result.cum_log_probs is None: - continue - batch_idx = request_ids.index(response.request_id) - clprobs = torch.tensor(response.result.cum_log_probs) - if is_beam_search: - cum_log_probs[batch_idx, :] = clprobs - else: - seq_idx = response.result.sequence_index - cum_log_probs[batch_idx, seq_idx] = clprobs - outputs['cum_log_probs'] = cum_log_probs.to(cuda_device) - - outputs = self._prepare_outputs(outputs, input_lengths) - else: - outputs = output_ids - - return outputs diff --git a/tensorrt_llm/runtime/multimodal_model_runner.py b/tensorrt_llm/runtime/multimodal_model_runner.py deleted file mode 100644 index 3ab5bb7ed824..000000000000 --- a/tensorrt_llm/runtime/multimodal_model_runner.py +++ /dev/null @@ -1,2744 +0,0 @@ -import json -import os -import sys -from io import BytesIO - -import requests - -# isort: off -import torch -import numpy as np -# isort: on -import math -from typing import Optional, Tuple - -import torch.nn.functional as F - -try: - from cuda.bindings import runtime as cudart -except ImportError: - from cuda import cudart - -from huggingface_hub import hf_hub_download -from PIL import Image, UnidentifiedImageError -from safetensors import safe_open -from torch import nn -from transformers import (AutoConfig, AutoModelForCausalLM, AutoProcessor, - AutoTokenizer) - -from .. import profiler -from .._deprecation import emit_engine_arch_deprecation -from .._utils import (get_hf_rope_theta, maybe_pin_memory, mpi_rank, - prefer_pinned, str_dtype_to_torch, str_dtype_to_trt, - supports_inflight_batching, torch_dtype_to_trt, - trt_dtype_to_torch) -from ..functional import RopeEmbeddingUtils, RotaryScalingType -from ..layers import MropeParams -from ..logger import logger -from .enc_dec_model_runner import EncDecModelRunner -from .model_runner import ModelRunner -from .session import Session, TensorInfo - -try: - import tensorrt_llm.bindings # NOQA - PYTHON_BINDINGS = True -except ImportError: - PYTHON_BINDINGS = False - -if PYTHON_BINDINGS: - from .model_runner_cpp import ModelRunnerCpp - - -class LlavaNextUtils: - # https://github.com/haotian-liu/LLaVA/blob/main/llava/mm_utils.py - - @staticmethod - def select_best_resolution(original_size, possible_resolutions): - """ - Selects the best resolution from a list of possible resolutions based on the original size. - - Args: - original_size (tuple): The original size of the image in the format (width, height). - possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...]. - - Returns: - tuple: The best fit resolution in the format (width, height). - """ - original_width, original_height = original_size - best_fit = None - max_effective_resolution = 0 - min_wasted_resolution = float('inf') - - for width, height in possible_resolutions: - scale = min(width / original_width, height / original_height) - downscaled_width, downscaled_height = int( - original_width * scale), int(original_height * scale) - effective_resolution = min(downscaled_width * downscaled_height, - original_width * original_height) - wasted_resolution = (width * height) - effective_resolution - - if effective_resolution > max_effective_resolution or ( - effective_resolution == max_effective_resolution - and wasted_resolution < min_wasted_resolution): - max_effective_resolution = effective_resolution - min_wasted_resolution = wasted_resolution - best_fit = (width, height) - - return best_fit - - @staticmethod - def get_anyres_image_grid_shape(image_size, - patch_size, - image_grid_pinpoints=None): - """ - Calculate the shape of the image patch grid after the preprocessing for images of any resolution. - - Args: - image_size (tuple): The size of the input image in the format (width, height). - patch_size (int): The size of each image patch. - - Returns: - tuple: The shape of the image patch grid in the format (width, height). - """ - if image_grid_pinpoints is None: - image_grid_pinpoints = [[336, 672], [672, 336], [672, 672], - [1008, 336], [336, 1008]] - width, height = LlavaNextUtils.select_best_resolution( - image_size, image_grid_pinpoints) - return width // patch_size, height // patch_size - - @staticmethod - def unpad_image(tensor, original_size): - """ - Unpads a PyTorch tensor of a padded and resized image. - - Args: - tensor (torch.Tensor): The image tensor, assumed to be in CxHxW format. - original_size (tuple): The original size of the image (width, height). - - Returns: - torch.Tensor: The unpadded image tensor. - """ - original_width, original_height = original_size - current_height, current_width = tensor.shape[1:] - - original_aspect_ratio = original_width / original_height - current_aspect_ratio = current_width / current_height - - if original_aspect_ratio > current_aspect_ratio: - scale_factor = current_width / original_width - new_height = int(original_height * scale_factor) - padding = (current_height - new_height) // 2 - unpadded_tensor = tensor[:, padding:current_height - padding, :] - else: - scale_factor = current_height / original_height - new_width = int(original_width * scale_factor) - padding = (current_width - new_width) // 2 - unpadded_tensor = tensor[:, :, padding:current_width - padding] - - return unpadded_tensor - - @staticmethod - def rearrange_image_features(image_feature, image_newline, image_size): - """ - Combine PyTorch feature grids from image patches. - - Args: - image_feature (torch.Tensor): The feature grids, assumed to be in NxCxHxW format. - image_newline (torch.Tensor): The newline embedding. - image_size (tuple): Size of the original image (width, height). - """ - CLIP_IMAGE_SIZE = 336 - CLIP_PATCH_SIZE = 14 - NUM_PATCHES_PER_SIDE = CLIP_IMAGE_SIZE // CLIP_PATCH_SIZE - if image_feature.shape[0] == 1: - return torch.cat((image_feature, image_newline[None]), dim=0) - - base_image_feature = image_feature[0] - image_feature = image_feature[1:] - height = width = NUM_PATCHES_PER_SIDE - assert height * width == base_image_feature.shape[0] - - num_patch_width, num_patch_height = LlavaNextUtils.get_anyres_image_grid_shape( - image_size, CLIP_IMAGE_SIZE) - image_feature = image_feature.view(num_patch_height, num_patch_width, - height, width, -1) - - image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() - image_feature = image_feature.flatten(1, 2).flatten(2, 3) - image_feature = LlavaNextUtils.unpad_image(image_feature, image_size) - image_feature = torch.cat( - (image_feature, image_newline[:, None, None].expand( - *image_feature.shape[:-1], 1)), - dim=-1) - image_feature = image_feature.flatten(1, 2).transpose(0, 1) - image_feature = torch.cat((base_image_feature, image_feature), dim=0) - return image_feature - - -class LlavaOnevisionUtils: - # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/modeling_llava_onevision.py - - @staticmethod - def pack_image_features(image_features, image_sizes, image_newline): - """ - Reshape, unpad and then pack each image_feature into a single image_features tensor containing all visual vectors. - - Args: - image_features (`torch.Tensor` of shape `(num_images, num_patches, image_length, embed_dim)`) - Image feature tensor, each contains all the visual feature of all patches. - image_sizes (`torch.Tensor` of shape `(num_images, 2)`) - Actual image size of each images (H, W). - image_newline (`torch.Tensor` of shape `(embed_dim)`) - New line embedding vector. - Returns: - image_features (`torch.Tensor` of shape `(all_feat_len, embed_dim)`) - """ - - IMAGE_SIZE = 384 - PATCH_SIZE = 14 - MAX_NUM_PATCHES = 9 - - new_image_features = [] - for image_idx, image_feature in enumerate(image_features): - if image_feature.shape[0] > 1: - base_image_feature = image_feature[0] - image_feature = image_feature[1:] - height = width = IMAGE_SIZE // PATCH_SIZE - if height * width != base_image_feature.shape[0]: - raise ValueError( - "The number of patches is not consistent with the image size." - ) - - IMAGE_GRID_PINPOINTS = [[384, 384], [384, 768], [384, 1152], - [384, 1536], [384, 1920], [384, 2304], - [768, 384], [768, 768], [768, 1152], - [768, 1536], [768, 1920], [768, 2304], - [1152, 384], [1152, 768], [1152, 1152], - [1152, 1536], - [1152, 1920], [1152, 2304], [1536, 384], - [1536, 768], [1536, 1152], [1536, 1536], - [1536, 1920], [1536, 2304], [1920, 384], - [1920, 768], [1920, 1152], [1920, 1536], - [1920, 1920], [1920, 2304], [2304, 384], - [2304, 768], [2304, 1152], [2304, 1536], - [2304, 1920], [2304, 2304]] - num_patch_width, num_patch_height = LlavaNextUtils.get_anyres_image_grid_shape( - image_sizes[image_idx][[1, 0]].tolist(), IMAGE_SIZE, - IMAGE_GRID_PINPOINTS) - image_feature = image_feature.view(num_patch_height, - num_patch_width, height, - width, -1) - image_feature = image_feature.permute(4, 0, 2, 1, - 3).contiguous() - image_feature = image_feature.flatten(1, 2).flatten(2, 3) - image_feature = LlavaNextUtils.unpad_image( - image_feature, image_sizes[image_idx][[1, 0]]) - - channels, curr_height, curr_width = image_feature.shape - ratio = math.sqrt(curr_height * curr_width / - (MAX_NUM_PATCHES * height**2)) - if ratio > 1.1: - image_feature = image_feature[None] - image_feature = nn.functional.interpolate( - image_feature, - [int(curr_height // ratio), - int(curr_width // ratio)], - mode="bilinear")[0] - - image_feature = torch.cat( - ( - image_feature, - image_newline[:, None, None].expand( - *image_feature.shape[:-1], 1).to( - image_feature.device, image_feature.dtype), - ), - dim=-1, - ) - image_feature = image_feature.flatten(1, 2).transpose(0, 1) - image_feature = torch.cat((base_image_feature, image_feature), - dim=0) - else: - image_feature = image_feature[0] - if image_newline is not None: - image_feature = torch.cat( - (image_feature, image_newline[None].to(image_feature)), - dim=0) - new_image_features.append(image_feature) - image_features = torch.stack(new_image_features) - return image_features - - @staticmethod - def apply_pooling(image_features): - IMAGE_SIZE = 384 - PATCH_SIZE = 14 - height = width = IMAGE_SIZE // PATCH_SIZE - batch_frames, seq_len, dim = image_features.shape - image_features = image_features.view(batch_frames, height, width, -1) - image_features = image_features.permute(0, 3, 1, 2).contiguous() - - height, width = image_features.shape[2:] - scaled_shape = [math.ceil(height / 2), math.ceil(width / 2)] - image_features = nn.functional.interpolate(image_features, - size=scaled_shape, - mode="bilinear") - - image_features = image_features.permute(0, 2, 3, 1) - image_features = image_features.view(batch_frames, -1, dim) - return image_features - - -class PhiMMUtils: - - @staticmethod - def add_image_newline(image_features, image_newline): - h, w, d = image_features.shape - image_newline = image_newline.expand(h, 1, -1) - image_features_newline = torch.cat([image_features, image_newline], - dim=1).flatten(0, 1) - return image_features_newline - - @staticmethod - def reshape_hd_patches(image_features, h_crop=1, w_crop=1): - n_crops, n_tokens, d = image_features.shape - assert n_crops == h_crop * w_crop - - h = w = int(n_tokens**0.5) - image_features = image_features.reshape( - h_crop, w_crop, h, w, - d).permute(0, 2, 1, 3, 4).reshape(h_crop * h, w_crop * w, d) - return image_features - - @staticmethod - def hd_feature_transform(image_features, - h_crop, - w_crop, - sub_GN, - glb_GN, - patch_mask=None): - glb_image_features = PhiMMUtils.add_image_newline( - PhiMMUtils.reshape_hd_patches(image_features[:1]), sub_GN) - - num_crops = h_crop * w_crop - sub_image_features = PhiMMUtils.reshape_hd_patches( - image_features[1:num_crops + 1], h_crop, w_crop) - if patch_mask is not None: - h, w = patch_mask.shape[1] // 2, patch_mask.shape[2] // 2 - sub_image_mask = (patch_mask[1:num_crops + 1, 0::2, - 0::2].bool().reshape( - h_crop, w_crop, h, - w).permute(0, 2, 1, 3).reshape( - h_crop * h, w_crop * w)) - hh = int(sub_image_mask[:, 0].sum().item()) - ww = int(sub_image_mask[0, :].sum().item()) - sub_image_features = sub_image_features[:hh, :ww] - sub_image_features = PhiMMUtils.add_image_newline( - sub_image_features, sub_GN) - - image_features = torch.cat( - [sub_image_features, - glb_GN.expand(1, -1), glb_image_features]) - return image_features - - @staticmethod - def reshape_audio_chunks(audio_features, chunk_mask=None): - audio_features = audio_features.flatten(0, 1) - if chunk_mask is not None: - # only the last chunk may include paddings - n_tokens = math.ceil(chunk_mask.flatten().sum().item() / 8) - audio_features = audio_features[:n_tokens] - return audio_features - - -class MultimodalModelRunner: - - def __init__(self, args): - emit_engine_arch_deprecation("MultimodalModelRunner") - self.args = args - self.use_trtllm_vision_engine = False - - self.runtime_rank = mpi_rank() - device_id = self.runtime_rank % torch.cuda.device_count() - torch.cuda.set_device(device_id) - self.device = "cuda:%d" % (device_id) - - self.stream = torch.cuda.Stream(torch.cuda.current_device()) - torch.cuda.set_stream(self.stream) - - if self.args.mm_embedding_offloading is None: - self.args.mm_embedding_offloading = self.args.enable_chunked_context - elif self.args.mm_embedding_offloading and not self.args.enable_chunked_context: - logger.warning( - "mm_embedding_offloading requires enable_chunked_context to be True. Setting mm_embedding_offloading to None." - ) - self.args.mm_embedding_offloading = None - - # parse model type from visual engine config - with open(os.path.join(self.visual_engine_dir, "config.json"), - "r") as f: - config = json.load(f) - if 'pretrained_config' in config: - if config['pretrained_config'][ - 'architecture'] == 'LlavaNextForConditionalGeneration': - self.model_type = 'llava_next' - self.vision_precision = config['pretrained_config']['dtype'] - self.use_trtllm_vision_engine = True - else: - logger.error( - "Currently only Llava-NeXT supports TRT-LLM vision engines." - ) - else: - self.model_type = config['builder_config']['model_type'] - self.vision_precision = config['builder_config']['precision'] - if self.model_type == 'pix2struct': - self.vision_precision = 'float16' - self.decoder_llm = not ( - 't5' in self.model_type - or self.model_type in ['nougat', 'pix2struct'] - ) # BLIP2-T5, pix2struct and Nougat are using encoder-decoder models as LLMs - - if self.model_type == 'video-neva': - self.num_frames = config['builder_config'].get('num_frames', None) - if self.model_type == "llava_next": - self.llm_name = AutoConfig.from_pretrained( - self.args.hf_model_dir).text_config._name_or_path - if 'internlm' in self.model_type: - self.args.lora_task_uids = ['0'] * args.batch_size - if self.model_type == "qwen2_vl": - hf_config = AutoConfig.from_pretrained(self.args.hf_model_dir) - self.vision_start_token_id = hf_config.vision_start_token_id - self.vision_end_token_id = hf_config.vision_end_token_id - self.vision_token_id = hf_config.vision_token_id - self.image_token_id = hf_config.image_token_id - self.video_token_id = hf_config.video_token_id - self.spatial_merge_size = hf_config.vision_config.spatial_merge_size - self.max_position_embeddings = hf_config.max_position_embeddings - self.hidden_size = hf_config.hidden_size - self.num_attention_heads = hf_config.num_attention_heads - self.rope_theta = get_hf_rope_theta(hf_config, 10000.0) - if self.model_type == 'llava_onevision': - self.num_frames = self.args.video_num_frames - if self.num_frames is None: - self.num_frames = 8 - assert self.args.video_path is None or self.args.image_path is None - if self.model_type == "pixtral": - hf_config = AutoConfig.from_pretrained(self.args.hf_model_dir) - self.image_size = hf_config.vision_config.image_size - self.patch_size = hf_config.vision_config.patch_size - self.vocab_size = hf_config.text_config.vocab_size - self.image_token_index = hf_config.image_token_index - self.spatial_merge_size = hf_config.spatial_merge_size - - self.audio_input_names = self.audio_output_names = None - if self.model_type == "mllama": - self.vision_input_names = [ - "pixel_values", - "aspect_ratio_ids", - "aspect_ratio_mask", - ] - self.vision_output_names = [ - "encoder_output", - ] - elif self.model_type == "llava_next" and self.use_trtllm_vision_engine: - self.vision_input_names = [ - "pixel_values", - ] - self.vision_output_names = [ - "image_features", - ] - elif self.model_type == "phi-4-multimodal": - self.vision_input_names = ["input", "attention_mask"] - self.audio_input_names = ["input", "attention_mask"] - self.audio_output_names = ["encoder_output"] - self.vision_output_names = ["encoder_output"] - else: - self.vision_input_names = ["input"] - self.vision_output_names = ["encoder_output"] - - self.session = args.session - if self.cpp_e2e: - self.visual_output_shape = config['builder_config'].get( - 'output_shape', None) - if self.decoder_llm: - if not supports_inflight_batching(self.llm_engine_dir): - logger.warning( - "The given engine does not support in-flight batching, both visual engine and LLM fallback to python session" - ) - self.session = 'python' - - if not PYTHON_BINDINGS and 'cpp' in args.session: - logger.warning( - "Python bindings of C++ session is unavailable, both visual engine and LLM fallback to Python session." - ) - self.session = 'python' - - args.debug_mode = False - if args.debug_mode and 'cpp' in args.session: - logger.warning( - "Debug mode is not supported in C++ session for now, both visual engine and LLM fallback to Python session." - ) - self.session = 'python' - - if self.model_type == 'qwen2_vl': - if self.args.session != "cpp_llm_only": - logger.warning( - "Qwen2-vl only support C++ session for now, fallback to C++ session." - ) - self.args.session = "cpp_llm_only" - - if (not (self.model_type in - ('llava', 'vila', 'blip2-opt', 'kosmos-2', 'fuyu', - 'cogvlm', 'neva', "internvl") or 'internlm' - in self.model_type)) and args.session == 'cpp': - logger.warning( - f'C++ end-to-end mode does not support {self.model_type}. Visual engine fallbacks to Python session. See support matrix in README.' - ) - args.session = 'cpp_llm_only' - self.session = args.session - - else: - self.session = 'cpp_llm_only' - - self.init_tokenizer() - self.init_processor() - self.init_image_encoder() - self.init_llm() - - if self.audio_input_names is not None: - with open(os.path.join(self.audio_engine_dir, "config.json"), - "r") as f: - config = json.load(f) - self.audio_precision = config['builder_config']['precision'] - self.init_audio_encoder() - else: - self.audio_encoder_session = self.audio_precision = None - - @property - def cpp_e2e(self): - return self.session == 'cpp' - - @property - def cpp_llm_only(self): - return self.session == 'cpp_llm_only' - - @property - def python_e2e(self): - return self.session == 'python' - - @property - def visual_engine_dir(self): - return os.path.join(self.args.engine_dir, 'vision') - - @property - def audio_engine_dir(self): - return os.path.join(self.args.engine_dir, 'audio') - - @property - def llm_engine_dir(self): - return os.path.join(self.args.engine_dir, 'llm') - - def init_tokenizer(self): - if self.model_type == 'nougat': - from transformers import NougatTokenizerFast - self.tokenizer = NougatTokenizerFast.from_pretrained( - self.args.hf_model_dir) - elif self.model_type == 'neva' or self.model_type == 'video-neva': - from sentencepiece import SentencePieceProcessor - - sp = SentencePieceProcessor( - os.path.join(self.args.hf_model_dir, 'tokenizer.model')) - - class return_obj: - - def __init__(self, input_ids): - self.input_ids = input_ids - - def __getitem__(self, name): - if name in "input_ids": - return self.input_ids - else: - raise AttributeError( - f"'return_obj' has no item '{name}'") - - # sentencepiece does not follow the same interface as HF - class HFTokenizerInterface(): - - def encode(self, x, return_tensors=None, **kwargs): - out = sp.encode(x) - if return_tensors == "pt": - out = torch.tensor(out) - return return_obj(out) - - def __call__(self, x, return_tensors=None, **kwargs): - return self.encode(x, return_tensors, **kwargs) - - def decode(self, x, **kwargs): - return sp.decode(x.tolist()) - - def batch_decode(self, x, **kwargs): - return self.decode(x, **kwargs) - - self.tokenizer = HFTokenizerInterface() - self.tokenizer.eos_token_id = sp.eos_id() - self.tokenizer.bos_token_id = sp.bos_id() - self.tokenizer.pad_token_id = sp.pad_id() - elif self.model_type == 'vila': - self.tokenizer = AutoTokenizer.from_pretrained( - self.args.hf_model_dir + "/llm", - use_fast=False, - use_legacy=False) - else: - use_fast = self.model_type in [ - "phi-3-vision", "phi-4-multimodal", "internvl" - ] - self.tokenizer = AutoTokenizer.from_pretrained( - self.args.hf_model_dir, - use_fast=use_fast, - use_legacy=False, - trust_remote_code=getattr(self.args, "trust_remote_code", - False)) - - self.tokenizer.padding_side = "right" - - def init_processor(self): - from torchvision import transforms - - if 'blip2' in self.model_type: - from transformers import Blip2Processor - self.processor = Blip2Processor.from_pretrained( - self.args.hf_model_dir) - - elif 'nougat' in self.model_type: - from transformers import NougatProcessor - self.processor = NougatProcessor.from_pretrained( - self.args.hf_model_dir) - - elif 'cogvlm' in self.model_type: - image_size = 490 - self.transform = transforms.Compose([ - transforms.Resize( - (image_size, image_size), - interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize((0.48145466, 0.4578275, 0.40821073), - (0.26862954, 0.26130258, 0.27577711)), - transforms.ConvertImageDtype(torch.bfloat16), - ]) - - elif self.model_type in [ - 'phi-3-vision', 'pix2struct', 'llava_next', 'llava', 'fuyu', - 'kosmos-2', 'mllama', 'llava_onevision', 'qwen2_vl', - 'phi-4-multimodal' - ]: - self.processor = AutoProcessor.from_pretrained( - self.args.hf_model_dir, - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - num_crops=16) - - elif 'pixtral' in self.model_type: - self.processor = AutoProcessor.from_pretrained( - self.args.hf_model_dir) - - elif 'internlm' in self.model_type: - image_size = 490 - self.processor = transforms.Compose([ - transforms.Resize( - (image_size, image_size), - interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize((0.48145466, 0.4578275, 0.40821073), - (0.26862954, 0.26130258, 0.27577711)), - ]) - - elif 'internvl' in self.model_type: - from transformers import CLIPImageProcessor - self.processor = CLIPImageProcessor.from_pretrained( - 'OpenGVLab/InternViT-300M-448px' - ) # You can change the InternViT model type according to your InternVL type - - elif self.model_type == "neva": - image_size = 384 - self.transform = transforms.Compose([ - transforms.Resize( - (image_size, image_size), - interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), - transforms.ConvertImageDtype(torch.float32), - ]) - - elif self.model_type == "video-neva": - pass - - elif self.model_type == "vila": - sys.path.append(self.args.hf_model_dir + "/../VILA") - from llava.mm_utils import process_images - from llava.model import LlavaLlamaConfig # noqa - from transformers import AutoModel - model = AutoModel.from_pretrained( - self.args.hf_model_dir, - device_map='auto', - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - ) - vision_tower = model.get_vision_tower() - vision_tower.image_processor - - def processor(raw_image): - return process_images(raw_image, vision_tower.image_processor, - model.config).to(model.device, - dtype=torch.float16) - - self.processor = processor - - if self.model_type == 'mllama': - from .processor_wrapper import MllamaProcessorWrapper - self.processor = MllamaProcessorWrapper(self.processor, logger) - - def init_image_encoder(self): - - # Phi-4-multimodal uses pytorch engine due to issues with creating TRT engine. - if self.model_type == "phi-4-multimodal": - model = AutoModelForCausalLM.from_pretrained( - self.args.hf_model_dir, - dtype=torch.float16, - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - device_map='cpu') - self.vision_model = model.model.embed_tokens_extend.image_embed.to( - self.device).eval() - self.image_newlines = {} - self.image_newlines['sub_GN'] = self.vision_model.img_projection( - self.vision_model.sub_GN).squeeze() - self.image_newlines['glb_GN'] = self.vision_model.img_projection( - self.vision_model.glb_GN).squeeze() - return - - if self.model_type == "phi-3-vision": - model = AutoModelForCausalLM.from_pretrained( - self.args.hf_model_dir, - dtype=torch.float16, - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - device_map='cpu') - self.vision_model = model.model.vision_embed_tokens.to( - self.device).eval() - - # Test run vision_model.get_img_features to pre-allocate memory for flash attention - image = self.processor(text="<|image_1|>", - images=Image.new('RGB', [10, 10]), - return_tensors="pt")['pixel_values'] - image = image.flatten(0, 1) - image = torch.rand(image.shape, - dtype=str_dtype_to_torch(self.vision_precision), - device=self.device) - self.vision_model.get_img_features(image) - return - - if self.cpp_e2e: - logger.info( - "Using C++ runtime for both visual engine and LLM decoder, skip loading visual engine in Python runtime." - ) - elif self.model_type == "llava_next" and self.use_trtllm_vision_engine: - cudart.cudaSetDevice(self.runtime_rank % torch.cuda.device_count()) - - vision_encoder_path = os.path.join( - self.visual_engine_dir, f"rank{self.runtime_rank}.engine") - logger.info(f'Loading engine from {vision_encoder_path}') - with open(vision_encoder_path, "rb") as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine {vision_encoder_path}') - assert engine_buffer is not None - - self.visual_encoder_session = Session.from_serialized_engine( - engine_buffer) - else: - vision_encoder_path = os.path.join(self.visual_engine_dir, - self.args.visual_engine_name) - logger.info(f'Loading engine from {vision_encoder_path}') - with open(vision_encoder_path, 'rb') as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine {vision_encoder_path}') - self.visual_encoder_session = Session.from_serialized_engine( - engine_buffer) - if self.model_type in ["llava_next", "llava_onevision"]: - self.image_newlines = {} - image_newlines_path = os.path.join(self.visual_engine_dir, - 'image_newlines.safetensors') - with safe_open(image_newlines_path, - framework="pt", - device=self.device) as f: - for k in f.keys(): - self.image_newlines[k] = f.get_tensor(k) - - def init_audio_encoder(self): - assert self.model_type == "phi-4-multimodal" - model = AutoModelForCausalLM.from_pretrained( - self.args.hf_model_dir, - dtype=torch.float16, - trust_remote_code=getattr(self.args, "trust_remote_code", False), - device_map='cpu') - self.audio_model = model.model.embed_tokens_extend.audio_embed.to( - self.device).eval() - - def init_llm(self): - if self.decoder_llm: - cross_kv_cache_fraction = None - if self.model_type == 'mllama': - cross_kv_cache_fraction = self.args.cross_kv_cache_fraction - if self.python_e2e: - logger.info(f'Running LLM with Python runner') - self.model = ModelRunner.from_dir( - self.llm_engine_dir, - rank=tensorrt_llm.mpi_rank(), - debug_mode=False, - stream=self.stream, - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc, - multi_block_mode=self.args.multi_block_mode, - ) - self.model_config = self.model.session._model_config - elif self.cpp_e2e: - logger.info( - f'Running both visual engine and LLM with Python runner') - self.model = ModelRunnerCpp.from_dir( - self.args.engine_dir, - rank=tensorrt_llm.mpi_rank(), - debug_mode=False, - is_enc_dec=True, # TODO: add a separate model variant here? - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc) - self.model_config = self.model.model_config - else: - logger.info(f'Running LLM with C++ runner') - self.model = ModelRunnerCpp.from_dir( - self.llm_engine_dir, - rank=tensorrt_llm.mpi_rank(), - debug_mode=False, - enable_chunked_context=self.args.enable_chunked_context, - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc, - kv_cache_free_gpu_memory_fraction=self.args. - kv_cache_free_gpu_memory_fraction, - cross_kv_cache_fraction=cross_kv_cache_fraction, - multi_block_mode=self.args.multi_block_mode, - mm_embedding_offloading=self.args.mm_embedding_offloading, - ) - self.model_config = self.model.model_config - self.runtime_mapping = self.model.mapping - else: - self.model = EncDecModelRunner.from_engine( - os.path.basename(self.args.hf_model_dir), - self.llm_engine_dir, - skip_encoder=self.model_type in ['nougat', 'pix2struct'], - debug_mode=False, - stream=self.stream, - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc) - if self.model_type in ['nougat', 'pix2struct']: - self.model_config = self.model.decoder_model_config - self.runtime_mapping = self.model.decoder_runtime_mapping - else: - self.model_config = self.model.encoder_model_config - self.runtime_mapping = self.model.encoder_runtime_mapping - - def video_preprocess(self, video_path): - from decord import VideoReader - if isinstance(video_path, str): - vr = VideoReader(video_path) - num_frames = self.num_frames - if num_frames == -1: - frames = [ - Image.fromarray(frame.asnumpy()[:, :, ::-1]).convert('RGB') - for frame in vr - ] - else: - # equally sliced frames into self.num_frames frames - # if self.num_frames is greater than the number of frames in the video, we will repeat the last frame - num_frames = min(num_frames, len(vr)) - indices = np.linspace(0, len(vr) - 1, num=num_frames, dtype=int) - frames = [ - Image.fromarray( - vr[idx].asnumpy()[:, :, ::-1]).convert('RGB') - for idx in indices - ] - if len(frames) < num_frames: - frames += [frames[-1]] * (num_frames - len(frames)) - else: - frames = self.video_path - - from transformers import CLIPImageProcessor - processor = CLIPImageProcessor.from_pretrained( - "openai/clip-vit-large-patch14", dtype=torch.bfloat16) - frames = processor.preprocess(frames, - return_tensors="pt")['pixel_values'] - # make dtype consistent with vision encoder - media_tensors = frames.to(str_dtype_to_torch( - self.vision_precision)) # [num_frames, 3, H, W] - return media_tensors.unsqueeze(0) #[1, num_frames, 3, H, W] - - def preprocess(self, pre_prompt, post_prompt, image, other_vision_inputs, - other_audio_inputs): - audio = None - # same prompt for single/multiple image(s) - n_prompts_n_images = False - if isinstance(post_prompt, - list) and len(post_prompt) > 1 and image is not None: - if hasattr(image, "pixel_values"): - if len(post_prompt) == image["pixel_values"].shape[0]: - n_prompts_n_images = True - # n prompts and n images - else: - if isinstance( - image, - torch.Tensor) and len(post_prompt) == image.shape[0]: - n_prompts_n_images = True - # n prompts and n images - - if self.model_type == 'kosmos-2': - input_ids = image['input_ids'].clone() - image_mask = image["image_embeds_position_mask"] - image = image['pixel_values'] - input_ids += image_mask * (self.model_config.vocab_size - 4) - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - length = input_ids.shape[1] - elif self.model_type == 'phi-3-vision': - input = image - image = input['pixel_values'] - image = image.flatten(0, 1) - elif self.model_type == 'phi-4-multimodal': - input = image - image = input['input_image_embeds'].flatten(0, 1) - other_vision_inputs['attention_mask'] = input[ - 'image_attention_mask'].flatten(0, 1).bool() - - audio = input['input_audio_embeds'] - l, d = audio.shape[1], audio.shape[2] - pad = 4000 - l % 4000 - audio = torch.cat([audio, audio.new_zeros(1, pad, d)], - dim=1).reshape(-1, 4000, d) - audio_mask = audio.new_ones(*audio.shape[:2]) - audio_mask[-1, -pad:] = 0 - other_audio_inputs['attention_mask'] = audio_mask.bool() - elif self.model_type == 'pixtral': - # Hold on to pixel_values and input_ids. - dtype = str_dtype_to_torch(self.vision_precision) - # Shape of pixel values from the processor varies with the raw image. - # So we create a new tensor with a fixed shape as expected by the vision - # encoder and create a corresponding attention mask. - image_size = self.image_size - patch_size = self.patch_size - d_min = torch.finfo(dtype).min - num_patches = (image_size // patch_size) - padded_image = torch.full( - (self.args.batch_size, 3, image_size, image_size), - fill_value=0, - dtype=dtype, - device="cuda") - padded_attention_mask = torch.full( - (self.args.batch_size, num_patches, num_patches), - fill_value=d_min, - dtype=dtype, - device="cuda") - h, w, input_ids = [], [], [] - for img_idx in range(self.args.batch_size): - pixel_values = image["pixel_values"][img_idx] - img_h, img_w = pixel_values.shape[-2:] - padded_image[img_idx, :, :img_h, :img_w] = pixel_values - padded_attention_mask[img_idx, :img_h // patch_size, :img_w // - patch_size] = 0 - input_ids.append(image["input_ids"][img_idx]) - h.append(img_h) - w.append(img_w) - - image = padded_image - other_vision_inputs = { - "attention_mask": padded_attention_mask, - } - elif self.model_type == 'llava_next': - input = image - image = input['pixel_values'] - image = image[0] - image_size = input['image_sizes'][0].cpu() - elif self.model_type == "qwen2_vl": - input = image - image = input['image'] - input_ids = input['input_ids'] - other_vision_inputs['image_grid_thw'].shape[0] - attention_mask = other_vision_inputs['attention_mask_llm'] - other_vision_inputs.pop('attention_mask_llm') - image_grid_thw = other_vision_inputs['image_grid_thw'] - other_vision_inputs.pop('image_grid_thw') - elif self.model_type == 'llava_onevision': - input = image - if self.args.video_path is None: - image = input['pixel_values'] - image = image[0].repeat(self.args.batch_size, 1, 1, 1) - image_size = input['image_sizes'][0] - image_size = image_size.repeat(self.args.batch_size, 1).cpu() - else: - image = input['pixel_values_videos'] - _, _, c, h, w = image.shape - image = image.repeat(self.args.batch_size, 1, 1, 1, 1) - image = image.view(-1, c, h, w) - elif self.model_type == "fuyu": - while len(image["image_patches"]) < self.args.batch_size: - image["image_patches"].append(image["image_patches"][0]) - - profiler.start("Vision encoder") - visual_features, visual_atts, model_runner_input = None, None, None - if image is not None: - model_runner_input = torch.stack( - image['image_patches'], - dim=0) if self.model_type == 'fuyu' else image - - if self.model_type == "phi-3-vision": - visual_features = self.vision_model.get_img_features( - image).reshape(1, image.shape[0], -1, - self.vision_model.image_dim_out) - visual_atts = None - elif self.model_type == "phi-4-multimodal": - visual_features = self.vision_model.get_img_features( - model_runner_input.to( - str_dtype_to_torch(self.vision_precision)), - other_vision_inputs['attention_mask']) - visual_features = self.vision_model.img_projection( - visual_features) - visual_atts = torch.ones(visual_features.size()[:-1], - dtype=torch.long).to( - model_runner_input.device) - else: - if self.cpp_e2e: - # If using E2E C++ runtime, visual_features will not be computed here in Python runtime. - # Instead, it only contains a shape read from the engine config, and is used for generating - # decoder prompt later - logger.info( - 'Skip running visual engine, get visual output shape from engine config.' - ) - model_runner_input = model_runner_input.to( - str_dtype_to_torch(self.vision_precision)) - batch_size = model_runner_input.shape[0] - output_shape = list(self.visual_output_shape) - output_shape[0] = batch_size - if self.model_type == 'fuyu': - output_shape[1] = model_runner_input.shape[ - 2] # fuyu's output patch number is not fixed, same as input patch number - visual_features = TensorInfo( - 'encoder_output', - str_dtype_to_trt(self.vision_precision), - tuple(output_shape)) - atts_shape = visual_features.shape[:-1] - visual_atts = TensorInfo('image_atts', None, - tuple(atts_shape)) - model_runner_input = torch.vsplit( - model_runner_input, model_runner_input.shape[0]) - else: - visual_features, visual_atts = self.get_visual_features( - model_runner_input, other_vision_inputs) - model_runner_input = None - profiler.stop("Vision encoder") - - profiler.start("Audio encoder") - audio_features = None - if audio is not None: - audio_features = self.get_audio_features(audio, other_audio_inputs) - profiler.stop("Audio encoder") - - if self.model_type == 'fuyu': - input_ids = image['input_ids'].to(torch.int32) - image_patches_indices = image['image_patches_indices'].to( - torch.int32) - - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - image_patches_indices = image_patches_indices.expand( - self.args.batch_size, *image_patches_indices.shape[1:]) - - input_ids = self.ptuning_setup_fuyu(input_ids, - image_patches_indices) - input_ids = torch.stack(input_ids, dim=0).to('cpu') - length = input_ids.shape[1] - if not self.cpp_e2e: # TODO: bs > 1 for C++ E2E Fuyu - visual_features = visual_features.repeat( - self.args.batch_size, 1, 1) - elif self.model_type == 'qwen2_vl': - length = input_ids.shape[1] - input_lengths = torch.IntTensor([length] * self.args.batch_size).to( - torch.int32) - input_ids, ptuning_args, mrope_args = self.setup_fake_prompts_qwen2vl( - visual_features, input_ids, image_grid_thw, attention_mask, - input_lengths) - return input_ids, input_lengths, ptuning_args, visual_features, mrope_args - - elif self.model_type == 'kosmos-2': - visual_features = visual_features.squeeze( - ) if visual_features is not None else None - elif self.model_type == 'vila': - if n_prompts_n_images: - input_ids = self.tokenizer_image_token(self.args.batch_size, - pre_prompt[0], - post_prompt, - self.tokenizer) - else: - input_ids = self.tokenizer_image_token(self.args.batch_size, - pre_prompt[0], - post_prompt[0], - self.tokenizer) - batch_split_prompts = self.split_prompt_by_images(input_ids) - if not n_prompts_n_images: - first_batch_split_prompts = batch_split_prompts[0] - # compute prompt length + visual length - length = sum( - [ids.shape[1] for ids in first_batch_split_prompts]) - if self.args.batch_size == 1 and len(image) > 1: - # mode 1: multiple image as a whole, flatten visual dims - length += visual_atts.shape[0] * visual_atts.shape[1] - else: - length += visual_atts.shape[1] - input_lengths = torch.IntTensor( - [length] * self.args.batch_size).to(torch.int32) - input_ids, ptuning_args = self.setup_fake_prompts_vila( - self.args.batch_size, visual_features, - first_batch_split_prompts, input_lengths) - else: - # mode 2: multiple different prompts corresponding to multiple images (1-1 correspondence) - length = [ - sum([ids.shape[1] for ids in batch_split_prompt]) - for batch_split_prompt in batch_split_prompts - ] - length = [l + visual_atts.shape[1] for l in length] - input_lengths = torch.IntTensor(length).to(torch.int32) - input_ids, ptuning_args = self.setup_fake_prompts_vila( - self.args.batch_size, visual_features, batch_split_prompts, - input_lengths) - return input_ids, input_lengths, ptuning_args, visual_features, model_runner_input - elif self.model_type == 'phi-3-vision': - image_sizes = input["image_sizes"] - profiler.start("Feature transform") - visual_features = self.vision_model.hd_feature_transform( - visual_features, image_sizes) - profiler.stop("Feature transform") - input_ids = input["input_ids"].clone() - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - num_img_tokens = [visual_features.shape[0]] - input_ids = self.ptuning_setup_phi3(visual_features=visual_features, - audio_features=None, - input_ids=input_ids, - num_img_tokens=num_img_tokens, - num_aud_tokens=None) - visual_features = visual_features.unsqueeze(0).repeat( - self.args.batch_size, 1, 1) - length = input_ids.shape[1] - elif self.model_type == 'phi-4-multimodal': - h, w = input["image_sizes"][0] - image_attention_mask = input.get("image_attention_mask") - if image_attention_mask is not None: - image_attention_mask = image_attention_mask[0].bool() - patch_size = 336 if self.model_type == 'phi-3-vision' else 448 - profiler.start("Feature transform") - visual_features = PhiMMUtils.hd_feature_transform( - visual_features, - h // patch_size, - w // patch_size, - self.image_newlines["sub_GN"], - self.image_newlines["glb_GN"], - patch_mask=image_attention_mask) - profiler.stop("Feature transform") - input_ids = input["input_ids"].clone() - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - num_img_tokens = [visual_features.shape[0]] - if audio_features is not None: - dim = audio_features.shape[-1] // 2 - audio_features = audio_features[..., -dim:] - audio_features = PhiMMUtils.reshape_audio_chunks( - audio_features, other_audio_inputs["attention_mask"]) - num_aud_tokens = [audio_features.shape[0]] - else: - num_aud_tokens = None - input_ids = self.ptuning_setup_phi3(visual_features=visual_features, - audio_features=audio_features, - input_ids=input_ids, - num_img_tokens=num_img_tokens, - num_aud_tokens=num_aud_tokens) - visual_features = visual_features.unsqueeze(0).repeat( - self.args.batch_size, 1, 1) - if audio_features is not None: - audio_features = audio_features.unsqueeze(0).repeat( - self.args.batch_size, 1, 1) - length = input_ids.shape[1] - - elif self.model_type == 'pixtral': - relevant_patch_size = self.patch_size * self.spatial_merge_size - output_img_size = self.image_size // relevant_patch_size - # Note: max_h * max_w shall serve as the `tokens_per_task` in ptuning prompt table. - max_h = max(h) // relevant_patch_size - max_w = max(w) // relevant_patch_size - visual_embed_dim = visual_features.shape[-1] - relevant_visual_features = torch.zeros(self.args.batch_size, - max_h * max_w, - visual_embed_dim) - for img_idx in range(self.args.batch_size): - complete_features = visual_features[img_idx] - complete_features = complete_features.reshape( - output_img_size, output_img_size, visual_embed_dim) - relevant_h = h[img_idx] // relevant_patch_size - relevant_w = w[img_idx] // relevant_patch_size - flattened_features = complete_features[:relevant_h, : - relevant_w, :].flatten( - 0, 1) - relevant_visual_features[img_idx, :relevant_h * - relevant_w, :] = flattened_features - visual_features = relevant_visual_features - input_ids = self.ptuning_setup_pixtral(input_ids=input_ids) - # Note: length is not used for pixtral model downstream. Setting it to a list - # of length of input_ids causes errors downstream. So, supplying a placeholder. - length = input_ids[0].shape[0] - - elif self.model_type == 'llava_next': - visual_features = LlavaNextUtils.rearrange_image_features( - visual_features, self.image_newlines["image_newline"], - image_size) - input_ids = self.ptuning_setup_llava_next(visual_features, - pre_prompt, post_prompt) - length = input_ids.shape[1] - elif self.model_type == 'mllama': - pre_input_ids = self.tokenizer(pre_prompt, - return_tensors="pt", - padding=True).input_ids - if n_prompts_n_images: - length = [pre_input_ids.shape[1]] * self.args.batch_size - else: - length = pre_input_ids.shape[1] - post_input_ids = None - elif self.model_type == 'llava_onevision': - if self.args.video_path is None: - visual_features = torch.split(visual_features, - visual_features.shape[0] // - self.args.batch_size, - dim=0) - visual_features = LlavaOnevisionUtils.pack_image_features( - visual_features, - image_size, - image_newline=self.image_newlines["image_newline"], - ) - else: - visual_features = LlavaOnevisionUtils.apply_pooling( - visual_features) - visual_features = visual_features.reshape( - self.args.batch_size, - self.num_frames * visual_features.shape[1], -1) - image_newline = self.image_newlines["image_newline"][ - None, None, :].repeat(self.args.batch_size, 1, - 1).to(visual_features.device) - visual_features = torch.cat((visual_features, image_newline), - dim=1) - - pre_input_ids = self.tokenizer(pre_prompt, - return_tensors="pt", - padding=True).input_ids - post_input_ids = self.tokenizer(post_prompt, - return_tensors="pt", - padding=True).input_ids - length = pre_input_ids.shape[1] + visual_features.shape[ - 1] + post_input_ids.shape[1] - else: - pre_input_ids = self.tokenizer(pre_prompt, - return_tensors="pt", - padding=True).input_ids - if post_prompt[0] is not None: - post_input_encoded = self.tokenizer(post_prompt, - return_tensors="pt", - padding=True) - post_input_ids = post_input_encoded.input_ids - if n_prompts_n_images and 'neva' not in self.model_type: - post_input_attention_mask = post_input_encoded.attention_mask - post_input_ids = [ - input_id[mask.bool()] for input_id, mask in zip( - post_input_ids, post_input_attention_mask) - ] - - if self.model_type == 'video-neva': - length = pre_input_ids.shape[1] + post_input_ids.shape[ - 1] + visual_atts.shape[2] * visual_atts.shape[1] - elif self.model_type == 'internvl': - length = pre_input_ids.shape[1] + post_input_ids.shape[ - 1] + visual_atts.shape[0] * visual_atts.shape[1] - else: - if n_prompts_n_images: - length = [ - pre_input_ids.shape[1] + visual_atts.shape[1] + - post_input_id.shape[0] - for post_input_id in post_input_ids - ] - else: - length = pre_input_ids.shape[1] + post_input_ids.shape[ - 1] + visual_atts.shape[1] - else: - post_input_ids = None - assert pre_input_ids.shape[0] == visual_atts.shape[0] - if visual_atts.shape[0] == 1: - length = pre_input_ids.shape[1] + visual_atts.shape[1] - else: - length = [ - pre_input_ids.shape[1] + visual_atts.shape[1] - for _ in range(visual_atts.shape[0]) - ] - - if n_prompts_n_images: - if isinstance(length, int): length = [length] - assert isinstance(length, list) - input_lengths = torch.IntTensor(length).to(torch.int32) - else: - assert isinstance(length, int) - input_lengths = torch.IntTensor([length] * self.args.batch_size).to( - torch.int32) - - if self.model_type in [ - 'fuyu', 'kosmos-2', 'phi-3-vision', 'llava_next', 'pixtral' - ]: - return input_ids, input_lengths, [ - visual_features - ], visual_features, model_runner_input - if self.model_type == 'phi-4-multimodal': - multimodal_features = torch.cat([visual_features, audio_features], - dim=1) - return input_ids, input_lengths, [ - multimodal_features - ], multimodal_features, model_runner_input - - input_ids, ptuning_args = self.setup_fake_prompts( - visual_features, pre_input_ids, post_input_ids, input_lengths) - - return input_ids, input_lengths, ptuning_args, visual_features, model_runner_input - - @staticmethod - def tokenizer_image_token(batch_size, - pre_prompt, - post_prompt, - tokenizer, - image_token_index=-200): - if isinstance(post_prompt, list): - prompts = [pre_prompt + item for item in post_prompt] - else: - prompts = [pre_prompt + post_prompt] - - def insert_separator(X, sep): - return [ - ele for sublist in zip(X, [sep] * len(X)) for ele in sublist - ][:-1] - - result = [] - for prompt in prompts: - prompt_chunks = [ - tokenizer(chunk).input_ids for chunk in prompt.split("") - ] - input_ids = [] - offset = 0 - if (len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 - and prompt_chunks[0][0] == tokenizer.bos_token_id): - offset = 1 - input_ids.append(prompt_chunks[0][0]) - - for x in insert_separator(prompt_chunks, - [image_token_index] * (offset + 1)): - input_ids.extend(x[offset:]) - - input_ids = torch.tensor(input_ids, dtype=torch.long) - input_ids[input_ids == image_token_index] = 0 - result.append(input_ids) - - if not isinstance(post_prompt, list): - result = result[0].unsqueeze(0).expand(batch_size, -1) - return result - - def split_prompt_by_images(self, tensor): - batch_splits = [] - for batch in tensor: - # Find indices where value is zero () - zero_indices = (batch == 0).nonzero(as_tuple=False).squeeze(0) - # Add starting point for slicing - start_idx = 0 - splits = [] - for idx in zero_indices: - if start_idx != idx: # Ensure not slicing zero-length tensors - splits.append(batch[start_idx:idx].unsqueeze(0)) - start_idx = idx + 1 # Move start index past the zero - if start_idx < len( - batch): # Handle last segment if it's not zero-ending - splits.append(batch[start_idx:].unsqueeze(0)) - # Remove empty tensors resulting from consecutive zeros - splits = [split for split in splits if split.numel() > 0] - batch_splits.append(splits) - - return batch_splits - - def prepare_position_ids_for_cogvlm(self, input_ids): - batch_size = len(input_ids) - position_ids = torch.arange(input_ids.shape[1]) - position_ids[2:1227] = 2 - position_ids[1227:] = torch.arange(3, input_ids.shape[1] + 1 - 1225) - - position_ids = position_ids.to(torch.int32).to('cuda') - input_position_ids = [] - for i in range(batch_size): - input_position_ids.append(position_ids) - - return input_position_ids - - def generate(self, - pre_prompt, - post_prompt, - image, - decoder_input_ids, - max_new_tokens, - other_vision_inputs={}, - other_audio_inputs={}, - other_decoder_inputs={}): - profiler.start("Generate") - profiler.start("Preprocess") - if 'qwen2_vl' in self.model_type: - input_ids, input_lengths, ptuning_args, visual_features, mrope_args = self.preprocess( - pre_prompt, post_prompt, image, other_vision_inputs, - other_audio_inputs) - mrope_params = MropeParams( - mrope_rotary_cos_sin=mrope_args[0], - mrope_position_deltas=mrope_args[1], - ) - else: - input_ids, input_lengths, ptuning_args, visual_features, model_runner_input = self.preprocess( - pre_prompt, post_prompt, image, other_vision_inputs, - other_audio_inputs) - profiler.stop("Preprocess") - - # use prompt tuning to pass multimodal features - # model.generate() expects the following params (see layers/embedding.py): - # args[0]: prompt embedding table, [batch_size, multimodal_len, hidden_size], later flattened to [batch_size * multimodal_len, hidden_size] - # args[1]: prompt task ids, [batch_size]. in multimodal case, arange(batch_size), i.e. in VILA batching mode 2, each image is treated separately in the batch instead of concated together (although the prompt embedding table has to be concated) - # args[2]: prompt task vocab size, [1]. assuming all table has the same length, which in multimodal case equals to multimodal_len - profiler.start("LLM") - if self.decoder_llm and self.model_type != "mllama": - end_id = self.tokenizer.eos_token_id - if 'opt' in self.model_type and 'blip2' in self.model_type: - # For BLIP2-OPT, model outputs a "\n" at the end. - # we avoid it by using newline as the end token - end_id = self.tokenizer.encode("\n", - add_special_tokens=False)[0] - - if self.model_type == 'cogvlm': - input_position_ids = self.prepare_position_ids_for_cogvlm( - input_ids) - - prompt_tasks = None - prompt_table = None - if not self.cpp_e2e: - batch_size = len(input_ids) - prompt_tasks = ",".join( - np.arange(batch_size, dtype=np.int32).astype(str)) - prompt_table = torch.stack([ptuning_args[0]]) - prompt_table = prompt_table.view(batch_size, -1, - prompt_table.shape[-1]) - - output_ids = self.model.generate( - input_ids, - input_position_ids=input_position_ids - if self.model_type == 'cogvlm' else None, - mrope_params=mrope_params - if self.model_type == 'qwen2_vl' else None, - encoder_input_features=model_runner_input - if self.cpp_e2e else None, - sampling_config=None, - prompt_table=prompt_table, - prompt_tasks=prompt_tasks, - max_new_tokens=max_new_tokens, - end_id=end_id, - pad_id=self.tokenizer.pad_token_id - if self.tokenizer.pad_token_id is not None else - self.tokenizer.all_special_ids[0], - top_k=self.args.top_k, - top_p=self.args.top_p, - temperature=self.args.temperature, - repetition_penalty=self.args.repetition_penalty, - num_beams=self.args.num_beams, - lora_uids=self.args.lora_task_uids, - output_sequence_lengths=False, - return_dict=False, - mm_embedding_offloading=self.args.mm_embedding_offloading) - elif self.model_type == "mllama": - # When image is passed: - # the shape of visual_features is [bs, 1, 4, 1025, hidden_size] - # the shape of cross_attention_mask is [bs, decode_input_len, 1, 4] - # When image is None, create dummy visual_features and cross_attention_mask - if visual_features is None: - visual_features = torch.zeros([ - self.args.batch_size, 1, 4, 1, - self.model_config.hidden_size * self.runtime_mapping.tp_size - ], - dtype=self.model.dtype, - device=self.device) - dummy_cross_attention_mask = torch.zeros( - [self.args.batch_size, input_ids.shape[1], 1, 4], - dtype=bool, - device=self.device) - skip_cross_attn_blocks = torch.ones([1], - dtype=torch.bool, - device='cpu') - else: - skip_cross_attn_blocks = torch.zeros([1], - dtype=torch.bool, - device='cpu') - - visual_features = visual_features.to(self.model.dtype).chunk( - self.args.batch_size, dim=0) - encoder_input_features = [] - cross_attention_masks = [] - encoder_output_lengths = [] - for batch_idx in range(self.args.batch_size): - visual_feature = visual_features[batch_idx] - num_vision_tokens = visual_feature.shape[3] - visual_feature = visual_feature.reshape( - [-1, visual_feature.shape[-1]]) - encoder_max_input_length = visual_feature.shape[0] - encoder_input_lengths = torch.IntTensor( - [encoder_max_input_length]).to(visual_feature.device) - - # prepare cross_attention_mask of context phase - if 'cross_attention_mask' in other_decoder_inputs: - cross_attention_mask = other_decoder_inputs[ - 'cross_attention_mask'][batch_idx] - else: - cross_attention_mask = dummy_cross_attention_mask[batch_idx] - text_total_length, *_ = cross_attention_mask.shape - cross_attention_mask = cross_attention_mask.repeat_interleave( - num_vision_tokens, dim=2) - cross_attention_mask = cross_attention_mask.view( - text_total_length, -1) - cross_attention_mask = cross_attention_mask.unsqueeze(1) - cross_attention_mask = cross_attention_mask.to( - visual_feature.device).to(torch.bool).reshape( - [-1, cross_attention_mask.shape[-1]]) - - # prepare cross_attention_mask for generation phase and concat them - tmp_mask = [cross_attention_mask] + [ - cross_attention_mask[-1:, :] for _ in range(max_new_tokens) - ] - cross_attention_mask = torch.concat(tmp_mask) - - encoder_input_features.append(visual_feature) - cross_attention_masks.append(cross_attention_mask) - encoder_output_lengths.append(encoder_input_lengths) - - outputs = self.model.generate( - batch_input_ids=input_ids, - encoder_input_ids=None, - encoder_input_features=encoder_input_features, - encoder_output_lengths=encoder_output_lengths, - cross_attention_masks=cross_attention_masks, - max_new_tokens=max_new_tokens, - # max_attention_window_size=args.max_attention_window_size, - # sink_token_length=args.sink_token_length, - end_id=self.tokenizer.eos_token_id, - pad_id=self.tokenizer.pad_token_id, - temperature=self.args.temperature, - top_k=self.args.top_k, - top_p=self.args.top_p, - num_beams=self.args.num_beams, - # length_penalty=args.length_penalty, - # early_stopping=args.early_stopping, - # beam_width_array=args.beam_width_array, - repetition_penalty=self.args.repetition_penalty, - # presence_penalty=args.presence_penalty, - # frequency_penalty=args.frequency_penalty, - # stop_words_list=stop_words_list, - # bad_words_list=bad_words_list, - # output_cum_log_probs=(args.output_cum_log_probs_npy != None), - # output_log_probs=(args.output_log_probs_npy != None), - # random_seed=args.random_seed, - # lora_uids=args.lora_task_uids, - # prompt_table=args.prompt_table_path, - # prompt_tasks=args.prompt_tasks, - # streaming=args.streaming, - output_sequence_lengths=True, - # no_repeat_ngram_size=self.args.no_repeat_ngram_size, - return_dict=True, - # medusa_choices=args.medusa_choices, - # return_all_generated_tokens=args.return_all_generated_tokens, - # input_token_extra_ids=input_token_extra_ids, - encoder_max_input_length=encoder_max_input_length, - skip_cross_attn_blocks=skip_cross_attn_blocks, - ) - if mpi_rank() == 0: - output_ids = outputs["output_ids"] - else: - if self.model_type in ['nougat', 'pix2struct']: - # Trim encoder input_ids to match visual features shape - ids_shape = (self.args.batch_size, visual_features.shape[1]) - if self.model_type == 'nougat': - input_ids = torch.zeros(ids_shape, dtype=torch.int32) - elif self.model_type == 'pix2struct': - input_ids = torch.ones(ids_shape, dtype=torch.int32) - - output_ids = self.model.generate( - input_ids, - decoder_input_ids, - max_new_tokens, - num_beams=self.args.num_beams, - bos_token_id=self.tokenizer.bos_token_id, - pad_token_id=self.tokenizer.pad_token_id, - eos_token_id=self.tokenizer.eos_token_id, - debug_mode=False, - prompt_embedding_table=ptuning_args[0], - prompt_tasks=ptuning_args[1], - prompt_vocab_size=ptuning_args[2]) - - # Reset input_lengths to match decoder_input_ids - input_lengths = torch.ones(input_lengths.shape, - dtype=input_lengths.dtype) - profiler.stop("LLM") - - if mpi_rank() == 0: - # Extract a list of tensors of shape beam_width x output_ids. - profiler.start("Tokenizer decode") - output_beams_list = [ - self.tokenizer.batch_decode( - output_ids[batch_idx, :, input_lengths[batch_idx]:], - skip_special_tokens=True) for batch_idx in range( - min(self.args.batch_size, input_lengths.shape[0])) - ] - - stripped_text = [[ - output_beams_list[batch_idx][beam_idx].strip() - for beam_idx in range(self.args.num_beams) - ] for batch_idx in range( - min(self.args.batch_size, input_lengths.shape[0]))] - profiler.stop("Tokenizer decode") - profiler.stop("Generate") - return stripped_text - else: - profiler.stop("Generate") - return None - - def get_visual_features(self, image, other_vision_inputs): - visual_features = { - self.vision_input_names[0]: - image.to(str_dtype_to_torch(self.vision_precision)), - } - if self.model_type == "qwen2_vl": - other_vision_inputs['attention_mask'] = other_vision_inputs[ - 'attention_mask'].to(str_dtype_to_torch(self.vision_precision)) - for key, tensor in other_vision_inputs.items(): - visual_features.update({key: tensor}) - - tensor_info = [ - TensorInfo(self.vision_input_names[0], - str_dtype_to_trt(self.vision_precision), image.shape), - ] - for key, tensor in other_vision_inputs.items(): - tensor_info.append( - TensorInfo(key, torch_dtype_to_trt(tensor.dtype), tensor.shape)) - - visual_output_info = self.visual_encoder_session.infer_shapes( - tensor_info) - self.visual_encoder_session.set_shapes(visual_features) - visual_outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device=image.device) - for t in visual_output_info - } - - ok = self.visual_encoder_session.run(visual_features, visual_outputs, - self.stream.cuda_stream) - assert ok, "Runtime execution failed for vision encoder session" - self.stream.synchronize() - - image_embeds = visual_outputs[self.vision_output_names[0]] - - if self.args.mm_embedding_offloading: - # CUDA Stream Overlapping Requirements: - # 1. Both memory copy stream and kernel execution stream must be non-default streams - # 2. For host<->device transfers (H2D/D2H), host memory MUST be page-locked (pinned) - # NOTE: pinning is skipped under Confidential Compute - # (see maybe_pin_memory() and prefer_pinned()) - pinned_embeds = torch.empty_like(image_embeds, - device='cpu', - pin_memory=prefer_pinned()) - pinned_embeds.copy_(image_embeds, non_blocking=True) - image_embeds = pinned_embeds - - image_atts = torch.ones(image_embeds.size()[:-1], - dtype=torch.long).to(image.device) - - return image_embeds, image_atts - - def get_audio_features(self, audio, other_audio_inputs): - tmp_features, _ = self.audio_model.encoder( - audio.to(str_dtype_to_torch(self.audio_precision)), - other_audio_inputs['attention_mask']) - speech_out = self.audio_model.audio_projection['speech'](tmp_features) - vision_out = self.audio_model.audio_projection['vision'](tmp_features) - return torch.cat((speech_out, vision_out), dim=-1) - - def setup_fake_prompts_vila(self, batch_size, visual_features, - split_input_ids, input_lengths): - # visual_features (num_images, feature_len, token_embed) - # Assemble fake prompts which points to image embedding actually - fake_prompt_counter = self.model_config.vocab_size - if batch_size == 1: - # only check for multi-image inference (mode 1) - assert len(visual_features) <= len( - split_input_ids - ), "Unexpected number of visual features. Please check # in prompt and the #image files." - - input_ids = [] - if batch_size == 1: - # mode 1: multiple image as a whole, concat all prompts together, 
...
-            input_ids = [split_input_ids[0]]
-            for idx in range(
-                    len(visual_features)
-            ):  # TODO:alternatively make TensorInfo iterable if this breaks others
-                fake_prompt_id = torch.arange(
-                    fake_prompt_counter,
-                    fake_prompt_counter + visual_features.shape[1])
-                fake_prompt_counter += visual_features.shape[1]
-                fake_prompt_id = fake_prompt_id.unsqueeze(0)
-                input_ids.append(fake_prompt_id)
-                # in case no inter or post prompt
-                if len(split_input_ids) > idx + 1:
-                    input_ids.append(split_input_ids[idx + 1])
-            input_ids = torch.cat(input_ids, dim=1).contiguous().to(torch.int32)
-            input_ids = input_ids.reshape(batch_size, -1)
-
-        elif batch_size > 1:
-            # mode 2: each image have individual prompt, 

-            for idx in range(len(visual_features)):
-                input_ids.append(split_input_ids[idx][0])
-                fake_prompt_id = torch.arange(
-                    fake_prompt_counter,
-                    fake_prompt_counter + visual_features.shape[1])
-                fake_prompt_id = fake_prompt_id.unsqueeze(0)
-                input_ids.append(fake_prompt_id)
-                if len(split_input_ids[idx]) > 1:
-                    input_ids.append(split_input_ids[idx][1])
-            result = []
-            for i in range(0, len(input_ids), 3):
-                # Concatenate every 3 items (
, , )
-                concatenated = torch.cat(input_ids[i:i + 3],
-                                         dim=1).to(torch.int32).squeeze(0)
-                result.append(concatenated)
-            input_ids = result
-
-        if (self.decoder_llm
-                or self.runtime_mapping.is_first_pp_rank()) and isinstance(
-                    visual_features, torch.Tensor):
-            ptuning_args = self.ptuning_setup(visual_features, input_ids,
-                                              input_lengths)
-        else:
-            ptuning_args = [None, None, None]
-
-        return input_ids, ptuning_args
-
-    def setup_fake_prompts(self, visual_features, pre_input_ids, post_input_ids,
-                           input_lengths):
-        # Assemble fake prompts which points to image embedding actually
-        if hasattr(self, 'num_frames') and (visual_features.shape[1]
-                                            == self.num_frames):
-            visual_features = visual_features.view(visual_features.shape[0], -1,
-                                                   visual_features.shape[-1])
-
-        if visual_features is not None:
-            if self.python_e2e:
-                # Non-IFB Mode(used in python session): All requests in a batch have their prompt_table concatenated in
-                # a shape of (bs*vision_embedding_len, vision_hidden). So only one fake_prompt_id is needed for the
-                # entire batch, with values from 0 to bs * vision_embedding_len-1.
-                fake_prompt_id = torch.arange(
-                    self.model_config.vocab_size, self.model_config.vocab_size +
-                    visual_features.shape[0] * visual_features.shape[1])
-                fake_prompt_id = fake_prompt_id.reshape(
-                    visual_features.shape[0], visual_features.shape[1])
-            else:
-                # IFB Mode(used in c++ session): Each request's prompt_table is independent and requires a fake_prompt_id
-                # for each request, with values ranging from 0 to vision_embedding_len-1.
-                fake_prompt_id = torch.arange(
-                    self.model_config.vocab_size,
-                    self.model_config.vocab_size + visual_features.shape[1])
-                fake_prompt_id = fake_prompt_id.repeat(visual_features.shape[0],
-                                                       1)
-
-        if 'internvl' in self.model_type:
-            fake_prompt_id = fake_prompt_id.reshape(1, -1)
-
-        if 'cogvlm' in self.model_type:
-            input_ids = torch.cat(
-                [pre_input_ids[:, 0:1], fake_prompt_id, pre_input_ids[:, 1:]],
-                dim=1).contiguous().to(torch.int32)
-        elif self.model_type == 'mllama':
-            input_ids = pre_input_ids.contiguous().to(torch.int32)
-        else:
-            if post_input_ids is not None:
-                if isinstance(post_input_ids, list):
-                    pre_input_fake_prompt_ids = [
-                        pre_input_ids[:len(fake_prompt_id)], fake_prompt_id
-                    ]
-                    pre_input_fake_prompt_ids = torch.cat(
-                        pre_input_fake_prompt_ids,
-                        dim=1).contiguous().to(torch.int32)
-                    input_ids = [
-                        torch.cat((pre_input_fake_prompt_id,
-                                   post_input_id)).contiguous().to(torch.int32)
-                        for pre_input_fake_prompt_id, post_input_id in zip(
-                            pre_input_fake_prompt_ids, post_input_ids)
-                    ]
-                else:
-                    input_ids = [pre_input_ids, fake_prompt_id, post_input_ids]
-                    input_ids = torch.cat(input_ids,
-                                          dim=1).contiguous().to(torch.int32)
-            else:
-                input_ids = [fake_prompt_id, pre_input_ids]
-                input_ids = torch.cat(input_ids,
-                                      dim=1).contiguous().to(torch.int32)
-
-        if (self.decoder_llm or self.runtime_mapping.is_first_pp_rank()
-            ) and self.model_type != "mllama" and isinstance(
-                visual_features, torch.Tensor):
-            ptuning_args = self.ptuning_setup(visual_features, input_ids,
-                                              input_lengths)
-        else:
-            ptuning_args = [None, None, None]
-
-        return input_ids, ptuning_args
-
-    def get_rope_index(
-        self,
-        input_ids: torch.IntTensor,
-        image_grid_thw: Optional[torch.LongTensor] = None,
-        video_grid_thw: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
-
-        Explanation:
-            Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
-
-            For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
-            Examples:
-                input_ids: [T T T T T], here T is for text.
-                temporal position_ids: [0, 1, 2, 3, 4]
-                height position_ids: [0, 1, 2, 3, 4]
-                width position_ids: [0, 1, 2, 3, 4]
-
-            For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
-            and 1D rotary position embedding for text part.
-            Examples:
-                Assume we have a video input with 3 temporal patches, 2 height patches and 2 width patches.
-                input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
-                vision temporal position_ids: [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
-                vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
-                vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
-                text temporal position_ids: [3, 4, 5, 6, 7]
-                text height position_ids: [3, 4, 5, 6, 7]
-                text width position_ids: [3, 4, 5, 6, 7]
-                Here we calculate the text start position_ids as the max vision position_ids plus 1.
-
-        Args:
-            input_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`):
-                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
-                it.
-            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
-                The temporal, height and width of feature shape of each image in LLM.
-            video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
-                The temporal, height and width of feature shape of each video in LLM.
-            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
-                - 1 for tokens that are **not masked**,
-                - 0 for tokens that are **masked**.
-
-        Returns:
-            position_ids (`torch.IntTensor` of shape `(3, batch_size, sequence_length)`)
-            mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
-        """
-        spatial_merge_size = self.spatial_merge_size
-        image_token_id = self.image_token_id
-        video_token_id = self.video_token_id
-        vision_start_token_id = self.vision_start_token_id
-        mrope_position_deltas = []
-        if image_grid_thw is not None or video_grid_thw is not None:
-            total_input_ids = input_ids
-            position_ids = torch.ones(3,
-                                      input_ids.shape[0],
-                                      input_ids.shape[1],
-                                      dtype=input_ids.dtype,
-                                      device=input_ids.device)
-            image_index, video_index = 0, 0
-            for i, input_ids in enumerate(total_input_ids):
-                if attention_mask is not None:
-                    input_ids = input_ids[attention_mask[i] == 1]
-                image_nums, video_nums = 0, 0
-                vision_start_indices = torch.argwhere(
-                    input_ids == vision_start_token_id).squeeze(1)
-                vision_tokens = input_ids[vision_start_indices + 1]
-                image_nums = (vision_tokens == image_token_id).sum()
-                video_nums = (vision_tokens == video_token_id).sum()
-                input_tokens = input_ids.tolist()
-                llm_pos_ids_list: list = []
-                st = 0
-                remain_images, remain_videos = image_nums, video_nums
-                for _ in range(image_nums + video_nums):
-                    if image_token_id in input_tokens and remain_images > 0:
-                        ed_image = input_tokens.index(image_token_id, st)
-                    else:
-                        ed_image = len(input_tokens) + 1
-                    if video_token_id in input_tokens and remain_videos > 0:
-                        ed_video = input_tokens.index(video_token_id, st)
-                    else:
-                        ed_video = len(input_tokens) + 1
-                    if ed_image < ed_video:
-                        t, h, w = (
-                            image_grid_thw[image_index][0],
-                            image_grid_thw[image_index][1],
-                            image_grid_thw[image_index][2],
-                        )
-                        image_index += 1
-                        remain_images -= 1
-                        ed = ed_image
-                    else:
-                        t, h, w = (
-                            video_grid_thw[video_index][0],
-                            video_grid_thw[video_index][1],
-                            video_grid_thw[video_index][2],
-                        )
-                        video_index += 1
-                        remain_videos -= 1
-                        ed = ed_video
-                    llm_grid_t, llm_grid_h, llm_grid_w = (
-                        t.item(),
-                        h.item() // spatial_merge_size,
-                        w.item() // spatial_merge_size,
-                    )
-                    text_len = ed - st
-
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                        llm_pos_ids_list) > 0 else 0
-                    llm_pos_ids_list.append(
-                        torch.arange(text_len).view(1, -1).expand(3, -1) +
-                        st_idx)
-
-                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
-                        -1, llm_grid_h * llm_grid_w).flatten()
-                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
-                        llm_grid_t, -1, llm_grid_w).flatten()
-                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
-                        llm_grid_t, llm_grid_h, -1).flatten()
-                    llm_pos_ids_list.append(
-                        torch.stack([t_index, h_index, w_index]) + text_len +
-                        st_idx)
-                    st = ed + llm_grid_t * llm_grid_h * llm_grid_w
-
-                if st < len(input_tokens):
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                        llm_pos_ids_list) > 0 else 0
-                    text_len = len(input_tokens) - st
-                    llm_pos_ids_list.append(
-                        torch.arange(text_len).view(1, -1).expand(3, -1) +
-                        st_idx)
-
-                llm_positions = torch.cat(llm_pos_ids_list,
-                                          dim=1).reshape(3, -1)
-                position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(
-                    position_ids.device)
-                mrope_position_deltas.append(llm_positions.max() + 1 -
-                                             len(total_input_ids[i]))
-            mrope_position_deltas = torch.tensor(
-                mrope_position_deltas, device=input_ids.device).unsqueeze(1)
-            return position_ids, mrope_position_deltas
-        else:
-            if attention_mask is not None:
-                position_ids = attention_mask.long().cumsum(-1) - 1
-                position_ids.masked_fill_(attention_mask == 0, 1)
-                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(
-                    input_ids.device)
-                max_position_ids = position_ids.max(0, keepdim=False)[0].max(
-                    -1, keepdim=True)[0]
-                mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[
-                    -1]
-            else:
-                position_ids = (torch.arange(input_ids.shape[1],
-                                             device=input_ids.device).view(
-                                                 1, 1, -1).expand(
-                                                     3, input_ids.shape[0], -1))
-                mrope_position_deltas = torch.zeros(
-                    [input_ids.shape[0], 1],
-                    device=input_ids.device,
-                    dtype=input_ids.dtype,
-                )
-
-            return position_ids, mrope_position_deltas
-
-    def setup_fake_prompts_qwen2vl(self, visual_features, input_ids,
-                                   vision_grid_thws, attention_mask,
-                                   input_lengths):
-
-        visual_features = torch.unsqueeze(visual_features, 0)
-
-        # Get the rope index
-        # From HF's preprocess code
-        mrope_position_ids, mrope_position_deltas = self.get_rope_index(
-            input_ids,
-            image_grid_thw=vision_grid_thws,
-            video_grid_thw=None,
-            attention_mask=attention_mask,
-        )
-
-        # This is where we convert input_ids of image features into fake_prompt_ids mapping for TRT-LLM engine.
-        masks = (input_ids == self.image_token_id) | (
-            input_ids == self.vision_token_id) | (input_ids
-                                                  == self.video_token_id)
-        cumulative_counts = masks.cumsum(dim=1)
-        values = (self.model_config.vocab_size - 1) + cumulative_counts
-        input_ids[masks] = values[masks]
-
-        if self.decoder_llm or self.runtime_mapping.is_first_pp_rank():
-            ptuning_args = self.ptuning_setup(visual_features, input_ids,
-                                              input_lengths)
-        else:
-            ptuning_args = [None, None, None]
-
-        # This does not have dependency on input.
-        # Switch to attributes to use across iterations.
-        if not hasattr(self, 'rotary_cos_sin'):
-            inv_freq, rotary_cos_sin = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin(
-                num_pos=self.max_position_embeddings,
-                dim=int(self.hidden_size / self.num_attention_heads),
-                theta=float(self.rope_theta),
-                scale_type=RotaryScalingType.mrope)
-            self.rotary_cos_sin = torch.from_numpy(rotary_cos_sin).to(
-                visual_features.device)
-            self.rotary_cos_sin = self.rotary_cos_sin.reshape(
-                self.max_position_embeddings,
-                int(self.hidden_size / self.num_attention_heads / 2), 2)
-            self.cos_ori = self.rotary_cos_sin[:, :, 0]
-            self.sin_ori = self.rotary_cos_sin[:, :, 1]
-
-        mrope_position_ids = mrope_position_ids.transpose(1, 0)
-        mrope_position_ids_padding = torch.zeros(
-            mrope_position_ids.shape[:-1] + (self.max_position_embeddings, ),
-            dtype=torch.int32,
-            device=visual_features.device)
-        mrope_position_ids_padding[:, :, :mrope_position_ids.
-                                   shape[-1]] = mrope_position_ids
-        cos = self.cos_ori[mrope_position_ids_padding]
-        sin = self.sin_ori[mrope_position_ids_padding]
-
-        mrope_section = [16, 24, 24]
-        cos = torch.cat([
-            m[:, i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))
-        ],
-                        dim=-1).unsqueeze(-1)
-        sin = torch.cat([
-            m[:, i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))
-        ],
-                        dim=-1).unsqueeze(-1)
-        concat_cos_sin = torch.concatenate((cos, sin), axis=-1)
-        concat_cos_sin = concat_cos_sin.reshape(concat_cos_sin.shape[0], -1)
-
-        mrope_args = [concat_cos_sin, mrope_position_deltas]
-        return input_ids, ptuning_args, mrope_args
-
-    def ptuning_setup_fuyu(self, input_ids, image_patches_indices):
-        res_input_ids = []
-        for cur_input_ids, cur_image_patches_indices in zip(
-                input_ids, image_patches_indices):
-            # Truncate input_ids to the length of image_patches_indices
-            cur_image_patches_indices = cur_image_patches_indices[:len(
-                cur_input_ids)]
-            # Get ids of the image_patches
-            non_zero_mask = cur_image_patches_indices != -1
-            # Replace input_ids with image_patches_indices values (where the patches are placed)
-            cur_input_ids = cur_input_ids.masked_scatter(
-                non_zero_mask,
-                cur_image_patches_indices[non_zero_mask] +
-                self.model_config.vocab_size,
-            )
-            res_input_ids.append(cur_input_ids)
-        return res_input_ids
-
-    def ptuning_setup_pixtral(self, input_ids):
-        # input_ids obtained from processor has token_ids for text as well as image tokens
-        # where each image token is represented by the same image_token_index.
-        image_token_index = self.image_token_index
-        vocab_size = self.vocab_size
-        # Replace all image tokens with a unique token_id > text_vacab_size.
-        # This shall be used to lookup the prompt table.
-        for img_idx in range(self.args.batch_size):
-            # Note: We reset replacer to text_vocab_size for each sample. This is as opposed to doing `replacer = vocab_size + img_idx * tokens_per_task`.
-            # That part of the look-up manipulation is done by the `task_ids` input to PromptEmbedding forward.
-            replacer = vocab_size
-            for token_idx in range(len(input_ids[img_idx])):
-                if input_ids[img_idx][token_idx] == image_token_index:
-                    input_ids[img_idx][token_idx] = replacer
-                    replacer += 1
-        return input_ids
-
-    def ptuning_setup_llava_next(self, visual_features, pre_prompt,
-                                 post_prompt):
-        input_ids = []
-        fake_prompt_ids = list(
-            range(self.model_config.vocab_size,
-                  self.model_config.vocab_size + visual_features.shape[0]))
-        input_ids = self.tokenizer.encode(
-            pre_prompt[0]) + fake_prompt_ids + self.tokenizer.encode(
-                post_prompt[0])[self.tokenizer.add_bos_token:]
-        input_ids = [input_ids] * len(pre_prompt)
-        input_ids = torch.tensor(input_ids)
-        return input_ids
-
-    def ptuning_setup_phi3(self, visual_features, audio_features, input_ids,
-                           num_img_tokens, num_aud_tokens):
-        fake_prompt_id = torch.arange(
-            self.model_config.vocab_size,
-            self.model_config.vocab_size + visual_features.shape[0])
-        if self.model_type == "phi-3-vision":
-            MAX_INPUT_ID = int(1e9)
-            positions = torch.nonzero(
-                (input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=False)
-        elif self.model_type == "phi-4-multimodal":
-            IMAGE_TOKEN_ID = 200010
-            positions = torch.nonzero(input_ids == IMAGE_TOKEN_ID,
-                                      as_tuple=False)
-        idx = 0
-        for _, cnt in enumerate(num_img_tokens):
-            input_ids[positions[idx, 0], positions[idx, 1]:positions[idx, 1] +
-                      cnt] = fake_prompt_id[idx:idx + cnt]
-            idx += cnt
-
-        if self.model_type == "phi-4-multimodal" and audio_features is not None:
-            prompt_id_offset = self.model_config.vocab_size + visual_features.shape[
-                0]
-            fake_prompt_id = torch.arange(
-                prompt_id_offset, prompt_id_offset + audio_features.shape[0])
-            AUDIO_TOKEN_ID = 200011
-            positions = torch.nonzero(input_ids == AUDIO_TOKEN_ID,
-                                      as_tuple=False)
-            idx = 0
-            for _, cnt in enumerate(num_aud_tokens):
-                input_ids[positions[idx, 0], positions[idx,
-                                                       1]:positions[idx, 1] +
-                          cnt] = fake_prompt_id[idx:idx + cnt]
-                idx += cnt
-        return input_ids
-
-    def ptuning_setup(self, prompt_table, input_ids, input_lengths):
-        hidden_size = self.model_config.hidden_size * self.runtime_mapping.tp_size
-        if prompt_table is not None:
-            task_vocab_size = torch.tensor(
-                [prompt_table.shape[1]],
-                dtype=torch.int32,
-            ).cuda()
-            prompt_table = prompt_table.view(
-                (prompt_table.shape[0] * prompt_table.shape[1],
-                 prompt_table.shape[2]))
-
-            assert prompt_table.shape[
-                1] == hidden_size, "Prompt table dimensions do not match hidden size"
-
-            if hasattr(self.model_config, 'dtype'):
-                prompt_table = prompt_table.cuda().to(
-                    dtype=str_dtype_to_torch(self.model_config.dtype))
-            else:
-                if self.args.mm_embedding_offloading:
-                    # CUDA Stream Overlapping Requirements:
-                    # 1. Both memory copy stream and kernel execution stream must be non-default streams
-                    # 2. For host<->device transfers (H2D/D2H), host memory MUST be page-locked (pinned)
-                    # NOTE: pinning is skipped under Confidential Compute
-                    # (see maybe_pin_memory() and prefer_pinned())
-                    prompt_table = maybe_pin_memory(prompt_table).to(
-                        dtype=self.model.dtype)
-                else:
-                    prompt_table = prompt_table.cuda().to(
-                        dtype=self.model.dtype)
-        else:
-            prompt_table = torch.empty([1, hidden_size]).cuda()
-            task_vocab_size = torch.zeros([1]).cuda()
-
-        remove_input_padding = self.model_config.remove_input_padding if hasattr(
-            self.model_config,
-            'remove_input_padding') else self.model_config.use_packed_input
-        if remove_input_padding:
-            tasks = torch.zeros([torch.sum(input_lengths)],
-                                dtype=torch.int32).cuda()
-            if self.decoder_llm: tasks = tasks.unsqueeze(0)
-        else:
-            if not isinstance(input_ids, list):
-                tasks = torch.zeros(input_ids.shape, dtype=torch.int32).cuda()
-            else:
-                max_length = max(input_id.size(-1) for input_id in input_ids)
-                tasks = torch.zeros((len(input_ids), max_length),
-                                    dtype=torch.int32).cuda()
-
-        return [prompt_table, tasks, task_vocab_size]
-
-    def load_test_data(self, image_path=None, video_path=None):
-
-        def load_images(image_paths):
-            if isinstance(image_paths, str):
-                image_paths = [image_paths]
-            images = []
-            for image_path in image_paths:
-                if image_path.startswith("http") or image_path.startswith(
-                        "https"):
-                    logger.info(f"downloading image from url {image_path}")
-                    try:
-                        response = requests.get(image_path, timeout=5)
-                        response.raise_for_status()
-                        if 'image' not in response.headers.get(
-                                'Content-Type', ''):
-                            raise Exception(
-                                f"URL does not point to an image: {image_path}."
-                            )
-                        image = Image.open(BytesIO(
-                            response.content)).convert("RGB")
-                    except (UnidentifiedImageError, IOError):
-                        raise Exception(
-                            f"Cannot identify image file at URL: {image_path}.")
-                    except Exception as e:
-                        raise Exception(
-                            f"Failed to download image from url {image_path}: {e}"
-                        )
-                else:
-                    image = Image.open(image_path).convert("RGB")
-                images.append(image)
-            return images if len(images) > 1 else images[0]
-
-        if "vila" in self.model_type:
-            if image_path is None:
-                img_urls = [
-                    'https://github.com/NVlabs/VILA/blob/6b941da19e31ddfdfaa60160908ccf0978d96615/demo_images/av.png?raw=true',
-                    'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
-                ] * 4
-
-                img_urls = img_urls[:self.args.batch_size]
-                self.args.image_path = ",".join(img_urls)
-                images = load_images(img_urls)
-            else:
-                if isinstance(image_path, str):
-                    image_path = image_path.split(self.args.path_sep)
-                images = load_images(image_path)
-        elif "pixtral" in self.model_type:
-            if image_path is None:
-                image_urls = [
-                    "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png",
-                    "https://www.ilankelman.org/stopsigns/australia.jpg",
-                    "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.png",
-                    "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
-                ]
-                while len(image_urls) < self.args.batch_size:
-                    image_urls *= 2
-                image_urls = image_urls[:self.args.batch_size]
-                self.args.image_path = ",".join(image_urls)
-                images = load_images(image_urls)
-            else:
-                if isinstance(image_path, str):
-                    image_path = image_path.split(self.args.path_sep)
-                images = load_images(image_path)
-            images = [images] if not isinstance(images, list) else images
-        elif "nougat" in self.model_type:
-            filepath = hf_hub_download(
-                repo_id="hf-internal-testing/fixtures_docvqa",
-                filename="nougat_paper.png",
-                revision="ec57bf8c8b1653a209c13f6e9ee66b12df0fc2db",
-                repo_type="dataset")
-            images = Image.open(filepath)
-        elif "fuyu" in self.model_type:
-            filepath = hf_hub_download(repo_id="adept/fuyu-8b",
-                                       filename="skateboard.png",
-                                       repo_type='model')
-            images = Image.open(filepath)
-        elif "kosmos" in self.model_type:
-            if image_path is None:
-                image_path = 'https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.png'
-            images = load_images(image_path)
-        elif "pix2struct" in self.model_type:
-            if image_path is None:
-                image_path = 'https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/val/png/multi_col_40963.png'
-            images = load_images(image_path)
-        elif "video-neva" in self.model_type:
-            images = video_path
-        elif "internlm" in self.model_type:
-            img_url = "https://huggingface.co/internlm/internlm-xcomposer2-vl-7b/resolve/main/image1.webp"
-            images = load_images(img_url)
-        elif "internvl" in self.model_type:
-            if image_path is None:
-                img_url = 'https://huggingface.co/OpenGVLab/InternVL2-4B/resolve/main/examples/image1.jpg'
-                images = Image.open(
-                    requests.get(img_url, stream=True,
-                                 timeout=5).raw).convert('RGB')
-            else:
-                images = Image.open(image_path).convert('RGB')
-        elif "qwen2_vl" in self.model_type:
-            images = []
-            if self.args.image_path is None:
-                img_url = 'https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg'
-                image = Image.open(
-                    requests.get(img_url, stream=True,
-                                 timeout=5).raw).convert('RGB')
-                image = image.resize((504, 504))
-                images.append(image)
-            else:
-                images = []
-                for image_path in self.args.image_path:
-                    image = Image.open(image_path).convert('RGB')
-                    image = image.resize((504, 504))
-                    images.append(image)
-        elif "llava_onevision" in self.model_type and self.args.video_path is not None:
-            if self.args.video_path == 'llava-onevision-accuracy':
-                self.args.video_path = hf_hub_download(
-                    repo_id="raushan-testing-hf/videos-test",
-                    filename="sample_demo_1.mp4",
-                    repo_type="dataset")
-            import av
-            with av.open(self.args.video_path) as container:
-                total_frames = container.streams.video[0].frames
-                assert total_frames >= self.num_frames
-                indices = np.arange(0, total_frames,
-                                    total_frames / self.num_frames).astype(int)
-                frames = []
-                container.seek(0)
-                start_index = indices[0]
-                end_index = indices[-1]
-                for i, frame in enumerate(container.decode(video=0)):
-                    if i > end_index:
-                        break
-                    if i >= start_index and i in indices:
-                        frames.append(frame)
-                images = np.stack(
-                    [x.to_ndarray(format="rgb24") for x in frames])
-            images = torch.tensor(images)
-        else:
-            if self.model_type != 'mllama':
-                if image_path is None:
-                    if self.model_type == "llava":
-                        image_path = [
-                            'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
-                        ] * 8
-                        image_path = image_path[:self.args.batch_size]
-                        self.args.image_path = ",".join(image_path)
-                    else:
-                        image_path = 'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
-                else:
-                    if isinstance(image_path, str):
-                        image_path = image_path.split(self.args.path_sep)
-
-            images = load_images(image_path) if image_path is not None else None
-        return images
-
-    def load_test_audio(self, audio_path):
-        if self.model_type != "phi-4-multimodal":
-            return None
-
-        assert audio_path is not None
-        import soundfile
-        audio = soundfile.read(audio_path)
-        return audio
-
-    def setup_inputs(self, input_text, raw_image, raw_audio=None):
-        from ..tools.multimodal_builder import compute_rotary_pos_emb
-        other_vision_inputs = {}
-        other_audio_inputs = {}
-        other_decoder_inputs = {}
-        if self.model_type not in ['qwen2_vl', 'vila', 'llava']:
-            input_text = input_text[0] if isinstance(input_text,
-                                                     list) else input_text
-
-        if 'blip2' in self.model_type:
-            image = self.processor(raw_image, input_text,
-                                   return_tensors="pt")['pixel_values']
-            if input_text is None:
-                input_text = "Question: which city is this? Answer:"
-            pre_prompt = input_text
-            post_prompt = None
-        elif 'internlm' in self.model_type:
-            #Feed the raw image into vis_processor, to get processed image
-            image = self.processor(raw_image).unsqueeze(0).cuda()
-            if input_text is None:
-                input_text = "Please describe this image in detail."
-            pre_prompt = ''
-            meta_instruction = 'You are an AI assistant whose name is InternLM-XComposer (浦语·灵笔).\n'
-            '- InternLM-XComposer (浦语·灵笔) is a multi-modality conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n'
-            '- InternLM-XComposer (浦语·灵笔) can understand and communicate fluently in the language chosen by the user such as English and 中文.\n'
-            '- InternLM-XComposer (浦语·灵笔) is capable of comprehending and articulating responses effectively based on the provided image.',
-            pre_prompt += f"""[UNUSED_TOKEN_146]system\n{meta_instruction}[UNUSED_TOKEN_145]\n"""
-            pre_prompt += f"""[UNUSED_TOKEN_146]user\n"""
-            post_prompt = f"""{input_text}[UNUSED_TOKEN_145]\n[UNUSED_TOKEN_146]assistant\n"""
-        elif 'qwen2_vl' in self.model_type:
-            from qwen_vl_utils import process_vision_info
-            from transformers.models.qwen2_vl.modeling_qwen2_vl import \
-                VisionRotaryEmbedding
-            hf_config = AutoConfig.from_pretrained(self.args.hf_model_dir)
-            if input_text is None:
-                input_text = ["Question: Describe this image. Answer:"
-                              ] * self.args.batch_size
-            messages = [[{
-                "role":
-                "user",
-                "content": [
-                    {
-                        "type": "image",
-                        "image": raw_image[idx],
-                    },
-                    {
-                        "type": "text",
-                        "text": input_text[idx],
-                    },
-                ],
-            }] for idx in range(self.args.batch_size)]
-
-            texts = [
-                self.processor.apply_chat_template(msg,
-                                                   tokenize=False,
-                                                   add_generation_prompt=True)
-                for msg in messages
-            ]
-            image_inputs, video_inputs = process_vision_info(messages)
-            inputs = self.processor(
-                text=texts,
-                images=image_inputs,
-                videos=video_inputs,
-                padding=True,
-                return_tensors="pt",
-            )
-            inputs = inputs.to(self.device)
-            image = inputs['pixel_values']
-            image_grid_thw = inputs['image_grid_thw']
-            input_ids = inputs['input_ids']
-            attention_mask = inputs['attention_mask']
-            cu_seqlens = torch.repeat_interleave(
-                image_grid_thw[:, 1] * image_grid_thw[:, 2],
-                image_grid_thw[:, 0]).cumsum(dim=0, dtype=torch.int32)
-            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
-
-            seq_length = image.shape[0]
-            # Create block indices using bucketing
-            block_indices = torch.bucketize(torch.arange(seq_length,
-                                                         device=image.device),
-                                            cu_seqlens,
-                                            right=True) - 1
-
-            # Generate block diagonal mask using matrix expansion
-            attention_mask_vit = torch.where(
-                block_indices.view(-1, 1) == block_indices.view(1, -1),
-                torch.zeros((), device=image.device, dtype=image.dtype),
-                torch.full((),
-                           torch.finfo(torch.float16).min,
-                           device=image.device,
-                           dtype=image.dtype)).unsqueeze(0)
-
-            decoder_input_ids = None
-            post_prompt = None
-            pre_prompt = None
-            images_qwenvl = {
-                "image": image,
-                "input_ids": input_ids,
-            }
-            rotary_pos_emb = compute_rotary_pos_emb(
-                image_grid_thw, hf_config, VisionRotaryEmbedding).to("cuda")
-            other_vision_inputs['attention_mask_llm'] = attention_mask
-            other_vision_inputs['image_grid_thw'] = image_grid_thw
-            other_vision_inputs['attention_mask'] = attention_mask_vit
-            other_vision_inputs['rotary_pos_emb'] = rotary_pos_emb
-            return input_text, pre_prompt, post_prompt, images_qwenvl, decoder_input_ids, other_vision_inputs, other_audio_inputs, other_decoder_inputs
-        elif 'nougat' in self.model_type:
-            image = self.processor(raw_image,
-                                   return_tensors="pt")['pixel_values']
-            # Nougat doesn't need text prompt (mBART use single token to start generation), just leave a dummy one here
-            if input_text is None:
-                input_text = "Question: which city is this? Answer:"
-            pre_prompt = input_text
-            post_prompt = None
-
-        elif 'cogvlm' in self.model_type:
-            image = self.transform(raw_image).unsqueeze(0)
-            if input_text is None:
-                input_text = " [INST] which city is this? [/INST] "
-            pre_prompt = input_text
-            post_prompt = None
-
-        elif 'phi-3-vision' in self.model_type:
-            pre_prompt = "<|user|>\n<|image_1|>\n"
-            if input_text is None:
-                input_text = "Which city is this?"
-            post_prompt = input_text + "<|end|>\n<|assistant|>\n"
-            prompt = pre_prompt + post_prompt
-            image = self.processor(text=prompt,
-                                   images=raw_image,
-                                   return_tensors="pt")
-        elif 'phi-4-multimodal' in self.model_type:
-            pre_prompt = "<|user|><|image_1|><|audio_1|>"
-            post_prompt = "<|end|><|assistant|>"
-            prompt = pre_prompt + post_prompt
-            image = self.processor(text=prompt,
-                                   images=[raw_image],
-                                   audios=[raw_audio],
-                                   return_tensors="pt")
-
-        elif 'pixtral' in self.model_type:
-            # Send image and text prompt to processor.
-            pre_prompt = "[INST][IMG]"
-            if input_text is None:
-                input_text = "What is in the image?"
-            post_prompt = "[/INST]"
-            prompt = pre_prompt + input_text + post_prompt
-            dtype = str_dtype_to_torch(self.vision_precision)
-            image = {'pixel_values': [], 'input_ids': []}
-            for img_idx in range(self.args.batch_size):
-                image_info = self.processor(text=prompt,
-                                            images=[raw_image[img_idx]],
-                                            return_tensors="pt").to(dtype)
-                image['pixel_values'].append(image_info['pixel_values'].to(
-                    self.device))
-                image['input_ids'].append(image_info['input_ids'][0].to(
-                    self.device))
-
-        elif 'internvl' in self.model_type:
-            pre_prompt = "<|system|>\n你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。<|end|><|user|>\n\n"
-            if input_text is None:
-                input_text = "Please describe the image shortly."
-            post_prompt = input_text + "<|end|><|assistant|>\n"
-            prompt = pre_prompt + post_prompt
-            image = self.processor(images=raw_image,
-                                   return_tensors='pt').pixel_values
-
-        elif self.model_type == "pix2struct":
-            if input_text is None:
-                input_text = ""
-            inputs = self.processor(
-                images=raw_image,
-                text=input_text,
-                return_tensors="pt",
-            )
-            image = inputs['flattened_patches']
-            image = image.expand(self.args.batch_size, -1, -1).contiguous()
-            pre_prompt = ""
-            post_prompt = None
-
-        elif self.model_type == "neva":
-            image = self.transform(raw_image).unsqueeze(0)
-
-            if input_text is None:
-                input_text = "Hi! What is in this image?"
-
-            pre_prompt = "System\n\nUser\n"
-            post_prompt = f"\n{input_text}\nAssistant\n"
-
-        elif self.model_type == "video-neva":
-
-            image = self.video_preprocess(
-                raw_image)  # shape (1, num_frames, 3, H, W)
-
-            if input_text is None:
-                input_text = "Hi! What is in this video?"
-
-            # SteerLM prompt template
-            pre_prompt = """System\nA chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n\nUser"""
-            post_prompt = f"\n{input_text}\nAssistant\nquality:4,toxicity:0,humor:0,creativity:0,helpfulness:4,correctness:4,coherence:4,complexity:4,verbosity:4\n" ""
-
-        elif self.model_type == "llava_next":
-            if self.llm_name == "mistralai/Mistral-7B-Instruct-v0.2":
-                pre_prompt = "[INST] "
-                if input_text is None:
-                    input_text = "Question: which city is this? Answer:"
-                post_prompt = f"\n{input_text} [/INST]"
-                prompt = pre_prompt + post_prompt
-
-            elif self.llm_name == "NousResearch/Nous-Hermes-2-Yi-34B":
-                pre_prompt = "<|im_start|>system\nAnswer the questions.<|im_end|><|im_start|>user\n"
-                if input_text is None:
-                    input_text = "Question: which city is this? Answer:"
-                post_prompt = f"\n{input_text}<|im_end|><|im_start|>assistant\n"
-                prompt = pre_prompt + post_prompt
-
-            else:
-                raise Exception(
-                    f"Prompt template for {self.llm_name} for not included currently"
-                )
-
-            image = self.processor(text=prompt,
-                                   images=raw_image,
-                                   return_tensors="pt")
-
-        elif self.model_type == 'vila':
-            if input_text is None:
-                input_text = "\n Please elaborate what you see in the images?"
-            if '8b' in self.args.hf_model_dir.lower():
-                pre_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"
-                post_prompt = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            elif '40b' in self.args.hf_model_dir.lower():
-                pre_prompt = "<|im_start|>system\nAnswer the questions.<|im_end|><|im_start|>user\n"
-                post_prompt = "<|im_end|><|im_start|>assistant\n"
-            else:
-                pre_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: "
-                post_prompt = " ASSISTANT:"
-            if isinstance(input_text, list):
-                post_prompt = [input + post_prompt for input in input_text]
-            else:
-                post_prompt = input_text + post_prompt
-            if not isinstance(raw_image, list):
-                raw_image = [raw_image]
-            image = self.processor(raw_image)
-
-        elif self.model_type in ['llava', 'fuyu', 'kosmos-2']:
-            if self.model_type == "llava":
-                pre_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: "
-                post_prompt = " ASSISTANT:"
-                if input_text is None:
-                    input_text = "\n Which city is this? Answer:"
-                if isinstance(input_text, list):
-                    post_prompt = [input + post_prompt for input in input_text]
-                else:
-                    post_prompt = input_text + post_prompt
-            elif self.model_type == 'fuyu':
-                pre_prompt = "Describe this image:"
-                post_prompt = None
-                if input_text is None:
-                    input_text = "Answer the following VQAv2 question based on the image: How many people are in the image?\n"
-            elif self.model_type == "kosmos-2":
-                pre_prompt = ""
-                post_prompt = None
-                if input_text is None:
-                    input_text = "An image of"
-
-            if self.model_type not in ['fuyu', 'kosmos-2']:
-                post_prompt = " ASSISTANT:"
-                if isinstance(input_text, list):
-                    post_prompt = [input + post_prompt for input in input_text]
-                else:
-                    post_prompt = input_text + post_prompt
-            else:
-                post_prompt = None
-            if self.model_type in ['fuyu', 'kosmos-2']:
-                image = self.processor(text=input_text,
-                                       images=raw_image,
-                                       return_tensors='pt')
-            else:
-                if isinstance(input_text, list):
-                    image = self.processor(text=input_text[0],
-                                           images=raw_image,
-                                           padding=True,
-                                           return_tensors="pt")['pixel_values']
-                else:
-                    image = self.processor(text=input_text,
-                                           images=raw_image,
-                                           return_tensors="pt")['pixel_values']
-
-        elif self.model_type in ['mllama']:
-            if raw_image is not None:
-                input_text = self.processor.apply_chat_template(
-                    images=raw_image, text=input_text)
-                inputs = self.processor(images=raw_image,
-                                        text=input_text,
-                                        return_tensors="pt")
-                other_vision_inputs = {
-                    "aspect_ratio_ids":
-                    inputs["aspect_ratio_ids"].to(self.device).expand(
-                        self.args.batch_size, -1).contiguous(),
-                    "aspect_ratio_mask":
-                    inputs["aspect_ratio_mask"].to(self.device).expand(
-                        self.args.batch_size, -1, -1).contiguous(),
-                }
-                other_decoder_inputs = {
-                    "cross_attention_mask":
-                    inputs["cross_attention_mask"].to(self.device).expand(
-                        self.args.batch_size, -1, -1, -1).contiguous(),
-                }
-                pre_prompt = input_text
-                post_prompt = None
-                image = inputs["pixel_values"]
-            else:
-                pre_prompt = input_text
-                post_prompt = None
-                image = None
-                logger.warning(
-                    "image_path is None. Will not pass image as input, skipping the vision encoder."
-                )
-                image = None
-        elif self.model_type in ['llava_onevision']:
-            pre_prompt = "<|im_start|>user " + "